From 22fa82f0bf0c69d5ee710bc605b20c64c82451e0 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 16 May 2026 01:46:17 +0800 Subject: [PATCH 01/78] 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 edcda233fcfe419fc4a6485852420c2fcf831d2c Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 16 May 2026 02:29:31 +0800 Subject: [PATCH 02/78] 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 492f0db75e702b7759bb6c7e134429ac6e819b33 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 16 May 2026 02:36:34 +0800 Subject: [PATCH 03/78] 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 288126368e5359a4ddf800195365aa9c8ae06281 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 16 May 2026 02:46:45 +0800 Subject: [PATCH 04/78] 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 917fb75ce627260eebef554a6c2e8a04e6e3b7c1 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 16 May 2026 02:50:26 +0800 Subject: [PATCH 05/78] 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 f1fef67e5abbcd1fe96f189cf93cddf1272ff144 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 16 May 2026 13:58:44 +0800 Subject: [PATCH 06/78] 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 69143f879974651396e12ef59d43dd18e9718d71 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 16 May 2026 14:13:53 +0800 Subject: [PATCH 07/78] 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 7c47d11803e1526ae1f91ece95e6d1fa937e9de0 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 16 May 2026 14:28:37 +0800 Subject: [PATCH 08/78] 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 25556d496a4261e84c5661d9be04eca99dd622c4 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 16 May 2026 15:20:17 +0800 Subject: [PATCH 09/78] 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 85ba461f4c1450a1ef15292601a48214dba9fe59 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 16 May 2026 15:50:19 +0800 Subject: [PATCH 10/78] 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 4d970c117dd7f5ce1fb540f875df6c17c76c1225 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 16 May 2026 16:32:59 +0800 Subject: [PATCH 11/78] 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 1d8f88f96d31a49d3b7a2f2ca14af4202c3a7ade Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sat, 16 May 2026 21:16:23 +0800 Subject: [PATCH 12/78] 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 459b5c8e739b593500dec79ba124c537bf337c9b Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 01:30:30 +0800 Subject: [PATCH 13/78] 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 53811a7f64162e856e21618092551c62b2558048 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 01:37:42 +0800 Subject: [PATCH 14/78] 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 a4be52bc8b39aeded35e4812599e9249a7732f23 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 01:48:29 +0800 Subject: [PATCH 15/78] 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 8c2985c8691898d106775f28224583dc58bfa835 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 02:28:51 +0800 Subject: [PATCH 16/78] 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 857351a03638c7d5d80fd54acbf8818fb8d497ce Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 02:43:41 +0800 Subject: [PATCH 17/78] 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 b15bf5097faae25e81672c9cb85d2ed927cfa595 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 02:48:29 +0800 Subject: [PATCH 18/78] 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 080821cbade12f6090bf9459784e0b9105382da7 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 13:18:55 +0800 Subject: [PATCH 19/78] 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 d6f737b111cd333caf3ddb6de292dfa9e2ba7ed9 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 13:50:30 +0800 Subject: [PATCH 20/78] 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 2e291032df6c8251f0372f01cb6de5ad8ff956ca Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 14:21:40 +0800 Subject: [PATCH 21/78] 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 f8c16e914ab40a6ff37e260a8380af09c54e0b60 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 14:48:49 +0800 Subject: [PATCH 22/78] 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 b3e51da818af1e158be6eea96e299604fa5ab66a Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 16:08:57 +0800 Subject: [PATCH 23/78] 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 fa8f553bdedc3d77acaf39c3f5c8fef1b89f221d Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 16:31:34 +0800 Subject: [PATCH 24/78] 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 f05e0e716b4c60707457de871774fbd46f7ed52e Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 16:42:02 +0800 Subject: [PATCH 25/78] 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 e00679af7f6d09b14f012837b146ed822d4d369f Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 16:54:13 +0800 Subject: [PATCH 26/78] 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 0c73c9c6173198f1a4094db103294c4845e3bda3 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 17:11:29 +0800 Subject: [PATCH 27/78] 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 d740fb12109398e2451835d2866aa2d4d68e6455 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 17:28:45 +0800 Subject: [PATCH 28/78] 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 dad7e460bdf1468320811a48f595ce79824bb30f Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 18:11:08 +0800 Subject: [PATCH 29/78] 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 defaff8acb640a2380890d3cf40f43aa890d0f4d Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 18:16:16 +0800 Subject: [PATCH 30/78] 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 4641a3518730f1cb97dcaf2f17aa36f61516e360 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 18:21:31 +0800 Subject: [PATCH 31/78] 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 ab803cfb7a33786b780c455696bb528669cb0b4d Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 18:25:04 +0800 Subject: [PATCH 32/78] 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 3dbfa96b0d1101a17c11e3a92b281390cf4b9595 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 18:29:18 +0800 Subject: [PATCH 33/78] 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 06cf7d635659c33b9daca7618876d15c979f09e2 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 18:33:24 +0800 Subject: [PATCH 34/78] 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 76c03168b14593b5471c7b87f296fda202e3825a Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 19:02:47 +0800 Subject: [PATCH 35/78] 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 0a6e55c55623a42ccb8cc99478636ecd51c4cf7e Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 22:55:21 +0800 Subject: [PATCH 36/78] 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 d146c6452f955d5e311e8c3bdb425878e0d07e20 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 16:02:59 +0000 Subject: [PATCH 37/78] 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 5dbf04824d4da5ac55d21934596b69950b621fad Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 16:55:17 +0000 Subject: [PATCH 38/78] 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 e4462ebef97b47fb9002c9112797e130a5341c36 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 17:34:10 +0000 Subject: [PATCH 39/78] 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 a6c99d735a340361201cd3f44c5125cbc22b3379 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 18:51:03 +0000 Subject: [PATCH 40/78] 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}") From d6de8c537355797de17a548d4df0a4e2c03414e3 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 01:52:03 +0000 Subject: [PATCH 41/78] refactor: deepen workflow architecture --- apps/api/app/api/v1/routes/s3_events.py | 3 +- .../app/services/s3_events/event_handlers.py | 20 +-- .../app/services/s3_events/intake_outcome.py | 43 +++++ .../s3_events/nested_message_decoder.py | 39 +++++ apps/api/app/services/s3_events/service.py | 7 +- .../tests/contract/test_retrieval_contract.py | 108 +++++++++++++ .../tests/contract/test_s3_event_contract.py | 34 ++++ .../orchestration/parse_pipeline.py | 27 ++++ .../services/document_parser/parse_service.py | 9 +- ...t_document_parser_architecture_contract.py | 57 +++++++ .../contract/test_parse_task_contract.py | 28 +++- .../shared/services/job_lifecycle_sync.py | 125 +++++++------- .../shared/services/job_result_sync.py | 10 +- .../shared/services/retrieval/app_service.py | 41 ++--- .../services/retrieval/execution_plan.py | 153 ++++++------------ .../services/retrieval/execution_routes.py | 35 +--- .../services/retrieval/query_request.py | 113 +++++++++++++ .../services/retrieval/reference_resolver.py | 100 ++++++++++++ 18 files changed, 718 insertions(+), 234 deletions(-) create mode 100644 apps/api/app/services/s3_events/intake_outcome.py create mode 100644 apps/api/app/services/s3_events/nested_message_decoder.py create mode 100644 apps/worker/app/services/document_parser/orchestration/parse_pipeline.py create mode 100644 packages/shared-python/shared/services/retrieval/query_request.py create mode 100644 packages/shared-python/shared/services/retrieval/reference_resolver.py diff --git a/apps/api/app/api/v1/routes/s3_events.py b/apps/api/app/api/v1/routes/s3_events.py index 56d4bdf58..ea3feff24 100644 --- a/apps/api/app/api/v1/routes/s3_events.py +++ b/apps/api/app/api/v1/routes/s3_events.py @@ -1,6 +1,7 @@ """S3 event webhook routes.""" from app.services.s3_events.service import safely_handle_s3_event_post +from app.services.s3_events.intake_outcome import sanitize_storage_event_headers from fastapi import APIRouter, Header, Request from loguru import logger @@ -40,7 +41,7 @@ async def handle_s3_events( ) -> 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)}" + f"S3 event Headers: {sanitize_storage_event_headers(dict(request.headers))}" ) if request.client: logger.info(f"Client IP: {request.client.host}") diff --git a/apps/api/app/services/s3_events/event_handlers.py b/apps/api/app/services/s3_events/event_handlers.py index 3d146954a..b3408012d 100644 --- a/apps/api/app/services/s3_events/event_handlers.py +++ b/apps/api/app/services/s3_events/event_handlers.py @@ -1,11 +1,11 @@ """Storage event protocol handlers.""" from __future__ import annotations -import base64 import json import os from typing import Any +from app.services.s3_events.nested_message_decoder import decode_nested_json_message from app.services.s3_events.signature_verification import ( verify_minio_signature, verify_oss_signature, @@ -169,24 +169,8 @@ def _unwrap_mns_message(event_data: dict[str, Any]) -> dict[str, Any]: if not isinstance(inner, str): return event_data - decoded = _decode_mns_inner_json(inner) + decoded = decode_nested_json_message(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/intake_outcome.py b/apps/api/app/services/s3_events/intake_outcome.py new file mode 100644 index 000000000..24a02189f --- /dev/null +++ b/apps/api/app/services/s3_events/intake_outcome.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +_SENSITIVE_HEADER_NAMES: frozenset[str] = frozenset( + { + "authorization", + "x-amz-sns-signature", + "x-amz-sns-signing-cert-url", + "x-minio-auth-token", + "x-oss-signature", + "x-oss-pub-key-url", + "x-mns-signature", + "x-mns-signing-cert-url", + } +) + + +@dataclass(frozen=True) +class StorageEventIntakeOutcome: + message: str + reason: str + + def to_response(self) -> dict[str, str]: + return {"message": self.message} + + +def build_storage_event_error_acknowledgement() -> StorageEventIntakeOutcome: + return StorageEventIntakeOutcome( + message="Event handling completed", + reason="acked_after_handler_error", + ) + + +def sanitize_storage_event_headers(headers: dict[str, str]) -> dict[str, str]: + sanitized: dict[str, str] = {} + for name, value in headers.items(): + if name.lower() in _SENSITIVE_HEADER_NAMES: + sanitized[name] = "" + else: + sanitized[name] = value + return sanitized diff --git a/apps/api/app/services/s3_events/nested_message_decoder.py b/apps/api/app/services/s3_events/nested_message_decoder.py new file mode 100644 index 000000000..eb3e29c0f --- /dev/null +++ b/apps/api/app/services/s3_events/nested_message_decoder.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +import base64 +import binascii +import json +from collections.abc import Callable +from typing import Any + +from loguru import logger + + +def decode_nested_json_message(inner: str) -> dict[str, Any] | None: + for decoder in (_decode_base64_json, _decode_plain_json): + decoded = _try_decode(inner, decoder) + if decoded is not None: + return decoded + return None + + +def _try_decode( + inner: str, + decoder: Callable[[str], object], +) -> dict[str, Any] | None: + try: + decoded = decoder(inner) + except (binascii.Error, json.JSONDecodeError, UnicodeDecodeError, ValueError) as exc: + logger.debug("Nested storage event message decode failed: {}", exc) + return None + return decoded if isinstance(decoded, dict) else None + + +def _decode_base64_json(inner: str) -> object: + decoded_bytes = base64.b64decode(inner, validate=True) + decoded_str = decoded_bytes.decode("utf-8") + return json.loads(decoded_str) + + +def _decode_plain_json(inner: str) -> object: + return json.loads(inner) diff --git a/apps/api/app/services/s3_events/service.py b/apps/api/app/services/s3_events/service.py index e06acedc3..9382c67bb 100644 --- a/apps/api/app/services/s3_events/service.py +++ b/apps/api/app/services/s3_events/service.py @@ -10,6 +10,9 @@ handle_sns_event, is_oss_event, ) +from app.services.s3_events.intake_outcome import ( + build_storage_event_error_acknowledgement, +) async def handle_s3_event_post( @@ -48,5 +51,7 @@ async def safely_handle_s3_event_post( minio_auth_token=minio_auth_token, ) except Exception as exc: + outcome = build_storage_event_error_acknowledgement() logger.error(f"Failed to handle S3 event: {exc}") - return {"message": "Event handling completed"} + logger.warning(f"S3 event intake outcome: reason={outcome.reason}") + return outcome.to_response() diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index 51c4ff48b..d3eecb67c 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -7,9 +7,11 @@ import pytest from httpx import AsyncClient from pytest import MonkeyPatch +from sqlalchemy.ext.asyncio import AsyncSession from tests.support.contract_database import ContractDatabase from shared.services.retrieval.agentic.types import AgenticResult +from shared.services.retrieval.query_request import RetrievalQuery from shared.services.retrieval.workflow.types import PlannedStep, QueryPlan, WorkflowResult @@ -139,6 +141,46 @@ def _result_source(result: dict[str, object]) -> dict[str, object]: return cast(dict[str, object], result["source"]) +def test_retrieval_query_owns_cache_and_route_request_shape() -> None: + query = RetrievalQuery.from_parameters( + db=cast(AsyncSession, object()), + user_id="user-1", + namespace="contract-retrieval", + query=" alpha ", + top_k=5, + exclude_document_ids=["doc-hidden"], + exclude_sections=[{"document_id": "doc-1", "section_path": "Root"}], + data_type=2, + signal_paths=["Root"], + filter_mode="keep", + channels=["content"], + channel_weights={"content": 1.5}, + rerank=True, + threshold=0.2, + internal_recall_k=17, + use_agentic=False, + ) + + route_context = query.build_route_context() + + assert query.query == "alpha" + assert query.build_cache_extra() == { + "data_type": 2, + "signal_paths": ["Root"], + "filter_mode": "keep", + "channels": ["content"], + "channel_weights": {"content": 1.5}, + "rerank": True, + "threshold": 0.2, + "internal_recall_k": 17, + "decomposition_enabled": True, + } + assert route_context.query == "alpha" + assert route_context.allowed_chunk_types == {"text"} + assert route_context.effective_recall_k == 17 + assert route_context.use_agentic is False + + @pytest.mark.asyncio async def test_should_return_seeded_retrieval_results_for_the_authenticated_user( developer_api_client_factory: Callable[ @@ -632,6 +674,72 @@ async def run(self, *_args: object, **kwargs: object) -> WorkflowResult: assert results == [] +@pytest.mark.asyncio +async def test_agentic_retrieval_should_drop_references_that_do_not_match_the_hydrated_section( + 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="mismatched section answer", + referenced_chunks=[ + { + "chunk_id": visible_chunk["chunk_id"], + "document_id": visible_chunk["document_id"], + "chunk_type": "text", + "section_path": "wrong/section", + "file_path": None, + "job_id": visible_chunk["job_id"], + } + ], + ) + + async with developer_api_client_factory() as api_client: + visible_chunk = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-reference-section-match", + source_file_name="visible.pdf", + section_path="right/section", + content="visible scoped content", + ) + await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-reference-section-match", + source_file_name="filler.pdf", + section_path="filler/section", + content="filler content", + ) + monkeypatch.setattr( + "shared.services.retrieval.workflow.orchestrator.WorkflowOrchestrator", + FakeWorkflowOrchestrator, + ) + + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": "contract-reference-section-match", + "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 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[ diff --git a/apps/api/tests/contract/test_s3_event_contract.py b/apps/api/tests/contract/test_s3_event_contract.py index b8e8561de..81ad263e9 100644 --- a/apps/api/tests/contract/test_s3_event_contract.py +++ b/apps/api/tests/contract/test_s3_event_contract.py @@ -1,4 +1,5 @@ import importlib +import base64 import json import socket from collections.abc import Callable @@ -48,6 +49,39 @@ async def _insert_waiting_file_job() -> tuple[str, str]: return user_id, job_id +def test_storage_event_intake_sanitizes_sensitive_headers() -> None: + from app.services.s3_events.intake_outcome import sanitize_storage_event_headers + + sanitized = sanitize_storage_event_headers( + { + "Authorization": "Bearer secret", + "x-amz-sns-signature": "signature", + "x-minio-auth-token": "token", + "content-type": "application/json", + } + ) + + assert sanitized == { + "Authorization": "", + "x-amz-sns-signature": "", + "x-minio-auth-token": "", + "content-type": "application/json", + } + + +def test_storage_event_nested_message_decoder_supports_plain_and_base64_json() -> None: + from app.services.s3_events.nested_message_decoder import decode_nested_json_message + + payload = {"Records": [{"eventName": "ObjectCreated:Put"}]} + encoded_payload = base64.b64encode(json.dumps(payload).encode("utf-8")).decode( + "utf-8" + ) + + assert decode_nested_json_message(json.dumps(payload)) == payload + assert decode_nested_json_message(encoded_payload) == payload + assert decode_nested_json_message("not-json") is None + + @pytest.mark.asyncio async def test_should_acknowledge_an_sns_subscription_confirmation_request( api_client_factory: Callable[[], AbstractAsyncContextManager[AsyncClient]], diff --git a/apps/worker/app/services/document_parser/orchestration/parse_pipeline.py b/apps/worker/app/services/document_parser/orchestration/parse_pipeline.py new file mode 100644 index 000000000..b31d36713 --- /dev/null +++ b/apps/worker/app/services/document_parser/orchestration/parse_pipeline.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import pandas as pd + +from app.services.document_parser.orchestration.parse_input import ParseInput +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 + + +@dataclass(frozen=True) +class ParsePipelineResult: + output_dir: str + parsed_df: pd.DataFrame | None + + def as_legacy_tuple(self) -> tuple[str, pd.DataFrame | None]: + return self.output_dir, self.parsed_df + + +def run_parse_pipeline(parse_input: ParseInput) -> ParsePipelineResult: + """Run parser session building, format routing, and output postprocessing.""" + session = build_parse_session(parse_input) + output_dir, parsed_df = route_document_parse(session) + processed_df = apply_parse_postprocess(output_dir, parsed_df) + return ParsePipelineResult(output_dir=output_dir, parsed_df=processed_df) diff --git a/apps/worker/app/services/document_parser/parse_service.py b/apps/worker/app/services/document_parser/parse_service.py index 000cef7bf..becfcd807 100644 --- a/apps/worker/app/services/document_parser/parse_service.py +++ b/apps/worker/app/services/document_parser/parse_service.py @@ -3,9 +3,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 +from app.services.document_parser.orchestration.parse_pipeline import run_parse_pipeline def checkerboard_inject_parse( @@ -49,7 +47,4 @@ def checkerboard_inject_parse( fragment_content=fragment_content, s3_key=s3_key, ) - 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 + return run_parse_pipeline(parse_input).as_legacy_tuple() 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 f0d5d381a..35ab7e531 100644 --- a/apps/worker/tests/contract/test_document_parser_architecture_contract.py +++ b/apps/worker/tests/contract/test_document_parser_architecture_contract.py @@ -115,6 +115,63 @@ def test_document_format_router_uses_adapters( assert actual_df is parsed_df +def test_parse_pipeline_owns_route_and_postprocess_contract( + worker_contract_environment: None, + monkeypatch: Any, + tmp_path: Path, +) -> None: + from app.services.document_parser.orchestration.parse_input import ParseInput + from app.services.document_parser.orchestration.parse_pipeline import ( + ParsePipelineResult, + run_parse_pipeline, + ) + + routed_df = pd.DataFrame([{"content": "routed"}]) + processed_df = pd.DataFrame([{"content": "processed"}]) + routed_output_dir = str(tmp_path / "routed") + + monkeypatch.setattr( + "app.services.document_parser.orchestration.parse_session.profile_document", + lambda *_args, **_kwargs: SimpleNamespace( + file_type="pdf", + page_count=1, + atlas_candidate=False, + doc_category="generic", + summary=lambda: "profile", + reasoning="test", + ), + ) + monkeypatch.setattr( + "app.services.document_parser.orchestration.parse_pipeline.route_document_parse", + lambda _session: (routed_output_dir, routed_df), + ) + + def fake_postprocess(output_dir: str, parsed_df: pd.DataFrame | None) -> pd.DataFrame: + assert output_dir == routed_output_dir + assert parsed_df is routed_df + return processed_df + + monkeypatch.setattr( + "app.services.document_parser.orchestration.parse_pipeline.apply_parse_postprocess", + fake_postprocess, + ) + + result = run_parse_pipeline( + ParseInput( + file_full_path=str(tmp_path / "contract.pdf"), + filename="contract.pdf", + output_dir=str(tmp_path), + internal_output_filename="contract.pdf", + ) + ) + + assert result == ParsePipelineResult( + output_dir=routed_output_dir, + parsed_df=processed_df, + ) + assert result.as_legacy_tuple() == (routed_output_dir, processed_df) + + def test_rendered_pdf_transform_centralizes_temporary_pdf_cleanup( worker_contract_environment: None, monkeypatch: Any, diff --git a/apps/worker/tests/contract/test_parse_task_contract.py b/apps/worker/tests/contract/test_parse_task_contract.py index b4c3aae49..816e04a15 100644 --- a/apps/worker/tests/contract/test_parse_task_contract.py +++ b/apps/worker/tests/contract/test_parse_task_contract.py @@ -433,7 +433,7 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: connection.execute( text( """ - SELECT delivery_mode, result_s3_key, result_size, inline_payload + SELECT id, delivery_mode, result_s3_key, result_size, inline_payload FROM job_results WHERE job_id = :job_id """ @@ -443,6 +443,21 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: .mappings() .one() ) + job_chunks = list( + connection.execute( + text( + """ + SELECT chunk_type, text, path, sort_order + FROM job_chunks + WHERE job_result_id = :job_result_id + ORDER BY sort_order + """ + ), + {"job_result_id": job_result_row["id"]}, + ) + .mappings() + .all() + ) document_row = ( connection.execute( text( @@ -546,6 +561,17 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: assert job_result_row["result_s3_key"] == f"results/{job_id}.zip" assert job_result_row["result_size"] > 0 assert dict(job_result_row["inline_payload"])["checksum"] + assert [chunk["chunk_type"] for chunk in job_chunks] == [ + "text", + "text", + "image", + "table", + ] + assert job_chunks[0]["text"] == text_content_with_refs + assert job_chunks[0]["path"].endswith( + "公司研究/自主可控加强，寒武纪或迎来营收快速放量周期" + ) + assert job_chunks[0]["sort_order"] == 0 assert document_row["namespace"] == "worker-contract" assert document_row["status"] == "active" assert document_row["source_file_name"] == source_file_name diff --git a/packages/shared-python/shared/services/job_lifecycle_sync.py b/packages/shared-python/shared/services/job_lifecycle_sync.py index 9d79b4ce0..d6d1d54d8 100644 --- a/packages/shared-python/shared/services/job_lifecycle_sync.py +++ b/packages/shared-python/shared/services/job_lifecycle_sync.py @@ -10,9 +10,11 @@ from __future__ import annotations import time -from typing import Any, Dict, List, Optional +from collections.abc import Callable +from typing import Any, Dict, List, Optional, TypeVar from loguru import logger +from sqlalchemy.orm import Session from shared.core.database_sync import get_sync_db_context from shared.services.job_failure_sync import SyncJobFailureFinalizer @@ -58,34 +60,25 @@ def finalize_job_success( """ logger.info(f"Finalizing job success: job_id={job_id}") - with get_sync_db_context() as db: - try: - finalization = self._success_finalizer.finalize( - db, - job_id=job_id, - 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 finalization.response.get("status") != "success": - db.rollback() - return finalization.response - - db.commit() - logger.info(f"Job {job_id} success transaction committed") - - self._success_finalizer.run_post_commit_actions(finalization) - - return finalization.response - - except Exception as exc: - logger.error(f"Failed to finalize job success {job_id}: {exc}") - db.rollback() - raise + return _run_lifecycle_transaction( + job_id=job_id, + label="success", + finalize=lambda db: self._success_finalizer.finalize( + db, + job_id=job_id, + 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, + ), + should_commit=lambda finalization: finalization.response.get("status") + == "success", + build_response=lambda finalization: finalization.response, + run_after_commit=self._success_finalizer.run_post_commit_actions, + ) def finalize_job_failure( self, @@ -106,31 +99,23 @@ def finalize_job_failure( """ logger.info(f"Finalizing job failure: job_id={job_id}") - with get_sync_db_context() as db: - try: - transition_ok, webhook_event = self._failure_finalizer.finalize( - db, - job_id=job_id, - error_message=error_message, - error_code=error_code, - error_details=error_details, - should_refund=should_refund, - ) - if not transition_ok: - db.rollback() - return False - - db.commit() - logger.info(f"Job {job_id} failure transaction committed") - - self._failure_finalizer.enqueue_webhook_after_commit(webhook_event) - - return True - - except Exception as exc: - logger.error(f"Failed to finalize job failure {job_id}: {exc}") - db.rollback() - raise + return _run_lifecycle_transaction( + job_id=job_id, + label="failure", + finalize=lambda db: self._failure_finalizer.finalize( + db, + job_id=job_id, + error_message=error_message, + error_code=error_code, + error_details=error_details, + should_refund=should_refund, + ), + should_commit=lambda finalization: finalization[0], + build_response=lambda finalization: finalization[0], + run_after_commit=lambda finalization: ( + self._failure_finalizer.enqueue_webhook_after_commit(finalization[1]) + ), + ) def update_progress( self, @@ -175,3 +160,35 @@ def get_sync_job_lifecycle_service() -> SyncJobLifecycleService: if _lifecycle_service is None: _lifecycle_service = SyncJobLifecycleService() return _lifecycle_service + + +_FinalizationT = TypeVar("_FinalizationT") +_ResponseT = TypeVar("_ResponseT") + + +def _run_lifecycle_transaction( + *, + job_id: str, + label: str, + finalize: Callable[[Session], _FinalizationT], + should_commit: Callable[[_FinalizationT], bool], + build_response: Callable[[_FinalizationT], _ResponseT], + run_after_commit: Callable[[_FinalizationT], None], +) -> _ResponseT: + with get_sync_db_context() as db: + try: + finalization = finalize(db) + if not should_commit(finalization): + db.rollback() + return build_response(finalization) + + db.commit() + logger.info(f"Job {job_id} {label} transaction committed") + + run_after_commit(finalization) + return build_response(finalization) + + except Exception as exc: + logger.error(f"Failed to finalize job {label} {job_id}: {exc}") + db.rollback() + raise diff --git a/packages/shared-python/shared/services/job_result_sync.py b/packages/shared-python/shared/services/job_result_sync.py index 2c7b2f5aa..22983a1ec 100644 --- a/packages/shared-python/shared/services/job_result_sync.py +++ b/packages/shared-python/shared/services/job_result_sync.py @@ -61,13 +61,19 @@ def replace_chunks( for index, chunk in enumerate(chunks): chunk_identifier = chunk.get("chunk_id") or str(uuid4()) metadata = chunk.get("metadata") + chunk_text = chunk.get("text") or chunk.get("content") + chunk_path = ( + metadata.get("path") + if isinstance(metadata, dict) and metadata.get("path") + else chunk.get("path") + ) 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, + text=str(chunk_text) if chunk_text is not None else None, + path=str(chunk_path) if chunk_path is not None else None, chunk_metadata=metadata, sort_order=chunk.get("order", index), ) diff --git a/packages/shared-python/shared/services/retrieval/app_service.py b/packages/shared-python/shared/services/retrieval/app_service.py index 13c60625e..0644c478e 100644 --- a/packages/shared-python/shared/services/retrieval/app_service.py +++ b/packages/shared-python/shared/services/retrieval/app_service.py @@ -4,7 +4,9 @@ from sqlalchemy.ext.asyncio import AsyncSession -from shared.services.retrieval.execution_plan import RetrievalExecutionPlan +from shared.services.retrieval.execution_plan import ( + run_retrieval_query as execute_retrieval_query, +) from shared.services.retrieval.scoring import merge_channels_rrf __all__ = ["merge_channels_rrf", "run_retrieval_query"] @@ -29,22 +31,21 @@ async def run_retrieval_query( internal_recall_k: int | None = None, use_agentic: bool | None = None, ) -> dict[str, Any]: - request: dict[str, Any] = { - "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, - } - return await RetrievalExecutionPlan(request).execute() + return await execute_retrieval_query( + 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, + ) diff --git a/packages/shared-python/shared/services/retrieval/execution_plan.py b/packages/shared-python/shared/services/retrieval/execution_plan.py index 71669f3ed..d94cc8f94 100644 --- a/packages/shared-python/shared/services/retrieval/execution_plan.py +++ b/packages/shared-python/shared/services/retrieval/execution_plan.py @@ -11,17 +11,13 @@ set_cached_retrieval_query_result, ) 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, ) from shared.services.retrieval.response_projection import ( project_public_retrieval_response, ) -from shared.services.retrieval.settings import ( - INTERNAL_RECALL_K_MULTIPLIER, - resolve_allowed_chunk_types, -) +from shared.services.retrieval.query_request import RetrievalQuery async def run_retrieval_query( @@ -45,100 +41,71 @@ async def run_retrieval_query( ) -> dict[str, Any]: """Run retrieval through the plan module.""" 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, - } + RetrievalQuery.from_parameters( + 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: + def __init__(self, request: RetrievalQuery) -> None: self.request = request async def execute(self) -> dict[str, Any]: - 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"] + request = self.request 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, + query=request.query, + user_id=request.user_id, + namespace=request.namespace, + top_k=request.top_k, + data_type=request.data_type, + exclude_document_ids=request.exclude_document_ids, + exclude_sections=request.exclude_sections, ) - if not query: + if not request.query: logger.info(" ⛔ Empty query filtered, skipping retrieval pipeline") return { - "namespace": namespace, - "query": query, + "namespace": request.namespace, + "query": request.query, "router_used": "empty_query_filtered", "results": [], } - 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 - ) + allowed_chunk_types = request.resolve_allowed_chunk_types() + effective_recall_k = request.resolve_effective_recall_k() 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}" + f"signal_paths={request.signal_paths} filter_mode={request.filter_mode} " + f"rerank={request.rerank} threshold={request.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_extra = request.build_cache_extra() 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, + user_id=request.user_id, + namespace=request.namespace, + query=request.query, + top_k=request.top_k, + exclude_document_ids=request.exclude_document_ids, + exclude_sections=request.exclude_sections, cache_extra=cache_extra, ) if cached_response is not None: @@ -146,42 +113,24 @@ async def execute(self) -> dict[str, Any]: 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, - threshold=threshold, - effective_recall_k=effective_recall_k, - use_agentic=use_agentic, - ) - outcome = await run_retrieval_route(route_context) + outcome = await run_retrieval_route(request.build_route_context()) if cache_version is not None: await _write_cached_response( - user_id=user_id, - namespace=namespace, + user_id=request.user_id, + namespace=request.namespace, version=cache_version, - query=query, - top_k=top_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, + query=request.query, + top_k=request.top_k, + exclude_document_ids=request.exclude_document_ids, + exclude_sections=request.exclude_sections, response=outcome.response, cache_extra=cache_extra, ) _schedule_hit_stats_update( - user_id=user_id, - namespace=namespace, + user_id=request.user_id, + namespace=request.namespace, results=outcome.hit_stats_results, ) _log_retrieval_complete( diff --git a/packages/shared-python/shared/services/retrieval/execution_routes.py b/packages/shared-python/shared/services/retrieval/execution_routes.py index ee5516159..cfd5984c7 100644 --- a/packages/shared-python/shared/services/retrieval/execution_routes.py +++ b/packages/shared-python/shared/services/retrieval/execution_routes.py @@ -5,11 +5,10 @@ from loguru import logger 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.reference_resolver import resolve_workflow_references from shared.services.retrieval.result_assembly import assemble_retrieval_results from shared.services.retrieval.response_projection import ( attach_citation, - enrich_referenced_chunks_with_asset_urls, ) from shared.services.retrieval.route_types import ( RetrievalRouteContext, @@ -122,41 +121,21 @@ async def _run_agentic_route( 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( + resolved_references = await resolve_workflow_references( db=context.db, user_id=context.user_id, namespace=context.namespace, - refs=enriched_refs, + refs=workflow_result.referenced_chunks, ) - 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, + rows=resolved_references.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["referenced_chunks"] = resolved_references.refs response["results"] = [attach_citation(row) for row in assembled_workflow_rows] completion_detail = ( @@ -165,8 +144,8 @@ async def _run_agentic_route( ) return RetrievalRouteOutcome( response=response, - hit_stats_results=enriched_refs, + hit_stats_results=resolved_references.refs, completion_label="AGENTIC RETRIEVAL", - completion_count=len(enriched_refs), + completion_count=len(resolved_references.refs), completion_detail=completion_detail, ) diff --git a/packages/shared-python/shared/services/retrieval/query_request.py b/packages/shared-python/shared/services/retrieval/query_request.py new file mode 100644 index 000000000..1877ae363 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/query_request.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.services.retrieval.route_types import RetrievalRouteContext +from shared.services.retrieval.settings import ( + INTERNAL_RECALL_K_MULTIPLIER, + resolve_allowed_chunk_types, +) + + +@dataclass(frozen=True) +class RetrievalQuery: + 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 + + @classmethod + def from_parameters( + cls, + *, + 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, + ) -> "RetrievalQuery": + return cls( + db=db, + user_id=user_id, + namespace=namespace, + query=str(query).strip(), + 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, + ) + + def build_cache_extra(self) -> dict[str, Any]: + return { + "data_type": self.data_type, + "signal_paths": self.signal_paths, + "filter_mode": self.filter_mode, + "channels": self.channels, + "channel_weights": self.channel_weights, + "rerank": self.rerank, + "threshold": self.threshold, + "internal_recall_k": self.internal_recall_k, + "decomposition_enabled": True, + } + + def resolve_allowed_chunk_types(self) -> set[str] | None: + return resolve_allowed_chunk_types(self.data_type) + + def resolve_effective_recall_k(self) -> int: + if self.internal_recall_k is not None: + return self.internal_recall_k + return self.top_k * INTERNAL_RECALL_K_MULTIPLIER + + def build_route_context(self) -> RetrievalRouteContext: + return RetrievalRouteContext( + db=self.db, + user_id=self.user_id, + namespace=self.namespace, + query=self.query, + top_k=self.top_k, + exclude_document_ids=self.exclude_document_ids, + exclude_sections=self.exclude_sections, + allowed_chunk_types=self.resolve_allowed_chunk_types(), + data_type=self.data_type, + signal_paths=self.signal_paths, + filter_mode=self.filter_mode, + channels=self.channels, + channel_weights=self.channel_weights, + threshold=self.threshold, + effective_recall_k=self.resolve_effective_recall_k(), + use_agentic=self.use_agentic, + ) diff --git a/packages/shared-python/shared/services/retrieval/reference_resolver.py b/packages/shared-python/shared/services/retrieval/reference_resolver.py new file mode 100644 index 000000000..7845ef177 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/reference_resolver.py @@ -0,0 +1,100 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.services.retrieval.reference_hydration import hydrate_referenced_chunk_rows +from shared.services.retrieval.response_projection import ( + enrich_referenced_chunks_with_asset_urls, +) +from shared.services.retrieval.row_utils import build_reference_lookup_key + + +@dataclass(frozen=True) +class ResolvedWorkflowReferences: + refs: list[dict[str, Any]] + rows: list[dict[str, Any]] + + +async def resolve_workflow_references( + *, + db: AsyncSession, + user_id: str, + namespace: str, + refs: list[dict[str, Any]], +) -> ResolvedWorkflowReferences: + enriched_refs = await enrich_referenced_chunks_with_asset_urls(refs) + hydrated_rows = await hydrate_referenced_chunk_rows( + db=db, + user_id=user_id, + namespace=namespace, + refs=enriched_refs, + ) + return _select_matching_references(enriched_refs, hydrated_rows) + + +def _select_matching_references( + refs: list[dict[str, Any]], + rows: list[dict[str, Any]], +) -> ResolvedWorkflowReferences: + selected_refs: list[dict[str, Any]] = [] + selected_rows: list[dict[str, Any]] = [] + seen_row_keys: set[tuple[str, str, str, str]] = set() + + for ref in refs: + matching_row = next( + ( + row + for row in rows + if _matches_reference(ref, row) + and _row_key(row) not in seen_row_keys + ), + None, + ) + if matching_row is None: + continue + + selected_refs.append(ref) + selected_rows.append(matching_row) + seen_row_keys.add(_row_key(matching_row)) + + return ResolvedWorkflowReferences(refs=selected_refs, rows=selected_rows) + + +def _matches_reference(ref: dict[str, Any], row: dict[str, Any]) -> bool: + ref_key = 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"), + ) + row_key = _row_key(row) + if ref_key[:2] != row_key[:2]: + return False + if ref_key[2] and ref_key[2] != row_key[2] and not _matches_root_alias(ref, row): + return False + if ref_key[3] and ref_key[3] != row_key[3]: + return False + return True + + +def _row_key(row: dict[str, Any]) -> tuple[str, str, str, str]: + return 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"), + ) + + +def _matches_root_alias(ref: dict[str, Any], row: dict[str, Any]) -> bool: + ref_section_path = str(ref.get("section_path") or "").strip() + row_section_path = str(row.get("section_path") or "").strip() + source_file_name = str(row.get("source_file_name") or "").strip() + return bool( + source_file_name + and row_section_path == "Root" + and ref_section_path == source_file_name + ) From 1e9c148b72dd2f89cffc3e82b3b0088363911094 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 02:37:57 +0000 Subject: [PATCH 42/78] refactor: deepen terminal workflow seams --- .../confirmation_service.py | 3 + .../document_ingestion/handoff_service.py | 29 ++- .../services/document_ingestion/service.py | 2 + .../contract/test_job_creation_contract.py | 81 +++++++ .../tests/contract/test_s3_event_contract.py | 39 ++++ .../app/core/tasks/stale_job_sweeper.py | 6 +- .../document_ingestion/job_state_gate.py | 7 +- .../document_ingestion/page_estimator.py | 210 +++++++++--------- .../parse_result_package.py | 92 ++++++++ .../document_ingestion/processing_billing.py | 13 +- .../document_ingestion/processing_run.py | 100 +++++---- .../contract/test_parse_task_contract.py | 31 ++- .../test_stale_job_sweeper_contract.py | 43 ++++ .../shared/core/state_machine/service.py | 140 +++++++++++- .../shared/core/state_machine/service_sync.py | 131 ++++++++++- .../core/state_machine/transition_outcome.py | 67 ++++++ .../shared/services/job_failure_sync.py | 9 +- .../shared/services/job_success_sync.py | 9 +- .../retrieval/workflow/orchestrator.py | 77 +++++-- .../retrieval/workflow/run_request.py | 50 +++++ .../retrieval/workflow/step_runner.py | 68 +++--- 21 files changed, 949 insertions(+), 258 deletions(-) create mode 100644 apps/worker/app/services/document_ingestion/parse_result_package.py create mode 100644 packages/shared-python/shared/core/state_machine/transition_outcome.py create mode 100644 packages/shared-python/shared/services/retrieval/workflow/run_request.py diff --git a/apps/api/app/services/document_ingestion/confirmation_service.py b/apps/api/app/services/document_ingestion/confirmation_service.py index cc513d22b..4cc3a7b8e 100644 --- a/apps/api/app/services/document_ingestion/confirmation_service.py +++ b/apps/api/app/services/document_ingestion/confirmation_service.py @@ -12,6 +12,7 @@ JobOperationException, NotFoundException, PermissionDeniedException, + UnavailableException, ValidationException, ) from shared.core.state_machine.states import JobStatus @@ -78,6 +79,8 @@ async def confirm_upload( raise except PermissionDeniedException: raise + except UnavailableException: + raise except ValidationException: raise except Exception as exc: diff --git a/apps/api/app/services/document_ingestion/handoff_service.py b/apps/api/app/services/document_ingestion/handoff_service.py index 9a14d92bb..68b1cd94a 100644 --- a/apps/api/app/services/document_ingestion/handoff_service.py +++ b/apps/api/app/services/document_ingestion/handoff_service.py @@ -3,9 +3,14 @@ from typing import Any from app.services.knowledge.kb_orchestrator import KBOrchestrator +from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession -from shared.core.exceptions.domain_exceptions import ValidationException +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import ( + UnavailableException, + ValidationException, +) from shared.core.state_machine.service import AsyncStateMachineService from shared.core.state_machine.states import JobStatus @@ -32,7 +37,7 @@ async def start_uploaded_file_workflow( user_id: str, trigger: str, ) -> None: - await self._state_machine.transition( + outcome = await self._state_machine.transition_outcome( db, job.job_id, JobStatus.PENDING.value, @@ -40,6 +45,19 @@ async def start_uploaded_file_workflow( None, "system", ) + if not outcome.succeeded: + logger.warning( + "Upload handoff transition rejected: " + f"job_id={job.job_id}, reason={outcome.reason}" + ) + raise UnavailableException( + internal_message=( + f"Could not advance uploaded job {job.job_id} to pending: " + f"{outcome.reason}" + ), + retry_after=settings.KB_TASK_RETRY_COUNTDOWN, + user_message="Job state is still settling. Retrying shortly.", + ) if job.job_type != _JOB_TYPE_KB_MANAGEMENT: raise ValidationException( @@ -64,9 +82,14 @@ async def start_uploaded_file_workflow( ) async def mark_upload_expired(self, db: AsyncSession, *, job: Any) -> None: - await self._state_machine.mark_failed( + outcome = await self._state_machine.mark_failed_outcome( db, job.job_id, "Upload expired: file was not uploaded within the allowed time window", error_code="UPLOAD_EXPIRED", ) + if not outcome.succeeded: + logger.warning( + "Upload expiry transition rejected: " + f"job_id={job.job_id}, reason={outcome.reason}" + ) diff --git a/apps/api/app/services/document_ingestion/service.py b/apps/api/app/services/document_ingestion/service.py index 4573f8051..defd2fbca 100644 --- a/apps/api/app/services/document_ingestion/service.py +++ b/apps/api/app/services/document_ingestion/service.py @@ -121,6 +121,8 @@ async def confirm_upload( raise except PermissionDeniedException: raise + except UnavailableException: + raise except ValidationException: raise 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 9a9d47926..edea174ec 100644 --- a/apps/api/tests/contract/test_job_creation_contract.py +++ b/apps/api/tests/contract/test_job_creation_contract.py @@ -1041,3 +1041,84 @@ async def _fake_start_workflow( "db_bound": True, } ] + + +@pytest.mark.asyncio +async def test_should_preserve_retryable_confirm_upload_transition_rejection( + monkeypatch: MonkeyPatch, + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + payload: dict[str, str] = { + "namespace": "contract-jobs", + "source_type": "file", + "file_name": "confirm-upload-retry.pdf", + "data_id": "contract-job-confirm-upload-retry", + } + + async def _fake_verify_s3_file_exists( + self: object, + s3_key: str, + bucket: str | None = None, + ) -> dict[str, object]: + assert bucket is None + return {"exists": True, "s3_key": s3_key} + + async def _fake_transition_outcome( + self: object, + db: object, + job_id: str, + to_state: str, + transition_reason: str = "normal_transition", + operator_id: str | None = None, + operator_type: str = "system", + metadata: dict[str, object] | None = None, + auto_commit: bool = True, + ) -> object: + del self, db, transition_reason, operator_id, operator_type, metadata, auto_commit + + from shared.core.state_machine.transition_outcome import JobTransitionOutcome + + return JobTransitionOutcome.rejected( + job_id=job_id, + to_state=to_state, + reason="cas_conflict", + attempts=3, + ) + + async with developer_api_client_factory() as api_client: + import shared.core.state_machine.service as state_machine_module + import shared.services.storage.file_upload_service as file_upload_service_module + + monkeypatch.setattr( + file_upload_service_module.FileUploadService, + "verify_s3_file_exists", + _fake_verify_s3_file_exists, + ) + monkeypatch.setattr( + state_machine_module.AsyncStateMachineService, + "transition_outcome", + _fake_transition_outcome, + ) + + create_response = await api_client.post("/api/v1/jobs", json=payload) + assert create_response.status_code == 200 + + create_response_json: dict[str, object] = create_response.json() + job_id = cast(str, create_response_json["job_id"]) + + confirm_response = await api_client.post(f"/api/v1/jobs/{job_id}/confirm-upload") + + assert confirm_response.status_code == 503 + assert confirm_response.headers["retry-after"] == "120" + assert confirm_response.headers["x-request-id"] + + response_json: dict[str, object] = confirm_response.json() + error = cast(dict[str, object], response_json["error"]) + details = cast(dict[str, object], error["details"]) + + assert response_json["success"] is False + assert error["code"] == "UNAVAILABLE" + assert error["message"] == "Job state is still settling. Retrying shortly." + assert details["retry_after"] == 120 diff --git a/apps/api/tests/contract/test_s3_event_contract.py b/apps/api/tests/contract/test_s3_event_contract.py index 81ad263e9..ae781eb36 100644 --- a/apps/api/tests/contract/test_s3_event_contract.py +++ b/apps/api/tests/contract/test_s3_event_contract.py @@ -10,6 +10,8 @@ import pytest from httpx import AsyncClient from pytest import MonkeyPatch +from sqlalchemy import text +from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from tests.support.contract_database import ContractDatabase @@ -207,6 +209,43 @@ async def start_workflow( assert job_row["status"] == "pending" +@pytest.mark.asyncio +async def test_should_expose_async_state_machine_rejection_reason( + api_client_factory: Callable[[], AbstractAsyncContextManager[AsyncClient]], +) -> None: + from shared.core.state_machine.service import AsyncStateMachineService + from shared.core.state_machine.states import JobStatus + from shared.testing.contract_runtime import get_contract_database_url + + async with api_client_factory(): + _, job_id = await _insert_waiting_file_job() + engine = create_async_engine(get_contract_database_url(), future=True) + session_factory = async_sessionmaker(engine, expire_on_commit=False) + + try: + async with engine.begin() as connection: + await connection.execute( + text("UPDATE jobs SET status = 'done' WHERE job_id = :job_id"), + {"job_id": job_id}, + ) + + async with session_factory() as db: + outcome = await AsyncStateMachineService().transition_outcome( + db, + job_id, + JobStatus.RUNNING.value, + transition_reason="contract_invalid_transition", + ) + finally: + await engine.dispose() + + assert outcome.succeeded is False + assert outcome.reason == "invalid_transition" + assert outcome.from_state == "done" + assert outcome.to_state == "running" + assert outcome.attempts == 1 + + @pytest.mark.asyncio async def test_should_reject_an_sns_subscription_confirmation_url_that_resolves_to_a_private_host( api_client_factory: Callable[[], AbstractAsyncContextManager[AsyncClient]], diff --git a/apps/worker/app/core/tasks/stale_job_sweeper.py b/apps/worker/app/core/tasks/stale_job_sweeper.py index 30142a614..235c5aabd 100644 --- a/apps/worker/app/core/tasks/stale_job_sweeper.py +++ b/apps/worker/app/core/tasks/stale_job_sweeper.py @@ -103,14 +103,14 @@ def expire_stale_jobs() -> dict: return {"status": "success", "expired": 0, "skipped": 0} for job in expired_jobs: - success = state_machine.mark_failed( + outcome = state_machine.mark_failed_outcome( db, job.job_id, error_message=JOB_EXPIRED_ERROR_MESSAGE, error_code=JOB_EXPIRED_ERROR_CODE, metadata={"sweeper": True, "stale_status": job.status}, ) - if success: + if outcome.succeeded: expired_count += 1 logger.info( f"Expired stale job {job.job_id} (was {job.status})" @@ -118,7 +118,7 @@ def expire_stale_jobs() -> dict: else: skipped_count += 1 logger.debug( - f"Job {job.job_id} already transitioned (CAS miss)" + f"Job {job.job_id} was not expired: reason={outcome.reason}" ) if expired_count > 0: diff --git a/apps/worker/app/services/document_ingestion/job_state_gate.py b/apps/worker/app/services/document_ingestion/job_state_gate.py index 5cd5d067c..c8b963105 100644 --- a/apps/worker/app/services/document_ingestion/job_state_gate.py +++ b/apps/worker/app/services/document_ingestion/job_state_gate.py @@ -72,17 +72,18 @@ def mark_job_running(job_id: str, redis_service: Any) -> bool: ), ) - if not state_machine.transition( + outcome = state_machine.transition_outcome( db, job_id, JobStatus.RUNNING.value, "start_processing", operator_type="system", - ): + ) + if not outcome.succeeded: raise UnavailableException( internal_message=( f"Failed to transition job {job_id} from pending to running; " - "the state may have changed concurrently" + f"reason={outcome.reason}" ), retry_after=settings.KB_TASK_RETRY_COUNTDOWN, user_message="Job state is still settling. Retrying shortly.", diff --git a/apps/worker/app/services/document_ingestion/page_estimator.py b/apps/worker/app/services/document_ingestion/page_estimator.py index ee681d4c4..bdc97a4d6 100644 --- a/apps/worker/app/services/document_ingestion/page_estimator.py +++ b/apps/worker/app/services/document_ingestion/page_estimator.py @@ -10,6 +10,8 @@ import math import tempfile +from collections.abc import Callable +from dataclasses import dataclass from pathlib import Path from shared.core.logging import logger @@ -19,153 +21,155 @@ ROWS_PER_PAGE = 50 +@dataclass(frozen=True) +class WorkloadEstimate: + page_count: int + method: str + fallback_reason: str | None = None + + @property + def used_fallback(self) -> bool: + return self.fallback_reason is not None + + class PageEstimator: """Estimate page count for billing purposes.""" @classmethod def estimate(cls, file_path: str) -> int: """Estimate the billable page count for a file.""" + return cls.estimate_workload(file_path).page_count + + @classmethod + def estimate_workload(cls, file_path: str) -> WorkloadEstimate: + """Estimate billable workload while keeping fallback policy explicit.""" path = Path(file_path) suffix = path.suffix.lower() - try: - if suffix == ".pdf": - return cls._estimate_pdf(file_path) - if suffix == ".pptx": - return cls._estimate_pptx(file_path) - if suffix == ".doc": - return cls._estimate_doc(file_path) - if suffix == ".docx": - return cls._estimate_docx(file_path) - if suffix == ".xls": - return cls._estimate_xls(file_path) - if suffix == ".xlsx": - return cls._estimate_xlsx(file_path) - if suffix in [".txt", ".md", ".json", ".fragment"]: - return cls._estimate_text(file_path) - if suffix in [".png", ".jpg", ".jpeg"]: - return 1 + if suffix == ".pdf": + return cls._estimate_with_fallback("pdf_metadata", cls._count_pdf, file_path) + if suffix == ".pptx": + return cls._estimate_with_fallback("pptx_slides", cls._count_pptx, file_path) + if suffix == ".doc": + return cls._estimate_with_fallback("doc_conversion", cls._count_doc, file_path) + if suffix == ".docx": + return cls._estimate_with_fallback("docx_words", cls._count_docx, file_path) + if suffix == ".xls": + return cls._estimate_with_fallback("xls_conversion", cls._count_xls, file_path) + if suffix == ".xlsx": + return cls._estimate_with_fallback("xlsx_rows", cls._count_xlsx, file_path) + if suffix in [".txt", ".md", ".json", ".fragment"]: + return cls._estimate_with_fallback("text_words", cls._count_text, file_path) + if suffix in [".png", ".jpg", ".jpeg"]: + return WorkloadEstimate(page_count=1, method="image_default") + + fallback_reason = f"unknown_file_type:{suffix or ''}" + logger.warning( + f"Unknown file type for billing: {suffix}, defaulting to 1 page" + ) + return WorkloadEstimate( + page_count=1, + method="unknown_file_type", + fallback_reason=fallback_reason, + ) + @classmethod + def _estimate_with_fallback( + cls, + method: str, + estimator: Callable[[str], int], + file_path: str, + ) -> WorkloadEstimate: + try: + page_count = max(1, estimator(file_path)) + return WorkloadEstimate(page_count=page_count, method=method) + except ImportError as exc: + fallback_reason = f"missing_dependency:{exc.name or exc}" logger.warning( - f"Unknown file type for billing: {suffix}, defaulting to 1 page" + f"Missing dependency for {method} page estimation, defaulting to 1 page: {exc}" ) - return 1 except Exception as exc: + fallback_reason = f"{method}_error:{type(exc).__name__}" logger.error(f"Error estimating pages for {file_path}: {exc}") - return 1 + + return WorkloadEstimate( + page_count=1, + method=method, + fallback_reason=fallback_reason, + ) @classmethod - def _estimate_pdf(cls, file_path: str) -> int: + def _count_pdf(cls, file_path: str) -> int: """Estimate pages for PDF using physical page count.""" - try: - from pypdf import PdfReader + from pypdf import PdfReader - reader = PdfReader(file_path) - return max(1, len(reader.pages)) - except ImportError: - logger.warning("pypdf not installed, defaulting to 1 page") - return 1 - except Exception as exc: - logger.error(f"PDF estimation error: {exc}") - return 1 + reader = PdfReader(file_path) + return len(reader.pages) @classmethod - def _estimate_pptx(cls, file_path: str) -> int: + def _count_pptx(cls, file_path: str) -> int: """Estimate pages for PPTX using slide count.""" - try: - from pptx import Presentation + from pptx import Presentation - 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 exc: - logger.error(f"PPTX estimation error: {exc}") - return 1 + presentation = Presentation(file_path) + return len(presentation.slides) @classmethod - def _estimate_docx(cls, file_path: str) -> int: + def _count_docx(cls, file_path: str) -> int: """Estimate pages for DOCX using word-based counting.""" - try: - from docx import Document + from docx import Document - document = Document(file_path) - total_text = "" + document = Document(file_path) + total_text = "" - for paragraph in document.paragraphs: - total_text += paragraph.text + " " + for paragraph in document.paragraphs: + total_text += paragraph.text + " " - for table in document.tables: - for row in table.rows: - for cell in row.cells: - total_text += cell.text + " " + 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 exc: - logger.error(f"DOCX estimation error: {exc}") - return 1 + word_count = count_cn_en(total_text) + return math.ceil(word_count / WORDS_PER_PAGE) @classmethod - def _estimate_doc(cls, file_path: str) -> int: + def _count_doc(cls, file_path: str) -> int: """Estimate pages for DOC by converting it to DOCX first.""" - try: - from app.services.document_parser.legacy_converter import doc_to_docx + 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 exc: - logger.error(f"DOC estimation error: {exc}") - return 1 + with tempfile.TemporaryDirectory(prefix="page-estimator-doc-") as temp_dir: + converted_path, _ = doc_to_docx(file_path, temp_dir) + return cls._count_docx(converted_path) @classmethod - def _estimate_xlsx(cls, file_path: str) -> int: + def _count_xlsx(cls, file_path: str) -> int: """Estimate pages for XLSX using row count.""" - try: - import pandas as pd + import pandas as pd - workbook = pd.ExcelFile(file_path) - total_rows = 0 + workbook = pd.ExcelFile(file_path) + total_rows = 0 - for sheet_name in workbook.sheet_names: - dataframe = pd.read_excel(workbook, sheet_name=sheet_name) - total_rows += len(dataframe) + 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 exc: - logger.error(f"XLSX estimation error: {exc}") - return 1 + return math.ceil(total_rows / ROWS_PER_PAGE) @classmethod - def _estimate_xls(cls, file_path: str) -> int: + def _count_xls(cls, file_path: str) -> int: """Estimate pages for XLS by converting it to XLSX first.""" - try: - from app.services.document_parser.legacy_converter import xls_to_xlsx + 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 exc: - logger.error(f"XLS estimation error: {exc}") - return 1 + with tempfile.TemporaryDirectory(prefix="page-estimator-xls-") as temp_dir: + converted_path, _ = xls_to_xlsx(file_path, temp_dir) + return cls._count_xlsx(converted_path) @classmethod - def _estimate_text(cls, file_path: str) -> int: + def _count_text(cls, file_path: str) -> int: """Estimate pages for text-like files using word-based counting.""" - try: - with open(file_path, "r", encoding="utf-8", errors="ignore") as file: - content = file.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 exc: - logger.error(f"Text estimation error: {exc}") - return 1 + word_count = count_cn_en(content) + return math.ceil(word_count / WORDS_PER_PAGE) diff --git a/apps/worker/app/services/document_ingestion/parse_result_package.py b/apps/worker/app/services/document_ingestion/parse_result_package.py new file mode 100644 index 000000000..f35c6414c --- /dev/null +++ b/apps/worker/app/services/document_ingestion/parse_result_package.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import pandas as pd +from loguru import logger + +from shared.core.exceptions.domain_exceptions import WorkerHandlingException +from shared.services.chunks.dataframe_chunk_converter import dataframe_to_chunks + + +@dataclass(frozen=True) +class ParseArtifact: + add_dir: str + dataframe: pd.DataFrame + + @property + def contents_count(self) -> int: + return len(self.dataframe) + + +@dataclass(frozen=True) +class ParseResultPackage: + artifact: ParseArtifact + chunks: list[dict[str, Any]] + + +@dataclass(frozen=True) +class GeneratedResultPackage: + zip_file_path: str + checksum_value: str + statistics: dict[str, Any] + zip_size: int + + @classmethod + def from_legacy_tuple( + cls, + zip_file_path: str, + checksum: dict[str, Any] | str | None, + statistics: dict[str, Any], + zip_size: int, + ) -> GeneratedResultPackage: + checksum_value = ( + str(checksum.get("value", "")) + if isinstance(checksum, dict) + else str(checksum or "") + ) + return cls( + zip_file_path=zip_file_path, + checksum_value=checksum_value, + statistics=statistics, + zip_size=zip_size, + ) + + +def build_parse_result_package( + *, + job_id: str, + filename: str, + add_dir: str, + parsed_contents_df: pd.DataFrame | None, +) -> ParseResultPackage: + artifact = _build_parse_artifact( + job_id=job_id, + filename=filename, + add_dir=add_dir, + parsed_contents_df=parsed_contents_df, + ) + chunks = dataframe_to_chunks(artifact.dataframe) + return ParseResultPackage(artifact=artifact, chunks=chunks) + + +def _build_parse_artifact( + *, + job_id: str, + filename: str, + add_dir: str, + parsed_contents_df: pd.DataFrame | None, +) -> ParseArtifact: + 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}" + ) + + return ParseArtifact(add_dir=add_dir, dataframe=parsed_contents_df) diff --git a/apps/worker/app/services/document_ingestion/processing_billing.py b/apps/worker/app/services/document_ingestion/processing_billing.py index 12aaa1159..e3d028208 100644 --- a/apps/worker/app/services/document_ingestion/processing_billing.py +++ b/apps/worker/app/services/document_ingestion/processing_billing.py @@ -3,6 +3,7 @@ from dataclasses import dataclass from datetime import datetime +from app.services.document_ingestion.page_estimator import WorkloadEstimate from app.services.document_ingestion.processing_context import ParseJobContext from loguru import logger from sqlalchemy import select @@ -28,8 +29,9 @@ def charge_parse_job_pages( job_id: str, filename: str | None, job_user_id: str | None, - page_count: int, + workload_estimate: WorkloadEstimate, ) -> ParseJobBillingSnapshot: + page_count = workload_estimate.page_count if not job_user_id: raise NotFoundException( resource="JobInfo", @@ -104,15 +106,20 @@ def record_processing_start( job_id: str, job_context: ParseJobContext, billing_snapshot: ParseJobBillingSnapshot, - page_count: int, processing_started_at: datetime, + workload_estimate: WorkloadEstimate, ) -> None: metadata_updates = { - "page_count": page_count, + "page_count": workload_estimate.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(), + "workload_estimate_method": workload_estimate.method, } + if workload_estimate.fallback_reason is not None: + metadata_updates["workload_estimate_fallback_reason"] = ( + workload_estimate.fallback_reason + ) 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_run.py b/apps/worker/app/services/document_ingestion/processing_run.py index 789b9ee21..88e63ef22 100644 --- a/apps/worker/app/services/document_ingestion/processing_run.py +++ b/apps/worker/app/services/document_ingestion/processing_run.py @@ -4,7 +4,6 @@ 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, @@ -13,6 +12,11 @@ ) 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.parse_result_package import ( + GeneratedResultPackage, + ParseArtifact, + build_parse_result_package, +) from app.services.document_ingestion.processing_billing import ( charge_parse_job_pages, record_processing_start, @@ -30,11 +34,7 @@ from app.services.document_parser.stage_profiler import stage_timer from loguru import logger -from shared.core.exceptions.domain_exceptions import ( - WorkerHandlingException, -) from shared.models.schemas.job_metadata import JobMetadataHelper -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 ( @@ -129,22 +129,28 @@ def _run_parse_job( 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}") + workload_estimate = PageEstimator.estimate_workload(local_temp_path) + page_count = workload_estimate.page_count + logger.info( + "Workload estimation: " + f"job_id={job_id}, page_count={page_count}, " + f"method={workload_estimate.method}, " + f"fallback_reason={workload_estimate.fallback_reason}" + ) 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, + workload_estimate=workload_estimate, ) 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, + workload_estimate=workload_estimate, ) doc_type = JobMetadataHelper.get_parsing_param( @@ -209,38 +215,33 @@ def _run_parse_job( 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) + result_package = build_parse_result_package( + job_id=job_id, + filename=filename, + add_dir=add_dir, + parsed_contents_df=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)}") + logger.info( + f"Chunks prepared: job_id={job_id}, count={len(result_package.chunks)}" + ) return _finalize_parse_job_success( - add_dir=add_dir, - chunks=chunks, + parse_artifact=result_package.artifact, + chunks=result_package.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, ) @@ -248,12 +249,11 @@ def _run_parse_job( def _finalize_parse_job_success( *, - add_dir: str, + parse_artifact: ParseArtifact, 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]: @@ -267,6 +267,8 @@ def _finalize_parse_job_success( document_top_summary = "" section_summaries: dict[str, str] = {} + add_dir = parse_artifact.add_dir + parsed_contents_df = parse_artifact.dataframe if add_dir and source_file_name: if "path" in parsed_contents_df.columns: ensure_doc_nav_json( @@ -317,22 +319,17 @@ def _finalize_parse_job_success( 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, - 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 "") + generated_package = GeneratedResultPackage.from_legacy_tuple( + *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, + ) ) lifecycle_service.update_progress( @@ -343,22 +340,29 @@ def _finalize_parse_job_success( 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, + zip_file_path=generated_package.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( + finalization_response = lifecycle_service.finalize_job_success( job_id=job_id, chunks=chunks, result_s3_key=result_s3_key, - checksum=checksum_value, - zip_size=zip_size, + checksum=generated_package.checksum_value, + zip_size=generated_package.zip_size, stored_count=stored_count, delivery_mode="url", section_summaries=section_summaries, ) + if finalization_response.get("status") != "success": + logger.error( + f"Worker processing finalization failed: job_id={job_id}, " + f"response={finalization_response}" + ) + return dict(finalization_response) + + lifecycle_service.update_progress(job_id, progress=100, message="Task complete!") logger.info( f"Worker processing complete: job_id={job_id}, result_s3_key={result_s3_key}" diff --git a/apps/worker/tests/contract/test_parse_task_contract.py b/apps/worker/tests/contract/test_parse_task_contract.py index 816e04a15..bf8bc61d2 100644 --- a/apps/worker/tests/contract/test_parse_task_contract.py +++ b/apps/worker/tests/contract/test_parse_task_contract.py @@ -397,6 +397,8 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: metadata = sync_job_metadata_service_cls(redis_service).get_metadata(job_id) assert metadata is not None assert metadata["page_count"] == 3 + assert metadata["workload_estimate_method"] == "pdf_metadata" + assert "workload_estimate_fallback_reason" not in metadata assert metadata["billing_status"] == expected_billing_status if billing_enabled: assert metadata["billing_amount_micro_dollars"] == expected_credits_charged @@ -1053,9 +1055,11 @@ def fake_download_s3_file_to_temp( billing_start_barrier = Barrier(len(job_ids)) - def fake_estimate_page_count(file_path: str) -> int: + def fake_estimate_workload(file_path: str) -> Any: + from app.services.document_ingestion.page_estimator import WorkloadEstimate + billing_start_barrier.wait(timeout=10) - return 1 + return WorkloadEstimate(page_count=1, method="contract_fake") def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: output_dir = ( @@ -1094,7 +1098,11 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: _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_job_service.PageEstimator, + "estimate_workload", + fake_estimate_workload, + ) monkeypatch.setattr(parse_service, "checkerboard_inject_parse", fake_checkerboard_inject_parse) monkeypatch.setattr(parse_job_service, "get_result_storage", lambda: FakeResultStorage()) @@ -1306,6 +1314,23 @@ def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: assert audit_transition_count == 0 +def test_should_make_workload_estimate_fallback_reason_explicit( + worker_contract_environment: None, + tmp_path: Path, +) -> None: + from app.services.document_ingestion.page_estimator import PageEstimator + + unknown_file = tmp_path / "payload.unknown" + unknown_file.write_text("content", encoding="utf-8") + + estimate = PageEstimator.estimate_workload(str(unknown_file)) + + assert estimate.page_count == 1 + assert estimate.method == "unknown_file_type" + assert estimate.fallback_reason == "unknown_file_type:.unknown" + assert PageEstimator.estimate(str(unknown_file)) == 1 + + def test_should_mark_the_job_failed_and_cleanup_the_workspace_when_parse_execution_raises( worker_contract_environment: None, monkeypatch: MonkeyPatch, 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 c5987ad8d..19a0f5858 100644 --- a/apps/worker/tests/contract/test_stale_job_sweeper_contract.py +++ b/apps/worker/tests/contract/test_stale_job_sweeper_contract.py @@ -224,3 +224,46 @@ def test_should_record_retry_transition_through_sync_state_machine( assert progress["status"] == "pending" assert progress["worker"] == "contract" assert progress["retry_count"] == 1 + + +def test_should_expose_sync_state_machine_rejection_reason( + worker_contract_environment: None, +) -> None: + from shared.core.database_sync import get_sync_db_context + from shared.core.state_machine.service_sync import SyncStateMachineService + from shared.core.state_machine.states import JobStatus + from shared.services.redis.redis_sync_service import SyncRedisServiceFactory + + job_id = f"job_outcome_{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="done", + source_type="file", + webhook_enabled=False, + job_metadata=_build_file_job_metadata(), + ) + + state_machine = SyncStateMachineService( + redis_service=SyncRedisServiceFactory.get_service() + ) + + with get_sync_db_context() as db: + outcome = state_machine.transition_outcome( + db, + job_id, + JobStatus.RUNNING.value, + transition_reason="contract_invalid_transition", + ) + + assert outcome.succeeded is False + assert outcome.reason == "invalid_transition" + assert outcome.from_state == "done" + assert outcome.to_state == "running" + assert outcome.attempts == 1 diff --git a/packages/shared-python/shared/core/state_machine/service.py b/packages/shared-python/shared/core/state_machine/service.py index 366649acb..8407b48e7 100644 --- a/packages/shared-python/shared/core/state_machine/service.py +++ b/packages/shared-python/shared/core/state_machine/service.py @@ -25,6 +25,7 @@ serialize_transition_metadata, utc_now_naive, ) +from shared.core.state_machine.transition_outcome import JobTransitionOutcome from shared.models.database.job import Job from shared.models.database.job_state_audit_log import JobStateAuditLog from shared.services.redis import RedisServiceFactory @@ -51,20 +52,56 @@ async def transition( auto_commit: bool = True, ) -> bool: """Execute an optimistic-lock state transition with up to 3 retries.""" + outcome = await self.transition_outcome( + db, + job_id, + to_state, + transition_reason, + operator_id, + operator_type, + metadata, + auto_commit=auto_commit, + ) + return outcome.as_bool() + + async def transition_outcome( + 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, + ) -> JobTransitionOutcome: + """Execute a state transition and preserve the reason if it is rejected.""" max_retries = 3 for attempt in range(max_retries): + attempts = attempt + 1 try: job = await self._get_job_with_version(db, job_id) if not job: logger.error(f"Job {job_id} does not exist") - return False + return JobTransitionOutcome.rejected( + job_id=job_id, + to_state=to_state, + reason="job_not_found", + attempts=attempts, + ) if not is_valid_transition(job.status, to_state): logger.warning( f"Job {job_id}: illegal transition {job.status} → {to_state}, rejected" ) - return False + return JobTransitionOutcome.rejected( + job_id=job_id, + from_state=job.status, + to_state=to_state, + reason="invalid_transition", + attempts=attempts, + ) old_state = job.status old_version = job.version @@ -99,7 +136,12 @@ async def transition( logger.info( f"Job {job_id} state transition: {old_state} → {to_state}" ) - return True + return JobTransitionOutcome.transitioned( + job_id=job_id, + from_state=old_state, + to_state=to_state, + attempts=attempts, + ) # CAS miss — retry with backoff if attempt < max_retries - 1: @@ -110,7 +152,13 @@ async def transition( continue else: logger.error(f"Job {job_id} CAS retries exhausted") - return False + return JobTransitionOutcome.rejected( + job_id=job_id, + from_state=old_state, + to_state=to_state, + reason="cas_conflict", + attempts=attempts, + ) except Exception as e: logger.error(f"Job {job_id} transition failed: {e}") @@ -119,9 +167,27 @@ async def transition( await db.rollback() except Exception as rollback_err: logger.warning(f"Job {job_id} rollback failed: {rollback_err}") - return False + return JobTransitionOutcome.rejected( + job_id=job_id, + to_state=to_state, + reason="rollback_exception", + attempts=attempts, + error_message=str(rollback_err), + ) + return JobTransitionOutcome.rejected( + job_id=job_id, + to_state=to_state, + reason="transition_exception", + attempts=attempts, + error_message=str(e), + ) - return False + return JobTransitionOutcome.rejected( + job_id=job_id, + to_state=to_state, + reason="cas_conflict", + attempts=max_retries, + ) async def mark_failed( self, @@ -135,6 +201,30 @@ async def mark_failed( auto_commit: bool = True, ) -> bool: """Mark a job as failed with error information.""" + outcome = await self.mark_failed_outcome( + db, + job_id, + error_message, + error_code=error_code, + error_details=error_details, + operator_id=operator_id, + metadata=metadata, + auto_commit=auto_commit, + ) + return outcome.as_bool() + + async def mark_failed_outcome( + 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, + ) -> JobTransitionOutcome: + """Mark a job as failed and expose why the transition was rejected.""" try: normalized_details, transition_metadata = ( build_failure_transition_metadata( @@ -152,7 +242,7 @@ async def mark_failed( normalized_details, ) - return await self.transition( + return await self.transition_outcome( db, job_id, JobStatus.FAILED.value, @@ -164,7 +254,13 @@ async def mark_failed( ) except Exception as e: logger.error(f"Failed to mark Job {job_id} as failed: {e}") - return False + return JobTransitionOutcome.rejected( + job_id=job_id, + to_state=JobStatus.FAILED.value, + reason="transition_exception", + attempts=1, + error_message=str(e), + ) async def mark_completed( self, @@ -175,8 +271,26 @@ async def mark_completed( auto_commit: bool = True, ) -> bool: """Mark a job as completed.""" + outcome = await self.mark_completed_outcome( + db, + job_id, + result_metadata=result_metadata, + operator_id=operator_id, + auto_commit=auto_commit, + ) + return outcome.as_bool() + + async def mark_completed_outcome( + self, + db: AsyncSession, + job_id: str, + result_metadata: Optional[Dict[str, Any]] = None, + operator_id: Optional[str] = None, + auto_commit: bool = True, + ) -> JobTransitionOutcome: + """Mark a job as completed and expose why the transition was rejected.""" try: - return await self.transition( + return await self.transition_outcome( db, job_id, JobStatus.DONE.value, @@ -188,7 +302,13 @@ async def mark_completed( ) except Exception as e: logger.error(f"Failed to mark Job {job_id} as completed: {e}") - return False + return JobTransitionOutcome.rejected( + job_id=job_id, + to_state=JobStatus.DONE.value, + reason="transition_exception", + attempts=1, + error_message=str(e), + ) async def handle_retry( self, 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 175802459..b527cb394 100644 --- a/packages/shared-python/shared/core/state_machine/service_sync.py +++ b/packages/shared-python/shared/core/state_machine/service_sync.py @@ -24,6 +24,7 @@ serialize_transition_metadata, utc_now_naive, ) +from shared.core.state_machine.transition_outcome import JobTransitionOutcome 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 @@ -49,20 +50,53 @@ def transition( metadata: Optional[Dict[str, Any]] = None, ) -> bool: """Execute an optimistic-lock state transition with up to 3 retries.""" + return self.transition_outcome( + db, + job_id, + to_state, + transition_reason, + operator_id, + operator_type, + metadata, + ).as_bool() + + def transition_outcome( + self, + db: Session, + 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, + ) -> JobTransitionOutcome: + """Execute a state transition and preserve the reason if it is rejected.""" max_retries = 3 for attempt in range(max_retries): + attempts = attempt + 1 try: job = self._get_job_with_version(db, job_id) if not job: logger.error(f"Job {job_id} does not exist") - return False + return JobTransitionOutcome.rejected( + job_id=job_id, + to_state=to_state, + reason="job_not_found", + attempts=attempts, + ) if not is_valid_transition(job.status, to_state): logger.warning( f"Job {job_id}: illegal transition {job.status} → {to_state}, rejected" ) - return False + return JobTransitionOutcome.rejected( + job_id=job_id, + from_state=job.status, + to_state=to_state, + reason="invalid_transition", + attempts=attempts, + ) old_state = job.status old_version = job.version @@ -85,7 +119,12 @@ def transition( logger.info( f"Job {job_id} state transition: {old_state} → {to_state}" ) - return True + return JobTransitionOutcome.transitioned( + job_id=job_id, + from_state=old_state, + to_state=to_state, + attempts=attempts, + ) if attempt < max_retries - 1: logger.warning( @@ -97,7 +136,13 @@ def transition( continue else: logger.error(f"Job {job_id} CAS retries exhausted") - return False + return JobTransitionOutcome.rejected( + job_id=job_id, + from_state=old_state, + to_state=to_state, + reason="cas_conflict", + attempts=attempts, + ) except Exception as e: logger.error(f"Job {job_id} transition failed: {e}") @@ -106,9 +151,27 @@ def transition( db.rollback() except Exception as rollback_err: logger.warning(f"Job {job_id} rollback failed: {rollback_err}") - return False + return JobTransitionOutcome.rejected( + job_id=job_id, + to_state=to_state, + reason="rollback_exception", + attempts=attempts, + error_message=str(rollback_err), + ) + return JobTransitionOutcome.rejected( + job_id=job_id, + to_state=to_state, + reason="transition_exception", + attempts=attempts, + error_message=str(e), + ) - return False + return JobTransitionOutcome.rejected( + job_id=job_id, + to_state=to_state, + reason="cas_conflict", + attempts=max_retries, + ) def mark_failed( self, @@ -121,6 +184,27 @@ def mark_failed( metadata: Optional[Dict[str, Any]] = None, ) -> bool: """Mark a job as failed with error information.""" + return self.mark_failed_outcome( + db, + job_id, + error_message, + error_code=error_code, + error_details=error_details, + operator_id=operator_id, + metadata=metadata, + ).as_bool() + + def mark_failed_outcome( + self, + db: Session, + 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, + ) -> JobTransitionOutcome: + """Mark a job as failed and expose why the transition was rejected.""" try: normalized_details, transition_metadata = ( build_failure_transition_metadata( @@ -138,7 +222,7 @@ def mark_failed( normalized_details, ) - return self.transition( + return self.transition_outcome( db, job_id, JobStatus.FAILED.value, @@ -149,7 +233,13 @@ def mark_failed( ) except Exception as e: logger.error(f"Failed to mark Job {job_id} as failed: {e}") - return False + return JobTransitionOutcome.rejected( + job_id=job_id, + to_state=JobStatus.FAILED.value, + reason="transition_exception", + attempts=1, + error_message=str(e), + ) def mark_completed( self, @@ -159,8 +249,23 @@ def mark_completed( operator_id: Optional[str] = None, ) -> bool: """Mark a job as completed.""" + return self.mark_completed_outcome( + db, + job_id, + result_metadata=result_metadata, + operator_id=operator_id, + ).as_bool() + + def mark_completed_outcome( + self, + db: Session, + job_id: str, + result_metadata: Optional[Dict[str, Any]] = None, + operator_id: Optional[str] = None, + ) -> JobTransitionOutcome: + """Mark a job as completed and expose why the transition was rejected.""" try: - return self.transition( + return self.transition_outcome( db, job_id, JobStatus.DONE.value, @@ -171,7 +276,13 @@ def mark_completed( ) except Exception as e: logger.error(f"Failed to mark Job {job_id} as completed: {e}") - return False + return JobTransitionOutcome.rejected( + job_id=job_id, + to_state=JobStatus.DONE.value, + reason="transition_exception", + attempts=1, + error_message=str(e), + ) def handle_retry( self, diff --git a/packages/shared-python/shared/core/state_machine/transition_outcome.py b/packages/shared-python/shared/core/state_machine/transition_outcome.py new file mode 100644 index 000000000..4caca8b28 --- /dev/null +++ b/packages/shared-python/shared/core/state_machine/transition_outcome.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal + +TransitionOutcomeReason = Literal[ + "transitioned", + "job_not_found", + "invalid_transition", + "cas_conflict", + "transition_exception", + "rollback_exception", +] + + +@dataclass(frozen=True) +class JobTransitionOutcome: + job_id: str + to_state: str + reason: TransitionOutcomeReason + attempts: int + from_state: str | None = None + error_message: str | None = None + + @property + def succeeded(self) -> bool: + return self.reason == "transitioned" + + def as_bool(self) -> bool: + return self.succeeded + + @classmethod + def transitioned( + cls, + *, + job_id: str, + from_state: str, + to_state: str, + attempts: int, + ) -> JobTransitionOutcome: + return cls( + job_id=job_id, + from_state=from_state, + to_state=to_state, + reason="transitioned", + attempts=attempts, + ) + + @classmethod + def rejected( + cls, + *, + job_id: str, + to_state: str, + reason: TransitionOutcomeReason, + attempts: int, + from_state: str | None = None, + error_message: str | None = None, + ) -> JobTransitionOutcome: + return cls( + job_id=job_id, + from_state=from_state, + to_state=to_state, + reason=reason, + attempts=attempts, + error_message=error_message, + ) diff --git a/packages/shared-python/shared/services/job_failure_sync.py b/packages/shared-python/shared/services/job_failure_sync.py index 0888b1c27..cb16775d2 100644 --- a/packages/shared-python/shared/services/job_failure_sync.py +++ b/packages/shared-python/shared/services/job_failure_sync.py @@ -39,15 +39,18 @@ def finalize( error_details: dict[str, Any] | None, should_refund: bool, ) -> tuple[bool, WebhookEvent | None]: - transition_ok = self._state_machine.mark_failed( + transition_outcome = self._state_machine.mark_failed_outcome( 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") + if not transition_outcome.succeeded: + logger.error( + f"Job {job_id} mark_failed transition failed: " + f"reason={transition_outcome.reason}" + ) return False, None if should_refund: diff --git a/packages/shared-python/shared/services/job_success_sync.py b/packages/shared-python/shared/services/job_success_sync.py index 110e46df3..fbdc36e92 100644 --- a/packages/shared-python/shared/services/job_success_sync.py +++ b/packages/shared-python/shared/services/job_success_sync.py @@ -68,7 +68,7 @@ def finalize( section_summaries=section_summaries, ) - transition_ok = self._state_machine.mark_completed( + transition_outcome = self._state_machine.mark_completed_outcome( db, job_id, result_metadata={ @@ -77,8 +77,11 @@ def finalize( "delivery_mode": delivery_mode, }, ) - if not transition_ok: - logger.error(f"Job {job_id} mark_completed transition failed") + if not transition_outcome.succeeded: + logger.error( + f"Job {job_id} mark_completed transition failed: " + f"reason={transition_outcome.reason}" + ) return JobSuccessFinalization( response={ "status": "failed", diff --git a/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py b/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py index 3fd425521..ad56e6ed7 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py +++ b/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py @@ -18,6 +18,7 @@ ) from shared.services.retrieval.workflow.plan_service import WorkflowPlanService from shared.services.retrieval.workflow.reference_projection import WorkflowReferenceProjection +from shared.services.retrieval.workflow.run_request import WorkflowRunRequest 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 @@ -25,14 +26,31 @@ from shared.services.retrieval.workflow.wallet import BudgetWallet DbSessionFactory = Callable[[], AbstractAsyncContextManager[AsyncSession]] +WorkflowStepRunnerFactory = Callable[[DbSessionFactory, str], WorkflowStepRunner] + + +def _create_workflow_step_runner( + db_factory: DbSessionFactory, + parent_run_id: str, +) -> WorkflowStepRunner: + return WorkflowStepRunner(db_factory=db_factory, parent_run_id=parent_run_id) class WorkflowOrchestrator: """Plan and execute a query workflow DAG.""" - def __init__(self, db_factory: DbSessionFactory | None = None) -> None: + def __init__( + self, + db_factory: DbSessionFactory | None = None, + plan_service: WorkflowPlanService | None = None, + step_runner_factory: WorkflowStepRunnerFactory | None = None, + ) -> None: self.parent_run_id = f'wret_{uuid4().hex[:12]}' self._db_factory = db_factory + self._plan_service = plan_service or WorkflowPlanService() + self._step_runner_factory = ( + step_runner_factory or _create_workflow_step_runner + ) def _get_db_factory(self) -> DbSessionFactory: if self._db_factory is not None: @@ -58,6 +76,28 @@ async def run( channels: list[str] | None = None, channel_weights: dict[str, float] | None = None, llm_fn=None, + ) -> WorkflowResult: + request = WorkflowRunRequest( + 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, + ) + return await self.run_request(db, request=request, llm_fn=llm_fn) + + async def run_request( + self, + db: AsyncSession, + *, + request: WorkflowRunRequest, + llm_fn=None, ) -> WorkflowResult: t0 = time.monotonic() config = WorkflowRuntimeConfig.from_env() @@ -72,16 +112,16 @@ async def run( ) total_chunks, total_docs, _chunks_count_by_doc = await _load_budget_inventory( db, - user_id=user_id, - namespace=namespace, - exclude_document_ids=exclude_document_ids, + user_id=request.user_id, + namespace=request.namespace, + exclude_document_ids=request.exclude_document_ids, ) planner_ledger.total_chunks = total_chunks planner_ledger.total_docs = total_docs - plan = await WorkflowPlanService().load_or_create( - user_id=user_id, - namespace=namespace, - query=query, + plan = await self._plan_service.load_or_create( + user_id=request.user_id, + namespace=request.namespace, + query=request.query, planner_llm=planner_llm, planner_ledger=planner_ledger, max_steps=config.max_steps, @@ -99,9 +139,9 @@ async def run( ledgers = await wallet.allocate(plan) results_by_id: dict[str, StepResult] = {} sem = asyncio.Semaphore(config.parallel_max) - step_runner = WorkflowStepRunner( - db_factory=self._get_db_factory(), - parent_run_id=self.parent_run_id, + step_runner = self._step_runner_factory( + self._get_db_factory(), + self.parent_run_id, ) for batch in plan.topological_batches(): @@ -112,16 +152,7 @@ async def run( ledger=ledgers[step.id], results_by_id=results_by_id, semaphore=sem, - user_id=user_id, - namespace=namespace, - top_k=step.top_k or top_k, - exclude_document_ids=exclude_document_ids, - exclude_sections=exclude_sections, - data_type=step.data_type or data_type, - signal_paths=signal_paths, - filter_mode=filter_mode, - channels=channels, - channel_weights=channel_weights, + request=request.for_step(step), llm_fn=llm_fn, ) for step in batch @@ -146,8 +177,8 @@ async def run( elapsed_ms, ) return WorkflowResult( - namespace=namespace, - query=query, + namespace=request.namespace, + query=request.query, router_used='workflow_decomposed' if len(plan.steps) > 1 else 'workflow_single_step', answer_text=answer_text, plan=plan, diff --git a/packages/shared-python/shared/services/retrieval/workflow/run_request.py b/packages/shared-python/shared/services/retrieval/workflow/run_request.py new file mode 100644 index 000000000..a695001af --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/workflow/run_request.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from shared.services.retrieval.workflow.types import PlannedStep + + +@dataclass(frozen=True) +class WorkflowRunRequest: + 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 + + def for_step(self, step: PlannedStep) -> WorkflowStepRequest: + return WorkflowStepRequest( + user_id=self.user_id, + namespace=self.namespace, + query=step.sub_query, + top_k=step.top_k or self.top_k, + exclude_document_ids=self.exclude_document_ids, + exclude_sections=self.exclude_sections, + data_type=step.data_type or self.data_type, + signal_paths=self.signal_paths, + filter_mode=self.filter_mode, + channels=self.channels, + channel_weights=self.channel_weights, + ) + + +@dataclass(frozen=True) +class WorkflowStepRequest: + 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 diff --git a/packages/shared-python/shared/services/retrieval/workflow/step_runner.py b/packages/shared-python/shared/services/retrieval/workflow/step_runner.py index 8a6be3d9d..a9e538f8e 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/step_runner.py +++ b/packages/shared-python/shared/services/retrieval/workflow/step_runner.py @@ -13,16 +13,25 @@ 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.run_request import WorkflowStepRequest from shared.services.retrieval.workflow.synthesizer import synthesize_step from shared.services.retrieval.workflow.types import PlannedStep, StepResult, StepStatus DbSessionFactory = Callable[[], AbstractAsyncContextManager[AsyncSession]] +RetrievalAgentFactory = Callable[[], RetrievalAgent] class WorkflowStepRunner: - def __init__(self, *, db_factory: DbSessionFactory, parent_run_id: str) -> None: + def __init__( + self, + *, + db_factory: DbSessionFactory, + parent_run_id: str, + agent_factory: RetrievalAgentFactory | None = None, + ) -> None: self._db_factory = db_factory self._parent_run_id = parent_run_id + self._agent_factory = agent_factory or RetrievalAgent self._references = WorkflowReferenceProjection() async def run_step( @@ -32,16 +41,7 @@ async def run_step( 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, + request: WorkflowStepRequest, llm_fn: LLMFn | None, ) -> None: async with semaphore: @@ -52,16 +52,7 @@ async def run_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, + request=request, llm_fn=llm_fn, ) @@ -71,34 +62,25 @@ async def _run_retrieve_step( 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, + request: WorkflowStepRequest, llm_fn: LLMFn | None, ) -> None: try: async with self._db_factory() as step_db: - agentic_result = await RetrievalAgent().run( + agentic_result = await self._agent_factory().run( step_db, - user_id=user_id, - namespace=namespace, - query=step.sub_query, - top_k=top_k, + user_id=request.user_id, + namespace=request.namespace, + query=request.query, + top_k=request.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, + exclude_document_ids=request.exclude_document_ids, + exclude_sections=request.exclude_sections, + data_type=request.data_type, + signal_paths=request.signal_paths, + filter_mode=request.filter_mode, + channels=request.channels, + channel_weights=request.channel_weights, ledger=ledger, parent_run_id=self._parent_run_id, workflow_step_id=step.id, From 26aa1036027b39d6a0d2c47f995ebe42b321ceac Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 03:38:38 +0000 Subject: [PATCH 43/78] refactor: deepen workflow policy seams --- AGENTS.md | 7 +- CONTEXT.md | 62 ++++- .../tests/contract/test_retrieval_contract.py | 224 +++++++++++++++++- .../orchestration/format_adapters.py | 51 ++-- .../orchestration/parse_output.py | 33 +++ .../orchestration/parse_pipeline.py | 22 +- .../orchestration/route_parse.py | 5 +- .../services/document_parser/parse_service.py | 54 ++++- ...t_document_parser_architecture_contract.py | 44 +++- ...test_job_lifecycle_post_commit_contract.py | 55 +++++ .../test_stale_job_sweeper_contract.py | 33 +++ ...eep-routes-and-worker-tasks-as-adapters.md | 25 ++ docs/adr/0002-use-typed-workflow-outcomes.md | 28 +++ ...keep-retrieval-workflow-policy-explicit.md | 29 +++ docs/adr/README.md | 12 + .../shared/core/state_machine/service.py | 78 +++--- .../shared/core/state_machine/service_sync.py | 78 +++--- .../core/state_machine/transition_runner.py | 120 ++++++++++ .../shared/services/job_failure_sync.py | 26 +- .../shared/services/job_lifecycle_sync.py | 22 +- .../services/job_post_commit_effects_sync.py | 62 +++++ .../shared/services/job_success_sync.py | 23 +- .../services/job_webhook_outbox_sync.py | 12 +- .../retrieval/agentic/discovery_phase.py | 2 + .../retrieval/agentic/discovery_tools.py | 6 +- .../retrieval/agentic/orchestrator.py | 3 + .../services/retrieval/execution_routes.py | 15 +- .../services/retrieval/query_request.py | 2 + .../shared/services/retrieval/route_types.py | 2 + .../retrieval/workflow/orchestrator.py | 6 + .../retrieval/workflow/run_request.py | 41 +++- .../retrieval/workflow/step_runner.py | 1 + 32 files changed, 1009 insertions(+), 174 deletions(-) create mode 100644 apps/worker/app/services/document_parser/orchestration/parse_output.py create mode 100644 apps/worker/tests/contract/test_job_lifecycle_post_commit_contract.py create mode 100644 docs/adr/0001-keep-routes-and-worker-tasks-as-adapters.md create mode 100644 docs/adr/0002-use-typed-workflow-outcomes.md create mode 100644 docs/adr/0003-keep-retrieval-workflow-policy-explicit.md create mode 100644 docs/adr/README.md create mode 100644 packages/shared-python/shared/core/state_machine/transition_runner.py create mode 100644 packages/shared-python/shared/services/job_post_commit_effects_sync.py diff --git a/AGENTS.md b/AGENTS.md index a47ed0712..eed149425 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -119,13 +119,16 @@ flowchart TB `apps/worker/app/services/document_parser/parse_service.py` → `checkerboard_inject_parse()` -This is the universal entry for all file types. It: +This is the legacy tuple-compatible entry for all file types. Internally, +`checkerboard_parse_output()` and parser adapters use the typed `ParseOutput` +contract. The parser flow: 1. **Profiles** the document via `doc_profiler.profile_document()` to detect file type, page count, and special categories (e.g. `atlas`). 2. **Routes** to the appropriate parser based on file extension. 3. **Post-processes**: cleans up unreferenced images, compresses PNG→JPG. -4. Returns `(output_dir, parsed_df)` — the parsed DataFrame. +4. Returns typed parse output internally, with `(output_dir, parsed_df)` kept as + the compatibility shape for older callers. ### Parser Routing Table diff --git a/CONTEXT.md b/CONTEXT.md index 44e050713..ba829e0c4 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -29,6 +29,18 @@ 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 Transition Outcome + +The typed result of a Job state-machine transition. It preserves whether the +transition succeeded, the target state, previous state when known, attempt +count, and rejection reason while keeping older boolean facades available. + +### Job Post-Commit Effect + +A post-transaction side effect planned during terminal Job finalization and run +only after the database commit succeeds. Current effects include retrieval cache +invalidation and outbound webhook publication. + ### Job Read The workflow that lists a User's Jobs and projects one Job into the public Job @@ -58,6 +70,26 @@ 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. +### Parse Output + +The stable parser adapter result with an output directory and optional parsed +DataFrame. It owns legacy tuple compatibility for older parser callers. + +### Parse Artifact + +The ingestion-side parsed content artifact. It validates parser output before +chunk conversion and result packaging. + +### Generated Result Package + +The generated ZIP bundle metadata used by terminal Job finalization, including +ZIP path, checksum, statistics, and byte size. + +### Workload Estimate + +The worker-side estimate used for billing and processing metadata. It records +page count, estimation method, and any fallback reason. + ### Parser Input The typed worker-side parse request assembled from Job metadata, parser options, @@ -106,6 +138,22 @@ Document Chunks, and document graph state. The query workflow that returns cited evidence from published documents. +### Retrieval Query + +The typed retrieval request that owns cache-shaping fields and route policy: +scope, filters, data type, channels, ranking options, and agentic toggle. + +### Workflow Run Request + +The agentic Retrieval request passed through planning and step execution. It +preserves user scope, filters, channel policy, internal recall, and explicit +ranking policy fields for the workflow path. + +### Workflow Step Request + +The per-step projection of a Workflow Run Request. It applies step-level query, +top-k, and data-type overrides while preserving the request policy. + ### Demo Source An API-owned canonical document shipped with the repository for demo and guest @@ -255,6 +303,11 @@ exceptions. - `app/api/v1/routes/s3_events.py` - `app/services/s3_events/*` +### Storage Event Intake + +The internal workflow that decodes S3-compatible storage events, sanitizes +headers, acknowledges malformed or unsafe events, and triggers upload handoff. + ### Async Callbacks - `app/api/v1/routes/qstash_callbacks.py` @@ -291,10 +344,15 @@ exceptions. 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. + legacy parser entrypoint; parser option shaping, format routing, rendered PDF + transforms, typed Parse Output, 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. +- Terminal Job finalization should plan post-commit effects with primitive + identifiers and run them after the database transaction commits. +- State-machine callers that need diagnostics should consume Job Transition + Outcome; boolean state-machine methods remain compatibility facades. - `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. diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index d3eecb67c..e0a4bcaf9 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -1,7 +1,7 @@ from collections.abc import AsyncIterator, Callable from contextlib import AbstractAsyncContextManager, asynccontextmanager from datetime import datetime, timezone -from typing import cast +from typing import Any, cast from uuid import uuid4 import pytest @@ -12,6 +12,7 @@ from tests.support.contract_database import ContractDatabase from shared.services.retrieval.agentic.types import AgenticResult from shared.services.retrieval.query_request import RetrievalQuery +from shared.services.retrieval.workflow.run_request import WorkflowRunRequest from shared.services.retrieval.workflow.types import PlannedStep, QueryPlan, WorkflowResult @@ -180,6 +181,205 @@ def test_retrieval_query_owns_cache_and_route_request_shape() -> None: assert route_context.effective_recall_k == 17 assert route_context.use_agentic is False + workflow_request = WorkflowRunRequest.from_route_context(route_context) + assert workflow_request.top_k == 5 + assert workflow_request.data_type == 2 + assert workflow_request.signal_paths == ["Root"] + assert workflow_request.filter_mode == "keep" + assert workflow_request.channels == ["content"] + assert workflow_request.channel_weights == {"content": 1.5} + assert workflow_request.internal_recall_k == 17 + assert workflow_request.threshold == 0.2 + assert workflow_request.rerank is True + assert ( + workflow_request.for_step( + PlannedStep(id="explicit-recall", sub_query="explicit", top_k=9) + ).internal_recall_k + == 17 + ) + + +def test_agentic_workflow_should_recompute_default_recall_for_step_top_k() -> None: + query = RetrievalQuery.from_parameters( + db=cast(AsyncSession, object()), + user_id="user-1", + namespace="contract-retrieval", + query="alpha", + top_k=1, + exclude_document_ids=[], + exclude_sections=[], + use_agentic=True, + ) + + workflow_request = WorkflowRunRequest.from_route_context(query.build_route_context()) + step_request = workflow_request.for_step( + PlannedStep(id="wide-step", sub_query="wide alpha", top_k=10) + ) + + assert workflow_request.internal_recall_k is None + assert step_request.top_k == 10 + assert step_request.internal_recall_k == 20 + + +@pytest.mark.asyncio +async def test_agentic_workflow_should_pass_full_request_policy_to_step_adapter( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + captured_requests: list[dict[str, object]] = [] + + 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="request-policy", sub_query="policy marker")], + final_strategy="concat_final_parts", + reasoning_summary=( + f"request policy contract for {kb_total_docs} docs " + f"and {kb_total_chunks} chunks" + ), + ) + + async def fake_retrieval_run( + self: object, + db: object, + **kwargs: object, + ) -> AgenticResult: + del self, db + captured_requests.append(kwargs) + return AgenticResult( + evidence_text="policy evidence", + answer_text="policy answer", + referenced_chunks=[ + { + "chunk_id": policy_document["chunk_id"], + "document_id": policy_document["document_id"], + "chunk_type": "text", + "section_path": policy_document["section_path"], + "file_path": "", + "job_id": policy_document["job_id"], + } + ], + router_used="contract_fake_agent", + ) + + async with developer_api_client_factory() as api_client: + policy_document = await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-agentic-request-policy", + source_file_name="policy.pdf", + section_path="Root/Policy", + content="policy marker content", + ) + await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-agentic-request-policy", + source_file_name="filler.pdf", + section_path="Root/Filler", + content="filler content", + ) + monkeypatch.setattr( + "shared.services.retrieval.workflow.planner.QueryPlanner.plan", + fake_plan, + ) + monkeypatch.setattr( + "shared.services.retrieval.workflow.step_runner.RetrievalAgent.run", + fake_retrieval_run, + ) + + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": "contract-agentic-request-policy", + "query": "policy marker", + "top_k": 1, + "data_type": 2, + "signal_paths": ["Root"], + "filter_mode": "keep", + "channels": ["content"], + "channel_weights": {"content": 2.0}, + "internal_recall_k": 23, + "threshold": 0.4, + "rerank": True, + "use_agentic": True, + }, + ) + + assert response.status_code == 200 + + assert len(captured_requests) == 1 + request = captured_requests[0] + assert request["user_id"] == "local-dev-user" + assert request["namespace"] == "contract-agentic-request-policy" + assert request["query"] == "policy marker" + assert request["top_k"] == 1 + assert request["exclude_document_ids"] == [] + assert request["exclude_sections"] == [] + assert request["data_type"] == 2 + assert request["signal_paths"] == ["Root"] + assert request["filter_mode"] == "keep" + assert request["channels"] == ["content"] + assert request["channel_weights"] == {"content": 2.0} + assert request["internal_recall_k"] == 23 + + +@pytest.mark.asyncio +async def test_agentic_discovery_should_use_internal_recall_for_fusion( + monkeypatch: MonkeyPatch, +) -> None: + from shared.services.retrieval.agentic import discovery_tools + + requested_top_k: list[int] = [] + + async def fake_empty_channel( + _db: AsyncSession, + **_kwargs: Any, + ) -> list[dict[str, Any]]: + return [] + + async def fake_content_channel( + _db: AsyncSession, + **kwargs: Any, + ) -> list[dict[str, Any]]: + requested_top_k.append(cast(int, kwargs["top_k"])) + return [ + { + "chunk_id": f"chunk-{index}", + "document_id": f"doc-{index}", + "section_path": f"Root/{index}", + "content": f"content {index}", + } + for index in range(5) + ] + + monkeypatch.setattr(discovery_tools, "path_channel", fake_empty_channel) + monkeypatch.setattr(discovery_tools, "content_channel", fake_content_channel) + monkeypatch.setattr(discovery_tools, "term_channel", fake_empty_channel) + + result = await discovery_tools.bottom_discovery( + cast(AsyncSession, None), + user_id="local-dev-user", + namespace="contract-agentic-recall-policy", + query="recall marker", + top_k=1, + exclude_document_ids=[], + exclude_sections=[], + channels=["content"], + internal_recall_k=4, + ) + + fused_rows = cast(list[dict[str, Any]], result.payload["fused_rows"]) + + assert requested_top_k == [4] + assert len(fused_rows) == 4 + @pytest.mark.asyncio async def test_should_return_seeded_retrieval_results_for_the_authenticated_user( @@ -608,10 +808,15 @@ async def test_agentic_retrieval_should_not_hydrate_references_outside_request_s monkeypatch: MonkeyPatch, ) -> None: class FakeWorkflowOrchestrator: - async def run(self, *_args: object, **kwargs: object) -> WorkflowResult: + async def run_request( + self, + _db: AsyncSession, + *, + request: WorkflowRunRequest, + ) -> WorkflowResult: return WorkflowResult( - namespace=str(kwargs["namespace"]), - query=str(kwargs["query"]), + namespace=request.namespace, + query=request.query, router_used="workflow_single_step", answer_text="foreign reference answer", referenced_chunks=[ @@ -682,10 +887,15 @@ async def test_agentic_retrieval_should_drop_references_that_do_not_match_the_hy monkeypatch: MonkeyPatch, ) -> None: class FakeWorkflowOrchestrator: - async def run(self, *_args: object, **kwargs: object) -> WorkflowResult: + async def run_request( + self, + _db: AsyncSession, + *, + request: WorkflowRunRequest, + ) -> WorkflowResult: return WorkflowResult( - namespace=str(kwargs["namespace"]), - query=str(kwargs["query"]), + namespace=request.namespace, + query=request.query, router_used="workflow_single_step", answer_text="mismatched section answer", referenced_chunks=[ 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 f10f9edd1..e9b5a8846 100644 --- a/apps/worker/app/services/document_parser/orchestration/format_adapters.py +++ b/apps/worker/app/services/document_parser/orchestration/format_adapters.py @@ -3,8 +3,7 @@ from dataclasses import dataclass from typing import Protocol -import pandas as pd - +from app.services.document_parser.orchestration.parse_output import ParseOutput from app.services.document_parser.orchestration.parse_session import ParseSession @@ -13,8 +12,8 @@ class DocumentParseAdapter(Protocol): 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.""" + def parse(self, session: ParseSession) -> ParseOutput: + """Parse a document session into a stable parser output object.""" ... @@ -22,7 +21,7 @@ def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: class FragmentParseAdapter: document_format: object - def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: + def parse(self, session: ParseSession) -> ParseOutput: from app.services.document_parser.fragment_parser import parse_fragment full_output_dir, _relative_root, parsed_df = parse_fragment( @@ -32,14 +31,14 @@ def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: kb_dir=session.kb_dir, base_llm_paras=session.base_llm_paras, ) - return full_output_dir, parsed_df + return ParseOutput(output_dir=full_output_dir, parsed_df=parsed_df) @dataclass(frozen=True) class TextParseAdapter: document_format: object - def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: + def parse(self, session: ParseSession) -> ParseOutput: from app.services.document_parser.md_parser import parse_md from app.services.document_parser.txt_parser import parse_texts @@ -51,14 +50,14 @@ def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: base_llm_paras=session.base_llm_paras, relative_root=session.relative_root, ) - return session.full_output_dir, parsed_df + return ParseOutput(output_dir=session.full_output_dir, parsed_df=parsed_df) @dataclass(frozen=True) class ImageParseAdapter: document_format: object - def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: + def parse(self, session: ParseSession) -> ParseOutput: from app.services.document_parser.image_parser import parse_image parsed_df = parse_image( @@ -69,14 +68,14 @@ def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: base_llm_paras=session.base_llm_paras, relative_root=session.relative_root, ) - return session.full_output_dir, parsed_df + return ParseOutput(output_dir=session.full_output_dir, parsed_df=parsed_df) @dataclass(frozen=True) class PdfParseAdapter: document_format: object - def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: + def parse(self, session: ParseSession) -> ParseOutput: from app.services.document_parser.pdf_parser import parse_pdfs parsed_df = parse_pdfs( @@ -88,14 +87,14 @@ def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: relative_root=session.relative_root, s3_key=session.s3_key, ) - return session.full_output_dir, parsed_df + return ParseOutput(output_dir=session.full_output_dir, parsed_df=parsed_df) @dataclass(frozen=True) class DocParseAdapter: document_format: object - def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: + def parse(self, session: ParseSession) -> ParseOutput: from app.services.document_parser.legacy_converter import doc_to_docx converted_docx_path, _ = doc_to_docx( @@ -109,7 +108,7 @@ def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: class DocxParseAdapter: document_format: object - def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: + def parse(self, session: ParseSession) -> ParseOutput: return _parse_docx_path(session.file_full_path, session) @@ -117,7 +116,7 @@ def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: class XlsParseAdapter: document_format: object - def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: + def parse(self, session: ParseSession) -> ParseOutput: from app.services.document_parser.legacy_converter import xls_to_xlsx converted_xlsx_path, _ = xls_to_xlsx( @@ -131,7 +130,7 @@ def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: class XlsxParseAdapter: document_format: object - def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: + def parse(self, session: ParseSession) -> ParseOutput: return _parse_xlsx_path(session.file_full_path, session) @@ -139,7 +138,7 @@ def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: class PptxParseAdapter: document_format: object - def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: + def parse(self, session: ParseSession) -> ParseOutput: from app.services.document_parser.pptx_parser import parse_pptx parsed_df = parse_pptx( @@ -152,14 +151,14 @@ def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: relative_root=session.relative_root, baseurl=session.base_url, ) - return session.full_output_dir, parsed_df + return ParseOutput(output_dir=session.full_output_dir, parsed_df=parsed_df) @dataclass(frozen=True) class MarkdownParseAdapter: document_format: object - def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: + def parse(self, session: ParseSession) -> ParseOutput: from app.services.document_parser.md_parser import parse_md parsed_df = parse_md( @@ -169,21 +168,21 @@ def parse(self, session: ParseSession) -> tuple[str, pd.DataFrame | None]: base_llm_paras=session.base_llm_paras, relative_root=session.relative_root, ) - return session.full_output_dir, parsed_df + return ParseOutput(output_dir=session.full_output_dir, parsed_df=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(self, session: ParseSession) -> ParseOutput: + return ParseOutput(output_dir=session.full_output_dir, parsed_df=None) def _parse_docx_path( docx_path: str, session: ParseSession, -) -> tuple[str, pd.DataFrame | None]: +) -> ParseOutput: from app.services.document_parser.doc_parser import convert_doc2dics, parse_docx parsed_structure, dataframe_list = parse_docx( @@ -201,13 +200,13 @@ def _parse_docx_path( base_llm_paras=session.base_llm_paras, relative_root=session.relative_root, ) - return session.full_output_dir, parsed_df + return ParseOutput(output_dir=session.full_output_dir, parsed_df=parsed_df) def _parse_xlsx_path( xlsx_path: str, session: ParseSession, -) -> tuple[str, pd.DataFrame | None]: +) -> ParseOutput: from app.services.document_parser.excel_table_parser import parse_xlsx parsed_df = parse_xlsx( @@ -218,4 +217,4 @@ def _parse_xlsx_path( base_llm_paras=session.base_llm_paras, relative_root=session.relative_root, ) - return session.full_output_dir, parsed_df + return ParseOutput(output_dir=session.full_output_dir, parsed_df=parsed_df) diff --git a/apps/worker/app/services/document_parser/orchestration/parse_output.py b/apps/worker/app/services/document_parser/orchestration/parse_output.py new file mode 100644 index 000000000..60084e6dd --- /dev/null +++ b/apps/worker/app/services/document_parser/orchestration/parse_output.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import pandas as pd + + +@dataclass(frozen=True) +class ParseOutput: + """Parser adapter output with a compatibility view for legacy callers.""" + + output_dir: str + parsed_df: pd.DataFrame | None + + @classmethod + def from_legacy_tuple( + cls, + output_dir: str, + parsed_df: pd.DataFrame | None, + ) -> ParseOutput: + return cls(output_dir=output_dir, parsed_df=parsed_df) + + @property + def rows_count(self) -> int: + if self.parsed_df is None: + return 0 + return len(self.parsed_df) + + def with_dataframe(self, parsed_df: pd.DataFrame | None) -> ParseOutput: + return ParseOutput(output_dir=self.output_dir, parsed_df=parsed_df) + + def as_legacy_tuple(self) -> tuple[str, pd.DataFrame | None]: + return self.output_dir, self.parsed_df diff --git a/apps/worker/app/services/document_parser/orchestration/parse_pipeline.py b/apps/worker/app/services/document_parser/orchestration/parse_pipeline.py index b31d36713..909a31197 100644 --- a/apps/worker/app/services/document_parser/orchestration/parse_pipeline.py +++ b/apps/worker/app/services/document_parser/orchestration/parse_pipeline.py @@ -1,27 +1,21 @@ from __future__ import annotations -from dataclasses import dataclass - -import pandas as pd - from app.services.document_parser.orchestration.parse_input import ParseInput +from app.services.document_parser.orchestration.parse_output import ParseOutput 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 -@dataclass(frozen=True) -class ParsePipelineResult: - output_dir: str - parsed_df: pd.DataFrame | None - - def as_legacy_tuple(self) -> tuple[str, pd.DataFrame | None]: - return self.output_dir, self.parsed_df +ParsePipelineResult = ParseOutput def run_parse_pipeline(parse_input: ParseInput) -> ParsePipelineResult: """Run parser session building, format routing, and output postprocessing.""" session = build_parse_session(parse_input) - output_dir, parsed_df = route_document_parse(session) - processed_df = apply_parse_postprocess(output_dir, parsed_df) - return ParsePipelineResult(output_dir=output_dir, parsed_df=processed_df) + parsed_output = route_document_parse(session) + processed_df = apply_parse_postprocess( + parsed_output.output_dir, + parsed_output.parsed_df, + ) + return parsed_output.with_dataframe(processed_df) 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 85402339e..68bb57e73 100644 --- a/apps/worker/app/services/document_parser/orchestration/route_parse.py +++ b/apps/worker/app/services/document_parser/orchestration/route_parse.py @@ -1,13 +1,12 @@ -import pandas as pd - from app.services.document_parser.orchestration.format_router import ( get_document_parse_adapter, resolve_document_format, ) +from app.services.document_parser.orchestration.parse_output import ParseOutput from app.services.document_parser.orchestration.parse_session import ParseSession -def route_document_parse(session: ParseSession) -> tuple[str, pd.DataFrame | None]: +def route_document_parse(session: ParseSession) -> ParseOutput: """Route a parser session to the correct adapter and return its output.""" document_format = resolve_document_format(session.file_full_path) adapter = get_document_parse_adapter(document_format) diff --git a/apps/worker/app/services/document_parser/parse_service.py b/apps/worker/app/services/document_parser/parse_service.py index becfcd807..6266d5990 100644 --- a/apps/worker/app/services/document_parser/parse_service.py +++ b/apps/worker/app/services/document_parser/parse_service.py @@ -3,10 +3,13 @@ import pandas as pd from app.services.document_parser.orchestration.parse_input import ParseInput, ParseOptions -from app.services.document_parser.orchestration.parse_pipeline import run_parse_pipeline +from app.services.document_parser.orchestration.parse_pipeline import ( + ParsePipelineResult, + run_parse_pipeline, +) -def checkerboard_inject_parse( +def checkerboard_parse_output( file_full_path: str, filename: str, output_dir: str, @@ -24,8 +27,8 @@ def checkerboard_inject_parse( base_url: str = "", fragment_content: str = "", s3_key: str | None = None, -) -> tuple[str, pd.DataFrame | None]: - """Run the stable parser seam using dedicated orchestration modules.""" +) -> ParsePipelineResult: + """Run the stable parser seam and return the parser output contract.""" parse_input = ParseInput( file_full_path=file_full_path, filename=filename, @@ -47,4 +50,45 @@ def checkerboard_inject_parse( fragment_content=fragment_content, s3_key=s3_key, ) - return run_parse_pipeline(parse_input).as_legacy_tuple() + return run_parse_pipeline(parse_input) + + +def checkerboard_inject_parse( + file_full_path: str, + filename: str, + output_dir: str, + internal_output_filename: str, + job_id: str | None = None, + kb_dir: str = "Default_Root", + 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 = "", + base_url: str = "", + fragment_content: str = "", + s3_key: str | None = None, +) -> tuple[str, pd.DataFrame | None]: + """Run the parser seam using the legacy tuple interface.""" + return checkerboard_parse_output( + file_full_path=file_full_path, + filename=filename, + output_dir=output_dir, + internal_output_filename=internal_output_filename, + job_id=job_id, + kb_dir=kb_dir, + 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, + add_frag_desc=add_frag_desc, + base_url=base_url, + fragment_content=fragment_content, + s3_key=s3_key, + ).as_legacy_tuple() 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 35ab7e531..c76f4413d 100644 --- a/apps/worker/tests/contract/test_document_parser_architecture_contract.py +++ b/apps/worker/tests/contract/test_document_parser_architecture_contract.py @@ -80,6 +80,7 @@ def test_document_format_router_uses_adapters( resolve_document_format, ) from app.services.document_parser.orchestration.parse_input import ParseInput + from app.services.document_parser.orchestration.parse_output import ParseOutput from app.services.document_parser.orchestration.parse_session import ParseSession from app.services.document_parser.orchestration.route_parse import route_document_parse @@ -109,10 +110,10 @@ def test_document_format_router_uses_adapters( relative_root="Default_Root/report.pdf", ) - output_dir, actual_df = route_document_parse(session) + output = route_document_parse(session) - assert output_dir == str(tmp_path) - assert actual_df is parsed_df + assert output == ParseOutput(output_dir=str(tmp_path), parsed_df=parsed_df) + assert output.as_legacy_tuple() == (str(tmp_path), parsed_df) def test_parse_pipeline_owns_route_and_postprocess_contract( @@ -121,6 +122,7 @@ def test_parse_pipeline_owns_route_and_postprocess_contract( tmp_path: Path, ) -> None: from app.services.document_parser.orchestration.parse_input import ParseInput + from app.services.document_parser.orchestration.parse_output import ParseOutput from app.services.document_parser.orchestration.parse_pipeline import ( ParsePipelineResult, run_parse_pipeline, @@ -143,7 +145,7 @@ def test_parse_pipeline_owns_route_and_postprocess_contract( ) monkeypatch.setattr( "app.services.document_parser.orchestration.parse_pipeline.route_document_parse", - lambda _session: (routed_output_dir, routed_df), + lambda _session: ParseOutput(output_dir=routed_output_dir, parsed_df=routed_df), ) def fake_postprocess(output_dir: str, parsed_df: pd.DataFrame | None) -> pd.DataFrame: @@ -172,6 +174,40 @@ def fake_postprocess(output_dir: str, parsed_df: pd.DataFrame | None) -> pd.Data assert result.as_legacy_tuple() == (routed_output_dir, processed_df) +def test_parser_service_keeps_legacy_tuple_compatibility( + worker_contract_environment: None, + monkeypatch: Any, + tmp_path: Path, +) -> None: + import app.services.document_parser.parse_service as parse_service + from app.services.document_parser.orchestration.parse_output import ParseOutput + + parsed_df = pd.DataFrame([{"content": "legacy"}]) + expected_output = ParseOutput(output_dir=str(tmp_path), parsed_df=parsed_df) + + monkeypatch.setattr( + parse_service, + "run_parse_pipeline", + lambda _parse_input: expected_output, + ) + + output = parse_service.checkerboard_parse_output( + file_full_path=str(tmp_path / "legacy.pdf"), + filename="legacy.pdf", + output_dir=str(tmp_path), + internal_output_filename="legacy.pdf", + ) + legacy_tuple = parse_service.checkerboard_inject_parse( + file_full_path=str(tmp_path / "legacy.pdf"), + filename="legacy.pdf", + output_dir=str(tmp_path), + internal_output_filename="legacy.pdf", + ) + + assert output is expected_output + assert legacy_tuple == (str(tmp_path), parsed_df) + + def test_rendered_pdf_transform_centralizes_temporary_pdf_cleanup( worker_contract_environment: None, monkeypatch: Any, diff --git a/apps/worker/tests/contract/test_job_lifecycle_post_commit_contract.py b/apps/worker/tests/contract/test_job_lifecycle_post_commit_contract.py new file mode 100644 index 000000000..a9b99fc26 --- /dev/null +++ b/apps/worker/tests/contract/test_job_lifecycle_post_commit_contract.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +from typing import Any, cast + + +def test_post_commit_effect_runner_uses_primitive_effect_plan( + worker_contract_environment: None, +) -> None: + from shared.services.job_post_commit_effects_sync import ( + PostCommitEffectPlan, + SyncJobPostCommitEffectRunner, + ) + + cache_invalidations: list[dict[str, Any]] = [] + webhook_event_ids: list[str] = [] + + class FakePublicationFinalizer: + def invalidate_cache_after_commit( + self, + cache_invalidation: dict[str, Any] | None, + ) -> None: + assert cache_invalidation is not None + cache_invalidations.append(cache_invalidation) + + class FakeWebhookOutbox: + def enqueue_event_id_after_commit( + self, + webhook_event_id: str | None, + ) -> None: + assert webhook_event_id is not None + webhook_event_ids.append(webhook_event_id) + + runner = SyncJobPostCommitEffectRunner( + publication_finalizer=cast(Any, FakePublicationFinalizer()), + webhook_outbox=cast(Any, FakeWebhookOutbox()), + ) + plan = PostCommitEffectPlan.from_success( + cache_invalidation={ + "job_id": "job_post_commit", + "user_id": "user-1", + "namespaces": ["default", "finance"], + }, + webhook_event_id="event-1", + ) + + runner.run(plan) + + assert cache_invalidations == [ + { + "job_id": "job_post_commit", + "user_id": "user-1", + "namespaces": ["default", "finance"], + } + ] + assert webhook_event_ids == ["event-1"] 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 19a0f5858..a6881d960 100644 --- a/apps/worker/tests/contract/test_stale_job_sweeper_contract.py +++ b/apps/worker/tests/contract/test_stale_job_sweeper_contract.py @@ -267,3 +267,36 @@ def test_should_expose_sync_state_machine_rejection_reason( assert outcome.from_state == "done" assert outcome.to_state == "running" assert outcome.attempts == 1 + + +def test_should_keep_sync_state_machine_bool_facade_for_missing_jobs( + worker_contract_environment: None, +) -> None: + from shared.core.database_sync import get_sync_db_context + from shared.core.state_machine.service_sync import SyncStateMachineService + from shared.core.state_machine.states import JobStatus + from shared.services.redis.redis_sync_service import SyncRedisServiceFactory + + missing_job_id = f"job_missing_{uuid4().hex[:12]}" + state_machine = SyncStateMachineService( + redis_service=SyncRedisServiceFactory.get_service() + ) + + with get_sync_db_context() as db: + outcome = state_machine.transition_outcome( + db, + missing_job_id, + JobStatus.RUNNING.value, + transition_reason="contract_missing_job", + ) + did_transition = state_machine.transition( + db, + missing_job_id, + JobStatus.RUNNING.value, + transition_reason="contract_missing_job_bool", + ) + + assert outcome.succeeded is False + assert outcome.reason == "job_not_found" + assert outcome.attempts == 1 + assert did_transition is False diff --git a/docs/adr/0001-keep-routes-and-worker-tasks-as-adapters.md b/docs/adr/0001-keep-routes-and-worker-tasks-as-adapters.md new file mode 100644 index 000000000..041ad3271 --- /dev/null +++ b/docs/adr/0001-keep-routes-and-worker-tasks-as-adapters.md @@ -0,0 +1,25 @@ +# 0001 Keep Routes And Worker Tasks As Adapters + +## Status + +Accepted + +## Context + +Knowhere has HTTP routes, internal callback routes, and Celery tasks that start +workflows. Bugs become harder to localize when these adapters also own parser, +publication, billing, retrieval, or webhook implementation details. + +## Decision + +Routes and worker tasks should translate external inputs into application or +shared workflow calls. Domain behavior should live behind workflow modules such +as Document Ingestion, Worker Document Parsing, Publication, Retrieval, Billing +Workflow, Storage Event Intake, and Webhook delivery. + +## Consequences + +Adapters stay thin and contract tests can assert public behavior at HTTP, +worker-task, database, storage, Redis, or event surfaces. New feature work +should prefer deep workflow modules over adding branch-heavy logic to routes or +Celery task functions. diff --git a/docs/adr/0002-use-typed-workflow-outcomes.md b/docs/adr/0002-use-typed-workflow-outcomes.md new file mode 100644 index 000000000..4e7b8ec50 --- /dev/null +++ b/docs/adr/0002-use-typed-workflow-outcomes.md @@ -0,0 +1,28 @@ +# 0002 Use Typed Workflow Outcomes + +## Status + +Accepted + +## Context + +Several Knowhere workflows used tuples, dictionaries, or booleans to represent +rich behavior. That made retry, fallback, and transition decisions harder to +reason about because callers lost the reason behind a result. + +## Decision + +New or deepened workflow seams should use typed outcomes when callers need more +than a yes/no answer. Current examples are `JobTransitionOutcome`, +`WorkloadEstimate`, `ParseOutput`, `ParseArtifact`, `GeneratedResultPackage`, +and `PostCommitEffectPlan`. + +Boolean and tuple facades may remain where they are public compatibility +interfaces, but the implementation should concentrate reason-bearing behavior +behind typed outcomes. + +## Consequences + +Contract tests can assert stable behavior and failure reasons without reaching +into private helper call order. Compatibility wrappers should be kept small and +named as wrappers, not used as the primary implementation seam. diff --git a/docs/adr/0003-keep-retrieval-workflow-policy-explicit.md b/docs/adr/0003-keep-retrieval-workflow-policy-explicit.md new file mode 100644 index 000000000..ccf97b846 --- /dev/null +++ b/docs/adr/0003-keep-retrieval-workflow-policy-explicit.md @@ -0,0 +1,29 @@ +# 0003 Keep Retrieval Workflow Policy Explicit + +## Status + +Accepted + +## Context + +Retrieval requests include scope, document/section exclusions, data type, +signal paths, channel selection, channel weights, internal recall, and ranking +policy. When these fields are passed as loose keyword arguments, the legacy and +agentic routes can drift. + +## Decision + +Retrieval should keep request policy in typed request modules. `RetrievalQuery` +owns cache and route policy. `WorkflowRunRequest` and `WorkflowStepRequest` +carry the agentic workflow projection of that policy through planning and step +execution. + +Fields that are accepted but not yet implemented in a route, such as reranking +or threshold semantics for multi-step workflow answers, must remain explicit in +the request type rather than disappearing silently. + +## Consequences + +The request type becomes the test surface for retrieval policy. Product +behavior changes such as enabling workflow reranking or threshold filtering +should be separate decisions with contract tests for the selected phase. diff --git a/docs/adr/README.md b/docs/adr/README.md new file mode 100644 index 000000000..3d0cb4fd0 --- /dev/null +++ b/docs/adr/README.md @@ -0,0 +1,12 @@ +# Architecture Decision Records + +This directory records durable architecture decisions for Knowhere. Keep ADRs +short, repo-facing, and focused on decisions future maintainers or agents might +otherwise re-litigate. + +Use this shape: + +- Status +- Context +- Decision +- Consequences diff --git a/packages/shared-python/shared/core/state_machine/service.py b/packages/shared-python/shared/core/state_machine/service.py index 8407b48e7..6457e3964 100644 --- a/packages/shared-python/shared/core/state_machine/service.py +++ b/packages/shared-python/shared/core/state_machine/service.py @@ -16,7 +16,6 @@ from shared.core.state_machine.states import ( JobStatus, - is_valid_transition, ) from shared.core.state_machine.transition_payloads import ( build_failure_transition_metadata, @@ -26,6 +25,16 @@ utc_now_naive, ) from shared.core.state_machine.transition_outcome import JobTransitionOutcome +from shared.core.state_machine.transition_runner import ( + MAX_TRANSITION_ATTEMPTS, + TransitionJobSnapshot, + build_cas_conflict_outcome, + build_rollback_exception_outcome, + build_transition_exception_outcome, + get_cas_retry_delay_seconds, + prepare_transition_attempt, + should_retry_cas_conflict, +) from shared.models.database.job import Job from shared.models.database.job_state_audit_log import JobStateAuditLog from shared.services.redis import RedisServiceFactory @@ -76,35 +85,45 @@ async def transition_outcome( auto_commit: bool = True, ) -> JobTransitionOutcome: """Execute a state transition and preserve the reason if it is rejected.""" - max_retries = 3 + max_retries = MAX_TRANSITION_ATTEMPTS for attempt in range(max_retries): attempts = attempt + 1 try: job = await self._get_job_with_version(db, job_id) - if not job: - logger.error(f"Job {job_id} does not exist") - return JobTransitionOutcome.rejected( - job_id=job_id, - to_state=to_state, - reason="job_not_found", - attempts=attempts, - ) + prepared = prepare_transition_attempt( + job_id=job_id, + to_state=to_state, + snapshot=( + TransitionJobSnapshot(status=job.status, version=job.version) + if job + else None + ), + attempts=attempts, + ) + if prepared.outcome is not None: + if prepared.outcome.reason == "job_not_found": + logger.error(f"Job {job_id} does not exist") + elif prepared.outcome.reason == "invalid_transition": + logger.warning( + f"Job {job_id}: illegal transition " + f"{prepared.outcome.from_state} → {to_state}, rejected" + ) + return prepared.outcome - if not is_valid_transition(job.status, to_state): - logger.warning( - f"Job {job_id}: illegal transition {job.status} → {to_state}, rejected" - ) - return JobTransitionOutcome.rejected( + if not prepared.can_write: + logger.error(f"Job {job_id} does not exist") + return build_transition_exception_outcome( job_id=job_id, - from_state=job.status, to_state=to_state, - reason="invalid_transition", attempts=attempts, + error="transition precondition did not provide state/version", ) - old_state = job.status - old_version = job.version + old_state = prepared.from_state + old_version = prepared.version + assert old_state is not None + assert old_version is not None # Record audit log (before CAS so the INSERT is within the same tx) await self._record_audit_log( @@ -144,19 +163,18 @@ async def transition_outcome( ) # CAS miss — retry with backoff - if attempt < max_retries - 1: + if should_retry_cas_conflict(attempt, max_attempts=max_retries): logger.warning( f"Job {job_id} CAS conflict, retry {attempt + 1}/{max_retries}" ) - await asyncio.sleep(0.1 * (2**attempt)) + await asyncio.sleep(get_cas_retry_delay_seconds(attempt)) continue else: logger.error(f"Job {job_id} CAS retries exhausted") - return JobTransitionOutcome.rejected( + return build_cas_conflict_outcome( job_id=job_id, from_state=old_state, to_state=to_state, - reason="cas_conflict", attempts=attempts, ) @@ -167,26 +185,24 @@ async def transition_outcome( await db.rollback() except Exception as rollback_err: logger.warning(f"Job {job_id} rollback failed: {rollback_err}") - return JobTransitionOutcome.rejected( + return build_rollback_exception_outcome( job_id=job_id, to_state=to_state, - reason="rollback_exception", attempts=attempts, - error_message=str(rollback_err), + error=rollback_err, ) - return JobTransitionOutcome.rejected( + return build_transition_exception_outcome( job_id=job_id, to_state=to_state, - reason="transition_exception", attempts=attempts, - error_message=str(e), + error=e, ) - return JobTransitionOutcome.rejected( + return build_cas_conflict_outcome( job_id=job_id, to_state=to_state, - reason="cas_conflict", attempts=max_retries, + from_state=None, ) async def mark_failed( 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 b527cb394..8935304e9 100644 --- a/packages/shared-python/shared/core/state_machine/service_sync.py +++ b/packages/shared-python/shared/core/state_machine/service_sync.py @@ -15,7 +15,6 @@ from shared.core.state_machine.states import ( JobStatus, - is_valid_transition, ) from shared.core.state_machine.transition_payloads import ( build_failure_transition_metadata, @@ -25,6 +24,16 @@ utc_now_naive, ) from shared.core.state_machine.transition_outcome import JobTransitionOutcome +from shared.core.state_machine.transition_runner import ( + MAX_TRANSITION_ATTEMPTS, + TransitionJobSnapshot, + build_cas_conflict_outcome, + build_rollback_exception_outcome, + build_transition_exception_outcome, + get_cas_retry_delay_seconds, + prepare_transition_attempt, + should_retry_cas_conflict, +) 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 @@ -71,35 +80,45 @@ def transition_outcome( metadata: Optional[Dict[str, Any]] = None, ) -> JobTransitionOutcome: """Execute a state transition and preserve the reason if it is rejected.""" - max_retries = 3 + max_retries = MAX_TRANSITION_ATTEMPTS for attempt in range(max_retries): attempts = attempt + 1 try: job = self._get_job_with_version(db, job_id) - if not job: - logger.error(f"Job {job_id} does not exist") - return JobTransitionOutcome.rejected( - job_id=job_id, - to_state=to_state, - reason="job_not_found", - attempts=attempts, - ) + prepared = prepare_transition_attempt( + job_id=job_id, + to_state=to_state, + snapshot=( + TransitionJobSnapshot(status=job.status, version=job.version) + if job + else None + ), + attempts=attempts, + ) + if prepared.outcome is not None: + if prepared.outcome.reason == "job_not_found": + logger.error(f"Job {job_id} does not exist") + elif prepared.outcome.reason == "invalid_transition": + logger.warning( + f"Job {job_id}: illegal transition " + f"{prepared.outcome.from_state} → {to_state}, rejected" + ) + return prepared.outcome - if not is_valid_transition(job.status, to_state): - logger.warning( - f"Job {job_id}: illegal transition {job.status} → {to_state}, rejected" - ) - return JobTransitionOutcome.rejected( + if not prepared.can_write: + logger.error(f"Job {job_id} does not exist") + return build_transition_exception_outcome( job_id=job_id, - from_state=job.status, to_state=to_state, - reason="invalid_transition", attempts=attempts, + error="transition precondition did not provide state/version", ) - old_state = job.status - old_version = job.version + old_state = prepared.from_state + old_version = prepared.version + assert old_state is not None + assert old_version is not None self._record_audit_log( db, @@ -126,21 +145,20 @@ def transition_outcome( attempts=attempts, ) - if attempt < max_retries - 1: + if should_retry_cas_conflict(attempt, max_attempts=max_retries): logger.warning( f"Job {job_id} CAS conflict, retry {attempt + 1}/{max_retries}" ) import gevent - gevent.sleep(0.1 * (2**attempt)) + gevent.sleep(get_cas_retry_delay_seconds(attempt)) continue else: logger.error(f"Job {job_id} CAS retries exhausted") - return JobTransitionOutcome.rejected( + return build_cas_conflict_outcome( job_id=job_id, from_state=old_state, to_state=to_state, - reason="cas_conflict", attempts=attempts, ) @@ -151,26 +169,24 @@ def transition_outcome( db.rollback() except Exception as rollback_err: logger.warning(f"Job {job_id} rollback failed: {rollback_err}") - return JobTransitionOutcome.rejected( + return build_rollback_exception_outcome( job_id=job_id, to_state=to_state, - reason="rollback_exception", attempts=attempts, - error_message=str(rollback_err), + error=rollback_err, ) - return JobTransitionOutcome.rejected( + return build_transition_exception_outcome( job_id=job_id, to_state=to_state, - reason="transition_exception", attempts=attempts, - error_message=str(e), + error=e, ) - return JobTransitionOutcome.rejected( + return build_cas_conflict_outcome( job_id=job_id, to_state=to_state, - reason="cas_conflict", attempts=max_retries, + from_state=None, ) def mark_failed( diff --git a/packages/shared-python/shared/core/state_machine/transition_runner.py b/packages/shared-python/shared/core/state_machine/transition_runner.py new file mode 100644 index 000000000..199ee7e59 --- /dev/null +++ b/packages/shared-python/shared/core/state_machine/transition_runner.py @@ -0,0 +1,120 @@ +from __future__ import annotations + +from dataclasses import dataclass + +from shared.core.state_machine.states import is_valid_transition +from shared.core.state_machine.transition_outcome import JobTransitionOutcome + +MAX_TRANSITION_ATTEMPTS = 3 + + +@dataclass(frozen=True) +class TransitionJobSnapshot: + status: str + version: int + + +@dataclass(frozen=True) +class PreparedTransitionAttempt: + outcome: JobTransitionOutcome | None + from_state: str | None = None + version: int | None = None + + @property + def can_write(self) -> bool: + return self.outcome is None and self.from_state is not None and self.version is not None + + +def prepare_transition_attempt( + *, + job_id: str, + to_state: str, + snapshot: TransitionJobSnapshot | None, + attempts: int, +) -> PreparedTransitionAttempt: + if snapshot is None: + return PreparedTransitionAttempt( + outcome=JobTransitionOutcome.rejected( + job_id=job_id, + to_state=to_state, + reason="job_not_found", + attempts=attempts, + ) + ) + + if not is_valid_transition(snapshot.status, to_state): + return PreparedTransitionAttempt( + outcome=JobTransitionOutcome.rejected( + job_id=job_id, + from_state=snapshot.status, + to_state=to_state, + reason="invalid_transition", + attempts=attempts, + ) + ) + + return PreparedTransitionAttempt( + outcome=None, + from_state=snapshot.status, + version=snapshot.version, + ) + + +def should_retry_cas_conflict( + attempt_index: int, + *, + max_attempts: int = MAX_TRANSITION_ATTEMPTS, +) -> bool: + return attempt_index < max_attempts - 1 + + +def get_cas_retry_delay_seconds(attempt_index: int) -> float: + return 0.1 * (2**attempt_index) + + +def build_cas_conflict_outcome( + *, + job_id: str, + from_state: str | None, + to_state: str, + attempts: int, +) -> JobTransitionOutcome: + return JobTransitionOutcome.rejected( + job_id=job_id, + from_state=from_state, + to_state=to_state, + reason="cas_conflict", + attempts=attempts, + ) + + +def build_transition_exception_outcome( + *, + job_id: str, + to_state: str, + attempts: int, + error: Exception | str, +) -> JobTransitionOutcome: + return JobTransitionOutcome.rejected( + job_id=job_id, + to_state=to_state, + reason="transition_exception", + attempts=attempts, + error_message=str(error), + ) + + +def build_rollback_exception_outcome( + *, + job_id: str, + to_state: str, + attempts: int, + error: Exception | str, +) -> JobTransitionOutcome: + return JobTransitionOutcome.rejected( + job_id=job_id, + to_state=to_state, + reason="rollback_exception", + attempts=attempts, + error_message=str(error), + ) diff --git a/packages/shared-python/shared/services/job_failure_sync.py b/packages/shared-python/shared/services/job_failure_sync.py index cb16775d2..ad3a313dd 100644 --- a/packages/shared-python/shared/services/job_failure_sync.py +++ b/packages/shared-python/shared/services/job_failure_sync.py @@ -1,5 +1,6 @@ from __future__ import annotations +from dataclasses import dataclass from typing import Any from loguru import logger @@ -9,12 +10,18 @@ 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_post_commit_effects_sync import PostCommitEffectPlan from shared.services.job_webhook_outbox_sync import SyncJobWebhookOutbox from shared.utils.error_details import normalize_error_details +@dataclass(frozen=True) +class JobFailureFinalization: + succeeded: bool + post_commit_effects: PostCommitEffectPlan + + class SyncJobFailureFinalizer: """Finalize failed Jobs inside the lifecycle transaction.""" @@ -38,7 +45,7 @@ def finalize( error_code: str, error_details: dict[str, Any] | None, should_refund: bool, - ) -> tuple[bool, WebhookEvent | None]: + ) -> JobFailureFinalization: transition_outcome = self._state_machine.mark_failed_outcome( db, job_id, @@ -51,7 +58,10 @@ def finalize( f"Job {job_id} mark_failed transition failed: " f"reason={transition_outcome.reason}" ) - return False, None + return JobFailureFinalization( + succeeded=False, + post_commit_effects=PostCommitEffectPlan.none(), + ) if should_refund: self._try_refund_credits(db, job_id) @@ -70,10 +80,12 @@ def finalize( ), }, ) - return True, webhook_event - - def enqueue_webhook_after_commit(self, webhook_event: WebhookEvent | None) -> None: - self._webhook_outbox.enqueue_after_commit(webhook_event) + return JobFailureFinalization( + succeeded=True, + post_commit_effects=PostCommitEffectPlan.from_failure( + webhook_event_id=webhook_event.id if webhook_event else None, + ), + ) def _try_refund_credits(self, db: Session, job_id: str) -> None: try: diff --git a/packages/shared-python/shared/services/job_lifecycle_sync.py b/packages/shared-python/shared/services/job_lifecycle_sync.py index d6d1d54d8..ad0dcfc70 100644 --- a/packages/shared-python/shared/services/job_lifecycle_sync.py +++ b/packages/shared-python/shared/services/job_lifecycle_sync.py @@ -18,6 +18,10 @@ from shared.core.database_sync import get_sync_db_context from shared.services.job_failure_sync import SyncJobFailureFinalizer +from shared.services.job_post_commit_effects_sync import ( + PostCommitEffectPlan, + SyncJobPostCommitEffectRunner, +) from shared.services.job_success_sync import SyncJobSuccessFinalizer from shared.services.redis.redis_sync_service import ( SyncRedisServiceFactory, @@ -34,6 +38,7 @@ class SyncJobLifecycleService: def __init__(self) -> None: self._success_finalizer = SyncJobSuccessFinalizer() self._failure_finalizer = SyncJobFailureFinalizer() + self._post_commit_effect_runner = SyncJobPostCommitEffectRunner() # ── Public API ────────────────────────────────────────────────────── @@ -77,7 +82,8 @@ def finalize_job_success( should_commit=lambda finalization: finalization.response.get("status") == "success", build_response=lambda finalization: finalization.response, - run_after_commit=self._success_finalizer.run_post_commit_actions, + build_effect_plan=lambda finalization: finalization.post_commit_effects, + run_after_commit_effects=self._post_commit_effect_runner.run, ) def finalize_job_failure( @@ -110,11 +116,10 @@ def finalize_job_failure( error_details=error_details, should_refund=should_refund, ), - should_commit=lambda finalization: finalization[0], - build_response=lambda finalization: finalization[0], - run_after_commit=lambda finalization: ( - self._failure_finalizer.enqueue_webhook_after_commit(finalization[1]) - ), + should_commit=lambda finalization: finalization.succeeded, + build_response=lambda finalization: finalization.succeeded, + build_effect_plan=lambda finalization: finalization.post_commit_effects, + run_after_commit_effects=self._post_commit_effect_runner.run, ) def update_progress( @@ -173,7 +178,8 @@ def _run_lifecycle_transaction( finalize: Callable[[Session], _FinalizationT], should_commit: Callable[[_FinalizationT], bool], build_response: Callable[[_FinalizationT], _ResponseT], - run_after_commit: Callable[[_FinalizationT], None], + build_effect_plan: Callable[[_FinalizationT], PostCommitEffectPlan], + run_after_commit_effects: Callable[[PostCommitEffectPlan], None], ) -> _ResponseT: with get_sync_db_context() as db: try: @@ -185,7 +191,7 @@ def _run_lifecycle_transaction( db.commit() logger.info(f"Job {job_id} {label} transaction committed") - run_after_commit(finalization) + run_after_commit_effects(build_effect_plan(finalization)) return build_response(finalization) except Exception as exc: diff --git a/packages/shared-python/shared/services/job_post_commit_effects_sync.py b/packages/shared-python/shared/services/job_post_commit_effects_sync.py new file mode 100644 index 000000000..35546ed71 --- /dev/null +++ b/packages/shared-python/shared/services/job_post_commit_effects_sync.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from shared.services.job_publication_sync import SyncJobPublicationFinalizer +from shared.services.job_webhook_outbox_sync import SyncJobWebhookOutbox + + +@dataclass(frozen=True) +class PostCommitEffectPlan: + retrieval_cache_invalidations: tuple[dict[str, Any], ...] = field( + default_factory=tuple + ) + webhook_event_ids: tuple[str, ...] = field(default_factory=tuple) + + @classmethod + def none(cls) -> PostCommitEffectPlan: + return cls() + + @classmethod + def from_success( + cls, + *, + cache_invalidation: dict[str, Any] | None, + webhook_event_id: str | None, + ) -> PostCommitEffectPlan: + return cls( + retrieval_cache_invalidations=( + (cache_invalidation,) if cache_invalidation else () + ), + webhook_event_ids=((webhook_event_id,) if webhook_event_id else ()), + ) + + @classmethod + def from_failure( + cls, + *, + webhook_event_id: str | None, + ) -> PostCommitEffectPlan: + return cls(webhook_event_ids=((webhook_event_id,) if webhook_event_id else ())) + + +class SyncJobPostCommitEffectRunner: + def __init__( + self, + *, + publication_finalizer: SyncJobPublicationFinalizer | None = None, + webhook_outbox: SyncJobWebhookOutbox | None = None, + ) -> None: + self._publication_finalizer = ( + publication_finalizer or SyncJobPublicationFinalizer() + ) + self._webhook_outbox = webhook_outbox or SyncJobWebhookOutbox() + + def run(self, plan: PostCommitEffectPlan) -> None: + for cache_invalidation in plan.retrieval_cache_invalidations: + self._publication_finalizer.invalidate_cache_after_commit( + cache_invalidation + ) + for webhook_event_id in plan.webhook_event_ids: + self._webhook_outbox.enqueue_event_id_after_commit(webhook_event_id) diff --git a/packages/shared-python/shared/services/job_success_sync.py b/packages/shared-python/shared/services/job_success_sync.py index fbdc36e92..1f7b8b20b 100644 --- a/packages/shared-python/shared/services/job_success_sync.py +++ b/packages/shared-python/shared/services/job_success_sync.py @@ -7,7 +7,7 @@ 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_post_commit_effects_sync import PostCommitEffectPlan 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 @@ -16,8 +16,7 @@ @dataclass(frozen=True) class JobSuccessFinalization: response: dict[str, Any] - cache_invalidation: dict[str, Any] | None - webhook_event: WebhookEvent | None + post_commit_effects: PostCommitEffectPlan class SyncJobSuccessFinalizer: @@ -88,8 +87,7 @@ def finalize( "job_id": job_id, "reason": "state_transition_failed", }, - cache_invalidation=None, - webhook_event=None, + post_commit_effects=PostCommitEffectPlan.none(), ) webhook_event = self._webhook_outbox.create_event( @@ -103,15 +101,8 @@ def finalize( "job_id": job_id, "stored_count": stored_count, }, - cache_invalidation=publication_outcome.cache_invalidation, - webhook_event=webhook_event, + post_commit_effects=PostCommitEffectPlan.from_success( + cache_invalidation=publication_outcome.cache_invalidation, + webhook_event_id=webhook_event.id if webhook_event else None, + ), ) - - 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 index d699a26d5..05ea84333 100644 --- a/packages/shared-python/shared/services/job_webhook_outbox_sync.py +++ b/packages/shared-python/shared/services/job_webhook_outbox_sync.py @@ -59,27 +59,31 @@ def create_event( def enqueue_after_commit(self, webhook_event: WebhookEvent | None) -> None: if not webhook_event: return + self.enqueue_event_id_after_commit(webhook_event.id) + def enqueue_event_id_after_commit(self, webhook_event_id: str | None) -> None: + if not webhook_event_id: + 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) + 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}" + 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"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}" + f"event_id={webhook_event_id}, error={exc}" ) diff --git a/packages/shared-python/shared/services/retrieval/agentic/discovery_phase.py b/packages/shared-python/shared/services/retrieval/agentic/discovery_phase.py index 2af9ce690..f244cadfa 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/discovery_phase.py +++ b/packages/shared-python/shared/services/retrieval/agentic/discovery_phase.py @@ -32,6 +32,7 @@ async def run_initial_discovery( filter_mode: str, channels: list[str] | None, channel_weights: dict[str, float] | None, + internal_recall_k: int | None, bootstrap_llm_fn: LLMFn | None, ) -> list[dict[str, Any]]: discovery_kwargs: dict[str, Any] = { @@ -46,6 +47,7 @@ async def run_initial_discovery( "filter_mode": filter_mode, "channels": channels, "channel_weights": channel_weights, + "internal_recall_k": internal_recall_k, } logger.info(" agentic: Phase 1 — discovery + document selection") diff --git a/packages/shared-python/shared/services/retrieval/agentic/discovery_tools.py b/packages/shared-python/shared/services/retrieval/agentic/discovery_tools.py index 5f32e36cb..44ac0f038 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/discovery_tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/discovery_tools.py @@ -131,7 +131,11 @@ async def bottom_discovery( 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 [] + ) fused_rows = merge_same_section_rows(fused_rows) if fused_rows: diff --git a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py index c8fdfa19a..3dfa91a80 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py +++ b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py @@ -85,6 +85,7 @@ async def run( filter_mode: str = 'delete', channels: list[str] | None = None, channel_weights: dict[str, float] | None = None, + internal_recall_k: int | None = None, config: AgentRunConfig | None = None, ledger: BudgetLedger | None = None, parent_run_id: str | None = None, @@ -133,6 +134,7 @@ async def run( 'exclude_document_ids': exclude_document_ids, 'exclude_sections': exclude_sections, 'signal_paths': signal_paths, + 'internal_recall_k': internal_recall_k, }, parent_run_id=parent_run_id, workflow_step_id=workflow_step_id, @@ -174,6 +176,7 @@ async def run( filter_mode=filter_mode, channels=channels, channel_weights=channel_weights, + internal_recall_k=internal_recall_k, bootstrap_llm_fn=bootstrap_llm_fn, ) diff --git a/packages/shared-python/shared/services/retrieval/execution_routes.py b/packages/shared-python/shared/services/retrieval/execution_routes.py index cfd5984c7..2cd4f5ea5 100644 --- a/packages/shared-python/shared/services/retrieval/execution_routes.py +++ b/packages/shared-python/shared/services/retrieval/execution_routes.py @@ -104,21 +104,12 @@ async def _run_agentic_route( context: RetrievalRouteContext, ) -> RetrievalRouteOutcome: from shared.services.retrieval.workflow.orchestrator import WorkflowOrchestrator + from shared.services.retrieval.workflow.run_request import WorkflowRunRequest workflow = WorkflowOrchestrator() - workflow_result = await workflow.run( + workflow_result = await workflow.run_request( 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, + request=WorkflowRunRequest.from_route_context(context), ) resolved_references = await resolve_workflow_references( diff --git a/packages/shared-python/shared/services/retrieval/query_request.py b/packages/shared-python/shared/services/retrieval/query_request.py index 1877ae363..253b2e7b2 100644 --- a/packages/shared-python/shared/services/retrieval/query_request.py +++ b/packages/shared-python/shared/services/retrieval/query_request.py @@ -107,7 +107,9 @@ def build_route_context(self) -> RetrievalRouteContext: filter_mode=self.filter_mode, channels=self.channels, channel_weights=self.channel_weights, + rerank=self.rerank, threshold=self.threshold, + internal_recall_k=self.internal_recall_k, effective_recall_k=self.resolve_effective_recall_k(), use_agentic=self.use_agentic, ) diff --git a/packages/shared-python/shared/services/retrieval/route_types.py b/packages/shared-python/shared/services/retrieval/route_types.py index f1aa94160..1e7e53d79 100644 --- a/packages/shared-python/shared/services/retrieval/route_types.py +++ b/packages/shared-python/shared/services/retrieval/route_types.py @@ -21,7 +21,9 @@ class RetrievalRouteContext: filter_mode: str channels: list[str] | None channel_weights: dict[str, float] | None + rerank: bool threshold: float + internal_recall_k: int | None effective_recall_k: int use_agentic: bool | None diff --git a/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py b/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py index ad56e6ed7..e830ee5fe 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py +++ b/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py @@ -75,6 +75,9 @@ async def run( filter_mode: str = 'delete', channels: list[str] | None = None, channel_weights: dict[str, float] | None = None, + internal_recall_k: int | None = None, + rerank: bool = False, + threshold: float = 0.0, llm_fn=None, ) -> WorkflowResult: request = WorkflowRunRequest( @@ -89,6 +92,9 @@ async def run( filter_mode=filter_mode, channels=channels, channel_weights=channel_weights, + internal_recall_k=internal_recall_k, + rerank=rerank, + threshold=threshold, ) return await self.run_request(db, request=request, llm_fn=llm_fn) diff --git a/packages/shared-python/shared/services/retrieval/workflow/run_request.py b/packages/shared-python/shared/services/retrieval/workflow/run_request.py index a695001af..4eff26d91 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/run_request.py +++ b/packages/shared-python/shared/services/retrieval/workflow/run_request.py @@ -2,6 +2,8 @@ from dataclasses import dataclass +from shared.services.retrieval.route_types import RetrievalRouteContext +from shared.services.retrieval.settings import INTERNAL_RECALL_K_MULTIPLIER from shared.services.retrieval.workflow.types import PlannedStep @@ -18,13 +20,39 @@ class WorkflowRunRequest: filter_mode: str = "delete" channels: list[str] | None = None channel_weights: dict[str, float] | None = None + internal_recall_k: int | None = None + rerank: bool = False + threshold: float = 0.0 + + @classmethod + def from_route_context( + cls, + context: RetrievalRouteContext, + ) -> WorkflowRunRequest: + return cls( + 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, + internal_recall_k=context.internal_recall_k, + rerank=context.rerank, + threshold=context.threshold, + ) def for_step(self, step: PlannedStep) -> WorkflowStepRequest: + step_top_k = step.top_k or self.top_k return WorkflowStepRequest( user_id=self.user_id, namespace=self.namespace, query=step.sub_query, - top_k=step.top_k or self.top_k, + top_k=step_top_k, exclude_document_ids=self.exclude_document_ids, exclude_sections=self.exclude_sections, data_type=step.data_type or self.data_type, @@ -32,8 +60,16 @@ def for_step(self, step: PlannedStep) -> WorkflowStepRequest: filter_mode=self.filter_mode, channels=self.channels, channel_weights=self.channel_weights, + internal_recall_k=self.resolve_effective_recall_k(step_top_k), + rerank=self.rerank, + threshold=self.threshold, ) + def resolve_effective_recall_k(self, top_k: int) -> int: + if self.internal_recall_k is not None: + return self.internal_recall_k + return top_k * INTERNAL_RECALL_K_MULTIPLIER + @dataclass(frozen=True) class WorkflowStepRequest: @@ -48,3 +84,6 @@ class WorkflowStepRequest: filter_mode: str channels: list[str] | None channel_weights: dict[str, float] | None + internal_recall_k: int | None + rerank: bool + threshold: float diff --git a/packages/shared-python/shared/services/retrieval/workflow/step_runner.py b/packages/shared-python/shared/services/retrieval/workflow/step_runner.py index a9e538f8e..772b334d3 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/step_runner.py +++ b/packages/shared-python/shared/services/retrieval/workflow/step_runner.py @@ -81,6 +81,7 @@ async def _run_retrieve_step( filter_mode=request.filter_mode, channels=request.channels, channel_weights=request.channel_weights, + internal_recall_k=request.internal_recall_k, ledger=ledger, parent_run_id=self._parent_run_id, workflow_step_id=step.id, From bbd2fe09d3353114e571ad7366ee7c2ec079555e Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 21:23:21 -0700 Subject: [PATCH 44/78] fix: resolve refactor review blockers --- .../app/services/demo_source_projection.py | 58 +- ...est_retrieval_asset_projection_contract.py | 124 -- .../document_ingestion/processing_billing.py | 6 +- .../document_parser/doc_profile_pdf.py | 5 - ...t_document_parser_architecture_contract.py | 1048 ----------------- 5 files changed, 40 insertions(+), 1201 deletions(-) delete mode 100644 apps/api/tests/contract/test_retrieval_asset_projection_contract.py delete mode 100644 apps/worker/tests/contract/test_document_parser_architecture_contract.py diff --git a/apps/api/app/services/demo_source_projection.py b/apps/api/app/services/demo_source_projection.py index 483e52de8..9f2817831 100644 --- a/apps/api/app/services/demo_source_projection.py +++ b/apps/api/app/services/demo_source_projection.py @@ -2,22 +2,38 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any +from typing import Any, Protocol from urllib.parse import quote -if TYPE_CHECKING: - from app.services.demo_source_catalog import ( - DemoCitationDefinition, - DemoExampleDefinition, - DemoSourceDefinition, - ) + +class _DemoCitationDefinition(Protocol): + section_path: str + description: str + content: str + + +class _DemoExampleDefinition(Protocol): + id: str + question: str + answer: str + citations: tuple[_DemoCitationDefinition, ...] + + +class _DemoSourceDefinition(Protocol): + demo_source_id: str + canonical_document_id: str + title: str + mime_type: str + size_bytes: int + chunk_count: int + examples: tuple[_DemoExampleDefinition, ...] class DemoSourceProjection: def source_catalog_payload( self, *, - source: DemoSourceDefinition, + source: _DemoSourceDefinition, chunks: tuple[dict[str, Any], ...], ) -> dict[str, Any]: return { @@ -43,7 +59,7 @@ def source_catalog_payload( def chunk_payload( self, *, - source: DemoSourceDefinition, + source: _DemoSourceDefinition, chunk: dict[str, Any], sort_order: int, ) -> dict[str, Any]: @@ -71,7 +87,7 @@ def chunk_payload( def publication_chunks( self, *, - source: DemoSourceDefinition, + source: _DemoSourceDefinition, chunks: tuple[dict[str, Any], ...], ) -> list[dict[str, Any]]: return [ @@ -82,7 +98,7 @@ def publication_chunks( def canonical_chunk_id( self, *, - source: DemoSourceDefinition, + source: _DemoSourceDefinition, chunk: dict[str, Any], ) -> str: return f"{source.demo_source_id}:{chunk['chunk_id']}" @@ -90,7 +106,7 @@ def canonical_chunk_id( def matches_chunk_id( self, *, - source: DemoSourceDefinition, + source: _DemoSourceDefinition, chunk: dict[str, Any], demo_chunk_id: str, ) -> bool: @@ -102,8 +118,8 @@ def matches_chunk_id( def _example_payload( self, *, - source: DemoSourceDefinition, - example: DemoExampleDefinition, + source: _DemoSourceDefinition, + example: _DemoExampleDefinition, chunks: tuple[dict[str, Any], ...], ) -> dict[str, Any]: return { @@ -119,8 +135,8 @@ def _example_payload( def _citation_payload( self, *, - source: DemoSourceDefinition, - citation: DemoCitationDefinition, + source: _DemoSourceDefinition, + citation: _DemoCitationDefinition, chunks: tuple[dict[str, Any], ...], ) -> dict[str, Any]: chunk = _resolve_citation_chunk(source=source, citation=citation, chunks=chunks) @@ -142,7 +158,7 @@ def _citation_payload( def _publication_chunk( *, - source: DemoSourceDefinition, + source: _DemoSourceDefinition, chunk: dict[str, Any], ) -> dict[str, Any]: materialized_chunk = dict(chunk) @@ -167,7 +183,7 @@ def _publication_chunk( def _publication_path( *, - source: DemoSourceDefinition, + source: _DemoSourceDefinition, raw_path: str | None, ) -> str: prefix = f"Default_Root/{source.title}" @@ -190,8 +206,8 @@ def _publication_path( def _resolve_citation_chunk( *, - source: DemoSourceDefinition, - citation: DemoCitationDefinition, + source: _DemoSourceDefinition, + citation: _DemoCitationDefinition, chunks: tuple[dict[str, Any], ...], ) -> dict[str, Any]: normalized_content = _normalize_text(citation.content) @@ -221,7 +237,7 @@ def _is_media_chunk(chunk: dict[str, Any]) -> bool: def _asset_url( *, - source: DemoSourceDefinition, + source: _DemoSourceDefinition, file_path: str | None, ) -> str | None: if not file_path: diff --git a/apps/api/tests/contract/test_retrieval_asset_projection_contract.py b/apps/api/tests/contract/test_retrieval_asset_projection_contract.py deleted file mode 100644 index fd23c106e..000000000 --- a/apps/api/tests/contract/test_retrieval_asset_projection_contract.py +++ /dev/null @@ -1,124 +0,0 @@ -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/apps/worker/app/services/document_ingestion/processing_billing.py b/apps/worker/app/services/document_ingestion/processing_billing.py index e3d028208..f0d313fa4 100644 --- a/apps/worker/app/services/document_ingestion/processing_billing.py +++ b/apps/worker/app/services/document_ingestion/processing_billing.py @@ -41,9 +41,9 @@ def charge_parse_job_pages( billing_service = WorkBillingService() billing_filename = filename or "" - billing_status = "skipped" - billing_amount_micro_dollars = 0 - billing_credits = 0.0 + billing_status: str + billing_amount_micro_dollars: int + billing_credits: float with get_sync_db_context() as db: job_result = db.execute(select(Job).where(Job.job_id == job_id).with_for_update()) diff --git a/apps/worker/app/services/document_parser/doc_profile_pdf.py b/apps/worker/app/services/document_parser/doc_profile_pdf.py index 216613372..4f9c447e6 100644 --- a/apps/worker/app/services/document_parser/doc_profile_pdf.py +++ b/apps/worker/app/services/document_parser/doc_profile_pdf.py @@ -436,11 +436,6 @@ def _profile_pdf_worker(queue, file_path: str) -> None: 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 diff --git a/apps/worker/tests/contract/test_document_parser_architecture_contract.py b/apps/worker/tests/contract/test_document_parser_architecture_contract.py deleted file mode 100644 index c76f4413d..000000000 --- a/apps/worker/tests/contract/test_document_parser_architecture_contract.py +++ /dev/null @@ -1,1048 +0,0 @@ -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_output import ParseOutput - 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 = route_document_parse(session) - - assert output == ParseOutput(output_dir=str(tmp_path), parsed_df=parsed_df) - assert output.as_legacy_tuple() == (str(tmp_path), parsed_df) - - -def test_parse_pipeline_owns_route_and_postprocess_contract( - worker_contract_environment: None, - monkeypatch: Any, - tmp_path: Path, -) -> None: - from app.services.document_parser.orchestration.parse_input import ParseInput - from app.services.document_parser.orchestration.parse_output import ParseOutput - from app.services.document_parser.orchestration.parse_pipeline import ( - ParsePipelineResult, - run_parse_pipeline, - ) - - routed_df = pd.DataFrame([{"content": "routed"}]) - processed_df = pd.DataFrame([{"content": "processed"}]) - routed_output_dir = str(tmp_path / "routed") - - monkeypatch.setattr( - "app.services.document_parser.orchestration.parse_session.profile_document", - lambda *_args, **_kwargs: SimpleNamespace( - file_type="pdf", - page_count=1, - atlas_candidate=False, - doc_category="generic", - summary=lambda: "profile", - reasoning="test", - ), - ) - monkeypatch.setattr( - "app.services.document_parser.orchestration.parse_pipeline.route_document_parse", - lambda _session: ParseOutput(output_dir=routed_output_dir, parsed_df=routed_df), - ) - - def fake_postprocess(output_dir: str, parsed_df: pd.DataFrame | None) -> pd.DataFrame: - assert output_dir == routed_output_dir - assert parsed_df is routed_df - return processed_df - - monkeypatch.setattr( - "app.services.document_parser.orchestration.parse_pipeline.apply_parse_postprocess", - fake_postprocess, - ) - - result = run_parse_pipeline( - ParseInput( - file_full_path=str(tmp_path / "contract.pdf"), - filename="contract.pdf", - output_dir=str(tmp_path), - internal_output_filename="contract.pdf", - ) - ) - - assert result == ParsePipelineResult( - output_dir=routed_output_dir, - parsed_df=processed_df, - ) - assert result.as_legacy_tuple() == (routed_output_dir, processed_df) - - -def test_parser_service_keeps_legacy_tuple_compatibility( - worker_contract_environment: None, - monkeypatch: Any, - tmp_path: Path, -) -> None: - import app.services.document_parser.parse_service as parse_service - from app.services.document_parser.orchestration.parse_output import ParseOutput - - parsed_df = pd.DataFrame([{"content": "legacy"}]) - expected_output = ParseOutput(output_dir=str(tmp_path), parsed_df=parsed_df) - - monkeypatch.setattr( - parse_service, - "run_parse_pipeline", - lambda _parse_input: expected_output, - ) - - output = parse_service.checkerboard_parse_output( - file_full_path=str(tmp_path / "legacy.pdf"), - filename="legacy.pdf", - output_dir=str(tmp_path), - internal_output_filename="legacy.pdf", - ) - legacy_tuple = parse_service.checkerboard_inject_parse( - file_full_path=str(tmp_path / "legacy.pdf"), - filename="legacy.pdf", - output_dir=str(tmp_path), - internal_output_filename="legacy.pdf", - ) - - assert output is expected_output - assert legacy_tuple == (str(tmp_path), 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 'North' in dataframe_html - 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: - 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, - ) - from app.services.document_parser.markdown_deferred_task import ( - ImageDeferredSummaryTask, - TableDeferredSummaryTask, - TextDeferredSummaryTask, - ) - - 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("
", 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 " 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 6fdf0da4d4f98d4089919a00dc8dc303684c4127 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 21:25:52 -0700 Subject: [PATCH 45/78] fix: satisfy demo source projection typing --- .../app/services/demo_source_projection.py | 53 ++++++++++++++----- 1 file changed, 39 insertions(+), 14 deletions(-) diff --git a/apps/api/app/services/demo_source_projection.py b/apps/api/app/services/demo_source_projection.py index 9f2817831..58d5ab0ef 100644 --- a/apps/api/app/services/demo_source_projection.py +++ b/apps/api/app/services/demo_source_projection.py @@ -7,26 +7,51 @@ class _DemoCitationDefinition(Protocol): - section_path: str - description: str - content: str + @property + def section_path(self) -> str: ... + + @property + def description(self) -> str: ... + + @property + def content(self) -> str: ... class _DemoExampleDefinition(Protocol): - id: str - question: str - answer: str - citations: tuple[_DemoCitationDefinition, ...] + @property + def id(self) -> str: ... + + @property + def question(self) -> str: ... + + @property + def answer(self) -> str: ... + + @property + def citations(self) -> tuple[_DemoCitationDefinition, ...]: ... class _DemoSourceDefinition(Protocol): - demo_source_id: str - canonical_document_id: str - title: str - mime_type: str - size_bytes: int - chunk_count: int - examples: tuple[_DemoExampleDefinition, ...] + @property + def demo_source_id(self) -> str: ... + + @property + def canonical_document_id(self) -> str: ... + + @property + def title(self) -> str: ... + + @property + def mime_type(self) -> str: ... + + @property + def size_bytes(self) -> int: ... + + @property + def chunk_count(self) -> int: ... + + @property + def examples(self) -> tuple[_DemoExampleDefinition, ...]: ... class DemoSourceProjection: From 788965f20b17effd62733811dcd76fad893fee6c Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 21:54:22 -0700 Subject: [PATCH 46/78] fix: use explicit protocol stubs --- .../app/services/demo_source_projection.py | 42 ++++++++++++------- 1 file changed, 28 insertions(+), 14 deletions(-) diff --git a/apps/api/app/services/demo_source_projection.py b/apps/api/app/services/demo_source_projection.py index 58d5ab0ef..c59511c17 100644 --- a/apps/api/app/services/demo_source_projection.py +++ b/apps/api/app/services/demo_source_projection.py @@ -8,50 +8,64 @@ class _DemoCitationDefinition(Protocol): @property - def section_path(self) -> str: ... + def section_path(self) -> str: + raise NotImplementedError @property - def description(self) -> str: ... + def description(self) -> str: + raise NotImplementedError @property - def content(self) -> str: ... + def content(self) -> str: + raise NotImplementedError class _DemoExampleDefinition(Protocol): @property - def id(self) -> str: ... + def id(self) -> str: + raise NotImplementedError @property - def question(self) -> str: ... + def question(self) -> str: + raise NotImplementedError @property - def answer(self) -> str: ... + def answer(self) -> str: + raise NotImplementedError @property - def citations(self) -> tuple[_DemoCitationDefinition, ...]: ... + def citations(self) -> tuple[_DemoCitationDefinition, ...]: + raise NotImplementedError class _DemoSourceDefinition(Protocol): @property - def demo_source_id(self) -> str: ... + def demo_source_id(self) -> str: + raise NotImplementedError @property - def canonical_document_id(self) -> str: ... + def canonical_document_id(self) -> str: + raise NotImplementedError @property - def title(self) -> str: ... + def title(self) -> str: + raise NotImplementedError @property - def mime_type(self) -> str: ... + def mime_type(self) -> str: + raise NotImplementedError @property - def size_bytes(self) -> int: ... + def size_bytes(self) -> int: + raise NotImplementedError @property - def chunk_count(self) -> int: ... + def chunk_count(self) -> int: + raise NotImplementedError @property - def examples(self) -> tuple[_DemoExampleDefinition, ...]: ... + def examples(self) -> tuple[_DemoExampleDefinition, ...]: + raise NotImplementedError class DemoSourceProjection: From 83c62f12b3662863d38b164e46b9bc44741a2966 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 22:22:09 -0700 Subject: [PATCH 47/78] test: remove unit-style contract tests --- .../test_agentic_answer_policy_contract.py | 26 ------ .../tests/contract/test_retrieval_contract.py | 81 ------------------- .../contract/test_excel_parser_contract.py | 61 -------------- 3 files changed, 168 deletions(-) delete mode 100644 apps/api/tests/contract/test_agentic_answer_policy_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 deleted file mode 100644 index 6ef4e07cd..000000000 --- a/apps/api/tests/contract/test_agentic_answer_policy_contract.py +++ /dev/null @@ -1,26 +0,0 @@ -from __future__ import annotations - -import pytest - -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: LLMFnInput) -> str: - return '{"status": "DONE", "answer": "truncated"' - - -@pytest.mark.asyncio -async def test_attempt_answer_should_not_expose_malformed_json_wrapper() -> None: - status, answer, reason = await attempt_answer( - _malformed_json_wrapper, - query="What changed?", - evidence_text="┈ evidence", - state=AgentState(), - config=AgentRunConfig(), - ) - - assert status == "NOT_FOUND" - assert answer == "" - assert reason == "attempt_answer returned malformed JSON" diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index e0a4bcaf9..2e9ca51d5 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -11,7 +11,6 @@ from tests.support.contract_database import ContractDatabase from shared.services.retrieval.agentic.types import AgenticResult -from shared.services.retrieval.query_request import RetrievalQuery from shared.services.retrieval.workflow.run_request import WorkflowRunRequest from shared.services.retrieval.workflow.types import PlannedStep, QueryPlan, WorkflowResult @@ -141,86 +140,6 @@ 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"]) - -def test_retrieval_query_owns_cache_and_route_request_shape() -> None: - query = RetrievalQuery.from_parameters( - db=cast(AsyncSession, object()), - user_id="user-1", - namespace="contract-retrieval", - query=" alpha ", - top_k=5, - exclude_document_ids=["doc-hidden"], - exclude_sections=[{"document_id": "doc-1", "section_path": "Root"}], - data_type=2, - signal_paths=["Root"], - filter_mode="keep", - channels=["content"], - channel_weights={"content": 1.5}, - rerank=True, - threshold=0.2, - internal_recall_k=17, - use_agentic=False, - ) - - route_context = query.build_route_context() - - assert query.query == "alpha" - assert query.build_cache_extra() == { - "data_type": 2, - "signal_paths": ["Root"], - "filter_mode": "keep", - "channels": ["content"], - "channel_weights": {"content": 1.5}, - "rerank": True, - "threshold": 0.2, - "internal_recall_k": 17, - "decomposition_enabled": True, - } - assert route_context.query == "alpha" - assert route_context.allowed_chunk_types == {"text"} - assert route_context.effective_recall_k == 17 - assert route_context.use_agentic is False - - workflow_request = WorkflowRunRequest.from_route_context(route_context) - assert workflow_request.top_k == 5 - assert workflow_request.data_type == 2 - assert workflow_request.signal_paths == ["Root"] - assert workflow_request.filter_mode == "keep" - assert workflow_request.channels == ["content"] - assert workflow_request.channel_weights == {"content": 1.5} - assert workflow_request.internal_recall_k == 17 - assert workflow_request.threshold == 0.2 - assert workflow_request.rerank is True - assert ( - workflow_request.for_step( - PlannedStep(id="explicit-recall", sub_query="explicit", top_k=9) - ).internal_recall_k - == 17 - ) - - -def test_agentic_workflow_should_recompute_default_recall_for_step_top_k() -> None: - query = RetrievalQuery.from_parameters( - db=cast(AsyncSession, object()), - user_id="user-1", - namespace="contract-retrieval", - query="alpha", - top_k=1, - exclude_document_ids=[], - exclude_sections=[], - use_agentic=True, - ) - - workflow_request = WorkflowRunRequest.from_route_context(query.build_route_context()) - step_request = workflow_request.for_step( - PlannedStep(id="wide-step", sub_query="wide alpha", top_k=10) - ) - - assert workflow_request.internal_recall_k is None - assert step_request.top_k == 10 - assert step_request.internal_recall_k == 20 - - @pytest.mark.asyncio async def test_agentic_workflow_should_pass_full_request_policy_to_step_adapter( developer_api_client_factory: Callable[ diff --git a/apps/worker/tests/contract/test_excel_parser_contract.py b/apps/worker/tests/contract/test_excel_parser_contract.py index 0f86e6e96..5d80dd3cf 100644 --- a/apps/worker/tests/contract/test_excel_parser_contract.py +++ b/apps/worker/tests/contract/test_excel_parser_contract.py @@ -2,8 +2,6 @@ from pathlib import Path -from pytest import MonkeyPatch - def _write_contract_workbook(workbook_path: Path) -> None: import openpyxl @@ -62,62 +60,3 @@ def test_xlsx_parser_contract_uses_stable_entrypoint_and_ignores_hidden_sheets( 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 From 9dbd5af97c27f7577a43185fe96b43fb409f900e Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 22:39:28 -0700 Subject: [PATCH 48/78] test: keep contract suite at behavior boundaries --- .../tests/contract/test_job_read_contract.py | 8 - .../tests/contract/test_retrieval_contract.py | 123 +------------ .../tests/contract/test_s3_event_contract.py | 73 -------- ...test_job_lifecycle_post_commit_contract.py | 55 ------ .../contract/test_parse_task_contract.py | 17 -- .../test_stale_job_sweeper_contract.py | 162 ------------------ .../test_worker_job_file_storage_contract.py | 105 ------------ 7 files changed, 2 insertions(+), 541 deletions(-) delete mode 100644 apps/worker/tests/contract/test_job_lifecycle_post_commit_contract.py delete mode 100644 apps/worker/tests/contract/test_worker_job_file_storage_contract.py diff --git a/apps/api/tests/contract/test_job_read_contract.py b/apps/api/tests/contract/test_job_read_contract.py index b913d2a6c..136d20f89 100644 --- a/apps/api/tests/contract/test_job_read_contract.py +++ b/apps/api/tests/contract/test_job_read_contract.py @@ -1,4 +1,3 @@ -import importlib from collections.abc import Callable from contextlib import AbstractAsyncContextManager from datetime import datetime, timedelta, timezone @@ -188,10 +187,3 @@ 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) diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index 2e9ca51d5..f438fec0e 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -1,5 +1,5 @@ -from collections.abc import AsyncIterator, Callable -from contextlib import AbstractAsyncContextManager, asynccontextmanager +from collections.abc import Callable +from contextlib import AbstractAsyncContextManager from datetime import datetime, timezone from typing import Any, cast from uuid import uuid4 @@ -249,57 +249,6 @@ async def fake_retrieval_run( assert request["internal_recall_k"] == 23 -@pytest.mark.asyncio -async def test_agentic_discovery_should_use_internal_recall_for_fusion( - monkeypatch: MonkeyPatch, -) -> None: - from shared.services.retrieval.agentic import discovery_tools - - requested_top_k: list[int] = [] - - async def fake_empty_channel( - _db: AsyncSession, - **_kwargs: Any, - ) -> list[dict[str, Any]]: - return [] - - async def fake_content_channel( - _db: AsyncSession, - **kwargs: Any, - ) -> list[dict[str, Any]]: - requested_top_k.append(cast(int, kwargs["top_k"])) - return [ - { - "chunk_id": f"chunk-{index}", - "document_id": f"doc-{index}", - "section_path": f"Root/{index}", - "content": f"content {index}", - } - for index in range(5) - ] - - monkeypatch.setattr(discovery_tools, "path_channel", fake_empty_channel) - monkeypatch.setattr(discovery_tools, "content_channel", fake_content_channel) - monkeypatch.setattr(discovery_tools, "term_channel", fake_empty_channel) - - result = await discovery_tools.bottom_discovery( - cast(AsyncSession, None), - user_id="local-dev-user", - namespace="contract-agentic-recall-policy", - query="recall marker", - top_k=1, - exclude_document_ids=[], - exclude_sections=[], - channels=["content"], - internal_recall_k=4, - ) - - fused_rows = cast(list[dict[str, Any]], result.payload["fused_rows"]) - - assert requested_top_k == [4] - assert len(fused_rows) == 4 - - @pytest.mark.asyncio async def test_should_return_seeded_retrieval_results_for_the_authenticated_user( developer_api_client_factory: Callable[ @@ -399,74 +348,6 @@ 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[ diff --git a/apps/api/tests/contract/test_s3_event_contract.py b/apps/api/tests/contract/test_s3_event_contract.py index ae781eb36..b8e8561de 100644 --- a/apps/api/tests/contract/test_s3_event_contract.py +++ b/apps/api/tests/contract/test_s3_event_contract.py @@ -1,5 +1,4 @@ import importlib -import base64 import json import socket from collections.abc import Callable @@ -10,8 +9,6 @@ import pytest from httpx import AsyncClient from pytest import MonkeyPatch -from sqlalchemy import text -from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from tests.support.contract_database import ContractDatabase @@ -51,39 +48,6 @@ async def _insert_waiting_file_job() -> tuple[str, str]: return user_id, job_id -def test_storage_event_intake_sanitizes_sensitive_headers() -> None: - from app.services.s3_events.intake_outcome import sanitize_storage_event_headers - - sanitized = sanitize_storage_event_headers( - { - "Authorization": "Bearer secret", - "x-amz-sns-signature": "signature", - "x-minio-auth-token": "token", - "content-type": "application/json", - } - ) - - assert sanitized == { - "Authorization": "", - "x-amz-sns-signature": "", - "x-minio-auth-token": "", - "content-type": "application/json", - } - - -def test_storage_event_nested_message_decoder_supports_plain_and_base64_json() -> None: - from app.services.s3_events.nested_message_decoder import decode_nested_json_message - - payload = {"Records": [{"eventName": "ObjectCreated:Put"}]} - encoded_payload = base64.b64encode(json.dumps(payload).encode("utf-8")).decode( - "utf-8" - ) - - assert decode_nested_json_message(json.dumps(payload)) == payload - assert decode_nested_json_message(encoded_payload) == payload - assert decode_nested_json_message("not-json") is None - - @pytest.mark.asyncio async def test_should_acknowledge_an_sns_subscription_confirmation_request( api_client_factory: Callable[[], AbstractAsyncContextManager[AsyncClient]], @@ -209,43 +173,6 @@ async def start_workflow( assert job_row["status"] == "pending" -@pytest.mark.asyncio -async def test_should_expose_async_state_machine_rejection_reason( - api_client_factory: Callable[[], AbstractAsyncContextManager[AsyncClient]], -) -> None: - from shared.core.state_machine.service import AsyncStateMachineService - from shared.core.state_machine.states import JobStatus - from shared.testing.contract_runtime import get_contract_database_url - - async with api_client_factory(): - _, job_id = await _insert_waiting_file_job() - engine = create_async_engine(get_contract_database_url(), future=True) - session_factory = async_sessionmaker(engine, expire_on_commit=False) - - try: - async with engine.begin() as connection: - await connection.execute( - text("UPDATE jobs SET status = 'done' WHERE job_id = :job_id"), - {"job_id": job_id}, - ) - - async with session_factory() as db: - outcome = await AsyncStateMachineService().transition_outcome( - db, - job_id, - JobStatus.RUNNING.value, - transition_reason="contract_invalid_transition", - ) - finally: - await engine.dispose() - - assert outcome.succeeded is False - assert outcome.reason == "invalid_transition" - assert outcome.from_state == "done" - assert outcome.to_state == "running" - assert outcome.attempts == 1 - - @pytest.mark.asyncio async def test_should_reject_an_sns_subscription_confirmation_url_that_resolves_to_a_private_host( api_client_factory: Callable[[], AbstractAsyncContextManager[AsyncClient]], diff --git a/apps/worker/tests/contract/test_job_lifecycle_post_commit_contract.py b/apps/worker/tests/contract/test_job_lifecycle_post_commit_contract.py deleted file mode 100644 index a9b99fc26..000000000 --- a/apps/worker/tests/contract/test_job_lifecycle_post_commit_contract.py +++ /dev/null @@ -1,55 +0,0 @@ -from __future__ import annotations - -from typing import Any, cast - - -def test_post_commit_effect_runner_uses_primitive_effect_plan( - worker_contract_environment: None, -) -> None: - from shared.services.job_post_commit_effects_sync import ( - PostCommitEffectPlan, - SyncJobPostCommitEffectRunner, - ) - - cache_invalidations: list[dict[str, Any]] = [] - webhook_event_ids: list[str] = [] - - class FakePublicationFinalizer: - def invalidate_cache_after_commit( - self, - cache_invalidation: dict[str, Any] | None, - ) -> None: - assert cache_invalidation is not None - cache_invalidations.append(cache_invalidation) - - class FakeWebhookOutbox: - def enqueue_event_id_after_commit( - self, - webhook_event_id: str | None, - ) -> None: - assert webhook_event_id is not None - webhook_event_ids.append(webhook_event_id) - - runner = SyncJobPostCommitEffectRunner( - publication_finalizer=cast(Any, FakePublicationFinalizer()), - webhook_outbox=cast(Any, FakeWebhookOutbox()), - ) - plan = PostCommitEffectPlan.from_success( - cache_invalidation={ - "job_id": "job_post_commit", - "user_id": "user-1", - "namespaces": ["default", "finance"], - }, - webhook_event_id="event-1", - ) - - runner.run(plan) - - assert cache_invalidations == [ - { - "job_id": "job_post_commit", - "user_id": "user-1", - "namespaces": ["default", "finance"], - } - ] - assert webhook_event_ids == ["event-1"] diff --git a/apps/worker/tests/contract/test_parse_task_contract.py b/apps/worker/tests/contract/test_parse_task_contract.py index bf8bc61d2..f0d45b529 100644 --- a/apps/worker/tests/contract/test_parse_task_contract.py +++ b/apps/worker/tests/contract/test_parse_task_contract.py @@ -1314,23 +1314,6 @@ def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: assert audit_transition_count == 0 -def test_should_make_workload_estimate_fallback_reason_explicit( - worker_contract_environment: None, - tmp_path: Path, -) -> None: - from app.services.document_ingestion.page_estimator import PageEstimator - - unknown_file = tmp_path / "payload.unknown" - unknown_file.write_text("content", encoding="utf-8") - - estimate = PageEstimator.estimate_workload(str(unknown_file)) - - assert estimate.page_count == 1 - assert estimate.method == "unknown_file_type" - assert estimate.fallback_reason == "unknown_file_type:.unknown" - assert PageEstimator.estimate(str(unknown_file)) == 1 - - def test_should_mark_the_job_failed_and_cleanup_the_workspace_when_parse_execution_raises( worker_contract_environment: None, monkeypatch: MonkeyPatch, 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 a6881d960..48cd82705 100644 --- a/apps/worker/tests/contract/test_stale_job_sweeper_contract.py +++ b/apps/worker/tests/contract/test_stale_job_sweeper_contract.py @@ -138,165 +138,3 @@ 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 - - -def test_should_expose_sync_state_machine_rejection_reason( - worker_contract_environment: None, -) -> None: - from shared.core.database_sync import get_sync_db_context - from shared.core.state_machine.service_sync import SyncStateMachineService - from shared.core.state_machine.states import JobStatus - from shared.services.redis.redis_sync_service import SyncRedisServiceFactory - - job_id = f"job_outcome_{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="done", - source_type="file", - webhook_enabled=False, - job_metadata=_build_file_job_metadata(), - ) - - state_machine = SyncStateMachineService( - redis_service=SyncRedisServiceFactory.get_service() - ) - - with get_sync_db_context() as db: - outcome = state_machine.transition_outcome( - db, - job_id, - JobStatus.RUNNING.value, - transition_reason="contract_invalid_transition", - ) - - assert outcome.succeeded is False - assert outcome.reason == "invalid_transition" - assert outcome.from_state == "done" - assert outcome.to_state == "running" - assert outcome.attempts == 1 - - -def test_should_keep_sync_state_machine_bool_facade_for_missing_jobs( - worker_contract_environment: None, -) -> None: - from shared.core.database_sync import get_sync_db_context - from shared.core.state_machine.service_sync import SyncStateMachineService - from shared.core.state_machine.states import JobStatus - from shared.services.redis.redis_sync_service import SyncRedisServiceFactory - - missing_job_id = f"job_missing_{uuid4().hex[:12]}" - state_machine = SyncStateMachineService( - redis_service=SyncRedisServiceFactory.get_service() - ) - - with get_sync_db_context() as db: - outcome = state_machine.transition_outcome( - db, - missing_job_id, - JobStatus.RUNNING.value, - transition_reason="contract_missing_job", - ) - did_transition = state_machine.transition( - db, - missing_job_id, - JobStatus.RUNNING.value, - transition_reason="contract_missing_job_bool", - ) - - assert outcome.succeeded is False - assert outcome.reason == "job_not_found" - assert outcome.attempts == 1 - assert did_transition is False 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 deleted file mode 100644 index b41b0bd9d..000000000 --- a/apps/worker/tests/contract/test_worker_job_file_storage_contract.py +++ /dev/null @@ -1,105 +0,0 @@ -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") - ] From 30bbbb16c44ad83f8bbb6524ba014eaf11acd6c0 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 22:41:49 -0700 Subject: [PATCH 49/78] test: fix contract test lint --- apps/api/tests/contract/test_retrieval_contract.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index f438fec0e..489e2c365 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -1,7 +1,7 @@ from collections.abc import Callable from contextlib import AbstractAsyncContextManager from datetime import datetime, timezone -from typing import Any, cast +from typing import cast from uuid import uuid4 import pytest From 7006127faf4ff254bfbdf9903467a89914a40622 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 17 May 2026 22:54:34 -0700 Subject: [PATCH 50/78] chore: remove empty service package modules --- apps/api/app/services/__init__.py | 0 apps/api/app/services/guest/__init__.py | 0 apps/worker/app/services/connect_builder/__init__.py | 0 apps/worker/app/services/document_parser/__init__.py | 0 apps/worker/app/services/storage/__init__.py | 0 5 files changed, 0 insertions(+), 0 deletions(-) delete mode 100644 apps/api/app/services/__init__.py delete mode 100644 apps/api/app/services/guest/__init__.py delete mode 100644 apps/worker/app/services/connect_builder/__init__.py delete mode 100644 apps/worker/app/services/document_parser/__init__.py delete mode 100644 apps/worker/app/services/storage/__init__.py diff --git a/apps/api/app/services/__init__.py b/apps/api/app/services/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/apps/api/app/services/guest/__init__.py b/apps/api/app/services/guest/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/apps/worker/app/services/connect_builder/__init__.py b/apps/worker/app/services/connect_builder/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/apps/worker/app/services/document_parser/__init__.py b/apps/worker/app/services/document_parser/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/apps/worker/app/services/storage/__init__.py b/apps/worker/app/services/storage/__init__.py deleted file mode 100644 index e69de29bb..000000000 From b6591a7d2fcdc7d1495099c5b9b462d4544d81f8 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 06:47:03 +0000 Subject: [PATCH 51/78] refactor: organize shared utilities by domain --- .github/codeql/codeql-config.yml | 2 +- AGENTS.md | 4 +++- apps/api/app/core/dependencies.py | 2 +- .../repositories/guest_device_repository.py | 2 +- .../auth/api_key_authentication_service.py | 2 +- .../auth/api_key_management_service.py | 2 +- .../stripe_credits_settlement_service.py | 2 +- .../stripe_refund_reconciliation_service.py | 2 +- .../document_ingestion/creation_service.py | 2 +- .../services/document_ingestion/service.py | 4 ++-- .../guest/guest_registration_service.py | 2 +- apps/api/app/services/jobs/job_read_model.py | 2 +- .../app/services/jobs/result_projection.py | 2 +- .../s3_events/subscription_service.py | 4 ++-- apps/api/main.py | 2 +- apps/api/scripts/add_credits.py | 2 +- apps/api/scripts/init_user.py | 2 +- .../tests/contract/test_billing_contract.py | 2 +- .../contract/test_job_creation_contract.py | 4 ++-- .../tests/contract/test_s3_event_contract.py | 2 +- apps/api/tests/support/contract_database.py | 2 +- apps/worker/app/core/worker_bootstrap.py | 2 +- apps/worker/app/services/common/__init__.py | 10 ++++----- .../app/services/common/device_checks.py | 0 .../app/services/common}/file_loading.py | 0 .../worker/app/services/common}/file_utils.py | 0 .../app/services/common/math_helpers.py | 0 .../app/services/common/resource_cleanup.py | 0 .../connect_builder/summary_builder.py | 2 +- .../tools/classify_special_pages.py | 2 +- .../document_agent/tools/probe_vlm_inspect.py | 2 +- .../tools/propose_shard_plan.py | 2 +- .../document_ingestion/page_estimator.py | 2 +- .../document_parser/atlas_classifier.py | 2 +- .../services/document_parser/atlas_parser.py | 4 ++-- .../services/document_parser/doc_parser.py | 8 +++---- .../document_parser/excel_table_parser.py | 8 +++---- .../document_parser/fragment_parser.py | 4 ++-- .../services/document_parser/html_parser.py | 2 +- .../services/document_parser/image_parser.py | 8 +++---- .../document_parser/internal_parse_name.py | 2 +- .../services/document_parser/layout_parser.py | 2 +- .../markdown_deferred_summary.py | 4 ++-- .../document_parser/markdown_image_asset.py | 4 ++-- .../document_parser/markdown_table_asset.py | 4 ++-- .../app/services/document_parser/md_parser.py | 6 ++--- .../document_parser/mineru_pdf_service.py | 2 +- .../document_parser/mineru_quota_manager.py | 2 +- .../document_parser/mineru_task_polling.py | 2 +- .../orchestration/parse_session.py | 2 +- .../services/document_parser/path_helpers.py | 4 ++-- .../services/document_parser/pptx_parser.py | 8 +++---- .../document_parser/table_frame_parser.py | 4 ++-- .../document_parser/table_text_parser.py | 2 +- .../services/document_parser/text_helpers.py | 2 +- .../services/document_parser/toc_parser.py | 2 +- .../services/document_parser/txt_parser.py | 6 ++--- .../services/workload/url_upload_transfer.py | 2 +- .../shared/{utils => core}/error_details.py | 0 .../shared-python/shared/core/security.py | 19 ++++++++++++++++ .../json_utils.py => core/serialization.py} | 0 .../shared/core/state_machine/service.py | 2 +- .../shared/core/state_machine/service_sync.py | 2 +- .../core/state_machine/transition_payloads.py | 4 ++-- .../shared/{utils/utc_now.py => core/time.py} | 0 .../shared/models/database/api_key.py | 2 +- .../models/database/credits_transaction.py | 2 +- .../models/database/demo_materialization.py | 2 +- .../shared/models/database/document.py | 2 +- .../shared/models/database/guest_device.py | 2 +- .../shared/models/database/job.py | 2 +- .../shared/models/database/job_result.py | 2 +- .../models/database/job_state_audit_log.py | 2 +- .../models/database/job_state_history.py | 2 +- .../shared/models/database/payment_record.py | 2 +- .../models/database/stripe_price_config.py | 2 +- .../shared/models/database/user_balance.py | 2 +- .../shared/models/database/webhook.py | 2 +- .../shared/models/database/webhook_log.py | 2 +- .../shared/models/database/webhook_secret.py | 2 +- .../shared/repositories/credits_repository.py | 2 +- .../repositories/credits_sync_repository.py | 2 +- .../ai}/ali_quota_manager.py | 2 +- .../ai}/iloveapi_quota_manager.py | 2 +- .../shared/{utils => services/ai}/llm_mock.py | 0 .../ai/openai_compatible_client_sync.py} | 8 +++---- .../{utils => services/ai}/quota_manager.py | 0 .../auth/api_key_tokens.py} | 0 .../services/billing/credits_service.py | 2 +- .../services/billing/credits_sync_service.py | 2 +- .../services/chunks/chunk_connections.py | 2 +- .../{utils => services/chunks}/chunk_refs.py | 0 .../shared/services/http/__init__.py | 1 + .../http/client_pool.py} | 0 .../http/pinned_outbound.py} | 2 +- .../{utils => services/http}/url_file_type.py | 6 ++--- .../{utils => services/http}/url_security.py | 0 .../{utils => services/http}/zip_download.py | 0 .../shared/services/job_failure_sync.py | 2 +- .../shared/services/job_lifecycle_sync.py | 2 +- .../shared/services/jobs/result_delivery.py | 2 +- .../shared/services/redis/distributed_lock.py | 2 +- .../services/redis/job_info_redis_service.py | 2 +- .../services/redis/job_metadata_service.py | 2 +- .../redis/key_builder.py} | 0 .../services/redis/rate_limit_service.py | 2 +- .../shared/services/redis/redis_monitor.py | 2 +- .../shared/services/redis/redis_service.py | 4 ++-- .../services/redis/redis_sync_service.py | 6 ++--- .../redis/retry_policy.py} | 0 .../services/redis/task_redis_service.py | 2 +- .../services/redis/user_redis_service.py | 2 +- .../services/retrieval/agentic/evidence.py | 2 +- .../services/retrieval/agentic/policy.py | 2 +- .../services/retrieval/agentic/runtime.py | 2 +- .../agentic/section_prompt_projection.py | 2 +- .../services/retrieval/lexical_ranker.py | 2 +- .../shared/services/retrieval/lexical_text.py | 2 +- .../shared/services/retrieval/llm_adapter.py | 6 ++--- .../retrieval}/token_estimate.py | 0 .../services/retrieval/workflow/planner.py | 2 +- .../retrieval/workflow/synthesizer.py | 2 +- .../services/storage/job_file_storage.py | 4 ++-- .../services/storage/zip_result_schema.py | 4 ++-- .../services/text_processing/__init__.py | 1 + .../text_processing/tokenization.py} | 2 +- .../services/webhook/delivery_client.py | 4 ++-- .../services/webhook/qstash_publisher.py | 2 +- .../shared-python/shared/utils/__init__.py | 0 .../shared/utils/security_utils.py | 22 ------------------- 130 files changed, 174 insertions(+), 173 deletions(-) rename packages/shared-python/shared/utils/device_utils.py => apps/worker/app/services/common/device_checks.py (100%) rename {packages/shared-python/shared/utils => apps/worker/app/services/common}/file_loading.py (100%) rename {packages/shared-python/shared/utils => apps/worker/app/services/common}/file_utils.py (100%) rename packages/shared-python/shared/utils/math_utils.py => apps/worker/app/services/common/math_helpers.py (100%) rename packages/shared-python/shared/utils/gc_utils.py => apps/worker/app/services/common/resource_cleanup.py (100%) rename packages/shared-python/shared/{utils => core}/error_details.py (100%) rename packages/shared-python/shared/{utils/json_utils.py => core/serialization.py} (100%) rename packages/shared-python/shared/{utils/utc_now.py => core/time.py} (100%) rename packages/shared-python/shared/{utils => services/ai}/ali_quota_manager.py (96%) rename packages/shared-python/shared/{utils => services/ai}/iloveapi_quota_manager.py (99%) rename packages/shared-python/shared/{utils => services/ai}/llm_mock.py (100%) rename packages/shared-python/shared/{utils/OpenAICompatibleClientSync.py => services/ai/openai_compatible_client_sync.py} (98%) rename packages/shared-python/shared/{utils => services/ai}/quota_manager.py (100%) rename packages/shared-python/shared/{utils/api_keys.py => services/auth/api_key_tokens.py} (100%) rename packages/shared-python/shared/{utils => services/chunks}/chunk_refs.py (100%) create mode 100644 packages/shared-python/shared/services/http/__init__.py rename packages/shared-python/shared/{utils/http_clients.py => services/http/client_pool.py} (100%) rename packages/shared-python/shared/{utils/pinned_outbound_http.py => services/http/pinned_outbound.py} (99%) rename packages/shared-python/shared/{utils => services/http}/url_file_type.py (97%) rename packages/shared-python/shared/{utils => services/http}/url_security.py (100%) rename packages/shared-python/shared/{utils => services/http}/zip_download.py (100%) rename packages/shared-python/shared/{utils/redis_key_builder.py => services/redis/key_builder.py} (100%) rename packages/shared-python/shared/{utils/redis_retry.py => services/redis/retry_policy.py} (100%) rename packages/shared-python/shared/{utils => services/retrieval}/token_estimate.py (100%) create mode 100644 packages/shared-python/shared/services/text_processing/__init__.py rename packages/shared-python/shared/{utils/text_utils.py => services/text_processing/tokenization.py} (99%) delete mode 100644 packages/shared-python/shared/utils/__init__.py delete mode 100644 packages/shared-python/shared/utils/security_utils.py diff --git a/.github/codeql/codeql-config.yml b/.github/codeql/codeql-config.yml index 705bc0dec..9ad24421d 100644 --- a/.github/codeql/codeql-config.yml +++ b/.github/codeql/codeql-config.yml @@ -2,4 +2,4 @@ name: "Knowhere CodeQL config" paths-ignore: - apps/api/scripts/** - - packages/shared-python/shared/utils/api_keys.py + - packages/shared-python/shared/services/auth/api_key_tokens.py diff --git a/AGENTS.md b/AGENTS.md index eed149425..421c133bd 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -47,10 +47,12 @@ knowhereapi-main/ │ └── shared-python/shared/ # Shared library (pip: knowhere-shared) │ ├── models/database/ # SQLAlchemy ORM models │ ├── models/schemas/ # Pydantic request/response schemas +│ ├── core/ # Shared primitives: time, serialization, security │ ├── services/retrieval/ # Core retrieval engine │ ├── services/chunks/ # DataFrame → ChunkPayload conversion │ ├── services/ai/ # LLM prompt service & AI client -│ └── utils/ # Text, file, and chunk utilities +│ ├── services/http/ # Outbound HTTP clients and URL safety +│ └── services/text_processing/ # Shared text tokenization helpers └── deploy/ # Docker Compose & deployment scripts ``` diff --git a/apps/api/app/core/dependencies.py b/apps/api/app/core/dependencies.py index 153177cd4..d14340151 100644 --- a/apps/api/app/core/dependencies.py +++ b/apps/api/app/core/dependencies.py @@ -18,7 +18,7 @@ AuthException, ) from shared.models.database.user import User -from shared.utils.api_keys import is_api_key_token +from shared.services.auth.api_key_tokens import is_api_key_token # Standard JWKS endpoint path (fixed, following OpenID Connect convention) JWKS_ENDPOINT_PATH = "/api/auth/jwks" diff --git a/apps/api/app/repositories/guest_device_repository.py b/apps/api/app/repositories/guest_device_repository.py index 1f0fbc247..ef2ca4f50 100644 --- a/apps/api/app/repositories/guest_device_repository.py +++ b/apps/api/app/repositories/guest_device_repository.py @@ -6,7 +6,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from shared.models.database.guest_device import GuestDevice -from shared.utils.utc_now import utc_now_naive +from shared.core.time import utc_now_naive class GuestDeviceRepository(BaseRepository[GuestDevice, dict, dict]): diff --git a/apps/api/app/services/auth/api_key_authentication_service.py b/apps/api/app/services/auth/api_key_authentication_service.py index 67e677247..a7083226e 100644 --- a/apps/api/app/services/auth/api_key_authentication_service.py +++ b/apps/api/app/services/auth/api_key_authentication_service.py @@ -13,7 +13,7 @@ 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 +from shared.services.auth.api_key_tokens import hash_api_key _API_KEY_USER_CACHE_TTL_SECONDS: int = 3600 diff --git a/apps/api/app/services/auth/api_key_management_service.py b/apps/api/app/services/auth/api_key_management_service.py index 0adbfee94..85189d77e 100644 --- a/apps/api/app/services/auth/api_key_management_service.py +++ b/apps/api/app/services/auth/api_key_management_service.py @@ -19,7 +19,7 @@ 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.auth.api_key_tokens import generate_api_key, hash_api_key, mask_api_key class APIKeyListItem(TypedDict): diff --git a/apps/api/app/services/billing/stripe_credits_settlement_service.py b/apps/api/app/services/billing/stripe_credits_settlement_service.py index 9091dd6e6..36bcd12b4 100644 --- a/apps/api/app/services/billing/stripe_credits_settlement_service.py +++ b/apps/api/app/services/billing/stripe_credits_settlement_service.py @@ -14,7 +14,7 @@ ) from shared.models.database.payment_record import PaymentRecord from shared.services.billing import CreditsService -from shared.utils.utc_now import utc_now_naive +from shared.core.time import utc_now_naive class StripeCreditsSettlementService: diff --git a/apps/api/app/services/billing/stripe_refund_reconciliation_service.py b/apps/api/app/services/billing/stripe_refund_reconciliation_service.py index d770195e1..1324402b8 100644 --- a/apps/api/app/services/billing/stripe_refund_reconciliation_service.py +++ b/apps/api/app/services/billing/stripe_refund_reconciliation_service.py @@ -11,7 +11,7 @@ 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 +from shared.core.time import utc_now_naive class StripeRefundReconciliationService: diff --git a/apps/api/app/services/document_ingestion/creation_service.py b/apps/api/app/services/document_ingestion/creation_service.py index a76977b65..acecb1974 100644 --- a/apps/api/app/services/document_ingestion/creation_service.py +++ b/apps/api/app/services/document_ingestion/creation_service.py @@ -29,7 +29,7 @@ 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.services.http.url_file_type import resolve_file_extension_async _JOB_TYPE_KB_MANAGEMENT = "kb_management" JobMetadata = dict[str, object] diff --git a/apps/api/app/services/document_ingestion/service.py b/apps/api/app/services/document_ingestion/service.py index defd2fbca..03f5da8fe 100644 --- a/apps/api/app/services/document_ingestion/service.py +++ b/apps/api/app/services/document_ingestion/service.py @@ -35,8 +35,8 @@ from shared.core.exceptions.webhook_exceptions import WebhookConfigException from shared.models.schemas.job import ConfirmUploadRequest, JobCreate, JobResponse from shared.models.schemas.job_metadata import JobMetadataHelper -from shared.utils.url_file_type import resolve_file_extension_async -from shared.utils.url_security import validate_http_url_and_resolve_ip_async +from shared.services.http.url_file_type import resolve_file_extension_async +from shared.services.http.url_security import validate_http_url_and_resolve_ip_async JobMetadata = dict[str, object] diff --git a/apps/api/app/services/guest/guest_registration_service.py b/apps/api/app/services/guest/guest_registration_service.py index 4f189807e..ae2cfbb68 100644 --- a/apps/api/app/services/guest/guest_registration_service.py +++ b/apps/api/app/services/guest/guest_registration_service.py @@ -24,7 +24,7 @@ GuestRegisterResponse, ) from shared.services.billing.credits_service import CreditsService -from shared.utils.api_keys import generate_api_key, hash_api_key, mask_api_key +from shared.services.auth.api_key_tokens import generate_api_key, hash_api_key, mask_api_key _GUEST_TIER: str = "guest" _GUEST_KEY_NAME_PREFIX: str = "guest-device" diff --git a/apps/api/app/services/jobs/job_read_model.py b/apps/api/app/services/jobs/job_read_model.py index 1d0a4ab39..94d25ba7f 100644 --- a/apps/api/app/services/jobs/job_read_model.py +++ b/apps/api/app/services/jobs/job_read_model.py @@ -20,7 +20,7 @@ ) from shared.models.schemas.job import JobList, JobResultResponse from shared.services.redis import RedisServiceFactory -from shared.utils.utc_now import utc_now_naive +from shared.core.time import utc_now_naive class JobReadModel: diff --git a/apps/api/app/services/jobs/result_projection.py b/apps/api/app/services/jobs/result_projection.py index 0a00ce2d2..710b92962 100644 --- a/apps/api/app/services/jobs/result_projection.py +++ b/apps/api/app/services/jobs/result_projection.py @@ -10,7 +10,7 @@ from shared.models.schemas.job import JobResultResponse, StandardErrorObject from shared.models.schemas.job_metadata import JobMetadataHelper from shared.services.jobs.result_delivery import JobResultDeliveryResolver -from shared.utils.error_details import normalize_error_details +from shared.core.error_details import normalize_error_details JobStatusValue = Literal[ "pending", "waiting-file", "running", "converting", "done", "failed" diff --git a/apps/api/app/services/s3_events/subscription_service.py b/apps/api/app/services/s3_events/subscription_service.py index 140b2eba3..c64ab78e5 100644 --- a/apps/api/app/services/s3_events/subscription_service.py +++ b/apps/api/app/services/s3_events/subscription_service.py @@ -3,8 +3,8 @@ 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 +from shared.services.http.pinned_outbound import send_pinned_outbound_request +from shared.services.http.url_security import validate_http_url_and_resolve_ip_async SNS_SUBSCRIPTION_TIMEOUT_SECONDS = 10 diff --git a/apps/api/main.py b/apps/api/main.py index bd7dec810..4b0dfde44 100644 --- a/apps/api/main.py +++ b/apps/api/main.py @@ -92,7 +92,7 @@ async def lifespan(app: FastAPI): logger.error(f"retrieval hit stats drain failed: {e}") try: - from shared.utils.http_clients import close_async_client + from shared.services.http.client_pool import close_async_client await close_async_client() except Exception as e: diff --git a/apps/api/scripts/add_credits.py b/apps/api/scripts/add_credits.py index 675f5ce26..dd0370d75 100644 --- a/apps/api/scripts/add_credits.py +++ b/apps/api/scripts/add_credits.py @@ -34,7 +34,7 @@ def _bootstrap_python_path() -> None: from shared.models.database.user import User from shared.models.database.user_balance import UserBalance from shared.services.billing import CreditsService -from shared.utils.utc_now import utc_now_naive +from shared.core.time import utc_now_naive def _build_parser() -> argparse.ArgumentParser: diff --git a/apps/api/scripts/init_user.py b/apps/api/scripts/init_user.py index de38fe61e..4847d5024 100644 --- a/apps/api/scripts/init_user.py +++ b/apps/api/scripts/init_user.py @@ -18,7 +18,7 @@ from shared.models.database.user_balance import UserBalance from shared.services.billing.credits_service import CreditsService from shared.services.auth.user_table_bootstrap import ensure_better_auth_user_table -from shared.utils.api_keys import generate_api_key, hash_api_key, mask_api_key +from shared.services.auth.api_key_tokens import generate_api_key, hash_api_key, mask_api_key _DEFAULT_API_KEY_NAME: str = "standalone-api-key" _DEFAULT_USER_TIER: str = "tier_5" diff --git a/apps/api/tests/contract/test_billing_contract.py b/apps/api/tests/contract/test_billing_contract.py index 6f5eb2be6..11a108ff5 100644 --- a/apps/api/tests/contract/test_billing_contract.py +++ b/apps/api/tests/contract/test_billing_contract.py @@ -12,7 +12,7 @@ from pytest import MonkeyPatch from tests.support.contract_database import ContractDatabase -from shared.utils.api_keys import hash_api_key +from shared.services.auth.api_key_tokens import hash_api_key def _utc_now() -> datetime: diff --git a/apps/api/tests/contract/test_job_creation_contract.py b/apps/api/tests/contract/test_job_creation_contract.py index edea174ec..d19932bc0 100644 --- a/apps/api/tests/contract/test_job_creation_contract.py +++ b/apps/api/tests/contract/test_job_creation_contract.py @@ -632,7 +632,7 @@ def resolve_public_address( 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 + import shared.services.http.client_pool as http_clients_module monkeypatch.setattr(socket, "getaddrinfo", resolve_public_address) monkeypatch.setattr( @@ -926,7 +926,7 @@ async def head( assert follow_redirects is False return _FakeHeadResponse() - import shared.utils.http_clients as http_clients_module + import shared.services.http.client_pool as http_clients_module monkeypatch.setattr( http_clients_module, "get_async_client", diff --git a/apps/api/tests/contract/test_s3_event_contract.py b/apps/api/tests/contract/test_s3_event_contract.py index b8e8561de..65c1c6378 100644 --- a/apps/api/tests/contract/test_s3_event_contract.py +++ b/apps/api/tests/contract/test_s3_event_contract.py @@ -210,7 +210,7 @@ def get(self, url: str, *args: object, **kwargs: object) -> object: async with api_client_factory() as api_client: monkeypatch.setattr(socket, "getaddrinfo", resolve_private_address) pinned_http_module = importlib.import_module( - "shared.utils.pinned_outbound_http" + "shared.services.http.pinned_outbound" ) monkeypatch.setattr( pinned_http_module.aiohttp, diff --git a/apps/api/tests/support/contract_database.py b/apps/api/tests/support/contract_database.py index 664fff00a..3358fff9e 100644 --- a/apps/api/tests/support/contract_database.py +++ b/apps/api/tests/support/contract_database.py @@ -9,7 +9,7 @@ from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine from shared.testing.contract_runtime import get_contract_database_url -from shared.utils.api_keys import hash_api_key +from shared.services.auth.api_key_tokens import hash_api_key async def _create_contract_engine() -> AsyncEngine: diff --git a/apps/worker/app/core/worker_bootstrap.py b/apps/worker/app/core/worker_bootstrap.py index 191f2fa27..2689bb71c 100644 --- a/apps/worker/app/core/worker_bootstrap.py +++ b/apps/worker/app/core/worker_bootstrap.py @@ -74,7 +74,7 @@ def shutdown_worker(**kwargs) -> None: logger.warning(f"Worker heartbeat cleanup failed: {exc}") try: - from shared.utils.http_clients import close_sync_client + from shared.services.http.client_pool import close_sync_client close_sync_client() logger.info("Worker sync HTTP client closed") diff --git a/apps/worker/app/services/common/__init__.py b/apps/worker/app/services/common/__init__.py index ffb53d823..fab06465a 100644 --- a/apps/worker/app/services/common/__init__.py +++ b/apps/worker/app/services/common/__init__.py @@ -1,10 +1,10 @@ """Common worker services.""" -from shared.utils.device_utils import check_internet -from shared.utils.file_utils import clean_file, path_handle -from shared.utils.gc_utils import gc_collect as _gc -from shared.utils.math_utils import min_max_normalize -from shared.utils.text_utils import ( +from app.services.common.device_checks import check_internet +from app.services.common.file_utils import clean_file, path_handle +from app.services.common.resource_cleanup import gc_collect as _gc +from app.services.common.math_helpers import min_max_normalize +from shared.services.text_processing.tokenization import ( count_cn_en, remove_duplicates_orderkept, tokenize2stw_remove, diff --git a/packages/shared-python/shared/utils/device_utils.py b/apps/worker/app/services/common/device_checks.py similarity index 100% rename from packages/shared-python/shared/utils/device_utils.py rename to apps/worker/app/services/common/device_checks.py diff --git a/packages/shared-python/shared/utils/file_loading.py b/apps/worker/app/services/common/file_loading.py similarity index 100% rename from packages/shared-python/shared/utils/file_loading.py rename to apps/worker/app/services/common/file_loading.py diff --git a/packages/shared-python/shared/utils/file_utils.py b/apps/worker/app/services/common/file_utils.py similarity index 100% rename from packages/shared-python/shared/utils/file_utils.py rename to apps/worker/app/services/common/file_utils.py diff --git a/packages/shared-python/shared/utils/math_utils.py b/apps/worker/app/services/common/math_helpers.py similarity index 100% rename from packages/shared-python/shared/utils/math_utils.py rename to apps/worker/app/services/common/math_helpers.py diff --git a/packages/shared-python/shared/utils/gc_utils.py b/apps/worker/app/services/common/resource_cleanup.py similarity index 100% rename from packages/shared-python/shared/utils/gc_utils.py rename to apps/worker/app/services/common/resource_cleanup.py diff --git a/apps/worker/app/services/connect_builder/summary_builder.py b/apps/worker/app/services/connect_builder/summary_builder.py index 7fcf8c6eb..44f472076 100644 --- a/apps/worker/app/services/connect_builder/summary_builder.py +++ b/apps/worker/app/services/connect_builder/summary_builder.py @@ -40,7 +40,7 @@ def _llm_summarize(snippets_text: str, node_name: str, max_tokens: int = 100) -> """ try: from shared.services.ai.prompt_service import build_prompt, _detect_text_language - from shared.utils.OpenAICompatibleClientSync import get_openai_client + from shared.services.ai.openai_compatible_client_sync import get_openai_client # Deterministic language lock — see prompt_service._language_directive detected_lang = _detect_text_language(snippets_text) diff --git a/apps/worker/app/services/document_agent/tools/classify_special_pages.py b/apps/worker/app/services/document_agent/tools/classify_special_pages.py index 564307b21..0b8efa6c8 100644 --- a/apps/worker/app/services/document_agent/tools/classify_special_pages.py +++ b/apps/worker/app/services/document_agent/tools/classify_special_pages.py @@ -141,7 +141,7 @@ def classify_special_pages( try: from shared.core.config import settings - from shared.utils.OpenAICompatibleClientSync import get_openai_client + from shared.services.ai.openai_compatible_client_sync import get_openai_client effective_model = model or settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL client = get_openai_client(model=effective_model) diff --git a/apps/worker/app/services/document_agent/tools/probe_vlm_inspect.py b/apps/worker/app/services/document_agent/tools/probe_vlm_inspect.py index 7b4de6396..f8949826d 100644 --- a/apps/worker/app/services/document_agent/tools/probe_vlm_inspect.py +++ b/apps/worker/app/services/document_agent/tools/probe_vlm_inspect.py @@ -71,7 +71,7 @@ def _call_vlm( max_tokens: int = 900, ) -> tuple[str, dict[str, int]]: from shared.core.config import settings - from shared.utils.OpenAICompatibleClientSync import get_openai_client + from shared.services.ai.openai_compatible_client_sync import get_openai_client effective_model = model or settings.IMAGE_MODEL or "qwen3.5-flash" client = get_openai_client(model=effective_model) diff --git a/apps/worker/app/services/document_agent/tools/propose_shard_plan.py b/apps/worker/app/services/document_agent/tools/propose_shard_plan.py index 2eb834ff4..1db72592b 100644 --- a/apps/worker/app/services/document_agent/tools/propose_shard_plan.py +++ b/apps/worker/app/services/document_agent/tools/propose_shard_plan.py @@ -295,7 +295,7 @@ def propose_shard_plan( if use_llm and page_count > max_pages_per_shard: try: from shared.core.config import settings - from shared.utils.OpenAICompatibleClientSync import get_openai_client + from shared.services.ai.openai_compatible_client_sync import get_openai_client effective_model = model or settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL client = get_openai_client(model=effective_model) diff --git a/apps/worker/app/services/document_ingestion/page_estimator.py b/apps/worker/app/services/document_ingestion/page_estimator.py index bdc97a4d6..4267cd233 100644 --- a/apps/worker/app/services/document_ingestion/page_estimator.py +++ b/apps/worker/app/services/document_ingestion/page_estimator.py @@ -15,7 +15,7 @@ from pathlib import Path from shared.core.logging import logger -from shared.utils.text_utils import count_cn_en +from shared.services.text_processing.tokenization import count_cn_en WORDS_PER_PAGE = 500 ROWS_PER_PAGE = 50 diff --git a/apps/worker/app/services/document_parser/atlas_classifier.py b/apps/worker/app/services/document_parser/atlas_classifier.py index 1ee8dc4ff..ed58ad49f 100644 --- a/apps/worker/app/services/document_parser/atlas_classifier.py +++ b/apps/worker/app/services/document_parser/atlas_classifier.py @@ -102,7 +102,7 @@ def _png_to_data_url(path: str) -> Optional[str]: def _call_vlm(image_data_urls: list[str]) -> bool: """Call VLM with preview images. Returns True if atlas, False otherwise.""" from shared.core.config import settings - from shared.utils.OpenAICompatibleClientSync import get_openai_client + from shared.services.ai.openai_compatible_client_sync import get_openai_client model = settings.IMAGE_MODEL or "qwen-vl-plus" client = get_openai_client(model=model) diff --git a/apps/worker/app/services/document_parser/atlas_parser.py b/apps/worker/app/services/document_parser/atlas_parser.py index fef806add..6e3666d0f 100644 --- a/apps/worker/app/services/document_parser/atlas_parser.py +++ b/apps/worker/app/services/document_parser/atlas_parser.py @@ -27,8 +27,8 @@ from loguru import logger from shared.core.config import settings -from shared.utils.chunk_refs import build_chunk_ref -from shared.utils.text_utils import tokenize2stw_remove +from shared.services.chunks.chunk_refs import build_chunk_ref +from shared.services.text_processing.tokenization import tokenize2stw_remove # ─── Config ─────────────────────────────────────────────────────────── VLM_CONCURRENCY = 3 # concurrent VLM calls diff --git a/apps/worker/app/services/document_parser/doc_parser.py b/apps/worker/app/services/document_parser/doc_parser.py index d405fed56..5eb56ab79 100755 --- a/apps/worker/app/services/document_parser/doc_parser.py +++ b/apps/worker/app/services/document_parser/doc_parser.py @@ -37,10 +37,10 @@ from shared.core.config import settings 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.file_loading import load_file_bytes -from shared.utils.file_utils import path_handle -from shared.utils.text_utils import tokenize2stw_remove +from shared.services.chunks.chunk_refs import build_chunk_ref, has_chunk_ref +from app.services.common.file_loading import load_file_bytes +from app.services.common.file_utils import path_handle +from shared.services.text_processing.tokenization import tokenize2stw_remove def get_leaf_dics(node, 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 index 1770a4163..2a6aa13ea 100644 --- a/apps/worker/app/services/document_parser/excel_table_parser.py +++ b/apps/worker/app/services/document_parser/excel_table_parser.py @@ -27,10 +27,10 @@ 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 +from shared.services.chunks.chunk_refs import build_chunk_ref +from app.services.common.file_loading import load_file_bytes +from app.services.common.file_utils import path_handle +from shared.services.text_processing.tokenization import tokenize2stw_remove @dataclass(frozen=True) diff --git a/apps/worker/app/services/document_parser/fragment_parser.py b/apps/worker/app/services/document_parser/fragment_parser.py index 5aef9f004..420d6b084 100644 --- a/apps/worker/app/services/document_parser/fragment_parser.py +++ b/apps/worker/app/services/document_parser/fragment_parser.py @@ -13,8 +13,8 @@ from openai.types.chat import ChatCompletionMessageParam from shared.core.config import settings -from shared.utils.file_utils import path_handle -from shared.utils.OpenAICompatibleClientSync import get_openai_client +from app.services.common.file_utils import path_handle +from shared.services.ai.openai_compatible_client_sync import get_openai_client def generate_fragment_title(content: str, max_tokens: int = 30) -> Optional[str]: diff --git a/apps/worker/app/services/document_parser/html_parser.py b/apps/worker/app/services/document_parser/html_parser.py index f412067d6..e326a709a 100644 --- a/apps/worker/app/services/document_parser/html_parser.py +++ b/apps/worker/app/services/document_parser/html_parser.py @@ -15,7 +15,7 @@ from bs4 import BeautifulSoup from shared.core.exceptions.domain_exceptions import TableParsingException -from shared.utils.text_utils import remove_duplicates_orderkept +from shared.services.text_processing.tokenization import remove_duplicates_orderkept class HTMLHeaderExpander: diff --git a/apps/worker/app/services/document_parser/image_parser.py b/apps/worker/app/services/document_parser/image_parser.py index cc9be0ab1..62b319450 100755 --- a/apps/worker/app/services/document_parser/image_parser.py +++ b/apps/worker/app/services/document_parser/image_parser.py @@ -22,10 +22,10 @@ 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 is_remote, load_file_bytes -from shared.utils.file_utils import path_handle -from shared.utils.OpenAICompatibleClientSync import ( +from shared.services.chunks.chunk_refs import build_chunk_ref +from app.services.common.file_loading import is_remote, load_file_bytes +from app.services.common.file_utils import path_handle +from shared.services.ai.openai_compatible_client_sync import ( OpenAICompatibleClientSync, get_openai_client, ) diff --git a/apps/worker/app/services/document_parser/internal_parse_name.py b/apps/worker/app/services/document_parser/internal_parse_name.py index 785fe6ff5..568dc13c6 100644 --- a/apps/worker/app/services/document_parser/internal_parse_name.py +++ b/apps/worker/app/services/document_parser/internal_parse_name.py @@ -4,7 +4,7 @@ import os from dataclasses import dataclass -from shared.utils.file_utils import path_handle +from app.services.common.file_utils import path_handle @dataclass(frozen=True) diff --git a/apps/worker/app/services/document_parser/layout_parser.py b/apps/worker/app/services/document_parser/layout_parser.py index b0490d576..ec7775904 100755 --- a/apps/worker/app/services/document_parser/layout_parser.py +++ b/apps/worker/app/services/document_parser/layout_parser.py @@ -35,7 +35,7 @@ from shared.services.ai.response_process_service import eval_response # ARQ dependency is removed, use Celery instead -from shared.utils.OpenAICompatibleClientSync import get_openai_client +from shared.services.ai.openai_compatible_client_sync import get_openai_client # ==================== Helper Functions ==================== 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 270303e90..ef5362ec8 100644 --- a/apps/worker/app/services/document_parser/markdown_deferred_summary.py +++ b/apps/worker/app/services/document_parser/markdown_deferred_summary.py @@ -23,8 +23,8 @@ 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 +from shared.services.chunks.chunk_refs import build_chunk_ref +from app.services.common.file_utils import path_handle DeferredResult = ( tuple[ diff --git a/apps/worker/app/services/document_parser/markdown_image_asset.py b/apps/worker/app/services/document_parser/markdown_image_asset.py index 10c93b5dc..406872c18 100644 --- a/apps/worker/app/services/document_parser/markdown_image_asset.py +++ b/apps/worker/app/services/document_parser/markdown_image_asset.py @@ -15,8 +15,8 @@ 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 +from shared.services.chunks.chunk_refs import build_chunk_ref +from app.services.common.file_utils import path_handle @dataclass(frozen=True) diff --git a/apps/worker/app/services/document_parser/markdown_table_asset.py b/apps/worker/app/services/document_parser/markdown_table_asset.py index a01e8ce97..c1ad0a3e7 100644 --- a/apps/worker/app/services/document_parser/markdown_table_asset.py +++ b/apps/worker/app/services/document_parser/markdown_table_asset.py @@ -14,8 +14,8 @@ 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 +from shared.services.chunks.chunk_refs import build_chunk_ref +from app.services.common.file_utils import path_handle @dataclass(frozen=True) diff --git a/apps/worker/app/services/document_parser/md_parser.py b/apps/worker/app/services/document_parser/md_parser.py index e284ea2c8..aded8f424 100755 --- a/apps/worker/app/services/document_parser/md_parser.py +++ b/apps/worker/app/services/document_parser/md_parser.py @@ -43,8 +43,8 @@ from loguru import logger from shared.core.config import settings -from shared.utils.chunk_refs import has_chunk_ref -from shared.utils.text_utils import tokenize2stw_remove +from shared.services.chunks.chunk_refs import has_chunk_ref +from shared.services.text_processing.tokenization import tokenize2stw_remove def find_surround_context(md_lines, lid): @@ -215,7 +215,7 @@ def parse_md( relative_root=None, ): if md_lines is None and file_path is not None: - from shared.utils.file_loading import is_remote, load_file_bytes + from app.services.common.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 bd4acdbc4..8be565518 100644 --- a/apps/worker/app/services/document_parser/mineru_pdf_service.py +++ b/apps/worker/app/services/document_parser/mineru_pdf_service.py @@ -23,7 +23,7 @@ UnavailableException, ) from shared.services.storage.job_file_storage import JobFileStorage -from shared.utils.file_loading import is_remote +from app.services.common.file_loading import is_remote MINERU_UPLOAD_TIMEOUT = ( settings.MINERU_UPLOAD_CONNECT_TIMEOUT, diff --git a/apps/worker/app/services/document_parser/mineru_quota_manager.py b/apps/worker/app/services/document_parser/mineru_quota_manager.py index 830330b09..8c2015959 100644 --- a/apps/worker/app/services/document_parser/mineru_quota_manager.py +++ b/apps/worker/app/services/document_parser/mineru_quota_manager.py @@ -14,7 +14,7 @@ SyncRedisService, SyncRedisServiceFactory, ) -from shared.utils.quota_manager import BaseQuotaManager, TokenConfig, TokenLease +from shared.services.ai.quota_manager import BaseQuotaManager, TokenConfig, TokenLease # Backward-compatible aliases so existing imports keep working MinerUTokenConfig = TokenConfig diff --git a/apps/worker/app/services/document_parser/mineru_task_polling.py b/apps/worker/app/services/document_parser/mineru_task_polling.py index c09c1ea86..746bf4cd7 100644 --- a/apps/worker/app/services/document_parser/mineru_task_polling.py +++ b/apps/worker/app/services/document_parser/mineru_task_polling.py @@ -22,7 +22,7 @@ UnavailableException, ) from shared.core.exceptions.knowhere_exception import KnowhereException -from shared.utils.zip_download import download_and_extract_zip +from shared.services.http.zip_download import download_and_extract_zip def get_batch_status(data: dict[str, Any]) -> Optional[dict[str, Any]]: 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 832e13969..9750a331d 100644 --- a/apps/worker/app/services/document_parser/orchestration/parse_session.py +++ b/apps/worker/app/services/document_parser/orchestration/parse_session.py @@ -12,7 +12,7 @@ from shared.core.config import settings from shared.core.exceptions.domain_exceptions import ValidationException -from shared.utils.file_utils import path_handle +from app.services.common.file_utils import path_handle PDF_PAGE_LIMIT = 600 diff --git a/apps/worker/app/services/document_parser/path_helpers.py b/apps/worker/app/services/document_parser/path_helpers.py index 41965e125..f337b6556 100644 --- a/apps/worker/app/services/document_parser/path_helpers.py +++ b/apps/worker/app/services/document_parser/path_helpers.py @@ -5,8 +5,8 @@ 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 +from shared.services.chunks.chunk_refs import extract_chunk_refs +from app.services.common.file_utils import path_handle SUMMARY_PATH_MARKERS: tuple[str, ...] = ("summary", "\u6458\u8981\u603b\u7ed3") diff --git a/apps/worker/app/services/document_parser/pptx_parser.py b/apps/worker/app/services/document_parser/pptx_parser.py index 19063937a..8329fb18a 100755 --- a/apps/worker/app/services/document_parser/pptx_parser.py +++ b/apps/worker/app/services/document_parser/pptx_parser.py @@ -26,8 +26,8 @@ FileSystemException, ) from shared.core.logging import LogEvent -from shared.utils.file_loading import load_file_bytes -from shared.utils.file_utils import path_handle +from app.services.common.file_loading import load_file_bytes +from app.services.common.file_utils import path_handle # ==================== LibreOffice conversion ==================== @@ -56,7 +56,7 @@ def pptx_to_pdf_libreoffice(pptx_path, outdir="."): def _get_iloveapi_token_lease(): """acquire iLoveAPI token lease from the quotas pool and generate a JWT token""" - from shared.utils.iloveapi_quota_manager import get_iloveapi_quota_manager + from shared.services.ai.iloveapi_quota_manager import get_iloveapi_quota_manager quota_manager = get_iloveapi_quota_manager() @@ -111,7 +111,7 @@ def _pptx_bytes_to_pdf_bytes(pptx_bytes: bytes, filename: str) -> bytes: Acquires an in-flight slot before starting. Slots are released only if the Redis-backed reservation actually succeeded. """ - from shared.utils.iloveapi_quota_manager import get_iloveapi_quota_manager + from shared.services.ai.iloveapi_quota_manager import get_iloveapi_quota_manager quota_manager = get_iloveapi_quota_manager() diff --git a/apps/worker/app/services/document_parser/table_frame_parser.py b/apps/worker/app/services/document_parser/table_frame_parser.py index 63f39feb8..24d76538c 100644 --- a/apps/worker/app/services/document_parser/table_frame_parser.py +++ b/apps/worker/app/services/document_parser/table_frame_parser.py @@ -15,8 +15,8 @@ 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 +from shared.services.ai.openai_compatible_client_sync import get_openai_client +from shared.services.text_processing.tokenization import remove_duplicates_orderkept def parse_headers( diff --git a/apps/worker/app/services/document_parser/table_text_parser.py b/apps/worker/app/services/document_parser/table_text_parser.py index 261bbe653..f83d3be0e 100644 --- a/apps/worker/app/services/document_parser/table_text_parser.py +++ b/apps/worker/app/services/document_parser/table_text_parser.py @@ -13,7 +13,7 @@ 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 + from shared.services.text_processing.tokenization import _is_meaningful_token if not raw_header_text: return "" diff --git a/apps/worker/app/services/document_parser/text_helpers.py b/apps/worker/app/services/document_parser/text_helpers.py index b11fdc772..bbacf4327 100644 --- a/apps/worker/app/services/document_parser/text_helpers.py +++ b/apps/worker/app/services/document_parser/text_helpers.py @@ -2,7 +2,7 @@ import re -from shared.utils.text_utils import _CN_EN_NUM_RE +from shared.services.text_processing.tokenization import _CN_EN_NUM_RE _CN_CHAR_RE = re.compile(r"[\u4e00-\u9fff]") EN_START_LIMIT = 15 diff --git a/apps/worker/app/services/document_parser/toc_parser.py b/apps/worker/app/services/document_parser/toc_parser.py index 2efb1a4cd..5c52ef5c4 100644 --- a/apps/worker/app/services/document_parser/toc_parser.py +++ b/apps/worker/app/services/document_parser/toc_parser.py @@ -21,7 +21,7 @@ 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.services.ai.openai_compatible_client_sync import get_openai_client # ==================== Markdown TOC Detection Functions ==================== diff --git a/apps/worker/app/services/document_parser/txt_parser.py b/apps/worker/app/services/document_parser/txt_parser.py index e1d41956b..b595bedbc 100755 --- a/apps/worker/app/services/document_parser/txt_parser.py +++ b/apps/worker/app/services/document_parser/txt_parser.py @@ -11,9 +11,9 @@ 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.chunk_refs import CHUNK_REF_PATTERN -from shared.utils.file_loading import load_file_bytes -from shared.utils.OpenAICompatibleClientSync import get_openai_client +from shared.services.chunks.chunk_refs import CHUNK_REF_PATTERN +from app.services.common.file_loading import load_file_bytes +from shared.services.ai.openai_compatible_client_sync import get_openai_client def clean_texts_by_form(text, form="html"): diff --git a/apps/worker/app/services/workload/url_upload_transfer.py b/apps/worker/app/services/workload/url_upload_transfer.py index 25f4faf13..c9da336a8 100644 --- a/apps/worker/app/services/workload/url_upload_transfer.py +++ b/apps/worker/app/services/workload/url_upload_transfer.py @@ -10,7 +10,7 @@ ValidationException, ) from shared.services.storage.job_file_storage import JobFileStorage -from shared.utils.url_file_type import resolve_file_extension_sync +from shared.services.http.url_file_type import resolve_file_extension_sync def resolve_supported_url_extension(source_url: str) -> str: diff --git a/packages/shared-python/shared/utils/error_details.py b/packages/shared-python/shared/core/error_details.py similarity index 100% rename from packages/shared-python/shared/utils/error_details.py rename to packages/shared-python/shared/core/error_details.py diff --git a/packages/shared-python/shared/core/security.py b/packages/shared-python/shared/core/security.py index c1384e008..8e93a8241 100644 --- a/packages/shared-python/shared/core/security.py +++ b/packages/shared-python/shared/core/security.py @@ -6,6 +6,25 @@ import bcrypt +def mask_api_key(api_key: str | None) -> str: + """Return a stable masked representation for logs and internal errors.""" + if not api_key: + return "" + + if len(api_key) <= 4: + return "*" * len(api_key) + + if len(api_key) <= 8: + prefix_length = 2 + suffix_length = 2 + else: + prefix_length = 4 + suffix_length = 4 + + masked_length = max(1, len(api_key) - prefix_length - suffix_length) + return f"{api_key[:prefix_length]}{'*' * masked_length}{api_key[-suffix_length:]}" + + def get_password_hash(password: str) -> str: """ Generate a password hash. diff --git a/packages/shared-python/shared/utils/json_utils.py b/packages/shared-python/shared/core/serialization.py similarity index 100% rename from packages/shared-python/shared/utils/json_utils.py rename to packages/shared-python/shared/core/serialization.py diff --git a/packages/shared-python/shared/core/state_machine/service.py b/packages/shared-python/shared/core/state_machine/service.py index 6457e3964..1206313b9 100644 --- a/packages/shared-python/shared/core/state_machine/service.py +++ b/packages/shared-python/shared/core/state_machine/service.py @@ -38,7 +38,7 @@ 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.redis_key_builder import RedisKeyType, redis_key_builder +from shared.services.redis.key_builder import RedisKeyType, redis_key_builder class AsyncStateMachineService: 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 8935304e9..6d8ba1cc1 100644 --- a/packages/shared-python/shared/core/state_machine/service_sync.py +++ b/packages/shared-python/shared/core/state_machine/service_sync.py @@ -37,7 +37,7 @@ 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.redis_key_builder import RedisKeyType, redis_key_builder +from shared.services.redis.key_builder import RedisKeyType, redis_key_builder class SyncStateMachineService: diff --git a/packages/shared-python/shared/core/state_machine/transition_payloads.py b/packages/shared-python/shared/core/state_machine/transition_payloads.py index c2a2cf58b..d7e6094de 100644 --- a/packages/shared-python/shared/core/state_machine/transition_payloads.py +++ b/packages/shared-python/shared/core/state_machine/transition_payloads.py @@ -4,8 +4,8 @@ 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 +from shared.core.error_details import normalize_error_details +from shared.core.serialization import make_json_safe def utc_now_naive() -> datetime: diff --git a/packages/shared-python/shared/utils/utc_now.py b/packages/shared-python/shared/core/time.py similarity index 100% rename from packages/shared-python/shared/utils/utc_now.py rename to packages/shared-python/shared/core/time.py diff --git a/packages/shared-python/shared/models/database/api_key.py b/packages/shared-python/shared/models/database/api_key.py index a70b24f9f..45b578867 100644 --- a/packages/shared-python/shared/models/database/api_key.py +++ b/packages/shared-python/shared/models/database/api_key.py @@ -12,7 +12,7 @@ from sqlalchemy.orm import Mapped, mapped_column from shared.core.database import Base -from shared.utils.utc_now import utc_now_naive +from shared.core.time import utc_now_naive class APIKey(Base): diff --git a/packages/shared-python/shared/models/database/credits_transaction.py b/packages/shared-python/shared/models/database/credits_transaction.py index 1b83d6bbc..61bcdbb7b 100644 --- a/packages/shared-python/shared/models/database/credits_transaction.py +++ b/packages/shared-python/shared/models/database/credits_transaction.py @@ -12,7 +12,7 @@ from sqlalchemy.orm import Mapped, mapped_column from shared.core.database import Base -from shared.utils.utc_now import utc_now_naive +from shared.core.time import utc_now_naive class CreditsTransaction(Base): diff --git a/packages/shared-python/shared/models/database/demo_materialization.py b/packages/shared-python/shared/models/database/demo_materialization.py index d4cc1a7a1..e63027c5a 100644 --- a/packages/shared-python/shared/models/database/demo_materialization.py +++ b/packages/shared-python/shared/models/database/demo_materialization.py @@ -9,7 +9,7 @@ from sqlalchemy.orm import Mapped, mapped_column from shared.core.database import Base -from shared.utils.utc_now import utc_now_naive +from shared.core.time import utc_now_naive class DemoMaterialization(Base): diff --git a/packages/shared-python/shared/models/database/document.py b/packages/shared-python/shared/models/database/document.py index 41bed45bb..d23ad235d 100644 --- a/packages/shared-python/shared/models/database/document.py +++ b/packages/shared-python/shared/models/database/document.py @@ -22,7 +22,7 @@ from sqlalchemy.orm import Mapped, mapped_column, relationship from shared.core.database import Base -from shared.utils.utc_now import utc_now_naive +from shared.core.time import utc_now_naive class Document(Base): diff --git a/packages/shared-python/shared/models/database/guest_device.py b/packages/shared-python/shared/models/database/guest_device.py index cf956bb26..2b62a1fc7 100644 --- a/packages/shared-python/shared/models/database/guest_device.py +++ b/packages/shared-python/shared/models/database/guest_device.py @@ -10,7 +10,7 @@ from sqlalchemy.orm import Mapped, mapped_column from shared.core.database import Base -from shared.utils.utc_now import utc_now_naive +from shared.core.time import utc_now_naive class GuestDevice(Base): diff --git a/packages/shared-python/shared/models/database/job.py b/packages/shared-python/shared/models/database/job.py index 76c9f1df1..4553ea655 100644 --- a/packages/shared-python/shared/models/database/job.py +++ b/packages/shared-python/shared/models/database/job.py @@ -25,7 +25,7 @@ from sqlalchemy.orm import Mapped, mapped_column, relationship from shared.core.database import Base -from shared.utils.utc_now import utc_now_naive +from shared.core.time import utc_now_naive if TYPE_CHECKING: diff --git a/packages/shared-python/shared/models/database/job_result.py b/packages/shared-python/shared/models/database/job_result.py index c6d4d9e73..f683c42d2 100644 --- a/packages/shared-python/shared/models/database/job_result.py +++ b/packages/shared-python/shared/models/database/job_result.py @@ -10,7 +10,7 @@ from sqlalchemy.orm import Mapped, mapped_column, relationship from shared.core.database import Base -from shared.utils.utc_now import utc_now_naive +from shared.core.time import utc_now_naive if TYPE_CHECKING: from shared.models.database.job import Job diff --git a/packages/shared-python/shared/models/database/job_state_audit_log.py b/packages/shared-python/shared/models/database/job_state_audit_log.py index 970b2af57..e10c4de58 100644 --- a/packages/shared-python/shared/models/database/job_state_audit_log.py +++ b/packages/shared-python/shared/models/database/job_state_audit_log.py @@ -9,7 +9,7 @@ from sqlalchemy.orm import Mapped, mapped_column, relationship from shared.core.database import Base -from shared.utils.utc_now import utc_now_naive +from shared.core.time import utc_now_naive if TYPE_CHECKING: from shared.models.database.job import Job diff --git a/packages/shared-python/shared/models/database/job_state_history.py b/packages/shared-python/shared/models/database/job_state_history.py index b480ae3f7..98ed383e4 100644 --- a/packages/shared-python/shared/models/database/job_state_history.py +++ b/packages/shared-python/shared/models/database/job_state_history.py @@ -10,7 +10,7 @@ from sqlalchemy.orm import Mapped, mapped_column, relationship from shared.core.database import Base -from shared.utils.utc_now import utc_now_naive +from shared.core.time import utc_now_naive if TYPE_CHECKING: from shared.models.database.job import Job diff --git a/packages/shared-python/shared/models/database/payment_record.py b/packages/shared-python/shared/models/database/payment_record.py index 69edb639e..923c7e1aa 100644 --- a/packages/shared-python/shared/models/database/payment_record.py +++ b/packages/shared-python/shared/models/database/payment_record.py @@ -21,7 +21,7 @@ from sqlalchemy.orm import Mapped, mapped_column from shared.core.database import Base -from shared.utils.utc_now import utc_now_naive +from shared.core.time import utc_now_naive class PaymentRecord(Base): """Payment Record Model (for idempotency guarantee)""" diff --git a/packages/shared-python/shared/models/database/stripe_price_config.py b/packages/shared-python/shared/models/database/stripe_price_config.py index 0fb1200f2..c820ead19 100644 --- a/packages/shared-python/shared/models/database/stripe_price_config.py +++ b/packages/shared-python/shared/models/database/stripe_price_config.py @@ -18,7 +18,7 @@ from sqlalchemy.orm import Mapped, mapped_column from shared.core.database import Base -from shared.utils.utc_now import utc_now_naive +from shared.core.time import utc_now_naive class StripePriceConfig(Base): diff --git a/packages/shared-python/shared/models/database/user_balance.py b/packages/shared-python/shared/models/database/user_balance.py index a67ea2565..77a380e6e 100644 --- a/packages/shared-python/shared/models/database/user_balance.py +++ b/packages/shared-python/shared/models/database/user_balance.py @@ -11,7 +11,7 @@ from sqlalchemy.orm import Mapped, mapped_column from shared.core.database import Base -from shared.utils.utc_now import utc_now_naive +from shared.core.time import utc_now_naive class UserBalance(Base): """User balance model — tracks credits balance and tier membership""" diff --git a/packages/shared-python/shared/models/database/webhook.py b/packages/shared-python/shared/models/database/webhook.py index 09f7f4dd1..c5da892d7 100644 --- a/packages/shared-python/shared/models/database/webhook.py +++ b/packages/shared-python/shared/models/database/webhook.py @@ -16,7 +16,7 @@ from sqlalchemy.orm import Mapped, mapped_column, relationship from shared.core.database import Base -from shared.utils.utc_now import utc_now_naive +from shared.core.time import utc_now_naive if TYPE_CHECKING: from shared.models.database.job import Job diff --git a/packages/shared-python/shared/models/database/webhook_log.py b/packages/shared-python/shared/models/database/webhook_log.py index 2beb6dca2..b41dff5a8 100644 --- a/packages/shared-python/shared/models/database/webhook_log.py +++ b/packages/shared-python/shared/models/database/webhook_log.py @@ -14,7 +14,7 @@ from sqlalchemy.orm import Mapped, mapped_column, relationship from shared.core.database import Base -from shared.utils.utc_now import utc_now_naive +from shared.core.time import utc_now_naive if TYPE_CHECKING: from shared.models.database.job import Job diff --git a/packages/shared-python/shared/models/database/webhook_secret.py b/packages/shared-python/shared/models/database/webhook_secret.py index c1d391b3f..50351ed8e 100644 --- a/packages/shared-python/shared/models/database/webhook_secret.py +++ b/packages/shared-python/shared/models/database/webhook_secret.py @@ -15,7 +15,7 @@ from sqlalchemy.orm import Mapped, mapped_column from shared.core.database import Base -from shared.utils.utc_now import utc_now_naive +from shared.core.time import utc_now_naive class WebhookSecretStatus: diff --git a/packages/shared-python/shared/repositories/credits_repository.py b/packages/shared-python/shared/repositories/credits_repository.py index 3e97e60c5..04fc1d452 100644 --- a/packages/shared-python/shared/repositories/credits_repository.py +++ b/packages/shared-python/shared/repositories/credits_repository.py @@ -12,7 +12,7 @@ from shared.models.database.credits_transaction import CreditsTransaction from shared.models.database.payment_record import PaymentRecord from shared.models.database.user_balance import UserBalance -from shared.utils.utc_now import utc_now_naive +from shared.core.time import utc_now_naive class CreditsRepository: diff --git a/packages/shared-python/shared/repositories/credits_sync_repository.py b/packages/shared-python/shared/repositories/credits_sync_repository.py index e0c198615..925b9857b 100644 --- a/packages/shared-python/shared/repositories/credits_sync_repository.py +++ b/packages/shared-python/shared/repositories/credits_sync_repository.py @@ -11,7 +11,7 @@ from shared.models.database.credits_transaction import CreditsTransaction from shared.models.database.payment_record import PaymentRecord from shared.models.database.user_balance import UserBalance -from shared.utils.utc_now import utc_now_naive +from shared.core.time import utc_now_naive class SyncCreditsRepository: diff --git a/packages/shared-python/shared/utils/ali_quota_manager.py b/packages/shared-python/shared/services/ai/ali_quota_manager.py similarity index 96% rename from packages/shared-python/shared/utils/ali_quota_manager.py rename to packages/shared-python/shared/services/ai/ali_quota_manager.py index 535b86c0c..764f3e7fb 100644 --- a/packages/shared-python/shared/utils/ali_quota_manager.py +++ b/packages/shared-python/shared/services/ai/ali_quota_manager.py @@ -14,7 +14,7 @@ SyncRedisService, SyncRedisServiceFactory, ) -from shared.utils.quota_manager import BaseQuotaManager, TokenConfig +from shared.services.ai.quota_manager import BaseQuotaManager, TokenConfig class AliQuotaManager(BaseQuotaManager): diff --git a/packages/shared-python/shared/utils/iloveapi_quota_manager.py b/packages/shared-python/shared/services/ai/iloveapi_quota_manager.py similarity index 99% rename from packages/shared-python/shared/utils/iloveapi_quota_manager.py rename to packages/shared-python/shared/services/ai/iloveapi_quota_manager.py index b3b923a41..7afb5aacf 100644 --- a/packages/shared-python/shared/utils/iloveapi_quota_manager.py +++ b/packages/shared-python/shared/services/ai/iloveapi_quota_manager.py @@ -26,7 +26,7 @@ SyncRedisService, SyncRedisServiceFactory, ) -from shared.utils.quota_manager import BaseQuotaManager, TokenConfig +from shared.services.ai.quota_manager import BaseQuotaManager, TokenConfig class ILoveApiQuotaManager(BaseQuotaManager): diff --git a/packages/shared-python/shared/utils/llm_mock.py b/packages/shared-python/shared/services/ai/llm_mock.py similarity index 100% rename from packages/shared-python/shared/utils/llm_mock.py rename to packages/shared-python/shared/services/ai/llm_mock.py diff --git a/packages/shared-python/shared/utils/OpenAICompatibleClientSync.py b/packages/shared-python/shared/services/ai/openai_compatible_client_sync.py similarity index 98% rename from packages/shared-python/shared/utils/OpenAICompatibleClientSync.py rename to packages/shared-python/shared/services/ai/openai_compatible_client_sync.py index 66da76a6d..d62718c79 100644 --- a/packages/shared-python/shared/utils/OpenAICompatibleClientSync.py +++ b/packages/shared-python/shared/services/ai/openai_compatible_client_sync.py @@ -17,9 +17,9 @@ from shared.core.config import settings from shared.core.exceptions.domain_exceptions import LLMServiceException -from shared.utils.http_clients import get_sync_client -from shared.utils.llm_mock import build_mock_chat_completion_response -from shared.utils.security_utils import mask_api_key +from shared.services.http.client_pool import get_sync_client +from shared.services.ai.llm_mock import build_mock_chat_completion_response +from shared.core.security import mask_api_key LOCAL_DEBUG = os.getenv("LOCAL_DEBUG", "0") == "1" LLMUsage = dict[str, int] @@ -185,7 +185,7 @@ def _make_ali_pool_call( api_kwargs: Dict[str, Any], ) -> tuple[str, LLMUsage]: """Acquire a token, make the call, and retry inline on 429.""" - from shared.utils.ali_quota_manager import get_ali_quota_manager + from shared.services.ai.ali_quota_manager import get_ali_quota_manager quota_manager = get_ali_quota_manager() base_url: Optional[str] = self._base_url diff --git a/packages/shared-python/shared/utils/quota_manager.py b/packages/shared-python/shared/services/ai/quota_manager.py similarity index 100% rename from packages/shared-python/shared/utils/quota_manager.py rename to packages/shared-python/shared/services/ai/quota_manager.py diff --git a/packages/shared-python/shared/utils/api_keys.py b/packages/shared-python/shared/services/auth/api_key_tokens.py similarity index 100% rename from packages/shared-python/shared/utils/api_keys.py rename to packages/shared-python/shared/services/auth/api_key_tokens.py diff --git a/packages/shared-python/shared/services/billing/credits_service.py b/packages/shared-python/shared/services/billing/credits_service.py index d0dc48fc4..3d690401f 100644 --- a/packages/shared-python/shared/services/billing/credits_service.py +++ b/packages/shared-python/shared/services/billing/credits_service.py @@ -24,7 +24,7 @@ from shared.models.database.payment_record import PaymentRecord from shared.models.database.user_balance import UserBalance from shared.repositories.credits_repository import CreditsRepository -from shared.utils.utc_now import utc_now_naive +from shared.core.time import utc_now_naive NON_EXPIRING_CREDIT_TRANSACTION_TYPES: tuple[str, ...] = ("refund",) diff --git a/packages/shared-python/shared/services/billing/credits_sync_service.py b/packages/shared-python/shared/services/billing/credits_sync_service.py index e6eb63446..d6b07d42b 100644 --- a/packages/shared-python/shared/services/billing/credits_sync_service.py +++ b/packages/shared-python/shared/services/billing/credits_sync_service.py @@ -19,7 +19,7 @@ from shared.models.database.payment_record import PaymentRecord from shared.models.database.user_balance import UserBalance from shared.repositories.credits_sync_repository import SyncCreditsRepository -from shared.utils.utc_now import utc_now_naive +from shared.core.time import utc_now_naive NON_EXPIRING_CREDIT_TRANSACTION_TYPES: tuple[str, ...] = ("refund",) diff --git a/packages/shared-python/shared/services/chunks/chunk_connections.py b/packages/shared-python/shared/services/chunks/chunk_connections.py index efb0804f1..13d4ab90c 100644 --- a/packages/shared-python/shared/services/chunks/chunk_connections.py +++ b/packages/shared-python/shared/services/chunks/chunk_connections.py @@ -5,7 +5,7 @@ from collections.abc import Mapping, Sequence from typing import Any, TypeAlias, TypedDict -from shared.utils.chunk_refs import ChunkRefSpan, extract_chunk_ref_spans +from shared.services.chunks.chunk_refs import ChunkRefSpan, extract_chunk_ref_spans class PositionPayload(TypedDict): diff --git a/packages/shared-python/shared/utils/chunk_refs.py b/packages/shared-python/shared/services/chunks/chunk_refs.py similarity index 100% rename from packages/shared-python/shared/utils/chunk_refs.py rename to packages/shared-python/shared/services/chunks/chunk_refs.py diff --git a/packages/shared-python/shared/services/http/__init__.py b/packages/shared-python/shared/services/http/__init__.py new file mode 100644 index 000000000..3bed643fe --- /dev/null +++ b/packages/shared-python/shared/services/http/__init__.py @@ -0,0 +1 @@ +"""Shared outbound HTTP adapters and URL safety helpers.""" diff --git a/packages/shared-python/shared/utils/http_clients.py b/packages/shared-python/shared/services/http/client_pool.py similarity index 100% rename from packages/shared-python/shared/utils/http_clients.py rename to packages/shared-python/shared/services/http/client_pool.py diff --git a/packages/shared-python/shared/utils/pinned_outbound_http.py b/packages/shared-python/shared/services/http/pinned_outbound.py similarity index 99% rename from packages/shared-python/shared/utils/pinned_outbound_http.py rename to packages/shared-python/shared/services/http/pinned_outbound.py index 808b219c8..c281dceb5 100644 --- a/packages/shared-python/shared/utils/pinned_outbound_http.py +++ b/packages/shared-python/shared/services/http/pinned_outbound.py @@ -19,7 +19,7 @@ from urllib3.util.connection import create_connection from shared.core.exceptions.domain_exceptions import ValidationException -from shared.utils.url_security import SafePublicHTTPURL +from shared.services.http.url_security import SafePublicHTTPURL @dataclass(frozen=True) diff --git a/packages/shared-python/shared/utils/url_file_type.py b/packages/shared-python/shared/services/http/url_file_type.py similarity index 97% rename from packages/shared-python/shared/utils/url_file_type.py rename to packages/shared-python/shared/services/http/url_file_type.py index 7314d7e23..a82f97a60 100644 --- a/packages/shared-python/shared/utils/url_file_type.py +++ b/packages/shared-python/shared/services/http/url_file_type.py @@ -12,7 +12,7 @@ from shared.core.config import settings from shared.core.exceptions.domain_exceptions import ValidationException -from shared.utils.url_security import ( +from shared.services.http.url_security import ( HTTPURLValidationResult, SafePublicHTTPURL, validate_http_url_and_resolve_ip, @@ -116,7 +116,7 @@ async def resolve_file_extension_async(url: str) -> str | None: return ext try: - from shared.utils.http_clients import get_async_client + from shared.services.http.client_pool import get_async_client client = get_async_client() response = None @@ -167,7 +167,7 @@ def resolve_file_extension_sync(url: str) -> str | None: return ext try: - from shared.utils.http_clients import get_sync_client + from shared.services.http.client_pool import get_sync_client client = get_sync_client() response = None diff --git a/packages/shared-python/shared/utils/url_security.py b/packages/shared-python/shared/services/http/url_security.py similarity index 100% rename from packages/shared-python/shared/utils/url_security.py rename to packages/shared-python/shared/services/http/url_security.py diff --git a/packages/shared-python/shared/utils/zip_download.py b/packages/shared-python/shared/services/http/zip_download.py similarity index 100% rename from packages/shared-python/shared/utils/zip_download.py rename to packages/shared-python/shared/services/http/zip_download.py diff --git a/packages/shared-python/shared/services/job_failure_sync.py b/packages/shared-python/shared/services/job_failure_sync.py index ad3a313dd..66c93cd1e 100644 --- a/packages/shared-python/shared/services/job_failure_sync.py +++ b/packages/shared-python/shared/services/job_failure_sync.py @@ -13,7 +13,7 @@ from shared.services.billing.credits_sync_service import SyncCreditsService from shared.services.job_post_commit_effects_sync import PostCommitEffectPlan from shared.services.job_webhook_outbox_sync import SyncJobWebhookOutbox -from shared.utils.error_details import normalize_error_details +from shared.core.error_details import normalize_error_details @dataclass(frozen=True) diff --git a/packages/shared-python/shared/services/job_lifecycle_sync.py b/packages/shared-python/shared/services/job_lifecycle_sync.py index ad0dcfc70..f160f41af 100644 --- a/packages/shared-python/shared/services/job_lifecycle_sync.py +++ b/packages/shared-python/shared/services/job_lifecycle_sync.py @@ -26,7 +26,7 @@ from shared.services.redis.redis_sync_service import ( SyncRedisServiceFactory, ) -from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder +from shared.services.redis.key_builder import RedisKeyType, redis_key_builder class SyncJobLifecycleService: diff --git a/packages/shared-python/shared/services/jobs/result_delivery.py b/packages/shared-python/shared/services/jobs/result_delivery.py index 0a5c9b0d6..406ba8627 100644 --- a/packages/shared-python/shared/services/jobs/result_delivery.py +++ b/packages/shared-python/shared/services/jobs/result_delivery.py @@ -5,7 +5,7 @@ from typing import Any from shared.services.storage.job_file_storage import JobFileStorage -from shared.utils.utc_now import utc_now_naive +from shared.core.time import utc_now_naive @dataclass(frozen=True) diff --git a/packages/shared-python/shared/services/redis/distributed_lock.py b/packages/shared-python/shared/services/redis/distributed_lock.py index a36a23bf9..8933737f1 100644 --- a/packages/shared-python/shared/services/redis/distributed_lock.py +++ b/packages/shared-python/shared/services/redis/distributed_lock.py @@ -18,7 +18,7 @@ from shared.core.config import settings from shared.core.exceptions.domain_exceptions import UnavailableException from shared.services.redis.redis_sync_service import SyncRedisService -from shared.utils.redis_key_builder import redis_key_builder +from shared.services.redis.key_builder import redis_key_builder # Lua script: atomically check owner token then delete. # Prevents a stale owner from accidentally deleting a lock diff --git a/packages/shared-python/shared/services/redis/job_info_redis_service.py b/packages/shared-python/shared/services/redis/job_info_redis_service.py index 89c1a2ec2..0ee34b0b7 100644 --- a/packages/shared-python/shared/services/redis/job_info_redis_service.py +++ b/packages/shared-python/shared/services/redis/job_info_redis_service.py @@ -8,7 +8,7 @@ from loguru import logger from shared.services.redis.redis_service import RedisService -from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder +from shared.services.redis.key_builder import RedisKeyType, redis_key_builder class JobInfoRedisService: diff --git a/packages/shared-python/shared/services/redis/job_metadata_service.py b/packages/shared-python/shared/services/redis/job_metadata_service.py index 62fb9c135..410a1c052 100644 --- a/packages/shared-python/shared/services/redis/job_metadata_service.py +++ b/packages/shared-python/shared/services/redis/job_metadata_service.py @@ -5,7 +5,7 @@ from loguru import logger from shared.services.redis.redis_service import RedisService -from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder +from shared.services.redis.key_builder import RedisKeyType, redis_key_builder class JobMetadataService: diff --git a/packages/shared-python/shared/utils/redis_key_builder.py b/packages/shared-python/shared/services/redis/key_builder.py similarity index 100% rename from packages/shared-python/shared/utils/redis_key_builder.py rename to packages/shared-python/shared/services/redis/key_builder.py diff --git a/packages/shared-python/shared/services/redis/rate_limit_service.py b/packages/shared-python/shared/services/redis/rate_limit_service.py index de7ef821b..c1cec343b 100644 --- a/packages/shared-python/shared/services/redis/rate_limit_service.py +++ b/packages/shared-python/shared/services/redis/rate_limit_service.py @@ -7,7 +7,7 @@ from loguru import logger from shared.services.redis.redis_service import RedisService -from shared.utils.redis_key_builder import redis_key_builder +from shared.services.redis.key_builder import redis_key_builder class RateLimitService: diff --git a/packages/shared-python/shared/services/redis/redis_monitor.py b/packages/shared-python/shared/services/redis/redis_monitor.py index 6ae090590..f6887576a 100644 --- a/packages/shared-python/shared/services/redis/redis_monitor.py +++ b/packages/shared-python/shared/services/redis/redis_monitor.py @@ -7,7 +7,7 @@ from loguru import logger from shared.services.redis.redis_service import RedisService -from shared.utils.redis_key_builder import redis_key_builder +from shared.services.redis.key_builder import redis_key_builder class RedisMonitor: diff --git a/packages/shared-python/shared/services/redis/redis_service.py b/packages/shared-python/shared/services/redis/redis_service.py index 99ed39c1d..fd735bcf8 100644 --- a/packages/shared-python/shared/services/redis/redis_service.py +++ b/packages/shared-python/shared/services/redis/redis_service.py @@ -13,7 +13,7 @@ RedisConnectionError, RedisOperationError, ) -from shared.utils.redis_retry import RedisHealthChecker, RedisRetry +from shared.services.redis.retry_policy import RedisHealthChecker, RedisRetry ResponseT = TypeVar("ResponseT") @@ -87,7 +87,7 @@ async def set( if isinstance(value, (dict, list)): # Use make_json_safe so complex types serialize consistently. - from shared.utils.json_utils import make_json_safe + from shared.core.serialization import make_json_safe safe_value = make_json_safe(value) value = json.dumps(safe_value, ensure_ascii=False) diff --git a/packages/shared-python/shared/services/redis/redis_sync_service.py b/packages/shared-python/shared/services/redis/redis_sync_service.py index aa725b314..2fe2cb7bf 100644 --- a/packages/shared-python/shared/services/redis/redis_sync_service.py +++ b/packages/shared-python/shared/services/redis/redis_sync_service.py @@ -12,7 +12,7 @@ from redis.connection import BlockingConnectionPool from shared.core.config.redis import RedisConfigManager -from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder +from shared.services.redis.key_builder import RedisKeyType, redis_key_builder class SyncRedisService: @@ -69,7 +69,7 @@ def set( client = self._get_client() full_key = self._build_key(key) if isinstance(value, (dict, list)): - from shared.utils.json_utils import make_json_safe + from shared.core.serialization import make_json_safe safe_value = make_json_safe(value) value = json.dumps(safe_value, ensure_ascii=False) @@ -407,7 +407,7 @@ def set_task_status(self, task_id: str, status: str) -> bool: def save_task_result(self, task_id: str, result: Dict[str, Any]) -> bool: try: - from shared.utils.json_utils import make_json_safe + from shared.core.serialization import make_json_safe task_ttl = redis_key_builder.get_key_ttl(RedisKeyType.TASK) result_key = self.redis._build_key(redis_key_builder.task_result(task_id)) diff --git a/packages/shared-python/shared/utils/redis_retry.py b/packages/shared-python/shared/services/redis/retry_policy.py similarity index 100% rename from packages/shared-python/shared/utils/redis_retry.py rename to packages/shared-python/shared/services/redis/retry_policy.py diff --git a/packages/shared-python/shared/services/redis/task_redis_service.py b/packages/shared-python/shared/services/redis/task_redis_service.py index 883d8a116..6c9444534 100644 --- a/packages/shared-python/shared/services/redis/task_redis_service.py +++ b/packages/shared-python/shared/services/redis/task_redis_service.py @@ -5,7 +5,7 @@ from loguru import logger from shared.services.redis.redis_service import RedisService -from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder +from shared.services.redis.key_builder import RedisKeyType, redis_key_builder class TaskRedisService: diff --git a/packages/shared-python/shared/services/redis/user_redis_service.py b/packages/shared-python/shared/services/redis/user_redis_service.py index 6db0a2acc..d28ee59d5 100644 --- a/packages/shared-python/shared/services/redis/user_redis_service.py +++ b/packages/shared-python/shared/services/redis/user_redis_service.py @@ -5,7 +5,7 @@ from loguru import logger from shared.services.redis.redis_service import RedisService -from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder +from shared.services.redis.key_builder import RedisKeyType, redis_key_builder class UserRedisService: diff --git a/packages/shared-python/shared/services/retrieval/agentic/evidence.py b/packages/shared-python/shared/services/retrieval/agentic/evidence.py index 3d54fb47b..5bd38f947 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/evidence.py +++ b/packages/shared-python/shared/services/retrieval/agentic/evidence.py @@ -11,7 +11,7 @@ from shared.services.retrieval.agentic.types import DocTreeNode 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 +from shared.services.retrieval.token_estimate import estimate_tokens def with_context_prompt_projection( diff --git a/packages/shared-python/shared/services/retrieval/agentic/policy.py b/packages/shared-python/shared/services/retrieval/agentic/policy.py index 25323580a..7a3bd3654 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/policy.py +++ b/packages/shared-python/shared/services/retrieval/agentic/policy.py @@ -21,7 +21,7 @@ from shared.services.retrieval.agentic.budget import BudgetExceeded from shared.services.retrieval.agentic.types import AgentRunConfig, AgentState from shared.services.retrieval.llm_adapter import LLMFn -from shared.utils.token_estimate import estimate_tokens +from shared.services.retrieval.token_estimate import estimate_tokens def _parse_answer_response(text: str) -> dict[str, Any] | None: diff --git a/packages/shared-python/shared/services/retrieval/agentic/runtime.py b/packages/shared-python/shared/services/retrieval/agentic/runtime.py index d658cd562..750541c25 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/runtime.py +++ b/packages/shared-python/shared/services/retrieval/agentic/runtime.py @@ -12,7 +12,7 @@ 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 +from shared.services.retrieval.token_estimate import estimate_tokens def build_config_from_env() -> AgentRunConfig: 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 index b7a251805..24d2d3ee9 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/section_prompt_projection.py +++ b/packages/shared-python/shared/services/retrieval/agentic/section_prompt_projection.py @@ -1,7 +1,7 @@ """Prompt projection for agentic section navigation.""" from __future__ import annotations -from shared.utils.text_utils import truncate_content_preview +from shared.services.text_processing.tokenization import truncate_content_preview def format_items_for_llm( diff --git a/packages/shared-python/shared/services/retrieval/lexical_ranker.py b/packages/shared-python/shared/services/retrieval/lexical_ranker.py index 8b50cd404..39fe446c0 100644 --- a/packages/shared-python/shared/services/retrieval/lexical_ranker.py +++ b/packages/shared-python/shared/services/retrieval/lexical_ranker.py @@ -4,7 +4,7 @@ from loguru import logger -from shared.utils.text_utils import tokenize_for_retrieval +from shared.services.text_processing.tokenization import tokenize_for_retrieval def tokenize_query_for_ranker(query: str) -> list[str]: diff --git a/packages/shared-python/shared/services/retrieval/lexical_text.py b/packages/shared-python/shared/services/retrieval/lexical_text.py index 42d15366f..7c98dc56c 100644 --- a/packages/shared-python/shared/services/retrieval/lexical_text.py +++ b/packages/shared-python/shared/services/retrieval/lexical_text.py @@ -6,7 +6,7 @@ from typing import Any, Optional -from shared.utils.text_utils import tokenize_contents_for_retrieval +from shared.services.text_processing.tokenization import tokenize_contents_for_retrieval def normalize_section_path(path: Optional[str]) -> str: diff --git a/packages/shared-python/shared/services/retrieval/llm_adapter.py b/packages/shared-python/shared/services/retrieval/llm_adapter.py index c5e592b5e..9c8bd9263 100644 --- a/packages/shared-python/shared/services/retrieval/llm_adapter.py +++ b/packages/shared-python/shared/services/retrieval/llm_adapter.py @@ -87,7 +87,7 @@ def create_retrieval_llm_fn( effective_model = model or _resolve_default_model() async def llm_fn(prompt: LLMFnInput) -> str: - from shared.utils.OpenAICompatibleClientSync import get_openai_client + from shared.services.ai.openai_compatible_client_sync import get_openai_client client = get_openai_client(model=effective_model) current_llm_usage.set(None) @@ -118,7 +118,7 @@ def create_retrieval_planner_fn( effective_model = model or _resolve_planner_model(thinking=thinking) async def llm_fn(prompt: LLMFnInput) -> str: - from shared.utils.OpenAICompatibleClientSync import get_openai_client + from shared.services.ai.openai_compatible_client_sync import get_openai_client client = get_openai_client(model=effective_model) current_llm_usage.set(None) @@ -159,7 +159,7 @@ def create_retrieval_vlm_fn( return None async def vlm_fn(prompt: LLMFnInput) -> str: - from shared.utils.OpenAICompatibleClientSync import get_openai_client + from shared.services.ai.openai_compatible_client_sync import get_openai_client client = get_openai_client(model=effective_model) current_llm_usage.set(None) diff --git a/packages/shared-python/shared/utils/token_estimate.py b/packages/shared-python/shared/services/retrieval/token_estimate.py similarity index 100% rename from packages/shared-python/shared/utils/token_estimate.py rename to packages/shared-python/shared/services/retrieval/token_estimate.py diff --git a/packages/shared-python/shared/services/retrieval/workflow/planner.py b/packages/shared-python/shared/services/retrieval/workflow/planner.py index e34713db6..463894879 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/planner.py +++ b/packages/shared-python/shared/services/retrieval/workflow/planner.py @@ -11,7 +11,7 @@ from shared.services.retrieval.agentic.budget import BudgetExceeded, BudgetLedger from shared.services.retrieval.llm_adapter import LLMFn, current_llm_usage from shared.services.retrieval.workflow.types import FinalStrategy, OutputRole, PlannedStep, QueryPlan, StepKind -from shared.utils.token_estimate import estimate_tokens +from shared.services.retrieval.token_estimate import estimate_tokens _PLAN_SCHEMA = { diff --git a/packages/shared-python/shared/services/retrieval/workflow/synthesizer.py b/packages/shared-python/shared/services/retrieval/workflow/synthesizer.py index 89d4ec031..2ee92bac9 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/synthesizer.py +++ b/packages/shared-python/shared/services/retrieval/workflow/synthesizer.py @@ -6,7 +6,7 @@ from shared.services.retrieval.agentic.budget import BudgetLedger from shared.services.retrieval.llm_adapter import LLMFn, current_llm_usage from shared.services.retrieval.workflow.types import PlannedStep, QueryPlan, StepResult -from shared.utils.token_estimate import estimate_tokens +from shared.services.retrieval.token_estimate import estimate_tokens _SYNTHESIZE_PROMPT = """\ 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 cc46f6bd3..77d49e2b4 100644 --- a/packages/shared-python/shared/services/storage/job_file_storage.py +++ b/packages/shared-python/shared/services/storage/job_file_storage.py @@ -8,8 +8,8 @@ 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 +from shared.services.http.pinned_outbound import download_pinned_outbound_file +from shared.services.http.url_security import validate_http_url_and_resolve_ip class JobFileStorage: diff --git a/packages/shared-python/shared/services/storage/zip_result_schema.py b/packages/shared-python/shared/services/storage/zip_result_schema.py index 410bfa111..72bb6b5a0 100644 --- a/packages/shared-python/shared/services/storage/zip_result_schema.py +++ b/packages/shared-python/shared/services/storage/zip_result_schema.py @@ -10,8 +10,8 @@ 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 +from shared.services.text_processing.tokenization import truncate_content_preview +from shared.core.time import utc_now_naive class ZipResultSchemaBuilder: diff --git a/packages/shared-python/shared/services/text_processing/__init__.py b/packages/shared-python/shared/services/text_processing/__init__.py new file mode 100644 index 000000000..7d77e9d64 --- /dev/null +++ b/packages/shared-python/shared/services/text_processing/__init__.py @@ -0,0 +1 @@ +"""Shared text processing and tokenization services.""" diff --git a/packages/shared-python/shared/utils/text_utils.py b/packages/shared-python/shared/services/text_processing/tokenization.py similarity index 99% rename from packages/shared-python/shared/utils/text_utils.py rename to packages/shared-python/shared/services/text_processing/tokenization.py index dd61bfdd2..8e00955d9 100644 --- a/packages/shared-python/shared/utils/text_utils.py +++ b/packages/shared-python/shared/services/text_processing/tokenization.py @@ -10,7 +10,7 @@ from syntok.tokenizer import Tokenizer as _SyntokTokenizer -from shared.utils.chunk_refs import CHUNK_REF_PATTERN +from shared.services.chunks.chunk_refs import CHUNK_REF_PATTERN try: from blingfire import text_to_words as _blingfire_text_to_words diff --git a/packages/shared-python/shared/services/webhook/delivery_client.py b/packages/shared-python/shared/services/webhook/delivery_client.py index 232f9a875..2e33aca20 100644 --- a/packages/shared-python/shared/services/webhook/delivery_client.py +++ b/packages/shared-python/shared/services/webhook/delivery_client.py @@ -8,8 +8,8 @@ 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 +from shared.services.http.pinned_outbound import send_pinned_outbound_request +from shared.services.http.url_security import validate_http_url_and_resolve_ip_async HTTP_TIMEOUT_SECONDS = 10 diff --git a/packages/shared-python/shared/services/webhook/qstash_publisher.py b/packages/shared-python/shared/services/webhook/qstash_publisher.py index 7a0580cb6..810f35def 100644 --- a/packages/shared-python/shared/services/webhook/qstash_publisher.py +++ b/packages/shared-python/shared/services/webhook/qstash_publisher.py @@ -26,7 +26,7 @@ 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 ( +from shared.services.http.url_security import ( validate_http_url_and_resolve_ip, ) diff --git a/packages/shared-python/shared/utils/__init__.py b/packages/shared-python/shared/utils/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/packages/shared-python/shared/utils/security_utils.py b/packages/shared-python/shared/utils/security_utils.py deleted file mode 100644 index 2702e79ef..000000000 --- a/packages/shared-python/shared/utils/security_utils.py +++ /dev/null @@ -1,22 +0,0 @@ -"""Security-related shared utilities.""" - -from typing import Optional - - -def mask_api_key(api_key: Optional[str]) -> str: - """Return a stable masked representation for logs and internal errors.""" - if not api_key: - return "" - - if len(api_key) <= 4: - return "*" * len(api_key) - - if len(api_key) <= 8: - prefix_length = 2 - suffix_length = 2 - else: - prefix_length = 4 - suffix_length = 4 - - masked_length = max(1, len(api_key) - prefix_length - suffix_length) - return f"{api_key[:prefix_length]}{'*' * masked_length}{api_key[-suffix_length:]}" From b7e344f2eaa3433d1a17ddbe7cf5ee2c5f670fdf Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 07:42:33 +0000 Subject: [PATCH 52/78] refactor: put shared quota base in neutral module --- .../worker/app/services/document_parser/mineru_quota_manager.py | 2 +- packages/shared-python/shared/services/ai/ali_quota_manager.py | 2 +- .../shared-python/shared/services/ai/iloveapi_quota_manager.py | 2 +- packages/shared-python/shared/services/quota/__init__.py | 1 + .../services/{ai/quota_manager.py => quota/token_pool.py} | 0 5 files changed, 4 insertions(+), 3 deletions(-) create mode 100644 packages/shared-python/shared/services/quota/__init__.py rename packages/shared-python/shared/services/{ai/quota_manager.py => quota/token_pool.py} (100%) diff --git a/apps/worker/app/services/document_parser/mineru_quota_manager.py b/apps/worker/app/services/document_parser/mineru_quota_manager.py index 8c2015959..f6fe88f3c 100644 --- a/apps/worker/app/services/document_parser/mineru_quota_manager.py +++ b/apps/worker/app/services/document_parser/mineru_quota_manager.py @@ -14,7 +14,7 @@ SyncRedisService, SyncRedisServiceFactory, ) -from shared.services.ai.quota_manager import BaseQuotaManager, TokenConfig, TokenLease +from shared.services.quota.token_pool import BaseQuotaManager, TokenConfig, TokenLease # Backward-compatible aliases so existing imports keep working MinerUTokenConfig = TokenConfig diff --git a/packages/shared-python/shared/services/ai/ali_quota_manager.py b/packages/shared-python/shared/services/ai/ali_quota_manager.py index 764f3e7fb..37bcc2e78 100644 --- a/packages/shared-python/shared/services/ai/ali_quota_manager.py +++ b/packages/shared-python/shared/services/ai/ali_quota_manager.py @@ -14,7 +14,7 @@ SyncRedisService, SyncRedisServiceFactory, ) -from shared.services.ai.quota_manager import BaseQuotaManager, TokenConfig +from shared.services.quota.token_pool import BaseQuotaManager, TokenConfig class AliQuotaManager(BaseQuotaManager): diff --git a/packages/shared-python/shared/services/ai/iloveapi_quota_manager.py b/packages/shared-python/shared/services/ai/iloveapi_quota_manager.py index 7afb5aacf..8cea78bb5 100644 --- a/packages/shared-python/shared/services/ai/iloveapi_quota_manager.py +++ b/packages/shared-python/shared/services/ai/iloveapi_quota_manager.py @@ -26,7 +26,7 @@ SyncRedisService, SyncRedisServiceFactory, ) -from shared.services.ai.quota_manager import BaseQuotaManager, TokenConfig +from shared.services.quota.token_pool import BaseQuotaManager, TokenConfig class ILoveApiQuotaManager(BaseQuotaManager): diff --git a/packages/shared-python/shared/services/quota/__init__.py b/packages/shared-python/shared/services/quota/__init__.py new file mode 100644 index 000000000..c6c7b4ef9 --- /dev/null +++ b/packages/shared-python/shared/services/quota/__init__.py @@ -0,0 +1 @@ +"""Shared quota and token-pool services.""" diff --git a/packages/shared-python/shared/services/ai/quota_manager.py b/packages/shared-python/shared/services/quota/token_pool.py similarity index 100% rename from packages/shared-python/shared/services/ai/quota_manager.py rename to packages/shared-python/shared/services/quota/token_pool.py From 3252cb8addb84a956620ffeef31fcb9b0157cfee Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 07:48:11 +0000 Subject: [PATCH 53/78] Revert "refactor: put shared quota base in neutral module" This reverts commit b15be8b8b85700ae01d1400db8ff7ede0b7f1a89. --- .../worker/app/services/document_parser/mineru_quota_manager.py | 2 +- packages/shared-python/shared/services/ai/ali_quota_manager.py | 2 +- .../shared-python/shared/services/ai/iloveapi_quota_manager.py | 2 +- .../services/{quota/token_pool.py => ai/quota_manager.py} | 0 packages/shared-python/shared/services/quota/__init__.py | 1 - 5 files changed, 3 insertions(+), 4 deletions(-) rename packages/shared-python/shared/services/{quota/token_pool.py => ai/quota_manager.py} (100%) delete mode 100644 packages/shared-python/shared/services/quota/__init__.py diff --git a/apps/worker/app/services/document_parser/mineru_quota_manager.py b/apps/worker/app/services/document_parser/mineru_quota_manager.py index f6fe88f3c..8c2015959 100644 --- a/apps/worker/app/services/document_parser/mineru_quota_manager.py +++ b/apps/worker/app/services/document_parser/mineru_quota_manager.py @@ -14,7 +14,7 @@ SyncRedisService, SyncRedisServiceFactory, ) -from shared.services.quota.token_pool import BaseQuotaManager, TokenConfig, TokenLease +from shared.services.ai.quota_manager import BaseQuotaManager, TokenConfig, TokenLease # Backward-compatible aliases so existing imports keep working MinerUTokenConfig = TokenConfig diff --git a/packages/shared-python/shared/services/ai/ali_quota_manager.py b/packages/shared-python/shared/services/ai/ali_quota_manager.py index 37bcc2e78..764f3e7fb 100644 --- a/packages/shared-python/shared/services/ai/ali_quota_manager.py +++ b/packages/shared-python/shared/services/ai/ali_quota_manager.py @@ -14,7 +14,7 @@ SyncRedisService, SyncRedisServiceFactory, ) -from shared.services.quota.token_pool import BaseQuotaManager, TokenConfig +from shared.services.ai.quota_manager import BaseQuotaManager, TokenConfig class AliQuotaManager(BaseQuotaManager): diff --git a/packages/shared-python/shared/services/ai/iloveapi_quota_manager.py b/packages/shared-python/shared/services/ai/iloveapi_quota_manager.py index 8cea78bb5..7afb5aacf 100644 --- a/packages/shared-python/shared/services/ai/iloveapi_quota_manager.py +++ b/packages/shared-python/shared/services/ai/iloveapi_quota_manager.py @@ -26,7 +26,7 @@ SyncRedisService, SyncRedisServiceFactory, ) -from shared.services.quota.token_pool import BaseQuotaManager, TokenConfig +from shared.services.ai.quota_manager import BaseQuotaManager, TokenConfig class ILoveApiQuotaManager(BaseQuotaManager): diff --git a/packages/shared-python/shared/services/quota/token_pool.py b/packages/shared-python/shared/services/ai/quota_manager.py similarity index 100% rename from packages/shared-python/shared/services/quota/token_pool.py rename to packages/shared-python/shared/services/ai/quota_manager.py diff --git a/packages/shared-python/shared/services/quota/__init__.py b/packages/shared-python/shared/services/quota/__init__.py deleted file mode 100644 index c6c7b4ef9..000000000 --- a/packages/shared-python/shared/services/quota/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Shared quota and token-pool services.""" From 1496638c7e5c3ea62f91d21f8b44d1b3e9001492 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 07:48:11 +0000 Subject: [PATCH 54/78] Revert "refactor: organize shared utilities by domain" This reverts commit 0a3252ad811f870080b9e0637b6899d1a44b5264. --- .github/codeql/codeql-config.yml | 2 +- AGENTS.md | 4 +--- apps/api/app/core/dependencies.py | 2 +- .../repositories/guest_device_repository.py | 2 +- .../auth/api_key_authentication_service.py | 2 +- .../auth/api_key_management_service.py | 2 +- .../stripe_credits_settlement_service.py | 2 +- .../stripe_refund_reconciliation_service.py | 2 +- .../document_ingestion/creation_service.py | 2 +- .../services/document_ingestion/service.py | 4 ++-- .../guest/guest_registration_service.py | 2 +- apps/api/app/services/jobs/job_read_model.py | 2 +- .../app/services/jobs/result_projection.py | 2 +- .../s3_events/subscription_service.py | 4 ++-- apps/api/main.py | 2 +- apps/api/scripts/add_credits.py | 2 +- apps/api/scripts/init_user.py | 2 +- .../tests/contract/test_billing_contract.py | 2 +- .../contract/test_job_creation_contract.py | 4 ++-- .../tests/contract/test_s3_event_contract.py | 2 +- apps/api/tests/support/contract_database.py | 2 +- apps/worker/app/core/worker_bootstrap.py | 2 +- apps/worker/app/services/common/__init__.py | 10 ++++----- .../connect_builder/summary_builder.py | 2 +- .../tools/classify_special_pages.py | 2 +- .../document_agent/tools/probe_vlm_inspect.py | 2 +- .../tools/propose_shard_plan.py | 2 +- .../document_ingestion/page_estimator.py | 2 +- .../document_parser/atlas_classifier.py | 2 +- .../services/document_parser/atlas_parser.py | 4 ++-- .../services/document_parser/doc_parser.py | 8 +++---- .../document_parser/excel_table_parser.py | 8 +++---- .../document_parser/fragment_parser.py | 4 ++-- .../services/document_parser/html_parser.py | 2 +- .../services/document_parser/image_parser.py | 8 +++---- .../document_parser/internal_parse_name.py | 2 +- .../services/document_parser/layout_parser.py | 2 +- .../markdown_deferred_summary.py | 4 ++-- .../document_parser/markdown_image_asset.py | 4 ++-- .../document_parser/markdown_table_asset.py | 4 ++-- .../app/services/document_parser/md_parser.py | 6 ++--- .../document_parser/mineru_pdf_service.py | 2 +- .../document_parser/mineru_quota_manager.py | 2 +- .../document_parser/mineru_task_polling.py | 2 +- .../orchestration/parse_session.py | 2 +- .../services/document_parser/path_helpers.py | 4 ++-- .../services/document_parser/pptx_parser.py | 8 +++---- .../document_parser/table_frame_parser.py | 4 ++-- .../document_parser/table_text_parser.py | 2 +- .../services/document_parser/text_helpers.py | 2 +- .../services/document_parser/toc_parser.py | 2 +- .../services/document_parser/txt_parser.py | 6 ++--- .../services/workload/url_upload_transfer.py | 2 +- .../shared-python/shared/core/security.py | 19 ---------------- .../shared/core/state_machine/service.py | 2 +- .../shared/core/state_machine/service_sync.py | 2 +- .../core/state_machine/transition_payloads.py | 4 ++-- .../shared/models/database/api_key.py | 2 +- .../models/database/credits_transaction.py | 2 +- .../models/database/demo_materialization.py | 2 +- .../shared/models/database/document.py | 2 +- .../shared/models/database/guest_device.py | 2 +- .../shared/models/database/job.py | 2 +- .../shared/models/database/job_result.py | 2 +- .../models/database/job_state_audit_log.py | 2 +- .../models/database/job_state_history.py | 2 +- .../shared/models/database/payment_record.py | 2 +- .../models/database/stripe_price_config.py | 2 +- .../shared/models/database/user_balance.py | 2 +- .../shared/models/database/webhook.py | 2 +- .../shared/models/database/webhook_log.py | 2 +- .../shared/models/database/webhook_secret.py | 2 +- .../shared/repositories/credits_repository.py | 2 +- .../repositories/credits_sync_repository.py | 2 +- .../services/billing/credits_service.py | 2 +- .../services/billing/credits_sync_service.py | 2 +- .../services/chunks/chunk_connections.py | 2 +- .../shared/services/http/__init__.py | 1 - .../shared/services/job_failure_sync.py | 2 +- .../shared/services/job_lifecycle_sync.py | 2 +- .../shared/services/jobs/result_delivery.py | 2 +- .../shared/services/redis/distributed_lock.py | 2 +- .../services/redis/job_info_redis_service.py | 2 +- .../services/redis/job_metadata_service.py | 2 +- .../services/redis/rate_limit_service.py | 2 +- .../shared/services/redis/redis_monitor.py | 2 +- .../shared/services/redis/redis_service.py | 4 ++-- .../services/redis/redis_sync_service.py | 6 ++--- .../services/redis/task_redis_service.py | 2 +- .../services/redis/user_redis_service.py | 2 +- .../services/retrieval/agentic/evidence.py | 2 +- .../services/retrieval/agentic/policy.py | 2 +- .../services/retrieval/agentic/runtime.py | 2 +- .../agentic/section_prompt_projection.py | 2 +- .../services/retrieval/lexical_ranker.py | 2 +- .../shared/services/retrieval/lexical_text.py | 2 +- .../shared/services/retrieval/llm_adapter.py | 6 ++--- .../services/retrieval/workflow/planner.py | 2 +- .../retrieval/workflow/synthesizer.py | 2 +- .../services/storage/job_file_storage.py | 4 ++-- .../services/storage/zip_result_schema.py | 4 ++-- .../services/text_processing/__init__.py | 1 - .../services/webhook/delivery_client.py | 4 ++-- .../services/webhook/qstash_publisher.py | 2 +- .../OpenAICompatibleClientSync.py} | 8 +++---- .../shared-python/shared/utils/__init__.py | 0 .../ai => utils}/ali_quota_manager.py | 2 +- .../api_key_tokens.py => utils/api_keys.py} | 0 .../{services/chunks => utils}/chunk_refs.py | 0 .../shared/utils/device_utils.py | 0 .../shared/{core => utils}/error_details.py | 0 .../shared/utils}/file_loading.py | 0 .../shared-python/shared/utils}/file_utils.py | 0 .../shared-python/shared/utils/gc_utils.py | 0 .../client_pool.py => utils/http_clients.py} | 0 .../ai => utils}/iloveapi_quota_manager.py | 2 +- .../serialization.py => utils/json_utils.py} | 0 .../shared/{services/ai => utils}/llm_mock.py | 0 .../shared-python/shared/utils/math_utils.py | 0 .../pinned_outbound_http.py} | 2 +- .../{services/ai => utils}/quota_manager.py | 0 .../redis_key_builder.py} | 0 .../retry_policy.py => utils/redis_retry.py} | 0 .../shared/utils/security_utils.py | 22 +++++++++++++++++++ .../tokenization.py => utils/text_utils.py} | 2 +- .../retrieval => utils}/token_estimate.py | 0 .../{services/http => utils}/url_file_type.py | 6 ++--- .../{services/http => utils}/url_security.py | 0 .../shared/{core/time.py => utils/utc_now.py} | 0 .../{services/http => utils}/zip_download.py | 0 130 files changed, 173 insertions(+), 174 deletions(-) delete mode 100644 packages/shared-python/shared/services/http/__init__.py delete mode 100644 packages/shared-python/shared/services/text_processing/__init__.py rename packages/shared-python/shared/{services/ai/openai_compatible_client_sync.py => utils/OpenAICompatibleClientSync.py} (98%) create mode 100644 packages/shared-python/shared/utils/__init__.py rename packages/shared-python/shared/{services/ai => utils}/ali_quota_manager.py (96%) rename packages/shared-python/shared/{services/auth/api_key_tokens.py => utils/api_keys.py} (100%) rename packages/shared-python/shared/{services/chunks => utils}/chunk_refs.py (100%) rename apps/worker/app/services/common/device_checks.py => packages/shared-python/shared/utils/device_utils.py (100%) rename packages/shared-python/shared/{core => utils}/error_details.py (100%) rename {apps/worker/app/services/common => packages/shared-python/shared/utils}/file_loading.py (100%) rename {apps/worker/app/services/common => packages/shared-python/shared/utils}/file_utils.py (100%) rename apps/worker/app/services/common/resource_cleanup.py => packages/shared-python/shared/utils/gc_utils.py (100%) rename packages/shared-python/shared/{services/http/client_pool.py => utils/http_clients.py} (100%) rename packages/shared-python/shared/{services/ai => utils}/iloveapi_quota_manager.py (99%) rename packages/shared-python/shared/{core/serialization.py => utils/json_utils.py} (100%) rename packages/shared-python/shared/{services/ai => utils}/llm_mock.py (100%) rename apps/worker/app/services/common/math_helpers.py => packages/shared-python/shared/utils/math_utils.py (100%) rename packages/shared-python/shared/{services/http/pinned_outbound.py => utils/pinned_outbound_http.py} (99%) rename packages/shared-python/shared/{services/ai => utils}/quota_manager.py (100%) rename packages/shared-python/shared/{services/redis/key_builder.py => utils/redis_key_builder.py} (100%) rename packages/shared-python/shared/{services/redis/retry_policy.py => utils/redis_retry.py} (100%) create mode 100644 packages/shared-python/shared/utils/security_utils.py rename packages/shared-python/shared/{services/text_processing/tokenization.py => utils/text_utils.py} (99%) rename packages/shared-python/shared/{services/retrieval => utils}/token_estimate.py (100%) rename packages/shared-python/shared/{services/http => utils}/url_file_type.py (97%) rename packages/shared-python/shared/{services/http => utils}/url_security.py (100%) rename packages/shared-python/shared/{core/time.py => utils/utc_now.py} (100%) rename packages/shared-python/shared/{services/http => utils}/zip_download.py (100%) diff --git a/.github/codeql/codeql-config.yml b/.github/codeql/codeql-config.yml index 9ad24421d..705bc0dec 100644 --- a/.github/codeql/codeql-config.yml +++ b/.github/codeql/codeql-config.yml @@ -2,4 +2,4 @@ name: "Knowhere CodeQL config" paths-ignore: - apps/api/scripts/** - - packages/shared-python/shared/services/auth/api_key_tokens.py + - packages/shared-python/shared/utils/api_keys.py diff --git a/AGENTS.md b/AGENTS.md index 421c133bd..eed149425 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -47,12 +47,10 @@ knowhereapi-main/ │ └── shared-python/shared/ # Shared library (pip: knowhere-shared) │ ├── models/database/ # SQLAlchemy ORM models │ ├── models/schemas/ # Pydantic request/response schemas -│ ├── core/ # Shared primitives: time, serialization, security │ ├── services/retrieval/ # Core retrieval engine │ ├── services/chunks/ # DataFrame → ChunkPayload conversion │ ├── services/ai/ # LLM prompt service & AI client -│ ├── services/http/ # Outbound HTTP clients and URL safety -│ └── services/text_processing/ # Shared text tokenization helpers +│ └── utils/ # Text, file, and chunk utilities └── deploy/ # Docker Compose & deployment scripts ``` diff --git a/apps/api/app/core/dependencies.py b/apps/api/app/core/dependencies.py index d14340151..153177cd4 100644 --- a/apps/api/app/core/dependencies.py +++ b/apps/api/app/core/dependencies.py @@ -18,7 +18,7 @@ AuthException, ) from shared.models.database.user import User -from shared.services.auth.api_key_tokens import is_api_key_token +from shared.utils.api_keys import is_api_key_token # Standard JWKS endpoint path (fixed, following OpenID Connect convention) JWKS_ENDPOINT_PATH = "/api/auth/jwks" diff --git a/apps/api/app/repositories/guest_device_repository.py b/apps/api/app/repositories/guest_device_repository.py index ef2ca4f50..1f0fbc247 100644 --- a/apps/api/app/repositories/guest_device_repository.py +++ b/apps/api/app/repositories/guest_device_repository.py @@ -6,7 +6,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from shared.models.database.guest_device import GuestDevice -from shared.core.time import utc_now_naive +from shared.utils.utc_now import utc_now_naive class GuestDeviceRepository(BaseRepository[GuestDevice, dict, dict]): diff --git a/apps/api/app/services/auth/api_key_authentication_service.py b/apps/api/app/services/auth/api_key_authentication_service.py index a7083226e..67e677247 100644 --- a/apps/api/app/services/auth/api_key_authentication_service.py +++ b/apps/api/app/services/auth/api_key_authentication_service.py @@ -13,7 +13,7 @@ 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.services.auth.api_key_tokens import hash_api_key +from shared.utils.api_keys import hash_api_key _API_KEY_USER_CACHE_TTL_SECONDS: int = 3600 diff --git a/apps/api/app/services/auth/api_key_management_service.py b/apps/api/app/services/auth/api_key_management_service.py index 85189d77e..0adbfee94 100644 --- a/apps/api/app/services/auth/api_key_management_service.py +++ b/apps/api/app/services/auth/api_key_management_service.py @@ -19,7 +19,7 @@ ValidationException, ) from shared.models.database.api_key import APIKey -from shared.services.auth.api_key_tokens import generate_api_key, hash_api_key, mask_api_key +from shared.utils.api_keys import generate_api_key, hash_api_key, mask_api_key class APIKeyListItem(TypedDict): diff --git a/apps/api/app/services/billing/stripe_credits_settlement_service.py b/apps/api/app/services/billing/stripe_credits_settlement_service.py index 36bcd12b4..9091dd6e6 100644 --- a/apps/api/app/services/billing/stripe_credits_settlement_service.py +++ b/apps/api/app/services/billing/stripe_credits_settlement_service.py @@ -14,7 +14,7 @@ ) from shared.models.database.payment_record import PaymentRecord from shared.services.billing import CreditsService -from shared.core.time import utc_now_naive +from shared.utils.utc_now import utc_now_naive class StripeCreditsSettlementService: diff --git a/apps/api/app/services/billing/stripe_refund_reconciliation_service.py b/apps/api/app/services/billing/stripe_refund_reconciliation_service.py index 1324402b8..d770195e1 100644 --- a/apps/api/app/services/billing/stripe_refund_reconciliation_service.py +++ b/apps/api/app/services/billing/stripe_refund_reconciliation_service.py @@ -11,7 +11,7 @@ from shared.core.logging import logger from shared.models.database.payment_record import PaymentRecord from shared.services.billing import CreditsService -from shared.core.time import utc_now_naive +from shared.utils.utc_now import utc_now_naive class StripeRefundReconciliationService: diff --git a/apps/api/app/services/document_ingestion/creation_service.py b/apps/api/app/services/document_ingestion/creation_service.py index acecb1974..a76977b65 100644 --- a/apps/api/app/services/document_ingestion/creation_service.py +++ b/apps/api/app/services/document_ingestion/creation_service.py @@ -29,7 +29,7 @@ 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.services.http.url_file_type import resolve_file_extension_async +from shared.utils.url_file_type import resolve_file_extension_async _JOB_TYPE_KB_MANAGEMENT = "kb_management" JobMetadata = dict[str, object] diff --git a/apps/api/app/services/document_ingestion/service.py b/apps/api/app/services/document_ingestion/service.py index 03f5da8fe..defd2fbca 100644 --- a/apps/api/app/services/document_ingestion/service.py +++ b/apps/api/app/services/document_ingestion/service.py @@ -35,8 +35,8 @@ from shared.core.exceptions.webhook_exceptions import WebhookConfigException from shared.models.schemas.job import ConfirmUploadRequest, JobCreate, JobResponse from shared.models.schemas.job_metadata import JobMetadataHelper -from shared.services.http.url_file_type import resolve_file_extension_async -from shared.services.http.url_security import validate_http_url_and_resolve_ip_async +from shared.utils.url_file_type import resolve_file_extension_async +from shared.utils.url_security import validate_http_url_and_resolve_ip_async JobMetadata = dict[str, object] diff --git a/apps/api/app/services/guest/guest_registration_service.py b/apps/api/app/services/guest/guest_registration_service.py index ae2cfbb68..4f189807e 100644 --- a/apps/api/app/services/guest/guest_registration_service.py +++ b/apps/api/app/services/guest/guest_registration_service.py @@ -24,7 +24,7 @@ GuestRegisterResponse, ) from shared.services.billing.credits_service import CreditsService -from shared.services.auth.api_key_tokens import generate_api_key, hash_api_key, mask_api_key +from shared.utils.api_keys import generate_api_key, hash_api_key, mask_api_key _GUEST_TIER: str = "guest" _GUEST_KEY_NAME_PREFIX: str = "guest-device" diff --git a/apps/api/app/services/jobs/job_read_model.py b/apps/api/app/services/jobs/job_read_model.py index 94d25ba7f..1d0a4ab39 100644 --- a/apps/api/app/services/jobs/job_read_model.py +++ b/apps/api/app/services/jobs/job_read_model.py @@ -20,7 +20,7 @@ ) from shared.models.schemas.job import JobList, JobResultResponse from shared.services.redis import RedisServiceFactory -from shared.core.time import utc_now_naive +from shared.utils.utc_now import utc_now_naive class JobReadModel: diff --git a/apps/api/app/services/jobs/result_projection.py b/apps/api/app/services/jobs/result_projection.py index 710b92962..0a00ce2d2 100644 --- a/apps/api/app/services/jobs/result_projection.py +++ b/apps/api/app/services/jobs/result_projection.py @@ -10,7 +10,7 @@ from shared.models.schemas.job import JobResultResponse, StandardErrorObject from shared.models.schemas.job_metadata import JobMetadataHelper from shared.services.jobs.result_delivery import JobResultDeliveryResolver -from shared.core.error_details import normalize_error_details +from shared.utils.error_details import normalize_error_details JobStatusValue = Literal[ "pending", "waiting-file", "running", "converting", "done", "failed" diff --git a/apps/api/app/services/s3_events/subscription_service.py b/apps/api/app/services/s3_events/subscription_service.py index c64ab78e5..140b2eba3 100644 --- a/apps/api/app/services/s3_events/subscription_service.py +++ b/apps/api/app/services/s3_events/subscription_service.py @@ -3,8 +3,8 @@ from loguru import logger -from shared.services.http.pinned_outbound import send_pinned_outbound_request -from shared.services.http.url_security import validate_http_url_and_resolve_ip_async +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 diff --git a/apps/api/main.py b/apps/api/main.py index 4b0dfde44..bd7dec810 100644 --- a/apps/api/main.py +++ b/apps/api/main.py @@ -92,7 +92,7 @@ async def lifespan(app: FastAPI): logger.error(f"retrieval hit stats drain failed: {e}") try: - from shared.services.http.client_pool import close_async_client + from shared.utils.http_clients import close_async_client await close_async_client() except Exception as e: diff --git a/apps/api/scripts/add_credits.py b/apps/api/scripts/add_credits.py index dd0370d75..675f5ce26 100644 --- a/apps/api/scripts/add_credits.py +++ b/apps/api/scripts/add_credits.py @@ -34,7 +34,7 @@ def _bootstrap_python_path() -> None: from shared.models.database.user import User from shared.models.database.user_balance import UserBalance from shared.services.billing import CreditsService -from shared.core.time import utc_now_naive +from shared.utils.utc_now import utc_now_naive def _build_parser() -> argparse.ArgumentParser: diff --git a/apps/api/scripts/init_user.py b/apps/api/scripts/init_user.py index 4847d5024..de38fe61e 100644 --- a/apps/api/scripts/init_user.py +++ b/apps/api/scripts/init_user.py @@ -18,7 +18,7 @@ from shared.models.database.user_balance import UserBalance from shared.services.billing.credits_service import CreditsService from shared.services.auth.user_table_bootstrap import ensure_better_auth_user_table -from shared.services.auth.api_key_tokens import generate_api_key, hash_api_key, mask_api_key +from shared.utils.api_keys import generate_api_key, hash_api_key, mask_api_key _DEFAULT_API_KEY_NAME: str = "standalone-api-key" _DEFAULT_USER_TIER: str = "tier_5" diff --git a/apps/api/tests/contract/test_billing_contract.py b/apps/api/tests/contract/test_billing_contract.py index 11a108ff5..6f5eb2be6 100644 --- a/apps/api/tests/contract/test_billing_contract.py +++ b/apps/api/tests/contract/test_billing_contract.py @@ -12,7 +12,7 @@ from pytest import MonkeyPatch from tests.support.contract_database import ContractDatabase -from shared.services.auth.api_key_tokens import hash_api_key +from shared.utils.api_keys import hash_api_key def _utc_now() -> datetime: diff --git a/apps/api/tests/contract/test_job_creation_contract.py b/apps/api/tests/contract/test_job_creation_contract.py index d19932bc0..edea174ec 100644 --- a/apps/api/tests/contract/test_job_creation_contract.py +++ b/apps/api/tests/contract/test_job_creation_contract.py @@ -632,7 +632,7 @@ def resolve_public_address( return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 0))] import shared.core.celery_app as celery_app_module - import shared.services.http.client_pool as http_clients_module + import shared.utils.http_clients as http_clients_module monkeypatch.setattr(socket, "getaddrinfo", resolve_public_address) monkeypatch.setattr( @@ -926,7 +926,7 @@ async def head( assert follow_redirects is False return _FakeHeadResponse() - import shared.services.http.client_pool as http_clients_module + import shared.utils.http_clients as http_clients_module monkeypatch.setattr( http_clients_module, "get_async_client", diff --git a/apps/api/tests/contract/test_s3_event_contract.py b/apps/api/tests/contract/test_s3_event_contract.py index 65c1c6378..b8e8561de 100644 --- a/apps/api/tests/contract/test_s3_event_contract.py +++ b/apps/api/tests/contract/test_s3_event_contract.py @@ -210,7 +210,7 @@ def get(self, url: str, *args: object, **kwargs: object) -> object: async with api_client_factory() as api_client: monkeypatch.setattr(socket, "getaddrinfo", resolve_private_address) pinned_http_module = importlib.import_module( - "shared.services.http.pinned_outbound" + "shared.utils.pinned_outbound_http" ) monkeypatch.setattr( pinned_http_module.aiohttp, diff --git a/apps/api/tests/support/contract_database.py b/apps/api/tests/support/contract_database.py index 3358fff9e..664fff00a 100644 --- a/apps/api/tests/support/contract_database.py +++ b/apps/api/tests/support/contract_database.py @@ -9,7 +9,7 @@ from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine from shared.testing.contract_runtime import get_contract_database_url -from shared.services.auth.api_key_tokens import hash_api_key +from shared.utils.api_keys import hash_api_key async def _create_contract_engine() -> AsyncEngine: diff --git a/apps/worker/app/core/worker_bootstrap.py b/apps/worker/app/core/worker_bootstrap.py index 2689bb71c..191f2fa27 100644 --- a/apps/worker/app/core/worker_bootstrap.py +++ b/apps/worker/app/core/worker_bootstrap.py @@ -74,7 +74,7 @@ def shutdown_worker(**kwargs) -> None: logger.warning(f"Worker heartbeat cleanup failed: {exc}") try: - from shared.services.http.client_pool import close_sync_client + from shared.utils.http_clients import close_sync_client close_sync_client() logger.info("Worker sync HTTP client closed") diff --git a/apps/worker/app/services/common/__init__.py b/apps/worker/app/services/common/__init__.py index fab06465a..ffb53d823 100644 --- a/apps/worker/app/services/common/__init__.py +++ b/apps/worker/app/services/common/__init__.py @@ -1,10 +1,10 @@ """Common worker services.""" -from app.services.common.device_checks import check_internet -from app.services.common.file_utils import clean_file, path_handle -from app.services.common.resource_cleanup import gc_collect as _gc -from app.services.common.math_helpers import min_max_normalize -from shared.services.text_processing.tokenization import ( +from shared.utils.device_utils import check_internet +from shared.utils.file_utils import clean_file, path_handle +from shared.utils.gc_utils import gc_collect as _gc +from shared.utils.math_utils import min_max_normalize +from shared.utils.text_utils import ( count_cn_en, remove_duplicates_orderkept, tokenize2stw_remove, diff --git a/apps/worker/app/services/connect_builder/summary_builder.py b/apps/worker/app/services/connect_builder/summary_builder.py index 44f472076..7fcf8c6eb 100644 --- a/apps/worker/app/services/connect_builder/summary_builder.py +++ b/apps/worker/app/services/connect_builder/summary_builder.py @@ -40,7 +40,7 @@ def _llm_summarize(snippets_text: str, node_name: str, max_tokens: int = 100) -> """ try: from shared.services.ai.prompt_service import build_prompt, _detect_text_language - from shared.services.ai.openai_compatible_client_sync import get_openai_client + from shared.utils.OpenAICompatibleClientSync import get_openai_client # Deterministic language lock — see prompt_service._language_directive detected_lang = _detect_text_language(snippets_text) diff --git a/apps/worker/app/services/document_agent/tools/classify_special_pages.py b/apps/worker/app/services/document_agent/tools/classify_special_pages.py index 0b8efa6c8..564307b21 100644 --- a/apps/worker/app/services/document_agent/tools/classify_special_pages.py +++ b/apps/worker/app/services/document_agent/tools/classify_special_pages.py @@ -141,7 +141,7 @@ def classify_special_pages( try: from shared.core.config import settings - from shared.services.ai.openai_compatible_client_sync import get_openai_client + from shared.utils.OpenAICompatibleClientSync import get_openai_client effective_model = model or settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL client = get_openai_client(model=effective_model) diff --git a/apps/worker/app/services/document_agent/tools/probe_vlm_inspect.py b/apps/worker/app/services/document_agent/tools/probe_vlm_inspect.py index f8949826d..7b4de6396 100644 --- a/apps/worker/app/services/document_agent/tools/probe_vlm_inspect.py +++ b/apps/worker/app/services/document_agent/tools/probe_vlm_inspect.py @@ -71,7 +71,7 @@ def _call_vlm( max_tokens: int = 900, ) -> tuple[str, dict[str, int]]: from shared.core.config import settings - from shared.services.ai.openai_compatible_client_sync import get_openai_client + from shared.utils.OpenAICompatibleClientSync import get_openai_client effective_model = model or settings.IMAGE_MODEL or "qwen3.5-flash" client = get_openai_client(model=effective_model) diff --git a/apps/worker/app/services/document_agent/tools/propose_shard_plan.py b/apps/worker/app/services/document_agent/tools/propose_shard_plan.py index 1db72592b..2eb834ff4 100644 --- a/apps/worker/app/services/document_agent/tools/propose_shard_plan.py +++ b/apps/worker/app/services/document_agent/tools/propose_shard_plan.py @@ -295,7 +295,7 @@ def propose_shard_plan( if use_llm and page_count > max_pages_per_shard: try: from shared.core.config import settings - from shared.services.ai.openai_compatible_client_sync import get_openai_client + from shared.utils.OpenAICompatibleClientSync import get_openai_client effective_model = model or settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL client = get_openai_client(model=effective_model) diff --git a/apps/worker/app/services/document_ingestion/page_estimator.py b/apps/worker/app/services/document_ingestion/page_estimator.py index 4267cd233..bdc97a4d6 100644 --- a/apps/worker/app/services/document_ingestion/page_estimator.py +++ b/apps/worker/app/services/document_ingestion/page_estimator.py @@ -15,7 +15,7 @@ from pathlib import Path from shared.core.logging import logger -from shared.services.text_processing.tokenization import count_cn_en +from shared.utils.text_utils import count_cn_en WORDS_PER_PAGE = 500 ROWS_PER_PAGE = 50 diff --git a/apps/worker/app/services/document_parser/atlas_classifier.py b/apps/worker/app/services/document_parser/atlas_classifier.py index ed58ad49f..1ee8dc4ff 100644 --- a/apps/worker/app/services/document_parser/atlas_classifier.py +++ b/apps/worker/app/services/document_parser/atlas_classifier.py @@ -102,7 +102,7 @@ def _png_to_data_url(path: str) -> Optional[str]: def _call_vlm(image_data_urls: list[str]) -> bool: """Call VLM with preview images. Returns True if atlas, False otherwise.""" from shared.core.config import settings - from shared.services.ai.openai_compatible_client_sync import get_openai_client + from shared.utils.OpenAICompatibleClientSync import get_openai_client model = settings.IMAGE_MODEL or "qwen-vl-plus" client = get_openai_client(model=model) diff --git a/apps/worker/app/services/document_parser/atlas_parser.py b/apps/worker/app/services/document_parser/atlas_parser.py index 6e3666d0f..fef806add 100644 --- a/apps/worker/app/services/document_parser/atlas_parser.py +++ b/apps/worker/app/services/document_parser/atlas_parser.py @@ -27,8 +27,8 @@ from loguru import logger from shared.core.config import settings -from shared.services.chunks.chunk_refs import build_chunk_ref -from shared.services.text_processing.tokenization import tokenize2stw_remove +from shared.utils.chunk_refs import build_chunk_ref +from shared.utils.text_utils import tokenize2stw_remove # ─── Config ─────────────────────────────────────────────────────────── VLM_CONCURRENCY = 3 # concurrent VLM calls diff --git a/apps/worker/app/services/document_parser/doc_parser.py b/apps/worker/app/services/document_parser/doc_parser.py index 5eb56ab79..d405fed56 100755 --- a/apps/worker/app/services/document_parser/doc_parser.py +++ b/apps/worker/app/services/document_parser/doc_parser.py @@ -37,10 +37,10 @@ from shared.core.config import settings from shared.core.exceptions.domain_exceptions import DocxParsingException from shared.core.exceptions.knowhere_exception import KnowhereException -from shared.services.chunks.chunk_refs import build_chunk_ref, has_chunk_ref -from app.services.common.file_loading import load_file_bytes -from app.services.common.file_utils import path_handle -from shared.services.text_processing.tokenization import tokenize2stw_remove +from shared.utils.chunk_refs import build_chunk_ref, has_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 def get_leaf_dics(node, 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 index 2a6aa13ea..1770a4163 100644 --- a/apps/worker/app/services/document_parser/excel_table_parser.py +++ b/apps/worker/app/services/document_parser/excel_table_parser.py @@ -27,10 +27,10 @@ from shared.core.exceptions.domain_exceptions import TableParsingException from shared.core.exceptions.knowhere_exception import KnowhereException -from shared.services.chunks.chunk_refs import build_chunk_ref -from app.services.common.file_loading import load_file_bytes -from app.services.common.file_utils import path_handle -from shared.services.text_processing.tokenization import tokenize2stw_remove +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) diff --git a/apps/worker/app/services/document_parser/fragment_parser.py b/apps/worker/app/services/document_parser/fragment_parser.py index 420d6b084..5aef9f004 100644 --- a/apps/worker/app/services/document_parser/fragment_parser.py +++ b/apps/worker/app/services/document_parser/fragment_parser.py @@ -13,8 +13,8 @@ from openai.types.chat import ChatCompletionMessageParam from shared.core.config import settings -from app.services.common.file_utils import path_handle -from shared.services.ai.openai_compatible_client_sync import get_openai_client +from shared.utils.file_utils import path_handle +from shared.utils.OpenAICompatibleClientSync import get_openai_client def generate_fragment_title(content: str, max_tokens: int = 30) -> Optional[str]: diff --git a/apps/worker/app/services/document_parser/html_parser.py b/apps/worker/app/services/document_parser/html_parser.py index e326a709a..f412067d6 100644 --- a/apps/worker/app/services/document_parser/html_parser.py +++ b/apps/worker/app/services/document_parser/html_parser.py @@ -15,7 +15,7 @@ from bs4 import BeautifulSoup from shared.core.exceptions.domain_exceptions import TableParsingException -from shared.services.text_processing.tokenization import remove_duplicates_orderkept +from shared.utils.text_utils import remove_duplicates_orderkept class HTMLHeaderExpander: diff --git a/apps/worker/app/services/document_parser/image_parser.py b/apps/worker/app/services/document_parser/image_parser.py index 62b319450..cc9be0ab1 100755 --- a/apps/worker/app/services/document_parser/image_parser.py +++ b/apps/worker/app/services/document_parser/image_parser.py @@ -22,10 +22,10 @@ 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.services.chunks.chunk_refs import build_chunk_ref -from app.services.common.file_loading import is_remote, load_file_bytes -from app.services.common.file_utils import path_handle -from shared.services.ai.openai_compatible_client_sync import ( +from shared.utils.chunk_refs import build_chunk_ref +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, get_openai_client, ) diff --git a/apps/worker/app/services/document_parser/internal_parse_name.py b/apps/worker/app/services/document_parser/internal_parse_name.py index 568dc13c6..785fe6ff5 100644 --- a/apps/worker/app/services/document_parser/internal_parse_name.py +++ b/apps/worker/app/services/document_parser/internal_parse_name.py @@ -4,7 +4,7 @@ import os from dataclasses import dataclass -from app.services.common.file_utils import path_handle +from shared.utils.file_utils import path_handle @dataclass(frozen=True) diff --git a/apps/worker/app/services/document_parser/layout_parser.py b/apps/worker/app/services/document_parser/layout_parser.py index ec7775904..b0490d576 100755 --- a/apps/worker/app/services/document_parser/layout_parser.py +++ b/apps/worker/app/services/document_parser/layout_parser.py @@ -35,7 +35,7 @@ from shared.services.ai.response_process_service import eval_response # ARQ dependency is removed, use Celery instead -from shared.services.ai.openai_compatible_client_sync import get_openai_client +from shared.utils.OpenAICompatibleClientSync import get_openai_client # ==================== Helper Functions ==================== 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 ef5362ec8..270303e90 100644 --- a/apps/worker/app/services/document_parser/markdown_deferred_summary.py +++ b/apps/worker/app/services/document_parser/markdown_deferred_summary.py @@ -23,8 +23,8 @@ from loguru import logger from shared.core.config import settings -from shared.services.chunks.chunk_refs import build_chunk_ref -from app.services.common.file_utils import path_handle +from shared.utils.chunk_refs import build_chunk_ref +from shared.utils.file_utils import path_handle DeferredResult = ( tuple[ diff --git a/apps/worker/app/services/document_parser/markdown_image_asset.py b/apps/worker/app/services/document_parser/markdown_image_asset.py index 406872c18..10c93b5dc 100644 --- a/apps/worker/app/services/document_parser/markdown_image_asset.py +++ b/apps/worker/app/services/document_parser/markdown_image_asset.py @@ -15,8 +15,8 @@ from app.services.document_parser.markdown_parse_state import ParserRowValues from loguru import logger -from shared.services.chunks.chunk_refs import build_chunk_ref -from app.services.common.file_utils import path_handle +from shared.utils.chunk_refs import build_chunk_ref +from shared.utils.file_utils import path_handle @dataclass(frozen=True) diff --git a/apps/worker/app/services/document_parser/markdown_table_asset.py b/apps/worker/app/services/document_parser/markdown_table_asset.py index c1ad0a3e7..a01e8ce97 100644 --- a/apps/worker/app/services/document_parser/markdown_table_asset.py +++ b/apps/worker/app/services/document_parser/markdown_table_asset.py @@ -14,8 +14,8 @@ 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.services.chunks.chunk_refs import build_chunk_ref -from app.services.common.file_utils import path_handle +from shared.utils.chunk_refs import build_chunk_ref +from shared.utils.file_utils import path_handle @dataclass(frozen=True) diff --git a/apps/worker/app/services/document_parser/md_parser.py b/apps/worker/app/services/document_parser/md_parser.py index aded8f424..e284ea2c8 100755 --- a/apps/worker/app/services/document_parser/md_parser.py +++ b/apps/worker/app/services/document_parser/md_parser.py @@ -43,8 +43,8 @@ from loguru import logger from shared.core.config import settings -from shared.services.chunks.chunk_refs import has_chunk_ref -from shared.services.text_processing.tokenization import tokenize2stw_remove +from shared.utils.chunk_refs import has_chunk_ref +from shared.utils.text_utils import tokenize2stw_remove def find_surround_context(md_lines, lid): @@ -215,7 +215,7 @@ def parse_md( relative_root=None, ): if md_lines is None and file_path is not None: - from app.services.common.file_loading 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 8be565518..bd4acdbc4 100644 --- a/apps/worker/app/services/document_parser/mineru_pdf_service.py +++ b/apps/worker/app/services/document_parser/mineru_pdf_service.py @@ -23,7 +23,7 @@ UnavailableException, ) from shared.services.storage.job_file_storage import JobFileStorage -from app.services.common.file_loading import is_remote +from shared.utils.file_loading import is_remote MINERU_UPLOAD_TIMEOUT = ( settings.MINERU_UPLOAD_CONNECT_TIMEOUT, diff --git a/apps/worker/app/services/document_parser/mineru_quota_manager.py b/apps/worker/app/services/document_parser/mineru_quota_manager.py index 8c2015959..830330b09 100644 --- a/apps/worker/app/services/document_parser/mineru_quota_manager.py +++ b/apps/worker/app/services/document_parser/mineru_quota_manager.py @@ -14,7 +14,7 @@ SyncRedisService, SyncRedisServiceFactory, ) -from shared.services.ai.quota_manager import BaseQuotaManager, TokenConfig, TokenLease +from shared.utils.quota_manager import BaseQuotaManager, TokenConfig, TokenLease # Backward-compatible aliases so existing imports keep working MinerUTokenConfig = TokenConfig diff --git a/apps/worker/app/services/document_parser/mineru_task_polling.py b/apps/worker/app/services/document_parser/mineru_task_polling.py index 746bf4cd7..c09c1ea86 100644 --- a/apps/worker/app/services/document_parser/mineru_task_polling.py +++ b/apps/worker/app/services/document_parser/mineru_task_polling.py @@ -22,7 +22,7 @@ UnavailableException, ) from shared.core.exceptions.knowhere_exception import KnowhereException -from shared.services.http.zip_download import download_and_extract_zip +from shared.utils.zip_download import download_and_extract_zip def get_batch_status(data: dict[str, Any]) -> Optional[dict[str, Any]]: 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 9750a331d..832e13969 100644 --- a/apps/worker/app/services/document_parser/orchestration/parse_session.py +++ b/apps/worker/app/services/document_parser/orchestration/parse_session.py @@ -12,7 +12,7 @@ from shared.core.config import settings from shared.core.exceptions.domain_exceptions import ValidationException -from app.services.common.file_utils import path_handle +from shared.utils.file_utils import path_handle PDF_PAGE_LIMIT = 600 diff --git a/apps/worker/app/services/document_parser/path_helpers.py b/apps/worker/app/services/document_parser/path_helpers.py index f337b6556..41965e125 100644 --- a/apps/worker/app/services/document_parser/path_helpers.py +++ b/apps/worker/app/services/document_parser/path_helpers.py @@ -5,8 +5,8 @@ from typing import Any from bs4 import BeautifulSoup -from shared.services.chunks.chunk_refs import extract_chunk_refs -from app.services.common.file_utils import path_handle +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") diff --git a/apps/worker/app/services/document_parser/pptx_parser.py b/apps/worker/app/services/document_parser/pptx_parser.py index 8329fb18a..19063937a 100755 --- a/apps/worker/app/services/document_parser/pptx_parser.py +++ b/apps/worker/app/services/document_parser/pptx_parser.py @@ -26,8 +26,8 @@ FileSystemException, ) from shared.core.logging import LogEvent -from app.services.common.file_loading import load_file_bytes -from app.services.common.file_utils import path_handle +from shared.utils.file_loading import load_file_bytes +from shared.utils.file_utils import path_handle # ==================== LibreOffice conversion ==================== @@ -56,7 +56,7 @@ def pptx_to_pdf_libreoffice(pptx_path, outdir="."): def _get_iloveapi_token_lease(): """acquire iLoveAPI token lease from the quotas pool and generate a JWT token""" - from shared.services.ai.iloveapi_quota_manager import get_iloveapi_quota_manager + from shared.utils.iloveapi_quota_manager import get_iloveapi_quota_manager quota_manager = get_iloveapi_quota_manager() @@ -111,7 +111,7 @@ def _pptx_bytes_to_pdf_bytes(pptx_bytes: bytes, filename: str) -> bytes: Acquires an in-flight slot before starting. Slots are released only if the Redis-backed reservation actually succeeded. """ - from shared.services.ai.iloveapi_quota_manager import get_iloveapi_quota_manager + from shared.utils.iloveapi_quota_manager import get_iloveapi_quota_manager quota_manager = get_iloveapi_quota_manager() diff --git a/apps/worker/app/services/document_parser/table_frame_parser.py b/apps/worker/app/services/document_parser/table_frame_parser.py index 24d76538c..63f39feb8 100644 --- a/apps/worker/app/services/document_parser/table_frame_parser.py +++ b/apps/worker/app/services/document_parser/table_frame_parser.py @@ -15,8 +15,8 @@ from shared.services.ai.prompt_service import build_prompt from shared.services.ai.response_process_service import eval_response -from shared.services.ai.openai_compatible_client_sync import get_openai_client -from shared.services.text_processing.tokenization import remove_duplicates_orderkept +from shared.utils.OpenAICompatibleClientSync import get_openai_client +from shared.utils.text_utils import remove_duplicates_orderkept def parse_headers( diff --git a/apps/worker/app/services/document_parser/table_text_parser.py b/apps/worker/app/services/document_parser/table_text_parser.py index f83d3be0e..261bbe653 100644 --- a/apps/worker/app/services/document_parser/table_text_parser.py +++ b/apps/worker/app/services/document_parser/table_text_parser.py @@ -13,7 +13,7 @@ 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.services.text_processing.tokenization import _is_meaningful_token + from shared.utils.text_utils import _is_meaningful_token if not raw_header_text: return "" diff --git a/apps/worker/app/services/document_parser/text_helpers.py b/apps/worker/app/services/document_parser/text_helpers.py index bbacf4327..b11fdc772 100644 --- a/apps/worker/app/services/document_parser/text_helpers.py +++ b/apps/worker/app/services/document_parser/text_helpers.py @@ -2,7 +2,7 @@ import re -from shared.services.text_processing.tokenization import _CN_EN_NUM_RE +from shared.utils.text_utils import _CN_EN_NUM_RE _CN_CHAR_RE = re.compile(r"[\u4e00-\u9fff]") EN_START_LIMIT = 15 diff --git a/apps/worker/app/services/document_parser/toc_parser.py b/apps/worker/app/services/document_parser/toc_parser.py index 5c52ef5c4..2efb1a4cd 100644 --- a/apps/worker/app/services/document_parser/toc_parser.py +++ b/apps/worker/app/services/document_parser/toc_parser.py @@ -21,7 +21,7 @@ from shared.services.ai.prompt_service import build_prompt from shared.services.ai.response_process_service import eval_response -from shared.services.ai.openai_compatible_client_sync import get_openai_client +from shared.utils.OpenAICompatibleClientSync import get_openai_client # ==================== Markdown TOC Detection Functions ==================== diff --git a/apps/worker/app/services/document_parser/txt_parser.py b/apps/worker/app/services/document_parser/txt_parser.py index b595bedbc..e1d41956b 100755 --- a/apps/worker/app/services/document_parser/txt_parser.py +++ b/apps/worker/app/services/document_parser/txt_parser.py @@ -11,9 +11,9 @@ 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.services.chunks.chunk_refs import CHUNK_REF_PATTERN -from app.services.common.file_loading import load_file_bytes -from shared.services.ai.openai_compatible_client_sync import get_openai_client +from shared.utils.chunk_refs import CHUNK_REF_PATTERN +from shared.utils.file_loading import load_file_bytes +from shared.utils.OpenAICompatibleClientSync import get_openai_client def clean_texts_by_form(text, form="html"): diff --git a/apps/worker/app/services/workload/url_upload_transfer.py b/apps/worker/app/services/workload/url_upload_transfer.py index c9da336a8..25f4faf13 100644 --- a/apps/worker/app/services/workload/url_upload_transfer.py +++ b/apps/worker/app/services/workload/url_upload_transfer.py @@ -10,7 +10,7 @@ ValidationException, ) from shared.services.storage.job_file_storage import JobFileStorage -from shared.services.http.url_file_type import resolve_file_extension_sync +from shared.utils.url_file_type import resolve_file_extension_sync def resolve_supported_url_extension(source_url: str) -> str: diff --git a/packages/shared-python/shared/core/security.py b/packages/shared-python/shared/core/security.py index 8e93a8241..c1384e008 100644 --- a/packages/shared-python/shared/core/security.py +++ b/packages/shared-python/shared/core/security.py @@ -6,25 +6,6 @@ import bcrypt -def mask_api_key(api_key: str | None) -> str: - """Return a stable masked representation for logs and internal errors.""" - if not api_key: - return "" - - if len(api_key) <= 4: - return "*" * len(api_key) - - if len(api_key) <= 8: - prefix_length = 2 - suffix_length = 2 - else: - prefix_length = 4 - suffix_length = 4 - - masked_length = max(1, len(api_key) - prefix_length - suffix_length) - return f"{api_key[:prefix_length]}{'*' * masked_length}{api_key[-suffix_length:]}" - - def get_password_hash(password: str) -> str: """ Generate a password hash. diff --git a/packages/shared-python/shared/core/state_machine/service.py b/packages/shared-python/shared/core/state_machine/service.py index 1206313b9..6457e3964 100644 --- a/packages/shared-python/shared/core/state_machine/service.py +++ b/packages/shared-python/shared/core/state_machine/service.py @@ -38,7 +38,7 @@ 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.services.redis.key_builder import RedisKeyType, redis_key_builder +from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder class AsyncStateMachineService: 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 6d8ba1cc1..8935304e9 100644 --- a/packages/shared-python/shared/core/state_machine/service_sync.py +++ b/packages/shared-python/shared/core/state_machine/service_sync.py @@ -37,7 +37,7 @@ 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.services.redis.key_builder import RedisKeyType, redis_key_builder +from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder class SyncStateMachineService: diff --git a/packages/shared-python/shared/core/state_machine/transition_payloads.py b/packages/shared-python/shared/core/state_machine/transition_payloads.py index d7e6094de..c2a2cf58b 100644 --- a/packages/shared-python/shared/core/state_machine/transition_payloads.py +++ b/packages/shared-python/shared/core/state_machine/transition_payloads.py @@ -4,8 +4,8 @@ from typing import Any from shared.core.state_machine.states import JobStatus -from shared.core.error_details import normalize_error_details -from shared.core.serialization import make_json_safe +from shared.utils.error_details import normalize_error_details +from shared.utils.json_utils import make_json_safe def utc_now_naive() -> datetime: diff --git a/packages/shared-python/shared/models/database/api_key.py b/packages/shared-python/shared/models/database/api_key.py index 45b578867..a70b24f9f 100644 --- a/packages/shared-python/shared/models/database/api_key.py +++ b/packages/shared-python/shared/models/database/api_key.py @@ -12,7 +12,7 @@ from sqlalchemy.orm import Mapped, mapped_column from shared.core.database import Base -from shared.core.time import utc_now_naive +from shared.utils.utc_now import utc_now_naive class APIKey(Base): diff --git a/packages/shared-python/shared/models/database/credits_transaction.py b/packages/shared-python/shared/models/database/credits_transaction.py index 61bcdbb7b..1b83d6bbc 100644 --- a/packages/shared-python/shared/models/database/credits_transaction.py +++ b/packages/shared-python/shared/models/database/credits_transaction.py @@ -12,7 +12,7 @@ from sqlalchemy.orm import Mapped, mapped_column from shared.core.database import Base -from shared.core.time import utc_now_naive +from shared.utils.utc_now import utc_now_naive class CreditsTransaction(Base): diff --git a/packages/shared-python/shared/models/database/demo_materialization.py b/packages/shared-python/shared/models/database/demo_materialization.py index e63027c5a..d4cc1a7a1 100644 --- a/packages/shared-python/shared/models/database/demo_materialization.py +++ b/packages/shared-python/shared/models/database/demo_materialization.py @@ -9,7 +9,7 @@ from sqlalchemy.orm import Mapped, mapped_column from shared.core.database import Base -from shared.core.time import utc_now_naive +from shared.utils.utc_now import utc_now_naive class DemoMaterialization(Base): diff --git a/packages/shared-python/shared/models/database/document.py b/packages/shared-python/shared/models/database/document.py index d23ad235d..41bed45bb 100644 --- a/packages/shared-python/shared/models/database/document.py +++ b/packages/shared-python/shared/models/database/document.py @@ -22,7 +22,7 @@ from sqlalchemy.orm import Mapped, mapped_column, relationship from shared.core.database import Base -from shared.core.time import utc_now_naive +from shared.utils.utc_now import utc_now_naive class Document(Base): diff --git a/packages/shared-python/shared/models/database/guest_device.py b/packages/shared-python/shared/models/database/guest_device.py index 2b62a1fc7..cf956bb26 100644 --- a/packages/shared-python/shared/models/database/guest_device.py +++ b/packages/shared-python/shared/models/database/guest_device.py @@ -10,7 +10,7 @@ from sqlalchemy.orm import Mapped, mapped_column from shared.core.database import Base -from shared.core.time import utc_now_naive +from shared.utils.utc_now import utc_now_naive class GuestDevice(Base): diff --git a/packages/shared-python/shared/models/database/job.py b/packages/shared-python/shared/models/database/job.py index 4553ea655..76c9f1df1 100644 --- a/packages/shared-python/shared/models/database/job.py +++ b/packages/shared-python/shared/models/database/job.py @@ -25,7 +25,7 @@ from sqlalchemy.orm import Mapped, mapped_column, relationship from shared.core.database import Base -from shared.core.time import utc_now_naive +from shared.utils.utc_now import utc_now_naive if TYPE_CHECKING: diff --git a/packages/shared-python/shared/models/database/job_result.py b/packages/shared-python/shared/models/database/job_result.py index f683c42d2..c6d4d9e73 100644 --- a/packages/shared-python/shared/models/database/job_result.py +++ b/packages/shared-python/shared/models/database/job_result.py @@ -10,7 +10,7 @@ from sqlalchemy.orm import Mapped, mapped_column, relationship from shared.core.database import Base -from shared.core.time import utc_now_naive +from shared.utils.utc_now import utc_now_naive if TYPE_CHECKING: from shared.models.database.job import Job diff --git a/packages/shared-python/shared/models/database/job_state_audit_log.py b/packages/shared-python/shared/models/database/job_state_audit_log.py index e10c4de58..970b2af57 100644 --- a/packages/shared-python/shared/models/database/job_state_audit_log.py +++ b/packages/shared-python/shared/models/database/job_state_audit_log.py @@ -9,7 +9,7 @@ from sqlalchemy.orm import Mapped, mapped_column, relationship from shared.core.database import Base -from shared.core.time import utc_now_naive +from shared.utils.utc_now import utc_now_naive if TYPE_CHECKING: from shared.models.database.job import Job diff --git a/packages/shared-python/shared/models/database/job_state_history.py b/packages/shared-python/shared/models/database/job_state_history.py index 98ed383e4..b480ae3f7 100644 --- a/packages/shared-python/shared/models/database/job_state_history.py +++ b/packages/shared-python/shared/models/database/job_state_history.py @@ -10,7 +10,7 @@ from sqlalchemy.orm import Mapped, mapped_column, relationship from shared.core.database import Base -from shared.core.time import utc_now_naive +from shared.utils.utc_now import utc_now_naive if TYPE_CHECKING: from shared.models.database.job import Job diff --git a/packages/shared-python/shared/models/database/payment_record.py b/packages/shared-python/shared/models/database/payment_record.py index 923c7e1aa..69edb639e 100644 --- a/packages/shared-python/shared/models/database/payment_record.py +++ b/packages/shared-python/shared/models/database/payment_record.py @@ -21,7 +21,7 @@ from sqlalchemy.orm import Mapped, mapped_column from shared.core.database import Base -from shared.core.time import utc_now_naive +from shared.utils.utc_now import utc_now_naive class PaymentRecord(Base): """Payment Record Model (for idempotency guarantee)""" diff --git a/packages/shared-python/shared/models/database/stripe_price_config.py b/packages/shared-python/shared/models/database/stripe_price_config.py index c820ead19..0fb1200f2 100644 --- a/packages/shared-python/shared/models/database/stripe_price_config.py +++ b/packages/shared-python/shared/models/database/stripe_price_config.py @@ -18,7 +18,7 @@ from sqlalchemy.orm import Mapped, mapped_column from shared.core.database import Base -from shared.core.time import utc_now_naive +from shared.utils.utc_now import utc_now_naive class StripePriceConfig(Base): diff --git a/packages/shared-python/shared/models/database/user_balance.py b/packages/shared-python/shared/models/database/user_balance.py index 77a380e6e..a67ea2565 100644 --- a/packages/shared-python/shared/models/database/user_balance.py +++ b/packages/shared-python/shared/models/database/user_balance.py @@ -11,7 +11,7 @@ from sqlalchemy.orm import Mapped, mapped_column from shared.core.database import Base -from shared.core.time import utc_now_naive +from shared.utils.utc_now import utc_now_naive class UserBalance(Base): """User balance model — tracks credits balance and tier membership""" diff --git a/packages/shared-python/shared/models/database/webhook.py b/packages/shared-python/shared/models/database/webhook.py index c5da892d7..09f7f4dd1 100644 --- a/packages/shared-python/shared/models/database/webhook.py +++ b/packages/shared-python/shared/models/database/webhook.py @@ -16,7 +16,7 @@ from sqlalchemy.orm import Mapped, mapped_column, relationship from shared.core.database import Base -from shared.core.time import utc_now_naive +from shared.utils.utc_now import utc_now_naive if TYPE_CHECKING: from shared.models.database.job import Job diff --git a/packages/shared-python/shared/models/database/webhook_log.py b/packages/shared-python/shared/models/database/webhook_log.py index b41dff5a8..2beb6dca2 100644 --- a/packages/shared-python/shared/models/database/webhook_log.py +++ b/packages/shared-python/shared/models/database/webhook_log.py @@ -14,7 +14,7 @@ from sqlalchemy.orm import Mapped, mapped_column, relationship from shared.core.database import Base -from shared.core.time import utc_now_naive +from shared.utils.utc_now import utc_now_naive if TYPE_CHECKING: from shared.models.database.job import Job diff --git a/packages/shared-python/shared/models/database/webhook_secret.py b/packages/shared-python/shared/models/database/webhook_secret.py index 50351ed8e..c1d391b3f 100644 --- a/packages/shared-python/shared/models/database/webhook_secret.py +++ b/packages/shared-python/shared/models/database/webhook_secret.py @@ -15,7 +15,7 @@ from sqlalchemy.orm import Mapped, mapped_column from shared.core.database import Base -from shared.core.time import utc_now_naive +from shared.utils.utc_now import utc_now_naive class WebhookSecretStatus: diff --git a/packages/shared-python/shared/repositories/credits_repository.py b/packages/shared-python/shared/repositories/credits_repository.py index 04fc1d452..3e97e60c5 100644 --- a/packages/shared-python/shared/repositories/credits_repository.py +++ b/packages/shared-python/shared/repositories/credits_repository.py @@ -12,7 +12,7 @@ from shared.models.database.credits_transaction import CreditsTransaction from shared.models.database.payment_record import PaymentRecord from shared.models.database.user_balance import UserBalance -from shared.core.time import utc_now_naive +from shared.utils.utc_now import utc_now_naive class CreditsRepository: diff --git a/packages/shared-python/shared/repositories/credits_sync_repository.py b/packages/shared-python/shared/repositories/credits_sync_repository.py index 925b9857b..e0c198615 100644 --- a/packages/shared-python/shared/repositories/credits_sync_repository.py +++ b/packages/shared-python/shared/repositories/credits_sync_repository.py @@ -11,7 +11,7 @@ from shared.models.database.credits_transaction import CreditsTransaction from shared.models.database.payment_record import PaymentRecord from shared.models.database.user_balance import UserBalance -from shared.core.time import utc_now_naive +from shared.utils.utc_now import utc_now_naive class SyncCreditsRepository: diff --git a/packages/shared-python/shared/services/billing/credits_service.py b/packages/shared-python/shared/services/billing/credits_service.py index 3d690401f..d0dc48fc4 100644 --- a/packages/shared-python/shared/services/billing/credits_service.py +++ b/packages/shared-python/shared/services/billing/credits_service.py @@ -24,7 +24,7 @@ from shared.models.database.payment_record import PaymentRecord from shared.models.database.user_balance import UserBalance from shared.repositories.credits_repository import CreditsRepository -from shared.core.time import utc_now_naive +from shared.utils.utc_now import utc_now_naive NON_EXPIRING_CREDIT_TRANSACTION_TYPES: tuple[str, ...] = ("refund",) diff --git a/packages/shared-python/shared/services/billing/credits_sync_service.py b/packages/shared-python/shared/services/billing/credits_sync_service.py index d6b07d42b..e6eb63446 100644 --- a/packages/shared-python/shared/services/billing/credits_sync_service.py +++ b/packages/shared-python/shared/services/billing/credits_sync_service.py @@ -19,7 +19,7 @@ from shared.models.database.payment_record import PaymentRecord from shared.models.database.user_balance import UserBalance from shared.repositories.credits_sync_repository import SyncCreditsRepository -from shared.core.time import utc_now_naive +from shared.utils.utc_now import utc_now_naive NON_EXPIRING_CREDIT_TRANSACTION_TYPES: tuple[str, ...] = ("refund",) diff --git a/packages/shared-python/shared/services/chunks/chunk_connections.py b/packages/shared-python/shared/services/chunks/chunk_connections.py index 13d4ab90c..efb0804f1 100644 --- a/packages/shared-python/shared/services/chunks/chunk_connections.py +++ b/packages/shared-python/shared/services/chunks/chunk_connections.py @@ -5,7 +5,7 @@ from collections.abc import Mapping, Sequence from typing import Any, TypeAlias, TypedDict -from shared.services.chunks.chunk_refs import ChunkRefSpan, extract_chunk_ref_spans +from shared.utils.chunk_refs import ChunkRefSpan, extract_chunk_ref_spans class PositionPayload(TypedDict): diff --git a/packages/shared-python/shared/services/http/__init__.py b/packages/shared-python/shared/services/http/__init__.py deleted file mode 100644 index 3bed643fe..000000000 --- a/packages/shared-python/shared/services/http/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Shared outbound HTTP adapters and URL safety helpers.""" diff --git a/packages/shared-python/shared/services/job_failure_sync.py b/packages/shared-python/shared/services/job_failure_sync.py index 66c93cd1e..ad3a313dd 100644 --- a/packages/shared-python/shared/services/job_failure_sync.py +++ b/packages/shared-python/shared/services/job_failure_sync.py @@ -13,7 +13,7 @@ from shared.services.billing.credits_sync_service import SyncCreditsService from shared.services.job_post_commit_effects_sync import PostCommitEffectPlan from shared.services.job_webhook_outbox_sync import SyncJobWebhookOutbox -from shared.core.error_details import normalize_error_details +from shared.utils.error_details import normalize_error_details @dataclass(frozen=True) diff --git a/packages/shared-python/shared/services/job_lifecycle_sync.py b/packages/shared-python/shared/services/job_lifecycle_sync.py index f160f41af..ad0dcfc70 100644 --- a/packages/shared-python/shared/services/job_lifecycle_sync.py +++ b/packages/shared-python/shared/services/job_lifecycle_sync.py @@ -26,7 +26,7 @@ from shared.services.redis.redis_sync_service import ( SyncRedisServiceFactory, ) -from shared.services.redis.key_builder import RedisKeyType, redis_key_builder +from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder class SyncJobLifecycleService: diff --git a/packages/shared-python/shared/services/jobs/result_delivery.py b/packages/shared-python/shared/services/jobs/result_delivery.py index 406ba8627..0a5c9b0d6 100644 --- a/packages/shared-python/shared/services/jobs/result_delivery.py +++ b/packages/shared-python/shared/services/jobs/result_delivery.py @@ -5,7 +5,7 @@ from typing import Any from shared.services.storage.job_file_storage import JobFileStorage -from shared.core.time import utc_now_naive +from shared.utils.utc_now import utc_now_naive @dataclass(frozen=True) diff --git a/packages/shared-python/shared/services/redis/distributed_lock.py b/packages/shared-python/shared/services/redis/distributed_lock.py index 8933737f1..a36a23bf9 100644 --- a/packages/shared-python/shared/services/redis/distributed_lock.py +++ b/packages/shared-python/shared/services/redis/distributed_lock.py @@ -18,7 +18,7 @@ from shared.core.config import settings from shared.core.exceptions.domain_exceptions import UnavailableException from shared.services.redis.redis_sync_service import SyncRedisService -from shared.services.redis.key_builder import redis_key_builder +from shared.utils.redis_key_builder import redis_key_builder # Lua script: atomically check owner token then delete. # Prevents a stale owner from accidentally deleting a lock diff --git a/packages/shared-python/shared/services/redis/job_info_redis_service.py b/packages/shared-python/shared/services/redis/job_info_redis_service.py index 0ee34b0b7..89c1a2ec2 100644 --- a/packages/shared-python/shared/services/redis/job_info_redis_service.py +++ b/packages/shared-python/shared/services/redis/job_info_redis_service.py @@ -8,7 +8,7 @@ from loguru import logger from shared.services.redis.redis_service import RedisService -from shared.services.redis.key_builder import RedisKeyType, redis_key_builder +from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder class JobInfoRedisService: diff --git a/packages/shared-python/shared/services/redis/job_metadata_service.py b/packages/shared-python/shared/services/redis/job_metadata_service.py index 410a1c052..62fb9c135 100644 --- a/packages/shared-python/shared/services/redis/job_metadata_service.py +++ b/packages/shared-python/shared/services/redis/job_metadata_service.py @@ -5,7 +5,7 @@ from loguru import logger from shared.services.redis.redis_service import RedisService -from shared.services.redis.key_builder import RedisKeyType, redis_key_builder +from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder class JobMetadataService: diff --git a/packages/shared-python/shared/services/redis/rate_limit_service.py b/packages/shared-python/shared/services/redis/rate_limit_service.py index c1cec343b..de7ef821b 100644 --- a/packages/shared-python/shared/services/redis/rate_limit_service.py +++ b/packages/shared-python/shared/services/redis/rate_limit_service.py @@ -7,7 +7,7 @@ from loguru import logger from shared.services.redis.redis_service import RedisService -from shared.services.redis.key_builder import redis_key_builder +from shared.utils.redis_key_builder import redis_key_builder class RateLimitService: diff --git a/packages/shared-python/shared/services/redis/redis_monitor.py b/packages/shared-python/shared/services/redis/redis_monitor.py index f6887576a..6ae090590 100644 --- a/packages/shared-python/shared/services/redis/redis_monitor.py +++ b/packages/shared-python/shared/services/redis/redis_monitor.py @@ -7,7 +7,7 @@ from loguru import logger from shared.services.redis.redis_service import RedisService -from shared.services.redis.key_builder import redis_key_builder +from shared.utils.redis_key_builder import redis_key_builder class RedisMonitor: diff --git a/packages/shared-python/shared/services/redis/redis_service.py b/packages/shared-python/shared/services/redis/redis_service.py index fd735bcf8..99ed39c1d 100644 --- a/packages/shared-python/shared/services/redis/redis_service.py +++ b/packages/shared-python/shared/services/redis/redis_service.py @@ -13,7 +13,7 @@ RedisConnectionError, RedisOperationError, ) -from shared.services.redis.retry_policy import RedisHealthChecker, RedisRetry +from shared.utils.redis_retry import RedisHealthChecker, RedisRetry ResponseT = TypeVar("ResponseT") @@ -87,7 +87,7 @@ async def set( if isinstance(value, (dict, list)): # Use make_json_safe so complex types serialize consistently. - from shared.core.serialization import make_json_safe + from shared.utils.json_utils import make_json_safe safe_value = make_json_safe(value) value = json.dumps(safe_value, ensure_ascii=False) diff --git a/packages/shared-python/shared/services/redis/redis_sync_service.py b/packages/shared-python/shared/services/redis/redis_sync_service.py index 2fe2cb7bf..aa725b314 100644 --- a/packages/shared-python/shared/services/redis/redis_sync_service.py +++ b/packages/shared-python/shared/services/redis/redis_sync_service.py @@ -12,7 +12,7 @@ from redis.connection import BlockingConnectionPool from shared.core.config.redis import RedisConfigManager -from shared.services.redis.key_builder import RedisKeyType, redis_key_builder +from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder class SyncRedisService: @@ -69,7 +69,7 @@ def set( client = self._get_client() full_key = self._build_key(key) if isinstance(value, (dict, list)): - from shared.core.serialization import make_json_safe + from shared.utils.json_utils import make_json_safe safe_value = make_json_safe(value) value = json.dumps(safe_value, ensure_ascii=False) @@ -407,7 +407,7 @@ def set_task_status(self, task_id: str, status: str) -> bool: def save_task_result(self, task_id: str, result: Dict[str, Any]) -> bool: try: - from shared.core.serialization import make_json_safe + from shared.utils.json_utils import make_json_safe task_ttl = redis_key_builder.get_key_ttl(RedisKeyType.TASK) result_key = self.redis._build_key(redis_key_builder.task_result(task_id)) diff --git a/packages/shared-python/shared/services/redis/task_redis_service.py b/packages/shared-python/shared/services/redis/task_redis_service.py index 6c9444534..883d8a116 100644 --- a/packages/shared-python/shared/services/redis/task_redis_service.py +++ b/packages/shared-python/shared/services/redis/task_redis_service.py @@ -5,7 +5,7 @@ from loguru import logger from shared.services.redis.redis_service import RedisService -from shared.services.redis.key_builder import RedisKeyType, redis_key_builder +from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder class TaskRedisService: diff --git a/packages/shared-python/shared/services/redis/user_redis_service.py b/packages/shared-python/shared/services/redis/user_redis_service.py index d28ee59d5..6db0a2acc 100644 --- a/packages/shared-python/shared/services/redis/user_redis_service.py +++ b/packages/shared-python/shared/services/redis/user_redis_service.py @@ -5,7 +5,7 @@ from loguru import logger from shared.services.redis.redis_service import RedisService -from shared.services.redis.key_builder import RedisKeyType, redis_key_builder +from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder class UserRedisService: diff --git a/packages/shared-python/shared/services/retrieval/agentic/evidence.py b/packages/shared-python/shared/services/retrieval/agentic/evidence.py index 5bd38f947..3d54fb47b 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/evidence.py +++ b/packages/shared-python/shared/services/retrieval/agentic/evidence.py @@ -11,7 +11,7 @@ from shared.services.retrieval.agentic.types import DocTreeNode from shared.services.retrieval.assets import build_retrieval_asset_url_map from shared.services.retrieval.hit_stats_service import compute_importance_score -from shared.services.retrieval.token_estimate import estimate_tokens +from shared.utils.token_estimate import estimate_tokens def with_context_prompt_projection( diff --git a/packages/shared-python/shared/services/retrieval/agentic/policy.py b/packages/shared-python/shared/services/retrieval/agentic/policy.py index 7a3bd3654..25323580a 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/policy.py +++ b/packages/shared-python/shared/services/retrieval/agentic/policy.py @@ -21,7 +21,7 @@ from shared.services.retrieval.agentic.budget import BudgetExceeded from shared.services.retrieval.agentic.types import AgentRunConfig, AgentState from shared.services.retrieval.llm_adapter import LLMFn -from shared.services.retrieval.token_estimate import estimate_tokens +from shared.utils.token_estimate import estimate_tokens def _parse_answer_response(text: str) -> dict[str, Any] | None: diff --git a/packages/shared-python/shared/services/retrieval/agentic/runtime.py b/packages/shared-python/shared/services/retrieval/agentic/runtime.py index 750541c25..d658cd562 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/runtime.py +++ b/packages/shared-python/shared/services/retrieval/agentic/runtime.py @@ -12,7 +12,7 @@ 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.services.retrieval.token_estimate import estimate_tokens +from shared.utils.token_estimate import estimate_tokens def build_config_from_env() -> AgentRunConfig: 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 index 24d2d3ee9..b7a251805 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/section_prompt_projection.py +++ b/packages/shared-python/shared/services/retrieval/agentic/section_prompt_projection.py @@ -1,7 +1,7 @@ """Prompt projection for agentic section navigation.""" from __future__ import annotations -from shared.services.text_processing.tokenization import truncate_content_preview +from shared.utils.text_utils import truncate_content_preview def format_items_for_llm( diff --git a/packages/shared-python/shared/services/retrieval/lexical_ranker.py b/packages/shared-python/shared/services/retrieval/lexical_ranker.py index 39fe446c0..8b50cd404 100644 --- a/packages/shared-python/shared/services/retrieval/lexical_ranker.py +++ b/packages/shared-python/shared/services/retrieval/lexical_ranker.py @@ -4,7 +4,7 @@ from loguru import logger -from shared.services.text_processing.tokenization import tokenize_for_retrieval +from shared.utils.text_utils import tokenize_for_retrieval def tokenize_query_for_ranker(query: str) -> list[str]: diff --git a/packages/shared-python/shared/services/retrieval/lexical_text.py b/packages/shared-python/shared/services/retrieval/lexical_text.py index 7c98dc56c..42d15366f 100644 --- a/packages/shared-python/shared/services/retrieval/lexical_text.py +++ b/packages/shared-python/shared/services/retrieval/lexical_text.py @@ -6,7 +6,7 @@ from typing import Any, Optional -from shared.services.text_processing.tokenization import tokenize_contents_for_retrieval +from shared.utils.text_utils import tokenize_contents_for_retrieval def normalize_section_path(path: Optional[str]) -> str: diff --git a/packages/shared-python/shared/services/retrieval/llm_adapter.py b/packages/shared-python/shared/services/retrieval/llm_adapter.py index 9c8bd9263..c5e592b5e 100644 --- a/packages/shared-python/shared/services/retrieval/llm_adapter.py +++ b/packages/shared-python/shared/services/retrieval/llm_adapter.py @@ -87,7 +87,7 @@ def create_retrieval_llm_fn( effective_model = model or _resolve_default_model() async def llm_fn(prompt: LLMFnInput) -> str: - from shared.services.ai.openai_compatible_client_sync import get_openai_client + from shared.utils.OpenAICompatibleClientSync import get_openai_client client = get_openai_client(model=effective_model) current_llm_usage.set(None) @@ -118,7 +118,7 @@ def create_retrieval_planner_fn( effective_model = model or _resolve_planner_model(thinking=thinking) async def llm_fn(prompt: LLMFnInput) -> str: - from shared.services.ai.openai_compatible_client_sync import get_openai_client + from shared.utils.OpenAICompatibleClientSync import get_openai_client client = get_openai_client(model=effective_model) current_llm_usage.set(None) @@ -159,7 +159,7 @@ def create_retrieval_vlm_fn( return None async def vlm_fn(prompt: LLMFnInput) -> str: - from shared.services.ai.openai_compatible_client_sync import get_openai_client + from shared.utils.OpenAICompatibleClientSync import get_openai_client client = get_openai_client(model=effective_model) current_llm_usage.set(None) diff --git a/packages/shared-python/shared/services/retrieval/workflow/planner.py b/packages/shared-python/shared/services/retrieval/workflow/planner.py index 463894879..e34713db6 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/planner.py +++ b/packages/shared-python/shared/services/retrieval/workflow/planner.py @@ -11,7 +11,7 @@ from shared.services.retrieval.agentic.budget import BudgetExceeded, BudgetLedger from shared.services.retrieval.llm_adapter import LLMFn, current_llm_usage from shared.services.retrieval.workflow.types import FinalStrategy, OutputRole, PlannedStep, QueryPlan, StepKind -from shared.services.retrieval.token_estimate import estimate_tokens +from shared.utils.token_estimate import estimate_tokens _PLAN_SCHEMA = { diff --git a/packages/shared-python/shared/services/retrieval/workflow/synthesizer.py b/packages/shared-python/shared/services/retrieval/workflow/synthesizer.py index 2ee92bac9..89d4ec031 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/synthesizer.py +++ b/packages/shared-python/shared/services/retrieval/workflow/synthesizer.py @@ -6,7 +6,7 @@ from shared.services.retrieval.agentic.budget import BudgetLedger from shared.services.retrieval.llm_adapter import LLMFn, current_llm_usage from shared.services.retrieval.workflow.types import PlannedStep, QueryPlan, StepResult -from shared.services.retrieval.token_estimate import estimate_tokens +from shared.utils.token_estimate import estimate_tokens _SYNTHESIZE_PROMPT = """\ 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 77d49e2b4..cc46f6bd3 100644 --- a/packages/shared-python/shared/services/storage/job_file_storage.py +++ b/packages/shared-python/shared/services/storage/job_file_storage.py @@ -8,8 +8,8 @@ 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.services.http.pinned_outbound import download_pinned_outbound_file -from shared.services.http.url_security import validate_http_url_and_resolve_ip +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: diff --git a/packages/shared-python/shared/services/storage/zip_result_schema.py b/packages/shared-python/shared/services/storage/zip_result_schema.py index 72bb6b5a0..410bfa111 100644 --- a/packages/shared-python/shared/services/storage/zip_result_schema.py +++ b/packages/shared-python/shared/services/storage/zip_result_schema.py @@ -10,8 +10,8 @@ normalize_connect_to_targets, parse_relationship_refs, ) -from shared.services.text_processing.tokenization import truncate_content_preview -from shared.core.time import utc_now_naive +from shared.utils.text_utils import truncate_content_preview +from shared.utils.utc_now import utc_now_naive class ZipResultSchemaBuilder: diff --git a/packages/shared-python/shared/services/text_processing/__init__.py b/packages/shared-python/shared/services/text_processing/__init__.py deleted file mode 100644 index 7d77e9d64..000000000 --- a/packages/shared-python/shared/services/text_processing/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Shared text processing and tokenization services.""" diff --git a/packages/shared-python/shared/services/webhook/delivery_client.py b/packages/shared-python/shared/services/webhook/delivery_client.py index 2e33aca20..232f9a875 100644 --- a/packages/shared-python/shared/services/webhook/delivery_client.py +++ b/packages/shared-python/shared/services/webhook/delivery_client.py @@ -8,8 +8,8 @@ from loguru import logger -from shared.services.http.pinned_outbound import send_pinned_outbound_request -from shared.services.http.url_security import validate_http_url_and_resolve_ip_async +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 diff --git a/packages/shared-python/shared/services/webhook/qstash_publisher.py b/packages/shared-python/shared/services/webhook/qstash_publisher.py index 810f35def..7a0580cb6 100644 --- a/packages/shared-python/shared/services/webhook/qstash_publisher.py +++ b/packages/shared-python/shared/services/webhook/qstash_publisher.py @@ -26,7 +26,7 @@ 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.services.http.url_security import ( +from shared.utils.url_security import ( validate_http_url_and_resolve_ip, ) diff --git a/packages/shared-python/shared/services/ai/openai_compatible_client_sync.py b/packages/shared-python/shared/utils/OpenAICompatibleClientSync.py similarity index 98% rename from packages/shared-python/shared/services/ai/openai_compatible_client_sync.py rename to packages/shared-python/shared/utils/OpenAICompatibleClientSync.py index d62718c79..66da76a6d 100644 --- a/packages/shared-python/shared/services/ai/openai_compatible_client_sync.py +++ b/packages/shared-python/shared/utils/OpenAICompatibleClientSync.py @@ -17,9 +17,9 @@ from shared.core.config import settings from shared.core.exceptions.domain_exceptions import LLMServiceException -from shared.services.http.client_pool import get_sync_client -from shared.services.ai.llm_mock import build_mock_chat_completion_response -from shared.core.security import mask_api_key +from shared.utils.http_clients import get_sync_client +from shared.utils.llm_mock import build_mock_chat_completion_response +from shared.utils.security_utils import mask_api_key LOCAL_DEBUG = os.getenv("LOCAL_DEBUG", "0") == "1" LLMUsage = dict[str, int] @@ -185,7 +185,7 @@ def _make_ali_pool_call( api_kwargs: Dict[str, Any], ) -> tuple[str, LLMUsage]: """Acquire a token, make the call, and retry inline on 429.""" - from shared.services.ai.ali_quota_manager import get_ali_quota_manager + from shared.utils.ali_quota_manager import get_ali_quota_manager quota_manager = get_ali_quota_manager() base_url: Optional[str] = self._base_url diff --git a/packages/shared-python/shared/utils/__init__.py b/packages/shared-python/shared/utils/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/packages/shared-python/shared/services/ai/ali_quota_manager.py b/packages/shared-python/shared/utils/ali_quota_manager.py similarity index 96% rename from packages/shared-python/shared/services/ai/ali_quota_manager.py rename to packages/shared-python/shared/utils/ali_quota_manager.py index 764f3e7fb..535b86c0c 100644 --- a/packages/shared-python/shared/services/ai/ali_quota_manager.py +++ b/packages/shared-python/shared/utils/ali_quota_manager.py @@ -14,7 +14,7 @@ SyncRedisService, SyncRedisServiceFactory, ) -from shared.services.ai.quota_manager import BaseQuotaManager, TokenConfig +from shared.utils.quota_manager import BaseQuotaManager, TokenConfig class AliQuotaManager(BaseQuotaManager): diff --git a/packages/shared-python/shared/services/auth/api_key_tokens.py b/packages/shared-python/shared/utils/api_keys.py similarity index 100% rename from packages/shared-python/shared/services/auth/api_key_tokens.py rename to packages/shared-python/shared/utils/api_keys.py diff --git a/packages/shared-python/shared/services/chunks/chunk_refs.py b/packages/shared-python/shared/utils/chunk_refs.py similarity index 100% rename from packages/shared-python/shared/services/chunks/chunk_refs.py rename to packages/shared-python/shared/utils/chunk_refs.py diff --git a/apps/worker/app/services/common/device_checks.py b/packages/shared-python/shared/utils/device_utils.py similarity index 100% rename from apps/worker/app/services/common/device_checks.py rename to packages/shared-python/shared/utils/device_utils.py diff --git a/packages/shared-python/shared/core/error_details.py b/packages/shared-python/shared/utils/error_details.py similarity index 100% rename from packages/shared-python/shared/core/error_details.py rename to packages/shared-python/shared/utils/error_details.py diff --git a/apps/worker/app/services/common/file_loading.py b/packages/shared-python/shared/utils/file_loading.py similarity index 100% rename from apps/worker/app/services/common/file_loading.py rename to packages/shared-python/shared/utils/file_loading.py diff --git a/apps/worker/app/services/common/file_utils.py b/packages/shared-python/shared/utils/file_utils.py similarity index 100% rename from apps/worker/app/services/common/file_utils.py rename to packages/shared-python/shared/utils/file_utils.py diff --git a/apps/worker/app/services/common/resource_cleanup.py b/packages/shared-python/shared/utils/gc_utils.py similarity index 100% rename from apps/worker/app/services/common/resource_cleanup.py rename to packages/shared-python/shared/utils/gc_utils.py diff --git a/packages/shared-python/shared/services/http/client_pool.py b/packages/shared-python/shared/utils/http_clients.py similarity index 100% rename from packages/shared-python/shared/services/http/client_pool.py rename to packages/shared-python/shared/utils/http_clients.py diff --git a/packages/shared-python/shared/services/ai/iloveapi_quota_manager.py b/packages/shared-python/shared/utils/iloveapi_quota_manager.py similarity index 99% rename from packages/shared-python/shared/services/ai/iloveapi_quota_manager.py rename to packages/shared-python/shared/utils/iloveapi_quota_manager.py index 7afb5aacf..b3b923a41 100644 --- a/packages/shared-python/shared/services/ai/iloveapi_quota_manager.py +++ b/packages/shared-python/shared/utils/iloveapi_quota_manager.py @@ -26,7 +26,7 @@ SyncRedisService, SyncRedisServiceFactory, ) -from shared.services.ai.quota_manager import BaseQuotaManager, TokenConfig +from shared.utils.quota_manager import BaseQuotaManager, TokenConfig class ILoveApiQuotaManager(BaseQuotaManager): diff --git a/packages/shared-python/shared/core/serialization.py b/packages/shared-python/shared/utils/json_utils.py similarity index 100% rename from packages/shared-python/shared/core/serialization.py rename to packages/shared-python/shared/utils/json_utils.py diff --git a/packages/shared-python/shared/services/ai/llm_mock.py b/packages/shared-python/shared/utils/llm_mock.py similarity index 100% rename from packages/shared-python/shared/services/ai/llm_mock.py rename to packages/shared-python/shared/utils/llm_mock.py diff --git a/apps/worker/app/services/common/math_helpers.py b/packages/shared-python/shared/utils/math_utils.py similarity index 100% rename from apps/worker/app/services/common/math_helpers.py rename to packages/shared-python/shared/utils/math_utils.py diff --git a/packages/shared-python/shared/services/http/pinned_outbound.py b/packages/shared-python/shared/utils/pinned_outbound_http.py similarity index 99% rename from packages/shared-python/shared/services/http/pinned_outbound.py rename to packages/shared-python/shared/utils/pinned_outbound_http.py index c281dceb5..808b219c8 100644 --- a/packages/shared-python/shared/services/http/pinned_outbound.py +++ b/packages/shared-python/shared/utils/pinned_outbound_http.py @@ -19,7 +19,7 @@ from urllib3.util.connection import create_connection from shared.core.exceptions.domain_exceptions import ValidationException -from shared.services.http.url_security import SafePublicHTTPURL +from shared.utils.url_security import SafePublicHTTPURL @dataclass(frozen=True) diff --git a/packages/shared-python/shared/services/ai/quota_manager.py b/packages/shared-python/shared/utils/quota_manager.py similarity index 100% rename from packages/shared-python/shared/services/ai/quota_manager.py rename to packages/shared-python/shared/utils/quota_manager.py diff --git a/packages/shared-python/shared/services/redis/key_builder.py b/packages/shared-python/shared/utils/redis_key_builder.py similarity index 100% rename from packages/shared-python/shared/services/redis/key_builder.py rename to packages/shared-python/shared/utils/redis_key_builder.py diff --git a/packages/shared-python/shared/services/redis/retry_policy.py b/packages/shared-python/shared/utils/redis_retry.py similarity index 100% rename from packages/shared-python/shared/services/redis/retry_policy.py rename to packages/shared-python/shared/utils/redis_retry.py diff --git a/packages/shared-python/shared/utils/security_utils.py b/packages/shared-python/shared/utils/security_utils.py new file mode 100644 index 000000000..2702e79ef --- /dev/null +++ b/packages/shared-python/shared/utils/security_utils.py @@ -0,0 +1,22 @@ +"""Security-related shared utilities.""" + +from typing import Optional + + +def mask_api_key(api_key: Optional[str]) -> str: + """Return a stable masked representation for logs and internal errors.""" + if not api_key: + return "" + + if len(api_key) <= 4: + return "*" * len(api_key) + + if len(api_key) <= 8: + prefix_length = 2 + suffix_length = 2 + else: + prefix_length = 4 + suffix_length = 4 + + masked_length = max(1, len(api_key) - prefix_length - suffix_length) + return f"{api_key[:prefix_length]}{'*' * masked_length}{api_key[-suffix_length:]}" diff --git a/packages/shared-python/shared/services/text_processing/tokenization.py b/packages/shared-python/shared/utils/text_utils.py similarity index 99% rename from packages/shared-python/shared/services/text_processing/tokenization.py rename to packages/shared-python/shared/utils/text_utils.py index 8e00955d9..dd61bfdd2 100644 --- a/packages/shared-python/shared/services/text_processing/tokenization.py +++ b/packages/shared-python/shared/utils/text_utils.py @@ -10,7 +10,7 @@ from syntok.tokenizer import Tokenizer as _SyntokTokenizer -from shared.services.chunks.chunk_refs import CHUNK_REF_PATTERN +from shared.utils.chunk_refs import CHUNK_REF_PATTERN try: from blingfire import text_to_words as _blingfire_text_to_words diff --git a/packages/shared-python/shared/services/retrieval/token_estimate.py b/packages/shared-python/shared/utils/token_estimate.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/token_estimate.py rename to packages/shared-python/shared/utils/token_estimate.py diff --git a/packages/shared-python/shared/services/http/url_file_type.py b/packages/shared-python/shared/utils/url_file_type.py similarity index 97% rename from packages/shared-python/shared/services/http/url_file_type.py rename to packages/shared-python/shared/utils/url_file_type.py index a82f97a60..7314d7e23 100644 --- a/packages/shared-python/shared/services/http/url_file_type.py +++ b/packages/shared-python/shared/utils/url_file_type.py @@ -12,7 +12,7 @@ from shared.core.config import settings from shared.core.exceptions.domain_exceptions import ValidationException -from shared.services.http.url_security import ( +from shared.utils.url_security import ( HTTPURLValidationResult, SafePublicHTTPURL, validate_http_url_and_resolve_ip, @@ -116,7 +116,7 @@ async def resolve_file_extension_async(url: str) -> str | None: return ext try: - from shared.services.http.client_pool import get_async_client + from shared.utils.http_clients import get_async_client client = get_async_client() response = None @@ -167,7 +167,7 @@ def resolve_file_extension_sync(url: str) -> str | None: return ext try: - from shared.services.http.client_pool import get_sync_client + from shared.utils.http_clients import get_sync_client client = get_sync_client() response = None diff --git a/packages/shared-python/shared/services/http/url_security.py b/packages/shared-python/shared/utils/url_security.py similarity index 100% rename from packages/shared-python/shared/services/http/url_security.py rename to packages/shared-python/shared/utils/url_security.py diff --git a/packages/shared-python/shared/core/time.py b/packages/shared-python/shared/utils/utc_now.py similarity index 100% rename from packages/shared-python/shared/core/time.py rename to packages/shared-python/shared/utils/utc_now.py diff --git a/packages/shared-python/shared/services/http/zip_download.py b/packages/shared-python/shared/utils/zip_download.py similarity index 100% rename from packages/shared-python/shared/services/http/zip_download.py rename to packages/shared-python/shared/utils/zip_download.py From bf7d0e9fbeb25f0d702909329bf90a34be5c48f1 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 09:27:58 +0000 Subject: [PATCH 55/78] refactor: organize shared utility ownership seams --- AGENTS.md | 5 +- CONTEXT.md | 18 + .../document_ingestion/creation_service.py | 2 +- .../services/document_ingestion/service.py | 4 +- .../s3_events/subscription_service.py | 4 +- apps/api/main.py | 2 +- .../contract/test_job_creation_contract.py | 8 +- .../tests/contract/test_s3_event_contract.py | 2 +- apps/worker/app/core/worker_bootstrap.py | 2 +- apps/worker/app/services/common/__init__.py | 8 +- .../app/services/common/device_checks.py | 0 .../app/services/common}/file_loading.py | 0 .../worker/app/services/common}/file_utils.py | 0 .../app/services/common/math_helpers.py | 0 .../app/services/common/resource_cleanup.py | 0 .../connect_builder/summary_builder.py | 2 +- .../tools/classify_special_pages.py | 2 +- .../document_agent/tools/probe_vlm_inspect.py | 2 +- .../tools/propose_shard_plan.py | 2 +- .../document_parser/atlas_classifier.py | 2 +- .../services/document_parser/doc_parser.py | 4 +- .../document_parser/excel_table_parser.py | 4 +- .../document_parser/fragment_parser.py | 4 +- .../services/document_parser/image_parser.py | 6 +- .../document_parser/internal_parse_name.py | 2 +- .../services/document_parser/layout_parser.py | 2 +- .../markdown_deferred_summary.py | 2 +- .../document_parser/markdown_image_asset.py | 2 +- .../document_parser/markdown_table_asset.py | 2 +- .../app/services/document_parser/md_parser.py | 2 +- .../document_parser/mineru_pdf_service.py | 2 +- .../document_parser/mineru_quota_manager.py | 2 +- .../orchestration/parse_session.py | 2 +- .../services/document_parser/path_helpers.py | 2 +- .../services/document_parser/pptx_parser.py | 8 +- .../document_parser/table_frame_parser.py | 2 +- .../services/document_parser/toc_parser.py | 2 +- .../services/document_parser/txt_parser.py | 4 +- .../services/workload/url_upload_transfer.py | 2 +- .../shared/core/state_machine/service.py | 2 +- .../shared/core/state_machine/service_sync.py | 2 +- .../shared/services/ai/ali_quota_manager.py | 65 +++ .../services/ai/iloveapi_quota_manager.py | 215 +++++++++ .../shared/services/ai/llm_mock.py | 157 +++++++ .../ai/openai_compatible_client_sync.py | 425 +++++++++++++++++ .../shared/services/http/__init__.py | 123 +++++ .../shared/services/http/client_pool.py | 77 ++++ .../shared/services/http/pinned_outbound.py | 313 +++++++++++++ .../shared/services/http/url_file_type.py | 195 ++++++++ .../shared/services/http/url_security.py | 243 ++++++++++ .../shared/services/job_lifecycle_sync.py | 2 +- .../shared/services/quota/__init__.py | 5 + .../shared/services/quota/token_pool.py | 368 +++++++++++++++ .../shared/services/redis/__init__.py | 64 ++- .../shared/services/redis/distributed_lock.py | 2 +- .../services/redis/job_info_redis_service.py | 2 +- .../services/redis/job_metadata_service.py | 2 +- .../shared/services/redis/key_builder.py | 287 ++++++++++++ .../services/redis/rate_limit_service.py | 2 +- .../shared/services/redis/redis_monitor.py | 2 +- .../shared/services/redis/redis_service.py | 2 +- .../services/redis/redis_sync_service.py | 2 +- .../shared/services/redis/retry_policy.py | 146 ++++++ .../services/redis/task_redis_service.py | 2 +- .../services/redis/user_redis_service.py | 2 +- .../shared/services/retrieval/llm_adapter.py | 6 +- .../services/storage/job_file_storage.py | 4 +- .../services/webhook/delivery_client.py | 4 +- .../services/webhook/qstash_publisher.py | 2 +- .../utils/OpenAICompatibleClientSync.py | 429 +----------------- .../shared/utils/ali_quota_manager.py | 66 +-- .../shared/utils/http_clients.py | 87 +--- .../shared/utils/iloveapi_quota_manager.py | 220 +-------- .../shared-python/shared/utils/llm_mock.py | 158 +------ .../shared/utils/pinned_outbound_http.py | 340 ++------------ .../shared/utils/quota_manager.py | 369 +-------------- .../shared/utils/redis_key_builder.py | 292 +----------- .../shared-python/shared/utils/redis_retry.py | 147 +----- .../shared/utils/url_file_type.py | 220 +-------- .../shared/utils/url_security.py | 262 +---------- 80 files changed, 2876 insertions(+), 2558 deletions(-) rename packages/shared-python/shared/utils/device_utils.py => apps/worker/app/services/common/device_checks.py (100%) rename {packages/shared-python/shared/utils => apps/worker/app/services/common}/file_loading.py (100%) rename {packages/shared-python/shared/utils => apps/worker/app/services/common}/file_utils.py (100%) rename packages/shared-python/shared/utils/math_utils.py => apps/worker/app/services/common/math_helpers.py (100%) rename packages/shared-python/shared/utils/gc_utils.py => apps/worker/app/services/common/resource_cleanup.py (100%) create mode 100644 packages/shared-python/shared/services/ai/ali_quota_manager.py create mode 100644 packages/shared-python/shared/services/ai/iloveapi_quota_manager.py create mode 100644 packages/shared-python/shared/services/ai/llm_mock.py create mode 100644 packages/shared-python/shared/services/ai/openai_compatible_client_sync.py create mode 100644 packages/shared-python/shared/services/http/__init__.py create mode 100644 packages/shared-python/shared/services/http/client_pool.py create mode 100644 packages/shared-python/shared/services/http/pinned_outbound.py create mode 100644 packages/shared-python/shared/services/http/url_file_type.py create mode 100644 packages/shared-python/shared/services/http/url_security.py create mode 100644 packages/shared-python/shared/services/quota/__init__.py create mode 100644 packages/shared-python/shared/services/quota/token_pool.py create mode 100644 packages/shared-python/shared/services/redis/key_builder.py create mode 100644 packages/shared-python/shared/services/redis/retry_policy.py diff --git a/AGENTS.md b/AGENTS.md index eed149425..a52642f81 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -50,7 +50,10 @@ knowhereapi-main/ │ ├── services/retrieval/ # Core retrieval engine │ ├── services/chunks/ # DataFrame → ChunkPayload conversion │ ├── services/ai/ # LLM prompt service & AI client -│ └── utils/ # Text, file, and chunk utilities +│ ├── services/http/ # Public URL validation and outbound HTTP +│ ├── services/redis/ # Redis state, key language, and retry policy +│ ├── services/quota/ # Shared token-pool quota primitives +│ └── utils/ # Compatibility and generic text/chunk utilities └── deploy/ # Docker Compose & deployment scripts ``` diff --git a/CONTEXT.md b/CONTEXT.md index ba829e0c4..9b69831a5 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -208,6 +208,24 @@ delivery logs. The verified async callback used to continue background work after external delivery. +### Public URL Policy + +The shared URL safety workflow used before Knowhere reaches user-provided or +third-party HTTP targets. It validates public HTTP/HTTPS URLs, pins resolved +addresses for outbound requests, blocks unsafe redirects, and detects URL file +types for Document Ingestion. + +### Redis State + +The shared Redis-backed runtime state used by background work, rate limits, +state-machine progress, distributed locks, and job metadata. It owns Redis key +language and Redis retry policy. + +### Quota Token Pool + +The shared Redis-backed token leasing workflow used by provider-specific quota +managers such as Ali, iLoveAPI, and MinerU. + ## apps/api Module Map ### HTTP Adapters diff --git a/apps/api/app/services/document_ingestion/creation_service.py b/apps/api/app/services/document_ingestion/creation_service.py index a76977b65..acecb1974 100644 --- a/apps/api/app/services/document_ingestion/creation_service.py +++ b/apps/api/app/services/document_ingestion/creation_service.py @@ -29,7 +29,7 @@ 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.services.http.url_file_type import resolve_file_extension_async _JOB_TYPE_KB_MANAGEMENT = "kb_management" JobMetadata = dict[str, object] diff --git a/apps/api/app/services/document_ingestion/service.py b/apps/api/app/services/document_ingestion/service.py index defd2fbca..03f5da8fe 100644 --- a/apps/api/app/services/document_ingestion/service.py +++ b/apps/api/app/services/document_ingestion/service.py @@ -35,8 +35,8 @@ from shared.core.exceptions.webhook_exceptions import WebhookConfigException from shared.models.schemas.job import ConfirmUploadRequest, JobCreate, JobResponse from shared.models.schemas.job_metadata import JobMetadataHelper -from shared.utils.url_file_type import resolve_file_extension_async -from shared.utils.url_security import validate_http_url_and_resolve_ip_async +from shared.services.http.url_file_type import resolve_file_extension_async +from shared.services.http.url_security import validate_http_url_and_resolve_ip_async JobMetadata = dict[str, object] diff --git a/apps/api/app/services/s3_events/subscription_service.py b/apps/api/app/services/s3_events/subscription_service.py index 140b2eba3..c64ab78e5 100644 --- a/apps/api/app/services/s3_events/subscription_service.py +++ b/apps/api/app/services/s3_events/subscription_service.py @@ -3,8 +3,8 @@ 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 +from shared.services.http.pinned_outbound import send_pinned_outbound_request +from shared.services.http.url_security import validate_http_url_and_resolve_ip_async SNS_SUBSCRIPTION_TIMEOUT_SECONDS = 10 diff --git a/apps/api/main.py b/apps/api/main.py index bd7dec810..4b0dfde44 100644 --- a/apps/api/main.py +++ b/apps/api/main.py @@ -92,7 +92,7 @@ async def lifespan(app: FastAPI): logger.error(f"retrieval hit stats drain failed: {e}") try: - from shared.utils.http_clients import close_async_client + from shared.services.http.client_pool import close_async_client await close_async_client() except Exception as e: diff --git a/apps/api/tests/contract/test_job_creation_contract.py b/apps/api/tests/contract/test_job_creation_contract.py index edea174ec..b56f482f6 100644 --- a/apps/api/tests/contract/test_job_creation_contract.py +++ b/apps/api/tests/contract/test_job_creation_contract.py @@ -632,11 +632,11 @@ def resolve_public_address( 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 + import shared.services.http.client_pool as client_pool_module monkeypatch.setattr(socket, "getaddrinfo", resolve_public_address) monkeypatch.setattr( - http_clients_module, + client_pool_module, "get_async_client", lambda: _FakeAsyncHttpClient(), ) @@ -926,9 +926,9 @@ async def head( assert follow_redirects is False return _FakeHeadResponse() - import shared.utils.http_clients as http_clients_module + import shared.services.http.client_pool as client_pool_module monkeypatch.setattr( - http_clients_module, + client_pool_module, "get_async_client", lambda: _FakeAsyncHttpClient(), ) diff --git a/apps/api/tests/contract/test_s3_event_contract.py b/apps/api/tests/contract/test_s3_event_contract.py index b8e8561de..65c1c6378 100644 --- a/apps/api/tests/contract/test_s3_event_contract.py +++ b/apps/api/tests/contract/test_s3_event_contract.py @@ -210,7 +210,7 @@ def get(self, url: str, *args: object, **kwargs: object) -> object: async with api_client_factory() as api_client: monkeypatch.setattr(socket, "getaddrinfo", resolve_private_address) pinned_http_module = importlib.import_module( - "shared.utils.pinned_outbound_http" + "shared.services.http.pinned_outbound" ) monkeypatch.setattr( pinned_http_module.aiohttp, diff --git a/apps/worker/app/core/worker_bootstrap.py b/apps/worker/app/core/worker_bootstrap.py index 191f2fa27..2689bb71c 100644 --- a/apps/worker/app/core/worker_bootstrap.py +++ b/apps/worker/app/core/worker_bootstrap.py @@ -74,7 +74,7 @@ def shutdown_worker(**kwargs) -> None: logger.warning(f"Worker heartbeat cleanup failed: {exc}") try: - from shared.utils.http_clients import close_sync_client + from shared.services.http.client_pool import close_sync_client close_sync_client() logger.info("Worker sync HTTP client closed") diff --git a/apps/worker/app/services/common/__init__.py b/apps/worker/app/services/common/__init__.py index ffb53d823..05787f01d 100644 --- a/apps/worker/app/services/common/__init__.py +++ b/apps/worker/app/services/common/__init__.py @@ -1,9 +1,9 @@ """Common worker services.""" -from shared.utils.device_utils import check_internet -from shared.utils.file_utils import clean_file, path_handle -from shared.utils.gc_utils import gc_collect as _gc -from shared.utils.math_utils import min_max_normalize +from app.services.common.device_checks import check_internet +from app.services.common.file_utils import clean_file, path_handle +from app.services.common.resource_cleanup import gc_collect as _gc +from app.services.common.math_helpers import min_max_normalize from shared.utils.text_utils import ( count_cn_en, remove_duplicates_orderkept, diff --git a/packages/shared-python/shared/utils/device_utils.py b/apps/worker/app/services/common/device_checks.py similarity index 100% rename from packages/shared-python/shared/utils/device_utils.py rename to apps/worker/app/services/common/device_checks.py diff --git a/packages/shared-python/shared/utils/file_loading.py b/apps/worker/app/services/common/file_loading.py similarity index 100% rename from packages/shared-python/shared/utils/file_loading.py rename to apps/worker/app/services/common/file_loading.py diff --git a/packages/shared-python/shared/utils/file_utils.py b/apps/worker/app/services/common/file_utils.py similarity index 100% rename from packages/shared-python/shared/utils/file_utils.py rename to apps/worker/app/services/common/file_utils.py diff --git a/packages/shared-python/shared/utils/math_utils.py b/apps/worker/app/services/common/math_helpers.py similarity index 100% rename from packages/shared-python/shared/utils/math_utils.py rename to apps/worker/app/services/common/math_helpers.py diff --git a/packages/shared-python/shared/utils/gc_utils.py b/apps/worker/app/services/common/resource_cleanup.py similarity index 100% rename from packages/shared-python/shared/utils/gc_utils.py rename to apps/worker/app/services/common/resource_cleanup.py diff --git a/apps/worker/app/services/connect_builder/summary_builder.py b/apps/worker/app/services/connect_builder/summary_builder.py index 7fcf8c6eb..44f472076 100644 --- a/apps/worker/app/services/connect_builder/summary_builder.py +++ b/apps/worker/app/services/connect_builder/summary_builder.py @@ -40,7 +40,7 @@ def _llm_summarize(snippets_text: str, node_name: str, max_tokens: int = 100) -> """ try: from shared.services.ai.prompt_service import build_prompt, _detect_text_language - from shared.utils.OpenAICompatibleClientSync import get_openai_client + from shared.services.ai.openai_compatible_client_sync import get_openai_client # Deterministic language lock — see prompt_service._language_directive detected_lang = _detect_text_language(snippets_text) diff --git a/apps/worker/app/services/document_agent/tools/classify_special_pages.py b/apps/worker/app/services/document_agent/tools/classify_special_pages.py index 564307b21..0b8efa6c8 100644 --- a/apps/worker/app/services/document_agent/tools/classify_special_pages.py +++ b/apps/worker/app/services/document_agent/tools/classify_special_pages.py @@ -141,7 +141,7 @@ def classify_special_pages( try: from shared.core.config import settings - from shared.utils.OpenAICompatibleClientSync import get_openai_client + from shared.services.ai.openai_compatible_client_sync import get_openai_client effective_model = model or settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL client = get_openai_client(model=effective_model) diff --git a/apps/worker/app/services/document_agent/tools/probe_vlm_inspect.py b/apps/worker/app/services/document_agent/tools/probe_vlm_inspect.py index 7b4de6396..f8949826d 100644 --- a/apps/worker/app/services/document_agent/tools/probe_vlm_inspect.py +++ b/apps/worker/app/services/document_agent/tools/probe_vlm_inspect.py @@ -71,7 +71,7 @@ def _call_vlm( max_tokens: int = 900, ) -> tuple[str, dict[str, int]]: from shared.core.config import settings - from shared.utils.OpenAICompatibleClientSync import get_openai_client + from shared.services.ai.openai_compatible_client_sync import get_openai_client effective_model = model or settings.IMAGE_MODEL or "qwen3.5-flash" client = get_openai_client(model=effective_model) diff --git a/apps/worker/app/services/document_agent/tools/propose_shard_plan.py b/apps/worker/app/services/document_agent/tools/propose_shard_plan.py index 2eb834ff4..1db72592b 100644 --- a/apps/worker/app/services/document_agent/tools/propose_shard_plan.py +++ b/apps/worker/app/services/document_agent/tools/propose_shard_plan.py @@ -295,7 +295,7 @@ def propose_shard_plan( if use_llm and page_count > max_pages_per_shard: try: from shared.core.config import settings - from shared.utils.OpenAICompatibleClientSync import get_openai_client + from shared.services.ai.openai_compatible_client_sync import get_openai_client effective_model = model or settings.HIERARCHY_LLM_MODEL or settings.NORMOL_MODEL client = get_openai_client(model=effective_model) diff --git a/apps/worker/app/services/document_parser/atlas_classifier.py b/apps/worker/app/services/document_parser/atlas_classifier.py index 1ee8dc4ff..ed58ad49f 100644 --- a/apps/worker/app/services/document_parser/atlas_classifier.py +++ b/apps/worker/app/services/document_parser/atlas_classifier.py @@ -102,7 +102,7 @@ def _png_to_data_url(path: str) -> Optional[str]: def _call_vlm(image_data_urls: list[str]) -> bool: """Call VLM with preview images. Returns True if atlas, False otherwise.""" from shared.core.config import settings - from shared.utils.OpenAICompatibleClientSync import get_openai_client + from shared.services.ai.openai_compatible_client_sync import get_openai_client model = settings.IMAGE_MODEL or "qwen-vl-plus" client = get_openai_client(model=model) diff --git a/apps/worker/app/services/document_parser/doc_parser.py b/apps/worker/app/services/document_parser/doc_parser.py index d405fed56..eadcf238f 100755 --- a/apps/worker/app/services/document_parser/doc_parser.py +++ b/apps/worker/app/services/document_parser/doc_parser.py @@ -38,8 +38,8 @@ 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.file_loading import load_file_bytes -from shared.utils.file_utils import path_handle +from app.services.common.file_loading import load_file_bytes +from app.services.common.file_utils import path_handle from shared.utils.text_utils import tokenize2stw_remove 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 1770a4163..589aba7a3 100644 --- a/apps/worker/app/services/document_parser/excel_table_parser.py +++ b/apps/worker/app/services/document_parser/excel_table_parser.py @@ -28,8 +28,8 @@ 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 app.services.common.file_loading import load_file_bytes +from app.services.common.file_utils import path_handle from shared.utils.text_utils import tokenize2stw_remove diff --git a/apps/worker/app/services/document_parser/fragment_parser.py b/apps/worker/app/services/document_parser/fragment_parser.py index 5aef9f004..420d6b084 100644 --- a/apps/worker/app/services/document_parser/fragment_parser.py +++ b/apps/worker/app/services/document_parser/fragment_parser.py @@ -13,8 +13,8 @@ from openai.types.chat import ChatCompletionMessageParam from shared.core.config import settings -from shared.utils.file_utils import path_handle -from shared.utils.OpenAICompatibleClientSync import get_openai_client +from app.services.common.file_utils import path_handle +from shared.services.ai.openai_compatible_client_sync import get_openai_client def generate_fragment_title(content: str, max_tokens: int = 30) -> Optional[str]: diff --git a/apps/worker/app/services/document_parser/image_parser.py b/apps/worker/app/services/document_parser/image_parser.py index cc9be0ab1..726fb6a99 100755 --- a/apps/worker/app/services/document_parser/image_parser.py +++ b/apps/worker/app/services/document_parser/image_parser.py @@ -23,9 +23,9 @@ 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 is_remote, load_file_bytes -from shared.utils.file_utils import path_handle -from shared.utils.OpenAICompatibleClientSync import ( +from app.services.common.file_loading import is_remote, load_file_bytes +from app.services.common.file_utils import path_handle +from shared.services.ai.openai_compatible_client_sync import ( OpenAICompatibleClientSync, get_openai_client, ) diff --git a/apps/worker/app/services/document_parser/internal_parse_name.py b/apps/worker/app/services/document_parser/internal_parse_name.py index 785fe6ff5..568dc13c6 100644 --- a/apps/worker/app/services/document_parser/internal_parse_name.py +++ b/apps/worker/app/services/document_parser/internal_parse_name.py @@ -4,7 +4,7 @@ import os from dataclasses import dataclass -from shared.utils.file_utils import path_handle +from app.services.common.file_utils import path_handle @dataclass(frozen=True) diff --git a/apps/worker/app/services/document_parser/layout_parser.py b/apps/worker/app/services/document_parser/layout_parser.py index b0490d576..ec7775904 100755 --- a/apps/worker/app/services/document_parser/layout_parser.py +++ b/apps/worker/app/services/document_parser/layout_parser.py @@ -35,7 +35,7 @@ from shared.services.ai.response_process_service import eval_response # ARQ dependency is removed, use Celery instead -from shared.utils.OpenAICompatibleClientSync import get_openai_client +from shared.services.ai.openai_compatible_client_sync import get_openai_client # ==================== Helper Functions ==================== 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 270303e90..49327122c 100644 --- a/apps/worker/app/services/document_parser/markdown_deferred_summary.py +++ b/apps/worker/app/services/document_parser/markdown_deferred_summary.py @@ -24,7 +24,7 @@ from shared.core.config import settings from shared.utils.chunk_refs import build_chunk_ref -from shared.utils.file_utils import path_handle +from app.services.common.file_utils import path_handle DeferredResult = ( tuple[ diff --git a/apps/worker/app/services/document_parser/markdown_image_asset.py b/apps/worker/app/services/document_parser/markdown_image_asset.py index 10c93b5dc..6672bc178 100644 --- a/apps/worker/app/services/document_parser/markdown_image_asset.py +++ b/apps/worker/app/services/document_parser/markdown_image_asset.py @@ -16,7 +16,7 @@ from loguru import logger from shared.utils.chunk_refs import build_chunk_ref -from shared.utils.file_utils import path_handle +from app.services.common.file_utils import path_handle @dataclass(frozen=True) diff --git a/apps/worker/app/services/document_parser/markdown_table_asset.py b/apps/worker/app/services/document_parser/markdown_table_asset.py index a01e8ce97..c0985b7a9 100644 --- a/apps/worker/app/services/document_parser/markdown_table_asset.py +++ b/apps/worker/app/services/document_parser/markdown_table_asset.py @@ -15,7 +15,7 @@ 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 +from app.services.common.file_utils import path_handle @dataclass(frozen=True) diff --git a/apps/worker/app/services/document_parser/md_parser.py b/apps/worker/app/services/document_parser/md_parser.py index e284ea2c8..7aea710fa 100755 --- a/apps/worker/app/services/document_parser/md_parser.py +++ b/apps/worker/app/services/document_parser/md_parser.py @@ -215,7 +215,7 @@ def parse_md( relative_root=None, ): if md_lines is None and file_path is not None: - from shared.utils.file_loading import is_remote, load_file_bytes + from app.services.common.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 bd4acdbc4..8be565518 100644 --- a/apps/worker/app/services/document_parser/mineru_pdf_service.py +++ b/apps/worker/app/services/document_parser/mineru_pdf_service.py @@ -23,7 +23,7 @@ UnavailableException, ) from shared.services.storage.job_file_storage import JobFileStorage -from shared.utils.file_loading import is_remote +from app.services.common.file_loading import is_remote MINERU_UPLOAD_TIMEOUT = ( settings.MINERU_UPLOAD_CONNECT_TIMEOUT, diff --git a/apps/worker/app/services/document_parser/mineru_quota_manager.py b/apps/worker/app/services/document_parser/mineru_quota_manager.py index 830330b09..f6fe88f3c 100644 --- a/apps/worker/app/services/document_parser/mineru_quota_manager.py +++ b/apps/worker/app/services/document_parser/mineru_quota_manager.py @@ -14,7 +14,7 @@ SyncRedisService, SyncRedisServiceFactory, ) -from shared.utils.quota_manager import BaseQuotaManager, TokenConfig, TokenLease +from shared.services.quota.token_pool import BaseQuotaManager, TokenConfig, TokenLease # Backward-compatible aliases so existing imports keep working MinerUTokenConfig = TokenConfig 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 832e13969..9750a331d 100644 --- a/apps/worker/app/services/document_parser/orchestration/parse_session.py +++ b/apps/worker/app/services/document_parser/orchestration/parse_session.py @@ -12,7 +12,7 @@ from shared.core.config import settings from shared.core.exceptions.domain_exceptions import ValidationException -from shared.utils.file_utils import path_handle +from app.services.common.file_utils import path_handle PDF_PAGE_LIMIT = 600 diff --git a/apps/worker/app/services/document_parser/path_helpers.py b/apps/worker/app/services/document_parser/path_helpers.py index 41965e125..4119ff53e 100644 --- a/apps/worker/app/services/document_parser/path_helpers.py +++ b/apps/worker/app/services/document_parser/path_helpers.py @@ -6,7 +6,7 @@ from bs4 import BeautifulSoup from shared.utils.chunk_refs import extract_chunk_refs -from shared.utils.file_utils import path_handle +from app.services.common.file_utils import path_handle SUMMARY_PATH_MARKERS: tuple[str, ...] = ("summary", "\u6458\u8981\u603b\u7ed3") diff --git a/apps/worker/app/services/document_parser/pptx_parser.py b/apps/worker/app/services/document_parser/pptx_parser.py index 19063937a..8329fb18a 100755 --- a/apps/worker/app/services/document_parser/pptx_parser.py +++ b/apps/worker/app/services/document_parser/pptx_parser.py @@ -26,8 +26,8 @@ FileSystemException, ) from shared.core.logging import LogEvent -from shared.utils.file_loading import load_file_bytes -from shared.utils.file_utils import path_handle +from app.services.common.file_loading import load_file_bytes +from app.services.common.file_utils import path_handle # ==================== LibreOffice conversion ==================== @@ -56,7 +56,7 @@ def pptx_to_pdf_libreoffice(pptx_path, outdir="."): def _get_iloveapi_token_lease(): """acquire iLoveAPI token lease from the quotas pool and generate a JWT token""" - from shared.utils.iloveapi_quota_manager import get_iloveapi_quota_manager + from shared.services.ai.iloveapi_quota_manager import get_iloveapi_quota_manager quota_manager = get_iloveapi_quota_manager() @@ -111,7 +111,7 @@ def _pptx_bytes_to_pdf_bytes(pptx_bytes: bytes, filename: str) -> bytes: Acquires an in-flight slot before starting. Slots are released only if the Redis-backed reservation actually succeeded. """ - from shared.utils.iloveapi_quota_manager import get_iloveapi_quota_manager + from shared.services.ai.iloveapi_quota_manager import get_iloveapi_quota_manager quota_manager = get_iloveapi_quota_manager() diff --git a/apps/worker/app/services/document_parser/table_frame_parser.py b/apps/worker/app/services/document_parser/table_frame_parser.py index 63f39feb8..a71cabf65 100644 --- a/apps/worker/app/services/document_parser/table_frame_parser.py +++ b/apps/worker/app/services/document_parser/table_frame_parser.py @@ -15,7 +15,7 @@ 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.services.ai.openai_compatible_client_sync import get_openai_client from shared.utils.text_utils import remove_duplicates_orderkept diff --git a/apps/worker/app/services/document_parser/toc_parser.py b/apps/worker/app/services/document_parser/toc_parser.py index 2efb1a4cd..5c52ef5c4 100644 --- a/apps/worker/app/services/document_parser/toc_parser.py +++ b/apps/worker/app/services/document_parser/toc_parser.py @@ -21,7 +21,7 @@ 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.services.ai.openai_compatible_client_sync import get_openai_client # ==================== Markdown TOC Detection Functions ==================== diff --git a/apps/worker/app/services/document_parser/txt_parser.py b/apps/worker/app/services/document_parser/txt_parser.py index e1d41956b..161e09864 100755 --- a/apps/worker/app/services/document_parser/txt_parser.py +++ b/apps/worker/app/services/document_parser/txt_parser.py @@ -12,8 +12,8 @@ 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.file_loading import load_file_bytes -from shared.utils.OpenAICompatibleClientSync import get_openai_client +from app.services.common.file_loading import load_file_bytes +from shared.services.ai.openai_compatible_client_sync import get_openai_client def clean_texts_by_form(text, form="html"): diff --git a/apps/worker/app/services/workload/url_upload_transfer.py b/apps/worker/app/services/workload/url_upload_transfer.py index 25f4faf13..c9da336a8 100644 --- a/apps/worker/app/services/workload/url_upload_transfer.py +++ b/apps/worker/app/services/workload/url_upload_transfer.py @@ -10,7 +10,7 @@ ValidationException, ) from shared.services.storage.job_file_storage import JobFileStorage -from shared.utils.url_file_type import resolve_file_extension_sync +from shared.services.http.url_file_type import resolve_file_extension_sync def resolve_supported_url_extension(source_url: str) -> str: diff --git a/packages/shared-python/shared/core/state_machine/service.py b/packages/shared-python/shared/core/state_machine/service.py index 6457e3964..1206313b9 100644 --- a/packages/shared-python/shared/core/state_machine/service.py +++ b/packages/shared-python/shared/core/state_machine/service.py @@ -38,7 +38,7 @@ 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.redis_key_builder import RedisKeyType, redis_key_builder +from shared.services.redis.key_builder import RedisKeyType, redis_key_builder class AsyncStateMachineService: 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 8935304e9..6d8ba1cc1 100644 --- a/packages/shared-python/shared/core/state_machine/service_sync.py +++ b/packages/shared-python/shared/core/state_machine/service_sync.py @@ -37,7 +37,7 @@ 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.redis_key_builder import RedisKeyType, redis_key_builder +from shared.services.redis.key_builder import RedisKeyType, redis_key_builder class SyncStateMachineService: diff --git a/packages/shared-python/shared/services/ai/ali_quota_manager.py b/packages/shared-python/shared/services/ai/ali_quota_manager.py new file mode 100644 index 000000000..37bcc2e78 --- /dev/null +++ b/packages/shared-python/shared/services/ai/ali_quota_manager.py @@ -0,0 +1,65 @@ +""" +AliQuotaManager — Redis-backed token pool for Aliyun DashScope API keys. + +Distributes requests across multiple ALI_API_KEYS using the shared +BaseQuotaManager, with per-token RPM / daily quotas and 429 cooldown. +""" + +from __future__ import annotations + +from typing import List, Optional + +from shared.core.config import settings +from shared.services.redis.redis_sync_service import ( + SyncRedisService, + SyncRedisServiceFactory, +) +from shared.services.quota.token_pool import BaseQuotaManager, TokenConfig + + +class AliQuotaManager(BaseQuotaManager): + """Quota-aware Ali DashScope token manager backed by Redis.""" + + SERVICE_PREFIX = "ali" + CURSOR_KEY = "ali:quota:cursor" + user_message = "AI service is busy right now. Please retry shortly." + + @classmethod + def from_settings( + cls, redis_service: Optional[SyncRedisService] = None + ) -> "AliQuotaManager": + """Create an AliQuotaManager from application settings.""" + tokens = cls.parse_tokens_from_settings() + instance = cls(redis_service or SyncRedisServiceFactory.get_service(), tokens) + instance.default_cooldown_seconds = settings.ALI_TOKEN_COOLDOWN_SECONDS + return instance + + @staticmethod + def parse_tokens_from_settings() -> List[TokenConfig]: + default_rpm_limit: int = settings.ALI_TOKEN_RPM_LIMIT + default_daily_limit: int = settings.ALI_TOKEN_DAILY_LIMIT + raw_pool: str = (settings.ALI_API_KEYS or "").strip() + + if raw_pool: + parsed = BaseQuotaManager.parse_token_specs( + raw_pool, + default_rpm_limit=default_rpm_limit, + default_daily_limit=default_daily_limit, + ) + if parsed: + return parsed + + raise ValueError( + "No Ali API keys configured. Set ALI_API_KEYS to a comma/newline-separated list or JSON array." + ) + + +_ali_quota_manager: Optional[AliQuotaManager] = None + + +def get_ali_quota_manager() -> AliQuotaManager: + """Return a singleton AliQuotaManager for worker processes.""" + global _ali_quota_manager + if _ali_quota_manager is None: + _ali_quota_manager = AliQuotaManager.from_settings() + return _ali_quota_manager diff --git a/packages/shared-python/shared/services/ai/iloveapi_quota_manager.py b/packages/shared-python/shared/services/ai/iloveapi_quota_manager.py new file mode 100644 index 000000000..8cea78bb5 --- /dev/null +++ b/packages/shared-python/shared/services/ai/iloveapi_quota_manager.py @@ -0,0 +1,215 @@ +""" +ILoveApiQuotaManager — Redis-backed project pool for iLoveAPI keys. + +Each iLoveAPI public_key+secret_key pair represents a separate project +with its own credit quota (free tier: 2,500 credits/month, officepdf +costs 10 credits/file = ~250 files/month). + +Distributes requests across multiple projects using the shared +BaseQuotaManager, with RPM burst protection, daily file limits, +429 cooldown, and a global in-flight concurrency limiter to prevent +worker pool starvation when iLoveAPI is slow or rate-limited. + +This is an optional enhancement path — all failures fail open to +LibreOffice local conversion. +""" + +from __future__ import annotations + +import json +from typing import List, Optional + +from loguru import logger + +from shared.core.config import settings +from shared.services.redis.redis_sync_service import ( + SyncRedisService, + SyncRedisServiceFactory, +) +from shared.services.quota.token_pool import BaseQuotaManager, TokenConfig + + +class ILoveApiQuotaManager(BaseQuotaManager): + """Quota-aware iLoveAPI token manager backed by Redis.""" + + SERVICE_PREFIX = "iloveapi" + CURSOR_KEY = "iloveapi:quota:cursor" + user_message = "Document parsing service is busy right now. Please retry shortly." + + INFLIGHT_KEY = "iloveapi:inflight" + INFLIGHT_TTL_SECONDS = 300 # 5-minute safety net; auto-expires stale counts + + _ACQUIRE_INFLIGHT_SCRIPT: str = """ +local key = KEYS[1] +local max_concurrent = tonumber(ARGV[1]) +local ttl_seconds = tonumber(ARGV[2]) + +local current = tonumber(redis.call('GET', key) or '0') +if current >= max_concurrent then + return {0, current} +end + +local new_count = redis.call('INCR', key) +-- Set TTL only if not already set (safety net for crash recovery) +if redis.call('TTL', key) < 0 then + redis.call('EXPIRE', key, ttl_seconds) +end + +return {1, new_count} +""" + + _RELEASE_INFLIGHT_SCRIPT: str = """ +local key = KEYS[1] +local current = tonumber(redis.call('GET', key) or '0') +if current <= 0 then + redis.call('DEL', key) + return 0 +end +return redis.call('DECR', key) +""" + + def __init__( + self, + redis_service: SyncRedisService, + tokens: List[TokenConfig], + max_concurrent: int = 5, + ) -> None: + super().__init__(redis_service, tokens) + self.max_concurrent = max_concurrent + + @classmethod + def from_settings( + cls, redis_service: Optional[SyncRedisService] = None + ) -> "ILoveApiQuotaManager": + """Create an ILoveApiQuotaManager from application settings.""" + tokens = cls.parse_tokens_from_settings() + max_concurrent: int = getattr(settings, "ILOVEAPI_MAX_CONCURRENT", 5) + instance = cls( + redis_service or SyncRedisServiceFactory.get_service(), + tokens, + max_concurrent=max_concurrent, + ) + instance.default_cooldown_seconds = 60 + return instance + + @staticmethod + def parse_tokens_from_settings() -> List[TokenConfig]: + default_rpm_limit: int = getattr(settings, "ILOVEAPI_TOKEN_RPM_LIMIT", 25) + default_daily_limit: int = getattr(settings, "ILOVEAPI_TOKEN_DAILY_LIMIT", 250) + raw_pool: str = (getattr(settings, "ILOVEAPI_KEYS", "") or "").strip() + + legacy_pub: str = (getattr(settings, "ILOVEAPI_PUBLIC_KEY", "") or "").strip() + legacy_sec: str = (getattr(settings, "ILOVEAPI_SECRET_KEY", "") or "").strip() + + specs: List[TokenConfig] = [] + if raw_pool: + try: + loaded = json.loads(raw_pool) + if isinstance(loaded, list): + for index, entry in enumerate(loaded): + pub = str(entry.get("public_key") or "").strip() + sec = str(entry.get("secret_key") or "").strip() + if pub and sec: + specs.append( + TokenConfig( + token_id=str( + entry.get("token_id") or f"iloveapi-{index + 1}" + ), + api_key=f"{pub}:{sec}", + rpm_limit=int( + entry.get("rpm_limit") or default_rpm_limit + ), + daily_limit=int( + entry.get("daily_limit") or default_daily_limit + ), + ) + ) + except json.JSONDecodeError: + pass + + if not specs and legacy_pub and legacy_sec: + specs.append( + TokenConfig( + token_id="iloveapi-default", + api_key=f"{legacy_pub}:{legacy_sec}", + rpm_limit=default_rpm_limit, + daily_limit=default_daily_limit, + ) + ) + + if not specs: + raise ValueError( + "No iLoveAPI keys configured. Set ILOVEAPI_KEYS to a JSON array containing public_key and secret_key." + ) + + return specs + + # ------------------------------------------------------------------ + # In-flight concurrency limiter + # ------------------------------------------------------------------ + + def acquire_inflight(self) -> Optional[bool]: + """Try to acquire an in-flight slot. + + Returns: + True: reserved a Redis-backed slot and must release it later. + False: at capacity; caller should fail open to local conversion. + None: Redis failed; caller may proceed, but must not release. + """ + try: + result = self.redis.eval( + self._ACQUIRE_INFLIGHT_SCRIPT, + keys=[self.INFLIGHT_KEY], + args=[self.max_concurrent, self.INFLIGHT_TTL_SECONDS], + ) + if isinstance(result, list) and len(result) >= 2: + acquired = int(result[0]) == 1 + current_count = int(result[1]) + if not acquired: + logger.bind( + service=self.SERVICE_PREFIX, + step="concurrency_exceeded", + current_inflight=current_count, + max_concurrent=self.max_concurrent, + ).warning( + f"iLoveAPI concurrency limit reached ({current_count}/{self.max_concurrent})" + ) + return acquired + except Exception: + logger.opt(exception=True).warning( + "Failed to check iLoveAPI in-flight counter; failing open to allow request" + ) + # Fail open on Redis errors — allow the request through without a lease. + return None + + def release_inflight(self) -> None: + """Release an in-flight slot after conversion completes (success or failure).""" + try: + self.redis.eval( + self._RELEASE_INFLIGHT_SCRIPT, + keys=[self.INFLIGHT_KEY], + args=[], + ) + except Exception: + logger.opt(exception=True).warning( + "Failed to release iLoveAPI in-flight slot" + ) + + def get_inflight_count(self) -> int: + """Return current in-flight count (for observability / testing).""" + try: + raw = self.redis.get(self.INFLIGHT_KEY, 0) + return max(0, int(raw or 0)) + except Exception: + return 0 + + +_iloveapi_quota_manager: Optional[ILoveApiQuotaManager] = None + + +def get_iloveapi_quota_manager() -> ILoveApiQuotaManager: + """Return a singleton ILoveApiQuotaManager for worker processes.""" + global _iloveapi_quota_manager + if _iloveapi_quota_manager is None: + _iloveapi_quota_manager = ILoveApiQuotaManager.from_settings() + return _iloveapi_quota_manager diff --git a/packages/shared-python/shared/services/ai/llm_mock.py b/packages/shared-python/shared/services/ai/llm_mock.py new file mode 100644 index 000000000..8bf155f21 --- /dev/null +++ b/packages/shared-python/shared/services/ai/llm_mock.py @@ -0,0 +1,157 @@ +"""Helpers for deterministic mock responses from OpenAI-compatible LLM calls.""" + +from typing import Any, Dict, List + +from loguru import logger +from openai.types.chat import ChatCompletionMessageParam + + +def build_mock_chat_completion_response( + messages: List[ChatCompletionMessageParam], + model_name: str, +) -> str: + """Return a canned response compatible with the inferred prompt contract.""" + prompt_text = _extract_prompt_text(messages) + task_name = _detect_mock_task(prompt_text) + logger.debug( + "LLM mock enabled, returning canned response: model={}, task={}", + model_name, + task_name, + ) + return _build_mock_response(task_name) + + +def _extract_message_text(content: Any) -> str: + """Flatten chat message content into plain text for mock-task detection.""" + if isinstance(content, str): + return content + + if not isinstance(content, list): + return "" + + parts: List[str] = [] + for item in content: + if isinstance(item, str): + parts.append(item) + continue + + if not isinstance(item, dict): + continue + + item_type = item.get("type") + if item_type == "text" and isinstance(item.get("text"), str): + parts.append(item["text"]) + continue + + if isinstance(item.get("content"), str): + parts.append(item["content"]) + + return "\n".join(parts) + + +def _extract_prompt_text(messages: List[ChatCompletionMessageParam]) -> str: + """Extract all text segments from a chat-completion message list.""" + parts: List[str] = [] + + for message in messages: + if not isinstance(message, dict): + continue + + text_part = _extract_message_text(message.get("content")) + if text_part: + parts.append(text_part) + + return "\n".join(parts) + + +def _detect_mock_task(prompt_text: str) -> str: + """Infer the prompt task so the mock can return a compatible response shape.""" + normalized_prompt = prompt_text.lower() + + if ( + "generate a concise title" in normalized_prompt + and "return only the title" in normalized_prompt + ): + return "fragment-title" + if ( + '"toc_start"' in normalized_prompt + and '"toc_end"' in normalized_prompt + and '"confidence"' in normalized_prompt + ): + return "detect-toc-range" + if ( + "row numbers where headers are located" in normalized_prompt + and '"answer"' in normalized_prompt + ): + return "detect-table-headers" + if '"answer"' in normalized_prompt and '"text" or "image"' in normalized_prompt: + return "judge-image-type" + if ( + '"title"' in normalized_prompt + and '"keywords"' in normalized_prompt + and '"summary"' in normalized_prompt + ): + return "summary-full" + if ( + 'json dictionary format with key "answer"' in normalized_prompt + and "keywords" in normalized_prompt + ): + return "summary-keywords" + if ( + "json array only" in normalized_prompt + and "toc, not body text" in normalized_prompt + ): + return "eval-toc-headings" + if ( + "json array" in normalized_prompt + and "only include rows that you judge to be headings" in normalized_prompt + ): + return "eval-headings" + if "scanned page from an engineering atlas" in normalized_prompt: + return "atlas-page-info" + if "perform ocr operation" in normalized_prompt: + return "ocr-image" + if ( + "you will receive an image" in normalized_prompt + and "line 1: output a short title" in normalized_prompt + ): + return "summary-images" + if ( + "you will receive one or more images and the user's current question" + in normalized_prompt + ): + return "ask-image" + if "summaries of sub-sections from a document section" in normalized_prompt: + return "file-summary" + if ( + "line 1: output a short title" in normalized_prompt + and "line 2 onward" in normalized_prompt + ): + return "summary-titled" + if "extract the main content of the material" in normalized_prompt: + return "summary" + + return "default" + + +def _build_mock_response(task_name: str) -> str: + """Return a canned response compatible with the inferred task contract.""" + response_by_task: Dict[str, str] = { + "fragment-title": "Mock Fragment Title", + "detect-toc-range": '{"toc_start": null, "toc_end": null, "confidence": "low"}', + "detect-table-headers": '{"answer": [0]}', + "judge-image-type": '{"answer": "image"}', + "summary-full": '{"title": "Mock Title", "keywords": "mock", "summary": "Mock summary"}', + "summary-keywords": '{"answer": "mock"}', + "eval-headings": "[]", + "eval-toc-headings": "[]", + "atlas-page-info": "Mock atlas page info", + "ocr-image": "Mock OCR text", + "summary-images": "Mock Image Title\nMock image summary", + "ask-image": "Mock image answer", + "file-summary": "Mock section summary", + "summary-titled": "Mock Title\nMock summary", + "summary": "Mock summary", + "default": "Mock LLM response", + } + return response_by_task.get(task_name, response_by_task["default"]) diff --git a/packages/shared-python/shared/services/ai/openai_compatible_client_sync.py b/packages/shared-python/shared/services/ai/openai_compatible_client_sync.py new file mode 100644 index 000000000..60ea29ae3 --- /dev/null +++ b/packages/shared-python/shared/services/ai/openai_compatible_client_sync.py @@ -0,0 +1,425 @@ +""" +Synchronous OpenAI-compatible client for gevent worker tasks. +Uses the official OpenAI SDK with shared httpx connection pooling +and built-in 429 retry with exponential backoff. + +For Aliyun (qwen) models, api keys are drawn from the AliQuotaManager +token pool and automatically rotated on 429 RateLimitError. +""" +import os +import threading +from typing import Any, Dict, List, Optional, Union + +import openai +from loguru import logger +from openai import OpenAI +from openai.types.chat import ChatCompletionMessageParam + +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import LLMServiceException +from shared.services.http.client_pool import get_sync_client +from shared.services.ai.llm_mock import build_mock_chat_completion_response +from shared.utils.security_utils import mask_api_key + +LOCAL_DEBUG = os.getenv("LOCAL_DEBUG", "0") == "1" +LLMUsage = dict[str, int] + +_client_cache: Dict[tuple, "OpenAICompatibleClientSync"] = {} +_client_cache_lock = threading.Lock() + + +def _is_ali_model(model_name: str) -> bool: + """Check whether a model name routes to Aliyun DashScope.""" + return "qwen" in (model_name or "").lower() + + +def _should_mock_llm_calls() -> bool: + """Whether all OpenAI-compatible LLM calls should short-circuit to mock responses.""" + return bool(getattr(settings, "LLM_MOCK_ENABLED", False)) + + +def _empty_usage() -> LLMUsage: + return {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + + +def _extract_usage(response: Any) -> LLMUsage: + usage = getattr(response, "usage", None) + if usage is None: + return _empty_usage() + return { + "prompt_tokens": int(getattr(usage, "prompt_tokens", 0) or 0), + "completion_tokens": int(getattr(usage, "completion_tokens", 0) or 0), + "total_tokens": int(getattr(usage, "total_tokens", 0) or 0), + } + + +def _summarize_exception_chain(exc: Exception, *, max_depth: int = 4) -> str: + parts: list[str] = [] + current: BaseException | None = exc + seen: set[int] = set() + + while current is not None and len(parts) < max_depth and id(current) not in seen: + seen.add(id(current)) + parts.append(f"{type(current).__name__}: {current}") + current = current.__cause__ or current.__context__ + + return " <- ".join(parts) + + +class OpenAICompatibleClientSync: + """Sync OpenAI-compatible client backed by the official OpenAI SDK.""" + + def __init__( + self, + api_key: Optional[str] = None, + api_url: Optional[str] = None, + default_model: Optional[str] = None, + timeout: int = 300, + max_retries: int = 2, + ): + self.default_model = default_model or getattr(settings, "NORMOL_MODEL", "deepseek-chat") + self._explicit_api_key = api_key + self._explicit_api_url = api_url + self._max_retries = max_retries + self._base_url: Optional[str] = self._resolve_base_url( + model_name=self.default_model, + api_url=api_url, + ) + + self.timeout = ( + getattr(settings, "OPENAI_CLIENT_TIMEOUT", 300) + if timeout is None + else timeout + ) + self._client: Optional[OpenAI] = None + if _should_mock_llm_calls(): + return + if not self._should_use_ali_pool(): + resolved_key: Optional[str] = self._resolve_direct_api_key( + model_name=self.default_model, + api_key=api_key, + ) + self._client = self._build_client( + api_key=resolved_key, + base_url=self._base_url, + max_retries=max_retries, + ) + + def _build_client( + self, + api_key: Optional[str], + base_url: Optional[str], + max_retries: int, + ) -> OpenAI: + """Build a direct OpenAI-compatible SDK client.""" + return OpenAI( + api_key=api_key, + base_url=base_url, + http_client=get_sync_client(), + max_retries=max_retries, + timeout=self.timeout, + ) + + @staticmethod + def _strip_chat_completions(url: Optional[str]) -> Optional[str]: + """Strip /chat/completions suffix so the OpenAI SDK can append it.""" + if url and url.rstrip("/").endswith("/chat/completions"): + return url.rstrip("/").removesuffix("/chat/completions") + return url + + def _resolve_base_url( + self, + model_name: str, + api_url: Optional[str], + ) -> Optional[str]: + if api_url: + return self._strip_chat_completions(api_url) + + model_lower = (model_name or "").lower() + if "qwen" in model_lower: + ali_base = getattr(settings, "ALI_URL", "https://dashscope.aliyuncs.com/compatible-mode/v1") + return self._strip_chat_completions(ali_base) + + if "glm" in model_lower: + glm_base = getattr(settings, "GLM_URL", "https://open.bigmodel.cn/api/paas/v4") + return self._strip_chat_completions(glm_base) + + if "doubao" in model_lower or model_lower.startswith("ep-"): + return self._strip_chat_completions(getattr(settings, "ARK_URL", None)) + + return self._strip_chat_completions(settings.DS_URL) + + def _resolve_direct_api_key( + self, + model_name: str, + api_key: Optional[str], + ) -> Optional[str]: + if api_key: + return api_key + + model_lower = (model_name or "").lower() + if "glm" in model_lower: + return getattr(settings, "GLM_API_KEY", None) + + if "doubao" in model_lower or model_lower.startswith("ep-"): + return getattr(settings, "ARK_API_KEY", None) + + return settings.DS_KEY + + # ------------------------------------------------------------------ + # Ali token-pool helpers + # ------------------------------------------------------------------ + + def _should_use_ali_pool(self) -> bool: + """Whether to route through the AliQuotaManager instead of a fixed key.""" + if self._explicit_api_key: + return False + return _is_ali_model(self.default_model) + + def _make_ali_pool_call( + self, + model: str, + all_messages: List[ChatCompletionMessageParam], + temperature: float, + max_tokens: int, + api_kwargs: Dict[str, Any], + ) -> tuple[str, LLMUsage]: + """Acquire a token, make the call, and retry inline on 429.""" + from shared.services.ai.ali_quota_manager import get_ali_quota_manager + + quota_manager = get_ali_quota_manager() + base_url: Optional[str] = self._base_url + + max_retries = settings.ALI_INLINE_MAX_RETRIES + for attempt in range(max_retries): + lease = quota_manager.acquire_request(operation="chat_completion") + try: + # Hybrid rate-limit control: SDK handles per-key backoff (exp. backoff + jitter), + # outer loop handles token rotation across the pool on persistent 429s. + client = self._build_client( + api_key=lease.api_key, + base_url=base_url, + max_retries=settings.ALI_SDK_MAX_RETRIES, + ) + response = client.chat.completions.create( + model=model, + messages=all_messages, + temperature=temperature, + max_tokens=max_tokens, + **api_kwargs, + ) + choices = response.choices + if not choices: + raise LLMServiceException( + internal_message="AI returned empty result", + provider=self.default_model, + ) + return choices[0].message.content or "", _extract_usage(response) + except openai.RateLimitError as exc: + retry_after = _parse_retry_after(exc) + quota_manager.mark_rate_limited(lease.token_id, retry_after) + logger.warning( + f"Ali token {lease.token_id} rate-limited (attempt {attempt + 1}/{max_retries}), " + f"cooling down {retry_after}s and retrying with next token" + ) + if attempt == max_retries - 1: + raise LLMServiceException( + internal_message=f"All Ali tokens exhausted after {max_retries} retries: {exc}", + provider=self.default_model, + original_exception=exc, + ) from exc + except LLMServiceException: + raise + except Exception as exc: + masked_api_key = mask_api_key(lease.api_key) + logger.error( + "LLM request failed (Ali pool): model={model}, token_id={token_id}, api_key={api_key}, error={error}", + model=model, + token_id=lease.token_id, + api_key=masked_api_key, + error=exc, + ) + raise LLMServiceException( + internal_message=( + "API request failed " + f"(Ali pool, token_id={lease.token_id}, api_key={masked_api_key}): {exc}" + ), + provider=self.default_model, + original_exception=exc, + ) from exc + # unreachable but satisfies the type checker + raise LLMServiceException( + internal_message="Ali token pool retry loop exited unexpectedly", + provider=self.default_model, + ) + + # ------------------------------------------------------------------ + + def chat_completion_with_usage( + self, + messages: Union[str, List[ChatCompletionMessageParam]], + model: Optional[str] = None, + temperature: float = 0.1, + max_tokens: int = 4096, + top_p: Optional[float] = None, + timeout: Optional[int] = None, + **kwargs, + ) -> tuple[str, LLMUsage]: + all_messages: List[ChatCompletionMessageParam] + if isinstance(messages, list): + all_messages = messages # type: ignore[assignment] + else: + all_messages = [{"role": "user", "content": str(messages)}] + + api_kwargs: Dict[str, Any] = {} + if top_p is not None: + api_kwargs["top_p"] = top_p + if timeout is not None: + api_kwargs["timeout"] = timeout + + allowed_api_params = { + "n", "stop", "presence_penalty", "frequency_penalty", + "logit_bias", "user", "seed", "tools", "tool_choice", + "response_format", "logprobs", "top_logprobs", + } + for key, value in kwargs.items(): + if key in allowed_api_params: + api_kwargs[key] = value + + # ── disable thinking mode ── + # Qwen3.5 by default enables thinking mode, which wastes tokens by outputting ... + # Explicitly disable it in all API calls + extra_body = api_kwargs.get("extra_body", {}) + if isinstance(extra_body, dict): + extra_body.setdefault("enable_thinking", False) + else: + extra_body = {"enable_thinking": False} + api_kwargs["extra_body"] = extra_body + + effective_model = model or self.default_model + if _should_mock_llm_calls(): + return build_mock_chat_completion_response( + messages=all_messages, + model_name=effective_model, + ), _empty_usage() + + # Route through Ali token pool when applicable + if self._should_use_ali_pool(): + try: + return self._make_ali_pool_call( + model=effective_model, + all_messages=all_messages, + temperature=temperature, + max_tokens=max_tokens, + api_kwargs=api_kwargs, + ) + except LLMServiceException: + raise + except Exception as exc: + logger.error(f"LLM request failed (Ali pool): model={effective_model}, error={exc}") + raise LLMServiceException( + internal_message=f"API request failed: {str(exc)}", + provider=self.default_model, + original_exception=exc, + ) from exc + + # Non-Ali path: use the single pre-configured client + client = self._client + if client is None: + raise LLMServiceException( + internal_message="OpenAI client is not initialized for direct provider requests", + provider=self.default_model, + ) + + try: + response = client.chat.completions.create( + model=effective_model, + messages=all_messages, + temperature=temperature, + max_tokens=max_tokens, + **api_kwargs, + ) + + choices = response.choices + if not choices: + raise LLMServiceException( + internal_message="AI returned empty result", + provider=self.default_model, + ) + + content = choices[0].message.content or "" + return content, _extract_usage(response) + except LLMServiceException: + raise + except Exception as exc: + logger.error( + "LLM request failed: model={model}, base_url={base_url}, error_chain={error_chain}", + model=effective_model, + base_url=client.base_url, + error_chain=_summarize_exception_chain(exc), + ) + raise LLMServiceException( + internal_message=f"API request failed: {_summarize_exception_chain(exc)}", + provider=self.default_model, + original_exception=exc, + ) from exc + + def chat_completion( + self, + messages: Union[str, List[ChatCompletionMessageParam]], + model: Optional[str] = None, + temperature: float = 0.1, + max_tokens: int = 4096, + top_p: Optional[float] = None, + timeout: Optional[int] = None, + **kwargs, + ) -> str: + content, _usage = self.chat_completion_with_usage( + messages=messages, + model=model, + temperature=temperature, + max_tokens=max_tokens, + top_p=top_p, + timeout=timeout, + **kwargs, + ) + return content + + +def _parse_retry_after(exc: openai.RateLimitError) -> int: + """Extract Retry-After seconds from a RateLimitError, with sane bounds.""" + try: + if hasattr(exc, "response") and exc.response is not None: + header_value = exc.response.headers.get("retry-after") or exc.response.headers.get("Retry-After") + if header_value: + return max(1, min(int(header_value), 120)) + except (ValueError, TypeError, AttributeError): + pass + return settings.ALI_TOKEN_COOLDOWN_SECONDS + + +def get_openai_client( + model: Optional[str] = None, + api_key: Optional[str] = None, + api_url: Optional[str] = None, + timeout: int = 300, + max_retries: int = 2, +) -> OpenAICompatibleClientSync: + """Return a cached OpenAICompatibleClientSync instance keyed by config.""" + cache_key = (model, api_key, api_url, timeout, max_retries) + client = _client_cache.get(cache_key) + if client is not None: + return client + with _client_cache_lock: + client = _client_cache.get(cache_key) + if client is not None: + return client + client = OpenAICompatibleClientSync( + api_key=api_key, + api_url=api_url, + default_model=model, + timeout=timeout, + max_retries=max_retries, + ) + _client_cache[cache_key] = client + return client diff --git a/packages/shared-python/shared/services/http/__init__.py b/packages/shared-python/shared/services/http/__init__.py new file mode 100644 index 000000000..e80b56b4c --- /dev/null +++ b/packages/shared-python/shared/services/http/__init__.py @@ -0,0 +1,123 @@ +"""Public URL and outbound HTTP policy service exports.""" + +from __future__ import annotations + +from importlib import import_module +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from shared.services.http.client_pool import ( + close_async_client, + close_sync_client, + get_async_client, + get_sync_client, + ) + from shared.services.http.pinned_outbound import ( + PinnedDownloadResult, + PinnedHTTPConnection, + PinnedHTTPConnectionPool, + PinnedHTTPSConnection, + PinnedHTTPSConnectionPool, + PinnedIPResolver, + PinnedOutboundResponse, + download_pinned_outbound_file, + download_pinned_outbound_file_async, + send_pinned_outbound_request, + ) + from shared.services.http.url_file_type import ( + CONTENT_TYPE_TO_EXTENSION, + MAX_SAFE_REDIRECTS, + REDIRECT_STATUS_CODES, + URL_VALIDATION_DESCRIPTIONS, + resolve_file_extension_async, + resolve_file_extension_sync, + ) + from shared.services.http.url_security import ( + BLOCKED_HOSTNAMES, + DEVELOPMENT_ENVIRONMENTS, + HTTP_SCHEMES, + HTTPURLValidationResult, + SafePublicHTTPURL, + URLValidationFailureReason, + validate_http_url_and_resolve_ip, + validate_http_url_and_resolve_ip_async, + ) + +__all__ = [ + "BLOCKED_HOSTNAMES", + "CONTENT_TYPE_TO_EXTENSION", + "DEVELOPMENT_ENVIRONMENTS", + "HTTP_SCHEMES", + "HTTPURLValidationResult", + "MAX_SAFE_REDIRECTS", + "PinnedDownloadResult", + "PinnedHTTPConnection", + "PinnedHTTPConnectionPool", + "PinnedHTTPSConnection", + "PinnedHTTPSConnectionPool", + "PinnedIPResolver", + "PinnedOutboundResponse", + "REDIRECT_STATUS_CODES", + "SafePublicHTTPURL", + "URLValidationFailureReason", + "URL_VALIDATION_DESCRIPTIONS", + "close_async_client", + "close_sync_client", + "download_pinned_outbound_file", + "download_pinned_outbound_file_async", + "get_async_client", + "get_sync_client", + "resolve_file_extension_async", + "resolve_file_extension_sync", + "send_pinned_outbound_request", + "validate_http_url_and_resolve_ip", + "validate_http_url_and_resolve_ip_async", +] + +_EXPORT_MODULES: dict[str, str] = { + "BLOCKED_HOSTNAMES": "shared.services.http.url_security", + "CONTENT_TYPE_TO_EXTENSION": "shared.services.http.url_file_type", + "DEVELOPMENT_ENVIRONMENTS": "shared.services.http.url_security", + "HTTP_SCHEMES": "shared.services.http.url_security", + "HTTPURLValidationResult": "shared.services.http.url_security", + "MAX_SAFE_REDIRECTS": "shared.services.http.url_file_type", + "PinnedDownloadResult": "shared.services.http.pinned_outbound", + "PinnedHTTPConnection": "shared.services.http.pinned_outbound", + "PinnedHTTPConnectionPool": "shared.services.http.pinned_outbound", + "PinnedHTTPSConnection": "shared.services.http.pinned_outbound", + "PinnedHTTPSConnectionPool": "shared.services.http.pinned_outbound", + "PinnedIPResolver": "shared.services.http.pinned_outbound", + "PinnedOutboundResponse": "shared.services.http.pinned_outbound", + "REDIRECT_STATUS_CODES": "shared.services.http.url_file_type", + "SafePublicHTTPURL": "shared.services.http.url_security", + "URLValidationFailureReason": "shared.services.http.url_security", + "URL_VALIDATION_DESCRIPTIONS": "shared.services.http.url_file_type", + "close_async_client": "shared.services.http.client_pool", + "close_sync_client": "shared.services.http.client_pool", + "download_pinned_outbound_file": "shared.services.http.pinned_outbound", + "download_pinned_outbound_file_async": "shared.services.http.pinned_outbound", + "get_async_client": "shared.services.http.client_pool", + "get_sync_client": "shared.services.http.client_pool", + "resolve_file_extension_async": "shared.services.http.url_file_type", + "resolve_file_extension_sync": "shared.services.http.url_file_type", + "send_pinned_outbound_request": "shared.services.http.pinned_outbound", + "validate_http_url_and_resolve_ip": "shared.services.http.url_security", + "validate_http_url_and_resolve_ip_async": "shared.services.http.url_security", +} + + +def __getattr__(name: str) -> Any: + """Load HTTP service exports on first access.""" + module_name = _EXPORT_MODULES.get(name) + if module_name is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + module = import_module(module_name) + value = getattr(module, name) + globals()[name] = value + return value + + +def __dir__() -> list[str]: + """Return public HTTP package exports.""" + return sorted([*globals(), *__all__]) diff --git a/packages/shared-python/shared/services/http/client_pool.py b/packages/shared-python/shared/services/http/client_pool.py new file mode 100644 index 000000000..1ffa03893 --- /dev/null +++ b/packages/shared-python/shared/services/http/client_pool.py @@ -0,0 +1,77 @@ +""" +Singleton HTTP clients with connection pooling. + +Reduces socket churn under high concurrency by reusing TCP and TLS connections +across requests. +""" + +import threading +from typing import Optional + +import httpx + +_DEFAULT_LIMITS = httpx.Limits( + max_keepalive_connections=20, + keepalive_expiry=30, +) + +_DEFAULT_TIMEOUT = httpx.Timeout( + connect=10.0, + read=300.0, + write=60.0, + pool=30.0, +) + +_sync_client: Optional[httpx.Client] = None +_sync_lock = threading.Lock() + + +def get_sync_client() -> httpx.Client: + """Return a shared sync httpx.Client with connection pooling.""" + global _sync_client + if _sync_client is None: + with _sync_lock: + if _sync_client is None: + _sync_client = httpx.Client( + limits=_DEFAULT_LIMITS, + timeout=_DEFAULT_TIMEOUT, + follow_redirects=True, + ) + return _sync_client + + +def close_sync_client() -> None: + """Close the shared sync client. Call on worker shutdown.""" + global _sync_client + with _sync_lock: + if _sync_client is not None: + _sync_client.close() + _sync_client = None + + +_async_client: Optional[httpx.AsyncClient] = None +_async_lock = threading.Lock() + + +def get_async_client() -> httpx.AsyncClient: + """Return a shared async httpx.AsyncClient with connection pooling.""" + global _async_client + if _async_client is None: + with _async_lock: + if _async_client is None: + _async_client = httpx.AsyncClient( + limits=_DEFAULT_LIMITS, + timeout=_DEFAULT_TIMEOUT, + follow_redirects=True, + ) + return _async_client + + +async def close_async_client() -> None: + """Close the shared async client. Call on API shutdown.""" + global _async_client + with _async_lock: + client = _async_client + _async_client = None + if client is not None: + await client.aclose() diff --git a/packages/shared-python/shared/services/http/pinned_outbound.py b/packages/shared-python/shared/services/http/pinned_outbound.py new file mode 100644 index 000000000..c281dceb5 --- /dev/null +++ b/packages/shared-python/shared/services/http/pinned_outbound.py @@ -0,0 +1,313 @@ +"""Pinned outbound HTTP helpers.""" + +from __future__ import annotations + +import os +import socket +import tempfile +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any +from urllib.parse import urlsplit + +import aiohttp +from aiohttp.abc import AbstractResolver +from urllib3 import Retry +from urllib3.connection import HTTPConnection, HTTPSConnection +from urllib3.connectionpool import HTTPConnectionPool, HTTPSConnectionPool +from urllib3.util import Timeout +from urllib3.util.connection import create_connection + +from shared.core.exceptions.domain_exceptions import ValidationException +from shared.services.http.url_security import SafePublicHTTPURL + + +@dataclass(frozen=True) +class PinnedOutboundResponse: + """Minimal response metadata for pinned outbound HTTP requests.""" + + status: int + + +@dataclass(frozen=True) +class PinnedDownloadResult: + """Metadata for a pinned HTTP download to a temporary file.""" + + status: int + temp_file_path: str + + +class PinnedIPResolver(AbstractResolver): + """Resolver that always returns the supplied pinned IP address.""" + + def __init__(self, pinned_ip: str) -> None: + self.pinned_ip = pinned_ip + + async def resolve( + self, + host: str, + port: int = 0, + family: int = socket.AF_INET, + ) -> list[dict[str, Any]]: + pinned_family: int = ( + socket.AF_INET6 if ":" in self.pinned_ip else socket.AF_INET + ) + return [ + { + "hostname": host, + "host": self.pinned_ip, + "port": port, + "family": pinned_family, + "proto": 0, + "flags": socket.AI_NUMERICHOST, + } + ] + + async def close(self) -> None: + pass + + +class PinnedHTTPConnection(HTTPConnection): + """HTTP connection that resolves the hostname to a pinned IP address.""" + + def __init__(self, *args: Any, pinned_ip: str, **kwargs: Any) -> None: + self._pinned_ip = pinned_ip + super().__init__(*args, **kwargs) + + def _new_conn(self) -> socket.socket: + return create_connection( + (self._pinned_ip, self.port), + self.timeout, + source_address=self.source_address, + socket_options=self.socket_options, + ) + + +class PinnedHTTPSConnection(HTTPSConnection): + """HTTPS connection that resolves the hostname to a pinned IP address.""" + + def __init__(self, *args: Any, pinned_ip: str, **kwargs: Any) -> None: + self._pinned_ip = pinned_ip + super().__init__(*args, **kwargs) + + def _new_conn(self) -> socket.socket: + return create_connection( + (self._pinned_ip, self.port), + self.timeout, + source_address=self.source_address, + socket_options=self.socket_options, + ) + + +class PinnedHTTPConnectionPool(HTTPConnectionPool): + """HTTP pool that pins DNS resolution to a fixed IP address.""" + + ConnectionCls = PinnedHTTPConnection # pyright: ignore[reportAssignmentType] + + def __init__(self, *args: Any, pinned_ip: str, **kwargs: Any) -> None: + kwargs["pinned_ip"] = pinned_ip + super().__init__(*args, **kwargs) + + +class PinnedHTTPSConnectionPool(HTTPSConnectionPool): + """HTTPS pool that pins DNS resolution to a fixed IP address.""" + + ConnectionCls = PinnedHTTPSConnection # pyright: ignore[reportAssignmentType] + + def __init__(self, *args: Any, pinned_ip: str, **kwargs: Any) -> None: + kwargs["pinned_ip"] = pinned_ip + super().__init__(*args, **kwargs) + + +def _build_host_header(parsed_url: Any) -> str: + hostname = parsed_url.hostname + if not hostname: + return "" + + if ":" in hostname and not hostname.startswith("["): + formatted_host = f"[{hostname}]" + else: + formatted_host = hostname + + if parsed_url.port is not None: + return f"{formatted_host}:{parsed_url.port}" + return formatted_host + + +def _validate_download_url(url: str, field: str) -> SafePublicHTTPURL: + parsed_url = urlsplit(url) + if parsed_url.scheme not in {"http", "https"}: + raise ValidationException( + user_message="Invalid URL", + violations=[ + {"field": field, "description": "URL must use http or https"} + ], + ) + if not parsed_url.hostname: + raise ValidationException( + user_message="Invalid URL", + violations=[{"field": field, "description": "URL must include a hostname"}], + ) + return SafePublicHTTPURL(url) + + +async def send_pinned_outbound_request( + *, + method: str, + url: str, + pinned_ip: str, + timeout_seconds: float, + headers: Mapping[str, str] | None = None, + json_body: Any | None = None, +) -> PinnedOutboundResponse: + """Send an outbound HTTP request through a resolver pinned to a validated IP.""" + validated_url = _validate_download_url(url, "url") + connector = aiohttp.TCPConnector(resolver=PinnedIPResolver(pinned_ip)) + timeout = aiohttp.ClientTimeout(total=timeout_seconds) + + async with aiohttp.ClientSession( + connector=connector, + timeout=timeout, + ) as session: + async with session.request( + method=method, + url=validated_url, + headers=headers, + json=json_body, + allow_redirects=False, + ) as response: + return PinnedOutboundResponse(status=response.status) + + +async def download_pinned_outbound_file_async( + *, + url: str, + pinned_ip: str, + timeout_seconds: float, + user_agent: str, + temp_dir: str | None = None, + field: str = "source_url", +) -> PinnedDownloadResult: + """Download a file with aiohttp while pinning DNS to a validated IP.""" + validated_url = _validate_download_url(url, field) + + temp_file_descriptor, temp_file_path = tempfile.mkstemp(dir=temp_dir) + os.close(temp_file_descriptor) + + connector = aiohttp.TCPConnector(resolver=PinnedIPResolver(pinned_ip)) + timeout = aiohttp.ClientTimeout(total=timeout_seconds) + + try: + async with aiohttp.ClientSession( + connector=connector, + timeout=timeout, + ) as session: + async with session.get( + validated_url, + allow_redirects=False, + headers={"User-Agent": user_agent}, + ) as response: + if not 200 <= response.status < 300: + raise ValidationException( + user_message="Invalid URL", + violations=[ + { + "field": field, + "description": f"URL request failed with status {response.status}", + } + ], + ) + + with open(temp_file_path, "wb") as output_file: + async for chunk in response.content.iter_chunked(65536): + if chunk: + output_file.write(chunk) + + return PinnedDownloadResult( + status=response.status, + temp_file_path=temp_file_path, + ) + except Exception: + if os.path.exists(temp_file_path): + os.remove(temp_file_path) + raise + + +def download_pinned_outbound_file( + *, + url: str, + pinned_ip: str, + timeout_seconds: float, + user_agent: str, + temp_dir: str | None = None, + field: str = "source_url", +) -> PinnedDownloadResult: + """ + Download a file through a resolver pinned to a validated IP. + + Redirects are blocked to prevent redirect-based SSRF. The request connects + to the pre-validated IP while preserving the original Host header. + """ + validated_url = _validate_download_url(url, field) + parsed_url = urlsplit(validated_url) + + temp_file_descriptor, temp_file_path = tempfile.mkstemp(dir=temp_dir) + os.close(temp_file_descriptor) + + try: + connection_pool: HTTPConnectionPool + request_path = parsed_url.path or "/" + if parsed_url.query: + request_path = f"{request_path}?{parsed_url.query}" + + if parsed_url.scheme == "https": + connection_pool = PinnedHTTPSConnectionPool( + parsed_url.hostname, + parsed_url.port or 443, + pinned_ip=pinned_ip, + retries=Retry(total=0, redirect=False), + ) + else: + connection_pool = PinnedHTTPConnectionPool( + parsed_url.hostname, + parsed_url.port or 80, + pinned_ip=pinned_ip, + retries=Retry(total=0, redirect=False), + ) + + response = connection_pool.urlopen( + "GET", + request_path, + timeout=Timeout.from_float(timeout_seconds), + preload_content=False, + redirect=False, + headers={"User-Agent": user_agent, "Host": _build_host_header(parsed_url)}, + ) + try: + if not 200 <= response.status < 300: + raise ValidationException( + user_message="Invalid URL", + violations=[ + { + "field": field, + "description": f"URL request failed with status {response.status}", + } + ], + ) + + with open(temp_file_path, "wb") as output_file: + for chunk in response.stream(65536): + if chunk: + output_file.write(chunk) + finally: + response.release_conn() + response.close() + + return PinnedDownloadResult( + status=response.status, + temp_file_path=temp_file_path, + ) + except Exception: + if os.path.exists(temp_file_path): + os.remove(temp_file_path) + raise diff --git a/packages/shared-python/shared/services/http/url_file_type.py b/packages/shared-python/shared/services/http/url_file_type.py new file mode 100644 index 000000000..14fd2a16d --- /dev/null +++ b/packages/shared-python/shared/services/http/url_file_type.py @@ -0,0 +1,195 @@ +""" +Resolve file extension from a URL. + +Tries the URL path first, then falls back to a HEAD request to read Content-Type. +Provides both async and sync variants. +""" + +import os +from urllib.parse import urljoin, urlparse + +from loguru import logger + +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import ValidationException +from shared.services.http.url_security import ( + HTTPURLValidationResult, + SafePublicHTTPURL, + validate_http_url_and_resolve_ip, +) + +CONTENT_TYPE_TO_EXTENSION: dict[str, str] = { + "application/pdf": ".pdf", + "application/msword": ".doc", + "application/vnd.openxmlformats-officedocument.wordprocessingml.document": ".docx", + "application/vnd.ms-excel": ".xls", + "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": ".xlsx", + "application/vnd.ms-powerpoint": ".ppt", + "application/vnd.openxmlformats-officedocument.presentationml.presentation": ".pptx", + "text/csv": ".csv", + "text/plain": ".txt", + "text/markdown": ".md", + "image/png": ".png", + "image/jpeg": ".jpg", + "image/gif": ".gif", + "image/bmp": ".bmp", + "image/tiff": ".tiff", + "image/svg+xml": ".svg", +} + +REDIRECT_STATUS_CODES: set[int] = {301, 302, 303, 307, 308} +MAX_SAFE_REDIRECTS: int = 5 +URL_VALIDATION_DESCRIPTIONS: dict[str, str] = { + "unsupported_scheme": "URL must use http or https", + "missing_hostname": "URL must include a hostname", + "hostname_resolution_failed": "URL hostname could not be resolved", + "invalid_resolved_address": "URL resolved to an invalid IP", + "hostname_not_allowed": "URL host is not allowed", +} + + +def _validate_source_url(url: str, field: str) -> SafePublicHTTPURL: + validation = validate_http_url_and_resolve_ip(url) + if not validation.is_valid: + _raise_url_validation_error(field, validation) + return SafePublicHTTPURL(url) + + +def _validate_redirect_url( + url: str, + redirect_url: str, + field: str, +) -> SafePublicHTTPURL: + return _validate_source_url(urljoin(url, redirect_url), field=field) + + +def _raise_url_validation_error( + field: str, + validation: HTTPURLValidationResult, +) -> None: + description = URL_VALIDATION_DESCRIPTIONS.get( + str(validation.failure_reason), + "URL is invalid", + ) + + raise ValidationException( + user_message="Invalid URL", + violations=[{"field": field, "description": description}], + internal_message=validation.error_message, + ) + + +def _extension_from_path(url: str) -> str | None: + """Extract a recognised file extension from the URL path.""" + path = urlparse(url).path + ext = os.path.splitext(path)[1].lower() + if ext and ext in settings.get_supported_extensions(): + return ext + return None + + +def _extension_from_content_type(content_type: str | None) -> str | None: + """Map a Content-Type header value to a supported file extension.""" + if not content_type: + return None + mime = content_type.split(";")[0].strip().lower() + ext = CONTENT_TYPE_TO_EXTENSION.get(mime) + if ext and ext in settings.get_supported_extensions(): + return ext + return None + + +async def resolve_file_extension_async(url: str) -> str | None: + """ + Resolve file extension from a URL. + + 1. Try extracting extension from URL path. + 2. If that fails, send a HEAD request and read Content-Type. + 3. Return None if neither method produces a supported extension. + """ + safe_url = _validate_source_url(url, field="source_url") + + ext = _extension_from_path(safe_url) + if ext: + return ext + + try: + from shared.services.http.client_pool import get_async_client + + client = get_async_client() + response = None + request_url: SafePublicHTTPURL = safe_url + for _ in range(MAX_SAFE_REDIRECTS + 1): + response = await client.head(request_url, follow_redirects=False) + if response.status_code not in REDIRECT_STATUS_CODES: + break + + location = response.headers.get("location") + if not location: + break + request_url = _validate_redirect_url( + request_url, + location, + field="source_url", + ) + + if response is None: + return None + content_type = response.headers.get("content-type") + ext = _extension_from_content_type(content_type) + if ext: + logger.info(f"Resolved file extension from Content-Type header: {ext}") + return ext + except ValidationException: + raise + except Exception as exc: + logger.warning(f"HEAD request failed for URL file type detection: {exc}") + + return None + + +def resolve_file_extension_sync(url: str) -> str | None: + """ + Resolve file extension from a URL. + + Same logic as async variant but uses the shared sync httpx client. + """ + safe_url = _validate_source_url(url, field="source_url") + + ext = _extension_from_path(safe_url) + if ext: + return ext + + try: + from shared.services.http.client_pool import get_sync_client + + client = get_sync_client() + response = None + request_url: SafePublicHTTPURL = safe_url + for _ in range(MAX_SAFE_REDIRECTS + 1): + response = client.head(request_url, follow_redirects=False) + if response.status_code not in REDIRECT_STATUS_CODES: + break + + location = response.headers.get("location") + if not location: + break + request_url = _validate_redirect_url( + request_url, + location, + field="source_url", + ) + + if response is None: + return None + content_type = response.headers.get("content-type") + ext = _extension_from_content_type(content_type) + if ext: + logger.info(f"Resolved file extension from Content-Type header: {ext}") + return ext + except ValidationException: + raise + except Exception as exc: + logger.warning(f"HEAD request failed for URL file type detection: {exc}") + + return None diff --git a/packages/shared-python/shared/services/http/url_security.py b/packages/shared-python/shared/services/http/url_security.py new file mode 100644 index 000000000..bc54c1107 --- /dev/null +++ b/packages/shared-python/shared/services/http/url_security.py @@ -0,0 +1,243 @@ +import asyncio +import ipaddress +import socket +from dataclasses import dataclass +from typing import Literal, TypeAlias, cast +from urllib.parse import urlparse + +SocketAddress: TypeAlias = tuple[object, ...] +AddressInfo: TypeAlias = tuple[int, int, int, str, SocketAddress] +AllowedIPAddress: TypeAlias = ipaddress.IPv4Address | ipaddress.IPv6Address +URLValidationFailureReason: TypeAlias = Literal[ + "unsupported_scheme", + "missing_hostname", + "hostname_resolution_failed", + "invalid_resolved_address", + "hostname_not_allowed", + "validation_failed", +] +HTTP_SCHEMES: frozenset[str] = frozenset({"http", "https"}) +DEVELOPMENT_ENVIRONMENTS: frozenset[str] = frozenset({"dev", "development", "local"}) +BLOCKED_HOSTNAMES: frozenset[str] = frozenset( + {"localhost", "ip6-localhost", "ip6-loopback"} +) + + +class SafePublicHTTPURL(str): + """A URL string that has passed public HTTP SSRF validation.""" + + +@dataclass(frozen=True) +class HTTPURLValidationResult: + """Validated HTTP URL and its pinned IP address.""" + + is_valid: bool + url: str + hostname: str | None = None + validated_ip: str | None = None + error_message: str | None = None + failure_reason: URLValidationFailureReason | None = None + + +def validate_http_url_and_resolve_ip( + url: str, +) -> HTTPURLValidationResult: + """Validate an HTTP URL policy and return the pinned resolved IP address.""" + try: + can_use_private_hosts = _can_use_private_hosts() + hostname_result = _validate_url_hostname(url) + if not hostname_result.is_valid: + return hostname_result + + hostname = cast(str, hostname_result.hostname) + if not can_use_private_hosts and _is_blocked_hostname(hostname): + return _build_url_validation_failure( + url, + f"Hostname {hostname} failed validation", + "hostname_not_allowed", + ) + + try: + address_infos = cast(list[AddressInfo], socket.getaddrinfo(hostname, None)) + except socket.gaierror as exc: + return _build_url_validation_failure( + url, + f"Unable to resolve hostname {hostname}: {exc}", + "hostname_resolution_failed", + ) + + return _validate_resolved_addresses( + url, + hostname, + address_infos, + can_use_private_hosts=can_use_private_hosts, + ) + except Exception as exc: + return _build_url_validation_failure( + url, + f"URL validation failed: {exc}", + "validation_failed", + ) + + +async def validate_http_url_and_resolve_ip_async( + url: str, +) -> HTTPURLValidationResult: + """Validate an HTTP URL policy asynchronously and return the pinned IP address.""" + try: + can_use_private_hosts = _can_use_private_hosts() + hostname_result = _validate_url_hostname(url) + if not hostname_result.is_valid: + return hostname_result + + hostname = cast(str, hostname_result.hostname) + if not can_use_private_hosts and _is_blocked_hostname(hostname): + return _build_url_validation_failure( + url, + f"Hostname {hostname} failed validation", + "hostname_not_allowed", + ) + + try: + loop = asyncio.get_running_loop() + address_infos = cast( + list[AddressInfo], + await loop.getaddrinfo(hostname, None), + ) + except socket.gaierror as exc: + return _build_url_validation_failure( + url, + f"Unable to resolve hostname {hostname}: {exc}", + "hostname_resolution_failed", + ) + + return _validate_resolved_addresses( + url, + hostname, + address_infos, + can_use_private_hosts=can_use_private_hosts, + ) + except Exception as exc: + return _build_url_validation_failure( + url, + f"URL validation failed: {exc}", + "validation_failed", + ) + + +def _can_use_private_hosts() -> bool: + from shared.core.config import app_config + + return app_config.ENVIRONMENT.lower() in DEVELOPMENT_ENVIRONMENTS + + +def _extract_resolved_addresses(address_infos: list[AddressInfo]) -> list[str]: + resolved_addresses: list[str] = [] + seen_addresses: set[str] = set() + for family, _, _, _, socket_address in address_infos: + if family not in (socket.AF_INET, socket.AF_INET6): + continue + if not socket_address: + continue + address = socket_address[0] + if not isinstance(address, str): + continue + if address and address not in seen_addresses: + resolved_addresses.append(address) + seen_addresses.add(address) + return resolved_addresses + + +def _is_blocked_hostname(hostname: str) -> bool: + normalized_hostname = hostname.rstrip(".").lower() + return normalized_hostname in BLOCKED_HOSTNAMES or normalized_hostname.endswith( + (".localhost", ".local") + ) + + +def _is_public_ip_address(ip_address: AllowedIPAddress) -> bool: + return ip_address.is_global and not _is_blocked_ip_address(ip_address) + + +def _is_blocked_ip_address(ip_address: AllowedIPAddress) -> bool: + return ( + ip_address.is_private + or ip_address.is_loopback + or ip_address.is_link_local + or ip_address.is_multicast + or ip_address.is_reserved + or ip_address.is_unspecified + ) + + +def _validate_url_hostname(url: str) -> HTTPURLValidationResult: + parsed_url = urlparse(url) + if parsed_url.scheme not in HTTP_SCHEMES: + return _build_url_validation_failure( + url, + f"Unsupported URL scheme: {parsed_url.scheme}", + "unsupported_scheme", + ) + + hostname = parsed_url.hostname + if not hostname: + return _build_url_validation_failure( + url, + "URL must include a hostname", + "missing_hostname", + ) + + return HTTPURLValidationResult(is_valid=True, url=url, hostname=hostname) + + +def _validate_resolved_addresses( + url: str, + hostname: str, + address_infos: list[AddressInfo], + *, + can_use_private_hosts: bool, +) -> HTTPURLValidationResult: + resolved_addresses = _extract_resolved_addresses(address_infos) + if not resolved_addresses: + return _build_url_validation_failure( + url, + f"Unable to resolve hostname {hostname}", + "hostname_resolution_failed", + ) + + for address in resolved_addresses: + try: + ip_address = ipaddress.ip_address(address) + except ValueError: + return _build_url_validation_failure( + url, + f"Hostname {hostname} resolved to invalid IP {address}", + "invalid_resolved_address", + ) + + if not can_use_private_hosts and not _is_public_ip_address(ip_address): + return _build_url_validation_failure( + url, + f"Hostname {hostname} failed validation", + "hostname_not_allowed", + ) + + return HTTPURLValidationResult( + is_valid=True, + url=url, + hostname=hostname, + validated_ip=resolved_addresses[0], + ) + + +def _build_url_validation_failure( + url: str, + error_message: str, + failure_reason: URLValidationFailureReason, +) -> HTTPURLValidationResult: + return HTTPURLValidationResult( + is_valid=False, + url=url, + error_message=error_message, + failure_reason=failure_reason, + ) diff --git a/packages/shared-python/shared/services/job_lifecycle_sync.py b/packages/shared-python/shared/services/job_lifecycle_sync.py index ad0dcfc70..f160f41af 100644 --- a/packages/shared-python/shared/services/job_lifecycle_sync.py +++ b/packages/shared-python/shared/services/job_lifecycle_sync.py @@ -26,7 +26,7 @@ from shared.services.redis.redis_sync_service import ( SyncRedisServiceFactory, ) -from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder +from shared.services.redis.key_builder import RedisKeyType, redis_key_builder class SyncJobLifecycleService: diff --git a/packages/shared-python/shared/services/quota/__init__.py b/packages/shared-python/shared/services/quota/__init__.py new file mode 100644 index 000000000..749a6d0ec --- /dev/null +++ b/packages/shared-python/shared/services/quota/__init__.py @@ -0,0 +1,5 @@ +"""Shared quota and token-pool services.""" + +from shared.services.quota.token_pool import BaseQuotaManager, TokenConfig, TokenLease + +__all__ = ["BaseQuotaManager", "TokenConfig", "TokenLease"] diff --git a/packages/shared-python/shared/services/quota/token_pool.py b/packages/shared-python/shared/services/quota/token_pool.py new file mode 100644 index 000000000..735f8fa87 --- /dev/null +++ b/packages/shared-python/shared/services/quota/token_pool.py @@ -0,0 +1,368 @@ +""" +BaseQuotaManager — Redis-backed quota management for any token-pool API. + +Design goals: +- Coordinate token usage across workers with Redis +- Enforce per-token RPM and daily quotas before calling the provider +- Keep selection fair with a quota-aware round-robin cursor +- Cool down tokens temporarily when the provider responds with 429 +""" + +from __future__ import annotations + +import json +import time +from dataclasses import dataclass +from typing import Any, Iterable, List, Optional + +from loguru import logger + +from shared.core.exceptions.domain_exceptions import UnavailableException +from shared.services.redis.redis_sync_service import SyncRedisService + + +@dataclass(frozen=True) +class TokenConfig: + """Generic configuration for a single API token.""" + + token_id: str + api_key: str + rpm_limit: int + daily_limit: int + + +@dataclass(frozen=True) +class TokenLease: + """A token reservation returned by acquire_request. + + Callers use ``api_key`` for the Authorization header, and pass + ``token_id`` back to ``mark_rate_limited`` when the provider + signals 429. + """ + + token_id: str + api_key: str + rpm_limit: int + daily_limit: int + + +class BaseQuotaManager: + """Abstract Redis-backed token pool with round-robin and cooldown. + + Subclasses only need to set: + - ``SERVICE_PREFIX`` — Redis key namespace (e.g. ``"ali"`` or ``"mineru"``) + - ``CURSOR_KEY`` — Redis key for the round-robin cursor + - ``default_cooldown_seconds`` — fallback cooldown when no Retry-After header + - ``user_message`` — user-facing error string on exhaustion + """ + + SERVICE_PREFIX: str = "base" + CURSOR_KEY: str = "base:quota:cursor" + + TOKEN_COOLDOWN_TTL_PADDING: int = 60 + MINUTE_WINDOW_SECONDS: int = 60 + DAY_WINDOW_SECONDS: int = 24 * 60 * 60 + + _RESERVE_REQUEST_SCRIPT: str = """ +local minute_limit = tonumber(ARGV[1]) +local day_limit = tonumber(ARGV[2]) +local minute_ttl = tonumber(ARGV[3]) +local day_ttl = tonumber(ARGV[4]) +local now_ts = tonumber(ARGV[5]) + +local cooldown_until = tonumber(redis.call('GET', KEYS[3]) or '0') +if cooldown_until > now_ts then + return {0, 'cooldown', cooldown_until - now_ts} +end + +local minute_count = tonumber(redis.call('GET', KEYS[1]) or '0') +if minute_count >= minute_limit then + local ttl = redis.call('TTL', KEYS[1]) + if ttl < 0 then + ttl = minute_ttl + redis.call('EXPIRE', KEYS[1], minute_ttl) + end + return {0, 'minute', ttl} +end + +local day_count = tonumber(redis.call('GET', KEYS[2]) or '0') +if day_count >= day_limit then + local ttl = redis.call('TTL', KEYS[2]) + if ttl < 0 then + ttl = day_ttl + redis.call('EXPIRE', KEYS[2], day_ttl) + end + return {0, 'day', ttl} +end + +minute_count = redis.call('INCR', KEYS[1]) +if minute_count == 1 then + redis.call('EXPIRE', KEYS[1], minute_ttl) +end + +day_count = redis.call('INCR', KEYS[2]) +if day_count == 1 then + redis.call('EXPIRE', KEYS[2], day_ttl) +end + +local minute_ttl_value = redis.call('TTL', KEYS[1]) +if minute_ttl_value < 0 then + minute_ttl_value = minute_ttl +end + +local day_ttl_value = redis.call('TTL', KEYS[2]) +if day_ttl_value < 0 then + day_ttl_value = day_ttl +end + +return {1, minute_count, day_count, minute_ttl_value, day_ttl_value} +""" + + # ---- overridable defaults ---- + default_cooldown_seconds: int = 60 + user_message: str = "Service is busy right now. Please retry shortly." + + def __init__( + self, redis_service: SyncRedisService, token_configs: Iterable[TokenConfig] + ) -> None: + self.redis = redis_service + self.tokens: List[TokenConfig] = list(token_configs) + if not self.tokens: + raise ValueError(f"{type(self).__name__} requires at least one token") + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def acquire_request( + self, + operation: str, + preferred_token_id: Optional[str] = None, + ) -> TokenLease: + """Reserve request capacity on one token before calling the provider.""" + now_ts = int(time.time()) + candidate_tokens = self._ordered_tokens(preferred_token_id) + earliest_retry_after: Optional[int] = None + retry_period = "minute" + + for token in candidate_tokens: + reserved, retry_after, quota_period = self._reserve_token_capacity( + token, now_ts + ) + if reserved: + logger.debug( + f"Reserved {self.SERVICE_PREFIX} token {token.token_id} " + f"for operation={operation}" + ) + return TokenLease( + token_id=token.token_id, + api_key=token.api_key, + rpm_limit=token.rpm_limit, + daily_limit=token.daily_limit, + ) + retry_after = max(1, retry_after) + if earliest_retry_after is None or retry_after < earliest_retry_after: + earliest_retry_after = retry_after + retry_period = quota_period + + logger.bind( + service=self.SERVICE_PREFIX, + step="quota_exhausted", + operation=operation, + preferred_token_id=preferred_token_id, + retry_after=earliest_retry_after or self.default_cooldown_seconds, + retry_period=retry_period, + token_count=len(self.tokens), + ).warning(f"{self.SERVICE_PREFIX} quota exhausted") + raise UnavailableException( + internal_message=( + f"{self.SERVICE_PREFIX} quota exhausted for operation={operation}; " + f"no token has available capacity right now" + ), + retry_after=earliest_retry_after or self.default_cooldown_seconds, + limit=max(token.rpm_limit for token in self.tokens), + period=retry_period, + user_message=self.user_message, + ) + + def mark_rate_limited( + self, token_id: str, retry_after: Optional[int] = None + ) -> None: + """Temporarily cool a token when the provider signals rate limiting.""" + retry_after_seconds = max(1, retry_after or self.default_cooldown_seconds) + cooldown_until = int(time.time()) + retry_after_seconds + cooldown_key = self._cooldown_key(token_id) + self.redis.set( + cooldown_key, + cooldown_until, + ttl=retry_after_seconds + self.TOKEN_COOLDOWN_TTL_PADDING, + ) + logger.warning( + f"{self.SERVICE_PREFIX} token {token_id} cooling down for " + f"{retry_after_seconds}s (until {cooldown_until})" + ) + + # ------------------------------------------------------------------ + # Token parsing utility (shared across subclasses) + # ------------------------------------------------------------------ + + @staticmethod + def parse_token_specs( + raw_value: str, + *, + default_rpm_limit: int, + default_daily_limit: int, + ) -> List[TokenConfig]: + """Parse token configs from JSON, newline, or comma-separated env values.""" + specs: List[TokenConfig] = [] + + def build_token(entry: Any, index: int) -> TokenConfig: + if isinstance(entry, dict): + api_key = str(entry.get("api_key") or entry.get("key") or "").strip() + if not api_key: + raise ValueError("Token entry missing api_key") + token_id = str( + entry.get("token_id") or entry.get("id") or f"token-{index + 1}" + ) + rpm_limit = int(entry.get("rpm_limit") or default_rpm_limit) + daily_limit = int(entry.get("daily_limit") or default_daily_limit) + return TokenConfig( + token_id=token_id, + api_key=api_key, + rpm_limit=rpm_limit, + daily_limit=daily_limit, + ) + + entry_text = str(entry).strip() + if not entry_text: + raise ValueError("Empty token entry") + + token_id = f"token-{index + 1}" + api_key = entry_text + if "=" in entry_text: + maybe_id, maybe_key = entry_text.split("=", 1) + if maybe_key.strip(): + token_id = maybe_id.strip() or token_id + api_key = maybe_key.strip() + if not api_key: + raise ValueError("Token entry missing api_key") + + return TokenConfig( + token_id=token_id, + api_key=api_key, + rpm_limit=default_rpm_limit, + daily_limit=default_daily_limit, + ) + + try: + loaded = json.loads(raw_value) + except json.JSONDecodeError: + loaded = None + + if isinstance(loaded, list): + for index, entry in enumerate(loaded): + specs.append(build_token(entry, index)) + return specs + + for index, line in enumerate(raw_value.replace(",", "\n").splitlines()): + line = line.strip() + if not line: + continue + specs.append(build_token(line, index)) + + return specs + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _ordered_tokens(self, preferred_token_id: Optional[str]) -> List[TokenConfig]: + if preferred_token_id: + preferred = [ + token for token in self.tokens if token.token_id == preferred_token_id + ] + remaining = [ + token + for token in self._ordered_tokens(None) + if token.token_id != preferred_token_id + ] + return preferred + remaining + + raw_cursor = self.redis.get(self.CURSOR_KEY, 0) + try: + cursor = int(raw_cursor or 0) + except (TypeError, ValueError): + cursor = 0 + + ordered = self.tokens[cursor:] + self.tokens[:cursor] + next_cursor = (cursor + 1) % len(self.tokens) + self.redis.set(self.CURSOR_KEY, next_cursor, ttl=self.DAY_WINDOW_SECONDS) + return ordered + + def _reserve_token_capacity( + self, token: TokenConfig, now_ts: int + ) -> tuple[bool, int, str]: + minute_key = self._minute_key(token.token_id, now_ts) + day_key = self._day_key(token.token_id, now_ts) + cooldown_key = self._cooldown_key(token.token_id) + + result = self.redis.eval( + self._RESERVE_REQUEST_SCRIPT, + keys=[minute_key, day_key, cooldown_key], + args=[ + token.rpm_limit, + token.daily_limit, + self._seconds_until_next_minute(now_ts), + self._seconds_until_next_day(now_ts), + now_ts, + ], + ) + + if isinstance(result, list) and result and int(result[0]) == 1: + return True, 0, "minute" + + if isinstance(result, list) and len(result) >= 3: + period = ( + str(result[1]) + if result[1] in {"minute", "day", "cooldown"} + else "minute" + ) + if period == "cooldown": + period = "minute" + return False, int(result[2]), period + + return False, self.default_cooldown_seconds, "minute" + + @staticmethod + def _seconds_until_next_minute(now_ts: int) -> int: + return max( + 1, + BaseQuotaManager.MINUTE_WINDOW_SECONDS + - (now_ts % BaseQuotaManager.MINUTE_WINDOW_SECONDS), + ) + + @staticmethod + def _seconds_until_next_day(now_ts: int) -> int: + return max( + 1, + BaseQuotaManager.DAY_WINDOW_SECONDS + - (now_ts % BaseQuotaManager.DAY_WINDOW_SECONDS), + ) + + def _minute_key(self, token_id: str, now_ts: int) -> str: + minute_bucket = now_ts // self.MINUTE_WINDOW_SECONDS + slot_tag = self._token_slot_tag(token_id) + return f"{self.SERVICE_PREFIX}:quota:{{{slot_tag}}}:minute:{minute_bucket}" + + def _day_key(self, token_id: str, now_ts: int) -> str: + day_bucket = now_ts // self.DAY_WINDOW_SECONDS + slot_tag = self._token_slot_tag(token_id) + return f"{self.SERVICE_PREFIX}:quota:{{{slot_tag}}}:day:{day_bucket}" + + def _cooldown_key(self, token_id: str) -> str: + slot_tag = self._token_slot_tag(token_id) + return f"{self.SERVICE_PREFIX}:quota:{{{slot_tag}}}:cooldown" + + @staticmethod + def _token_slot_tag(token_id: str) -> str: + return str(token_id).replace("{", "_").replace("}", "_") diff --git a/packages/shared-python/shared/services/redis/__init__.py b/packages/shared-python/shared/services/redis/__init__.py index b4186f21f..a0cb3f46e 100644 --- a/packages/shared-python/shared/services/redis/__init__.py +++ b/packages/shared-python/shared/services/redis/__init__.py @@ -1,13 +1,21 @@ """Redis service exports.""" -from .job_info_redis_service import JobInfoRedisService -from .job_metadata_service import JobMetadataService -from .redis_alerts import AlertRule, RedisAlertManager, RedisAlertNotifier -from .redis_monitor import RedisMonitor -from .redis_service import RedisService -from .redis_service_factory import RedisServiceFactory -from .task_redis_service import TaskRedisService -from .user_redis_service import UserRedisService +from __future__ import annotations + +from importlib import import_module +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from .job_info_redis_service import JobInfoRedisService + from .job_metadata_service import JobMetadataService + from .key_builder import RedisKeyBuilder, RedisKeyType, redis_key_builder + from .redis_alerts import AlertRule, RedisAlertManager, RedisAlertNotifier + from .redis_monitor import RedisMonitor + from .redis_service import RedisService + from .redis_service_factory import RedisServiceFactory + from .retry_policy import RedisHealthChecker, RedisRetry + from .task_redis_service import TaskRedisService + from .user_redis_service import UserRedisService __all__ = [ "RedisService", @@ -20,4 +28,44 @@ "UserRedisService", "JobInfoRedisService", "JobMetadataService", + "RedisKeyBuilder", + "RedisKeyType", + "redis_key_builder", + "RedisHealthChecker", + "RedisRetry", ] + +_EXPORT_MODULES: dict[str, str] = { + "RedisService": "shared.services.redis.redis_service", + "RedisServiceFactory": "shared.services.redis.redis_service_factory", + "RedisMonitor": "shared.services.redis.redis_monitor", + "RedisAlertManager": "shared.services.redis.redis_alerts", + "RedisAlertNotifier": "shared.services.redis.redis_alerts", + "AlertRule": "shared.services.redis.redis_alerts", + "TaskRedisService": "shared.services.redis.task_redis_service", + "UserRedisService": "shared.services.redis.user_redis_service", + "JobInfoRedisService": "shared.services.redis.job_info_redis_service", + "JobMetadataService": "shared.services.redis.job_metadata_service", + "RedisKeyBuilder": "shared.services.redis.key_builder", + "RedisKeyType": "shared.services.redis.key_builder", + "redis_key_builder": "shared.services.redis.key_builder", + "RedisHealthChecker": "shared.services.redis.retry_policy", + "RedisRetry": "shared.services.redis.retry_policy", +} + + +def __getattr__(name: str) -> Any: + """Load Redis package exports on first access.""" + module_name = _EXPORT_MODULES.get(name) + if module_name is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + module = import_module(module_name) + value = getattr(module, name) + globals()[name] = value + return value + + +def __dir__() -> list[str]: + """Return public Redis package exports.""" + return sorted([*globals(), *__all__]) diff --git a/packages/shared-python/shared/services/redis/distributed_lock.py b/packages/shared-python/shared/services/redis/distributed_lock.py index a36a23bf9..8933737f1 100644 --- a/packages/shared-python/shared/services/redis/distributed_lock.py +++ b/packages/shared-python/shared/services/redis/distributed_lock.py @@ -18,7 +18,7 @@ from shared.core.config import settings from shared.core.exceptions.domain_exceptions import UnavailableException from shared.services.redis.redis_sync_service import SyncRedisService -from shared.utils.redis_key_builder import redis_key_builder +from shared.services.redis.key_builder import redis_key_builder # Lua script: atomically check owner token then delete. # Prevents a stale owner from accidentally deleting a lock diff --git a/packages/shared-python/shared/services/redis/job_info_redis_service.py b/packages/shared-python/shared/services/redis/job_info_redis_service.py index 89c1a2ec2..0ee34b0b7 100644 --- a/packages/shared-python/shared/services/redis/job_info_redis_service.py +++ b/packages/shared-python/shared/services/redis/job_info_redis_service.py @@ -8,7 +8,7 @@ from loguru import logger from shared.services.redis.redis_service import RedisService -from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder +from shared.services.redis.key_builder import RedisKeyType, redis_key_builder class JobInfoRedisService: diff --git a/packages/shared-python/shared/services/redis/job_metadata_service.py b/packages/shared-python/shared/services/redis/job_metadata_service.py index 62fb9c135..410a1c052 100644 --- a/packages/shared-python/shared/services/redis/job_metadata_service.py +++ b/packages/shared-python/shared/services/redis/job_metadata_service.py @@ -5,7 +5,7 @@ from loguru import logger from shared.services.redis.redis_service import RedisService -from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder +from shared.services.redis.key_builder import RedisKeyType, redis_key_builder class JobMetadataService: diff --git a/packages/shared-python/shared/services/redis/key_builder.py b/packages/shared-python/shared/services/redis/key_builder.py new file mode 100644 index 000000000..62d3d5b28 --- /dev/null +++ b/packages/shared-python/shared/services/redis/key_builder.py @@ -0,0 +1,287 @@ +"""Helpers for standardized Redis key naming.""" + +from enum import Enum +from typing import Any, Dict + + +class RedisKeyType(Enum): + """Enumeration of Redis key types.""" + + USER = "user" + TASK = "task" + CONVERSATION = "conversation" + KNOWLEDGE_BASE = "kb" + SESSION = "session" + CACHE = "cache" + QUEUE = "queue" + COUNTER = "counter" + SET = "set" + HASH = "hash" + LIST = "list" + + +class RedisKeyBuilder: + """Redis key builder.""" + + def __init__(self, prefix: str = "knowhere-api"): + self.prefix = prefix + + def build_key(self, key_type: RedisKeyType, *parts: str, **kwargs) -> str: + """ + Build a Redis key. + + Args: + key_type: Key type. + *parts: Positional key segments. + **kwargs: Additional named segments. + + Returns: + Full Redis key. + """ + key_parts = [self.prefix, key_type.value] + key_parts.extend(parts) + + # Append any extra key-value segments. + if kwargs: + for key, value in sorted(kwargs.items()): + key_parts.append(f"{key}:{value}") + + return ":".join(key_parts) + + # ==================== User Keys ==================== + + def user_config(self, username: str) -> str: + """User config key.""" + return self.build_key(RedisKeyType.USER, username, "config") + + def user_session(self, username: str) -> str: + """User session key.""" + return self.build_key(RedisKeyType.USER, username, "session") + + def user_permissions(self, username: str) -> str: + """User permissions key.""" + return self.build_key(RedisKeyType.USER, username, "permissions") + + def user_activity(self, username: str) -> str: + """User activity key.""" + return self.build_key(RedisKeyType.USER, username, "activity") + + # ==================== Task Keys ==================== + + def task_status(self, task_id: str) -> str: + """Task status key.""" + return self.build_key(RedisKeyType.TASK, task_id, "status") + + def task_result(self, task_id: str) -> str: + """Task result key.""" + return self.build_key(RedisKeyType.TASK, task_id, "result") + + def task_metadata(self, task_id: str) -> str: + """Task metadata key.""" + return self.build_key(RedisKeyType.TASK, task_id, "metadata") + + def task_progress(self, task_id: str) -> str: + """Task progress key.""" + return self.build_key(RedisKeyType.TASK, task_id, "progress") + + def task_info(self, task_id: str) -> str: + """Task info key.""" + return self.build_key(RedisKeyType.TASK, task_id, "info") + + def task_queue(self, queue_name: str) -> str: + """Task queue key.""" + return self.build_key(RedisKeyType.QUEUE, queue_name) + + # ==================== Conversation Keys ==================== + + def conversation_state(self, conversation_id: str) -> str: + """Conversation state key.""" + return self.build_key(RedisKeyType.CONVERSATION, conversation_id, "state") + + def conversation_history(self, conversation_id: str) -> str: + """Conversation history key.""" + return self.build_key(RedisKeyType.CONVERSATION, conversation_id, "history") + + def conversation_context(self, conversation_id: str) -> str: + """Conversation context key.""" + return self.build_key(RedisKeyType.CONVERSATION, conversation_id, "context") + + # ==================== Knowledge Base Keys ==================== + + def kb_status(self, user_id: str) -> str: + """Knowledge-base status key.""" + return self.build_key(RedisKeyType.KNOWLEDGE_BASE, user_id, "status") + + def kb_vectors(self, user_id: str) -> str: + """Knowledge-base vectors key.""" + return self.build_key(RedisKeyType.KNOWLEDGE_BASE, user_id, "vectors") + + def kb_metadata(self, user_id: str) -> str: + """Knowledge-base metadata key.""" + return self.build_key(RedisKeyType.KNOWLEDGE_BASE, user_id, "metadata") + + def kb_index(self, user_id: str) -> str: + """Knowledge-base index key.""" + return self.build_key(RedisKeyType.KNOWLEDGE_BASE, user_id, "index") + + # ==================== Session Keys ==================== + + def session_data(self, session_id: str) -> str: + """Session data key.""" + return self.build_key(RedisKeyType.SESSION, session_id, "data") + + def session_expiry(self, session_id: str) -> str: + """Session expiry key.""" + return self.build_key(RedisKeyType.SESSION, session_id, "expiry") + + # ==================== Cache Keys ==================== + + def cache_data(self, cache_key: str) -> str: + """Cache data key.""" + return self.build_key(RedisKeyType.CACHE, cache_key) + + def cache_metadata(self, cache_key: str) -> str: + """Cache metadata key.""" + return self.build_key(RedisKeyType.CACHE, cache_key, "metadata") + + # ==================== Counter Keys ==================== + + def counter_user_requests(self, username: str) -> str: + """User-request counter key.""" + return self.build_key(RedisKeyType.COUNTER, "user_requests", username) + + def counter_api_calls(self, api_name: str) -> str: + """API-call counter key.""" + return self.build_key(RedisKeyType.COUNTER, "api_calls", api_name) + + def counter_task_completed(self, task_type: str) -> str: + """Task-completed counter key.""" + return self.build_key(RedisKeyType.COUNTER, "task_completed", task_type) + + def rate_limit_api(self, user_id: str, api_name: str) -> str: + """API rate-limit key.""" + return self.build_key(RedisKeyType.COUNTER, "rate_limit", api_name, user_id) + + # ==================== Set Keys ==================== + + def set_active_users(self) -> str: + """Active-users set key.""" + return self.build_key(RedisKeyType.SET, "active_users") + + def set_online_users(self) -> str: + """Online-users set key.""" + return self.build_key(RedisKeyType.SET, "online_users") + + def set_processing_tasks(self) -> str: + """Processing-tasks set key.""" + return self.build_key(RedisKeyType.SET, "processing_tasks") + + # ==================== Hash Keys ==================== + + def hash_user_stats(self, username: str) -> str: + """User-stats hash key.""" + return self.build_key(RedisKeyType.HASH, "user_stats", username) + + def hash_system_stats(self) -> str: + """System-stats hash key.""" + return self.build_key(RedisKeyType.HASH, "system_stats") + + def hash_task_stats(self, task_id: str) -> str: + """Task-stats hash key.""" + return self.build_key(RedisKeyType.HASH, "task_stats", task_id) + + # ==================== List Keys ==================== + + def list_recent_activities(self, username: str) -> str: + """Recent-activities list key.""" + return self.build_key(RedisKeyType.LIST, "recent_activities", username) + + def list_error_logs(self) -> str: + """Error-logs list key.""" + return self.build_key(RedisKeyType.LIST, "error_logs") + + def list_audit_logs(self) -> str: + """Audit-logs list key.""" + return self.build_key(RedisKeyType.LIST, "audit_logs") + + # ==================== Lock Keys ==================== + + def lock_job_processing(self, job_id: str) -> str: + """Distributed lock key for exclusive job processing.""" + return f"lock:job_processing:{job_id}" + + # ==================== Utility Methods ==================== + + def parse_key(self, key: str) -> Dict[str, Any]: + """ + Parse a Redis key into its components. + + Args: + key: Redis key. + + Returns: + Parsed key information. + """ + if not key.startswith(self.prefix): + raise ValueError(f"Key '{key}' is not a valid Redis key format") + + parts = key[len(self.prefix) :].lstrip(":").split(":") + + if len(parts) < 2: + raise ValueError(f"Key '{key}' has an invalid format") + + result = { + "prefix": self.prefix, + "type": parts[0], + "parts": parts[1:], + "original_key": key, + } + + return result + + def is_key_type(self, key: str, key_type: RedisKeyType) -> bool: + """ + Check whether a key belongs to the requested type. + + Args: + key: Redis key. + key_type: Key type. + + Returns: + Whether the key matches the requested type. + """ + try: + parsed = self.parse_key(key) + return parsed["type"] == key_type.value + except ValueError: + return False + + def get_key_ttl(self, key_type: RedisKeyType) -> int: + """ + Get the default TTL for a key type. + + Args: + key_type: Key type. + + Returns: + TTL in seconds. + """ + ttl_mapping = { + RedisKeyType.USER: 86400, # 1 day (user_config cache). + RedisKeyType.TASK: 86400, # 1 day. + RedisKeyType.CONVERSATION: 3600 * 2, # 2 hours. + RedisKeyType.KNOWLEDGE_BASE: 86400 * 30, # 30 days. + RedisKeyType.SESSION: 3600, # 1 hour. + RedisKeyType.CACHE: 3600, # 1 hour. + RedisKeyType.QUEUE: 86400, # 1 day. + RedisKeyType.COUNTER: 86400, # 1 day. + RedisKeyType.SET: 86400, # 1 day. + RedisKeyType.HASH: 86400, # 1 day. + RedisKeyType.LIST: 86400, # 1 day. + } + + return ttl_mapping.get(key_type, 3600) # Default to 1 hour. + + +# Global key-builder instance. +redis_key_builder = RedisKeyBuilder() diff --git a/packages/shared-python/shared/services/redis/rate_limit_service.py b/packages/shared-python/shared/services/redis/rate_limit_service.py index de7ef821b..c1cec343b 100644 --- a/packages/shared-python/shared/services/redis/rate_limit_service.py +++ b/packages/shared-python/shared/services/redis/rate_limit_service.py @@ -7,7 +7,7 @@ from loguru import logger from shared.services.redis.redis_service import RedisService -from shared.utils.redis_key_builder import redis_key_builder +from shared.services.redis.key_builder import redis_key_builder class RateLimitService: diff --git a/packages/shared-python/shared/services/redis/redis_monitor.py b/packages/shared-python/shared/services/redis/redis_monitor.py index 6ae090590..f6887576a 100644 --- a/packages/shared-python/shared/services/redis/redis_monitor.py +++ b/packages/shared-python/shared/services/redis/redis_monitor.py @@ -7,7 +7,7 @@ from loguru import logger from shared.services.redis.redis_service import RedisService -from shared.utils.redis_key_builder import redis_key_builder +from shared.services.redis.key_builder import redis_key_builder class RedisMonitor: diff --git a/packages/shared-python/shared/services/redis/redis_service.py b/packages/shared-python/shared/services/redis/redis_service.py index 99ed39c1d..fa9d39fa9 100644 --- a/packages/shared-python/shared/services/redis/redis_service.py +++ b/packages/shared-python/shared/services/redis/redis_service.py @@ -13,7 +13,7 @@ RedisConnectionError, RedisOperationError, ) -from shared.utils.redis_retry import RedisHealthChecker, RedisRetry +from shared.services.redis.retry_policy import RedisHealthChecker, RedisRetry ResponseT = TypeVar("ResponseT") diff --git a/packages/shared-python/shared/services/redis/redis_sync_service.py b/packages/shared-python/shared/services/redis/redis_sync_service.py index aa725b314..90e363aa1 100644 --- a/packages/shared-python/shared/services/redis/redis_sync_service.py +++ b/packages/shared-python/shared/services/redis/redis_sync_service.py @@ -12,7 +12,7 @@ from redis.connection import BlockingConnectionPool from shared.core.config.redis import RedisConfigManager -from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder +from shared.services.redis.key_builder import RedisKeyType, redis_key_builder class SyncRedisService: diff --git a/packages/shared-python/shared/services/redis/retry_policy.py b/packages/shared-python/shared/services/redis/retry_policy.py new file mode 100644 index 000000000..e1915605a --- /dev/null +++ b/packages/shared-python/shared/services/redis/retry_policy.py @@ -0,0 +1,146 @@ +"""Retry helpers for Redis operations.""" + +import asyncio +import random +from typing import Any, Callable + +from loguru import logger + +from shared.core.exceptions.redis_exceptions import ( + RedisConnectionError, + RedisOperationError, + RedisTimeoutError, +) + + +class RedisRetry: + """Retry helper for Redis operations.""" + + @staticmethod + async def with_retry( + operation: Callable, + max_retries: int = 3, + base_delay: float = 1.0, + max_delay: float = 60.0, + exponential_base: float = 2.0, + jitter: bool = True, + ) -> Any: + """ + Execute a Redis operation with retry support. + + Args: + operation: Operation callback to execute. + max_retries: Maximum retry count. + base_delay: Base delay in seconds. + max_delay: Maximum delay in seconds. + exponential_base: Exponential backoff base. + jitter: Whether to add random jitter. + + Returns: + Operation result. + + Raises: + RedisOperationError: Raised when all retries are exhausted. + """ + last_exception = None + + for attempt in range(max_retries + 1): + try: + return await operation() + except (RedisConnectionError, RedisOperationError, RedisTimeoutError) as e: + last_exception = e + + if attempt == max_retries: + logger.error( + f"Redis operation failed after {max_retries} retries: {e}" + ) + raise RedisOperationError( + internal_message=f"Redis operation failed after {max_retries} retries: {str(e)}", + original_exception=e, + ) + + # Compute the retry delay. + delay = min(base_delay * (exponential_base**attempt), max_delay) + + # Add random jitter when enabled. + if jitter: + delay *= 0.5 + random.random() * 0.5 + + logger.warning( + f"Redis operation failed, retrying in {delay:.2f}s (attempt {attempt + 1}/{max_retries + 1}): {e}" + ) + await asyncio.sleep(delay) + + # Should not reach here + raise RedisOperationError( + internal_message=f"Redis operation failed: {str(last_exception)}", + original_exception=last_exception, + ) + + @staticmethod + async def with_circuit_breaker( + operation: Callable, + failure_threshold: int = 5, + recovery_timeout: float = 60.0, + expected_exception: type = Exception, + ) -> Any: + """ + Execute a Redis operation with a circuit-breaker placeholder. + + Args: + operation: Operation callback to execute. + failure_threshold: Failure threshold. + recovery_timeout: Recovery timeout. + expected_exception: Expected exception type. + + Returns: + Operation result. + """ + # Circuit-breaker logic can be implemented here later. + # For now, execute the operation directly. + return await operation() + + +class RedisHealthChecker: + """Health checker for Redis connections.""" + + def __init__(self, redis_client): + self.redis_client = redis_client + self._is_healthy = True + self._last_check = 0 + self._check_interval = 30 # Check every 30 seconds. + + async def is_healthy(self) -> bool: + """Check whether Redis is healthy.""" + import time + + current_time = time.time() + + # Return the cached result if the last check is still fresh. + if current_time - self._last_check < self._check_interval: + return self._is_healthy + + try: + # Use PING to validate the connection. + await self.redis_client.ping() + self._is_healthy = True + self._last_check = current_time + return True + except Exception as e: + logger.error(f"Redis health check failed: {e}") + self._is_healthy = False + self._last_check = current_time + return False + + async def wait_for_healthy(self, timeout: float = 30.0) -> bool: + """Wait for Redis to become healthy again.""" + import time + + start_time = time.time() + + while time.time() - start_time < timeout: + if await self.is_healthy(): + return True + await asyncio.sleep(1) + + return False diff --git a/packages/shared-python/shared/services/redis/task_redis_service.py b/packages/shared-python/shared/services/redis/task_redis_service.py index 883d8a116..6c9444534 100644 --- a/packages/shared-python/shared/services/redis/task_redis_service.py +++ b/packages/shared-python/shared/services/redis/task_redis_service.py @@ -5,7 +5,7 @@ from loguru import logger from shared.services.redis.redis_service import RedisService -from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder +from shared.services.redis.key_builder import RedisKeyType, redis_key_builder class TaskRedisService: diff --git a/packages/shared-python/shared/services/redis/user_redis_service.py b/packages/shared-python/shared/services/redis/user_redis_service.py index 6db0a2acc..d28ee59d5 100644 --- a/packages/shared-python/shared/services/redis/user_redis_service.py +++ b/packages/shared-python/shared/services/redis/user_redis_service.py @@ -5,7 +5,7 @@ from loguru import logger from shared.services.redis.redis_service import RedisService -from shared.utils.redis_key_builder import RedisKeyType, redis_key_builder +from shared.services.redis.key_builder import RedisKeyType, redis_key_builder class UserRedisService: diff --git a/packages/shared-python/shared/services/retrieval/llm_adapter.py b/packages/shared-python/shared/services/retrieval/llm_adapter.py index c5e592b5e..9c8bd9263 100644 --- a/packages/shared-python/shared/services/retrieval/llm_adapter.py +++ b/packages/shared-python/shared/services/retrieval/llm_adapter.py @@ -87,7 +87,7 @@ def create_retrieval_llm_fn( effective_model = model or _resolve_default_model() async def llm_fn(prompt: LLMFnInput) -> str: - from shared.utils.OpenAICompatibleClientSync import get_openai_client + from shared.services.ai.openai_compatible_client_sync import get_openai_client client = get_openai_client(model=effective_model) current_llm_usage.set(None) @@ -118,7 +118,7 @@ def create_retrieval_planner_fn( effective_model = model or _resolve_planner_model(thinking=thinking) async def llm_fn(prompt: LLMFnInput) -> str: - from shared.utils.OpenAICompatibleClientSync import get_openai_client + from shared.services.ai.openai_compatible_client_sync import get_openai_client client = get_openai_client(model=effective_model) current_llm_usage.set(None) @@ -159,7 +159,7 @@ def create_retrieval_vlm_fn( return None async def vlm_fn(prompt: LLMFnInput) -> str: - from shared.utils.OpenAICompatibleClientSync import get_openai_client + from shared.services.ai.openai_compatible_client_sync import get_openai_client client = get_openai_client(model=effective_model) current_llm_usage.set(None) 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 cc46f6bd3..77d49e2b4 100644 --- a/packages/shared-python/shared/services/storage/job_file_storage.py +++ b/packages/shared-python/shared/services/storage/job_file_storage.py @@ -8,8 +8,8 @@ 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 +from shared.services.http.pinned_outbound import download_pinned_outbound_file +from shared.services.http.url_security import validate_http_url_and_resolve_ip class JobFileStorage: diff --git a/packages/shared-python/shared/services/webhook/delivery_client.py b/packages/shared-python/shared/services/webhook/delivery_client.py index 232f9a875..2e33aca20 100644 --- a/packages/shared-python/shared/services/webhook/delivery_client.py +++ b/packages/shared-python/shared/services/webhook/delivery_client.py @@ -8,8 +8,8 @@ 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 +from shared.services.http.pinned_outbound import send_pinned_outbound_request +from shared.services.http.url_security import validate_http_url_and_resolve_ip_async HTTP_TIMEOUT_SECONDS = 10 diff --git a/packages/shared-python/shared/services/webhook/qstash_publisher.py b/packages/shared-python/shared/services/webhook/qstash_publisher.py index 7a0580cb6..810f35def 100644 --- a/packages/shared-python/shared/services/webhook/qstash_publisher.py +++ b/packages/shared-python/shared/services/webhook/qstash_publisher.py @@ -26,7 +26,7 @@ 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 ( +from shared.services.http.url_security import ( validate_http_url_and_resolve_ip, ) diff --git a/packages/shared-python/shared/utils/OpenAICompatibleClientSync.py b/packages/shared-python/shared/utils/OpenAICompatibleClientSync.py index 66da76a6d..c9b31d6b8 100644 --- a/packages/shared-python/shared/utils/OpenAICompatibleClientSync.py +++ b/packages/shared-python/shared/utils/OpenAICompatibleClientSync.py @@ -1,425 +1,8 @@ -""" -Synchronous OpenAI-compatible client for gevent worker tasks. -Uses the official OpenAI SDK with shared httpx connection pooling -and built-in 429 retry with exponential backoff. +"""Compatibility wrapper for the OpenAI-compatible sync client.""" -For Aliyun (qwen) models, api keys are drawn from the AliQuotaManager -token pool and automatically rotated on 429 RateLimitError. -""" -import os -import threading -from typing import Any, Dict, List, Optional, Union +from shared.services.ai.openai_compatible_client_sync import ( + OpenAICompatibleClientSync, + get_openai_client, +) -import openai -from loguru import logger -from openai import OpenAI -from openai.types.chat import ChatCompletionMessageParam - -from shared.core.config import settings -from shared.core.exceptions.domain_exceptions import LLMServiceException -from shared.utils.http_clients import get_sync_client -from shared.utils.llm_mock import build_mock_chat_completion_response -from shared.utils.security_utils import mask_api_key - -LOCAL_DEBUG = os.getenv("LOCAL_DEBUG", "0") == "1" -LLMUsage = dict[str, int] - -_client_cache: Dict[tuple, "OpenAICompatibleClientSync"] = {} -_client_cache_lock = threading.Lock() - - -def _is_ali_model(model_name: str) -> bool: - """Check whether a model name routes to Aliyun DashScope.""" - return "qwen" in (model_name or "").lower() - - -def _should_mock_llm_calls() -> bool: - """Whether all OpenAI-compatible LLM calls should short-circuit to mock responses.""" - return bool(getattr(settings, "LLM_MOCK_ENABLED", False)) - - -def _empty_usage() -> LLMUsage: - return {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} - - -def _extract_usage(response: Any) -> LLMUsage: - usage = getattr(response, "usage", None) - if usage is None: - return _empty_usage() - return { - "prompt_tokens": int(getattr(usage, "prompt_tokens", 0) or 0), - "completion_tokens": int(getattr(usage, "completion_tokens", 0) or 0), - "total_tokens": int(getattr(usage, "total_tokens", 0) or 0), - } - - -def _summarize_exception_chain(exc: Exception, *, max_depth: int = 4) -> str: - parts: list[str] = [] - current: BaseException | None = exc - seen: set[int] = set() - - while current is not None and len(parts) < max_depth and id(current) not in seen: - seen.add(id(current)) - parts.append(f"{type(current).__name__}: {current}") - current = current.__cause__ or current.__context__ - - return " <- ".join(parts) - - -class OpenAICompatibleClientSync: - """Sync OpenAI-compatible client backed by the official OpenAI SDK.""" - - def __init__( - self, - api_key: Optional[str] = None, - api_url: Optional[str] = None, - default_model: Optional[str] = None, - timeout: int = 300, - max_retries: int = 2, - ): - self.default_model = default_model or getattr(settings, "NORMOL_MODEL", "deepseek-chat") - self._explicit_api_key = api_key - self._explicit_api_url = api_url - self._max_retries = max_retries - self._base_url: Optional[str] = self._resolve_base_url( - model_name=self.default_model, - api_url=api_url, - ) - - self.timeout = ( - getattr(settings, "OPENAI_CLIENT_TIMEOUT", 300) - if timeout is None - else timeout - ) - self._client: Optional[OpenAI] = None - if _should_mock_llm_calls(): - return - if not self._should_use_ali_pool(): - resolved_key: Optional[str] = self._resolve_direct_api_key( - model_name=self.default_model, - api_key=api_key, - ) - self._client = self._build_client( - api_key=resolved_key, - base_url=self._base_url, - max_retries=max_retries, - ) - - def _build_client( - self, - api_key: Optional[str], - base_url: Optional[str], - max_retries: int, - ) -> OpenAI: - """Build a direct OpenAI-compatible SDK client.""" - return OpenAI( - api_key=api_key, - base_url=base_url, - http_client=get_sync_client(), - max_retries=max_retries, - timeout=self.timeout, - ) - - @staticmethod - def _strip_chat_completions(url: Optional[str]) -> Optional[str]: - """Strip /chat/completions suffix so the OpenAI SDK can append it.""" - if url and url.rstrip("/").endswith("/chat/completions"): - return url.rstrip("/").removesuffix("/chat/completions") - return url - - def _resolve_base_url( - self, - model_name: str, - api_url: Optional[str], - ) -> Optional[str]: - if api_url: - return self._strip_chat_completions(api_url) - - model_lower = (model_name or "").lower() - if "qwen" in model_lower: - ali_base = getattr(settings, "ALI_URL", "https://dashscope.aliyuncs.com/compatible-mode/v1") - return self._strip_chat_completions(ali_base) - - if "glm" in model_lower: - glm_base = getattr(settings, "GLM_URL", "https://open.bigmodel.cn/api/paas/v4") - return self._strip_chat_completions(glm_base) - - if "doubao" in model_lower or model_lower.startswith("ep-"): - return self._strip_chat_completions(getattr(settings, "ARK_URL", None)) - - return self._strip_chat_completions(settings.DS_URL) - - def _resolve_direct_api_key( - self, - model_name: str, - api_key: Optional[str], - ) -> Optional[str]: - if api_key: - return api_key - - model_lower = (model_name or "").lower() - if "glm" in model_lower: - return getattr(settings, "GLM_API_KEY", None) - - if "doubao" in model_lower or model_lower.startswith("ep-"): - return getattr(settings, "ARK_API_KEY", None) - - return settings.DS_KEY - - # ------------------------------------------------------------------ - # Ali token-pool helpers - # ------------------------------------------------------------------ - - def _should_use_ali_pool(self) -> bool: - """Whether to route through the AliQuotaManager instead of a fixed key.""" - if self._explicit_api_key: - return False - return _is_ali_model(self.default_model) - - def _make_ali_pool_call( - self, - model: str, - all_messages: List[ChatCompletionMessageParam], - temperature: float, - max_tokens: int, - api_kwargs: Dict[str, Any], - ) -> tuple[str, LLMUsage]: - """Acquire a token, make the call, and retry inline on 429.""" - from shared.utils.ali_quota_manager import get_ali_quota_manager - - quota_manager = get_ali_quota_manager() - base_url: Optional[str] = self._base_url - - max_retries = settings.ALI_INLINE_MAX_RETRIES - for attempt in range(max_retries): - lease = quota_manager.acquire_request(operation="chat_completion") - try: - # Hybrid rate-limit control: SDK handles per-key backoff (exp. backoff + jitter), - # outer loop handles token rotation across the pool on persistent 429s. - client = self._build_client( - api_key=lease.api_key, - base_url=base_url, - max_retries=settings.ALI_SDK_MAX_RETRIES, - ) - response = client.chat.completions.create( - model=model, - messages=all_messages, - temperature=temperature, - max_tokens=max_tokens, - **api_kwargs, - ) - choices = response.choices - if not choices: - raise LLMServiceException( - internal_message="AI returned empty result", - provider=self.default_model, - ) - return choices[0].message.content or "", _extract_usage(response) - except openai.RateLimitError as exc: - retry_after = _parse_retry_after(exc) - quota_manager.mark_rate_limited(lease.token_id, retry_after) - logger.warning( - f"Ali token {lease.token_id} rate-limited (attempt {attempt + 1}/{max_retries}), " - f"cooling down {retry_after}s and retrying with next token" - ) - if attempt == max_retries - 1: - raise LLMServiceException( - internal_message=f"All Ali tokens exhausted after {max_retries} retries: {exc}", - provider=self.default_model, - original_exception=exc, - ) from exc - except LLMServiceException: - raise - except Exception as exc: - masked_api_key = mask_api_key(lease.api_key) - logger.error( - "LLM request failed (Ali pool): model={model}, token_id={token_id}, api_key={api_key}, error={error}", - model=model, - token_id=lease.token_id, - api_key=masked_api_key, - error=exc, - ) - raise LLMServiceException( - internal_message=( - "API request failed " - f"(Ali pool, token_id={lease.token_id}, api_key={masked_api_key}): {exc}" - ), - provider=self.default_model, - original_exception=exc, - ) from exc - # unreachable but satisfies the type checker - raise LLMServiceException( - internal_message="Ali token pool retry loop exited unexpectedly", - provider=self.default_model, - ) - - # ------------------------------------------------------------------ - - def chat_completion_with_usage( - self, - messages: Union[str, List[ChatCompletionMessageParam]], - model: Optional[str] = None, - temperature: float = 0.1, - max_tokens: int = 4096, - top_p: Optional[float] = None, - timeout: Optional[int] = None, - **kwargs, - ) -> tuple[str, LLMUsage]: - all_messages: List[ChatCompletionMessageParam] - if isinstance(messages, list): - all_messages = messages # type: ignore[assignment] - else: - all_messages = [{"role": "user", "content": str(messages)}] - - api_kwargs: Dict[str, Any] = {} - if top_p is not None: - api_kwargs["top_p"] = top_p - if timeout is not None: - api_kwargs["timeout"] = timeout - - allowed_api_params = { - "n", "stop", "presence_penalty", "frequency_penalty", - "logit_bias", "user", "seed", "tools", "tool_choice", - "response_format", "logprobs", "top_logprobs", - } - for key, value in kwargs.items(): - if key in allowed_api_params: - api_kwargs[key] = value - - # ── disable thinking mode ── - # Qwen3.5 by default enables thinking mode, which wastes tokens by outputting ... - # Explicitly disable it in all API calls - extra_body = api_kwargs.get("extra_body", {}) - if isinstance(extra_body, dict): - extra_body.setdefault("enable_thinking", False) - else: - extra_body = {"enable_thinking": False} - api_kwargs["extra_body"] = extra_body - - effective_model = model or self.default_model - if _should_mock_llm_calls(): - return build_mock_chat_completion_response( - messages=all_messages, - model_name=effective_model, - ), _empty_usage() - - # Route through Ali token pool when applicable - if self._should_use_ali_pool(): - try: - return self._make_ali_pool_call( - model=effective_model, - all_messages=all_messages, - temperature=temperature, - max_tokens=max_tokens, - api_kwargs=api_kwargs, - ) - except LLMServiceException: - raise - except Exception as exc: - logger.error(f"LLM request failed (Ali pool): model={effective_model}, error={exc}") - raise LLMServiceException( - internal_message=f"API request failed: {str(exc)}", - provider=self.default_model, - original_exception=exc, - ) from exc - - # Non-Ali path: use the single pre-configured client - client = self._client - if client is None: - raise LLMServiceException( - internal_message="OpenAI client is not initialized for direct provider requests", - provider=self.default_model, - ) - - try: - response = client.chat.completions.create( - model=effective_model, - messages=all_messages, - temperature=temperature, - max_tokens=max_tokens, - **api_kwargs, - ) - - choices = response.choices - if not choices: - raise LLMServiceException( - internal_message="AI returned empty result", - provider=self.default_model, - ) - - content = choices[0].message.content or "" - return content, _extract_usage(response) - except LLMServiceException: - raise - except Exception as exc: - logger.error( - "LLM request failed: model={model}, base_url={base_url}, error_chain={error_chain}", - model=effective_model, - base_url=client.base_url, - error_chain=_summarize_exception_chain(exc), - ) - raise LLMServiceException( - internal_message=f"API request failed: {_summarize_exception_chain(exc)}", - provider=self.default_model, - original_exception=exc, - ) from exc - - def chat_completion( - self, - messages: Union[str, List[ChatCompletionMessageParam]], - model: Optional[str] = None, - temperature: float = 0.1, - max_tokens: int = 4096, - top_p: Optional[float] = None, - timeout: Optional[int] = None, - **kwargs, - ) -> str: - content, _usage = self.chat_completion_with_usage( - messages=messages, - model=model, - temperature=temperature, - max_tokens=max_tokens, - top_p=top_p, - timeout=timeout, - **kwargs, - ) - return content - - -def _parse_retry_after(exc: openai.RateLimitError) -> int: - """Extract Retry-After seconds from a RateLimitError, with sane bounds.""" - try: - if hasattr(exc, "response") and exc.response is not None: - header_value = exc.response.headers.get("retry-after") or exc.response.headers.get("Retry-After") - if header_value: - return max(1, min(int(header_value), 120)) - except (ValueError, TypeError, AttributeError): - pass - return settings.ALI_TOKEN_COOLDOWN_SECONDS - - -def get_openai_client( - model: Optional[str] = None, - api_key: Optional[str] = None, - api_url: Optional[str] = None, - timeout: int = 300, - max_retries: int = 2, -) -> OpenAICompatibleClientSync: - """Return a cached OpenAICompatibleClientSync instance keyed by config.""" - cache_key = (model, api_key, api_url, timeout, max_retries) - client = _client_cache.get(cache_key) - if client is not None: - return client - with _client_cache_lock: - client = _client_cache.get(cache_key) - if client is not None: - return client - client = OpenAICompatibleClientSync( - api_key=api_key, - api_url=api_url, - default_model=model, - timeout=timeout, - max_retries=max_retries, - ) - _client_cache[cache_key] = client - return client +__all__ = ["OpenAICompatibleClientSync", "get_openai_client"] diff --git a/packages/shared-python/shared/utils/ali_quota_manager.py b/packages/shared-python/shared/utils/ali_quota_manager.py index 535b86c0c..304a36e8e 100644 --- a/packages/shared-python/shared/utils/ali_quota_manager.py +++ b/packages/shared-python/shared/utils/ali_quota_manager.py @@ -1,65 +1,5 @@ -""" -AliQuotaManager — Redis-backed token pool for Aliyun DashScope API keys. +"""Compatibility wrapper for the Ali quota manager.""" -Distributes requests across multiple ALI_API_KEYS using the shared -BaseQuotaManager, with per-token RPM / daily quotas and 429 cooldown. -""" +from shared.services.ai.ali_quota_manager import AliQuotaManager, get_ali_quota_manager -from __future__ import annotations - -from typing import List, Optional - -from shared.core.config import settings -from shared.services.redis.redis_sync_service import ( - SyncRedisService, - SyncRedisServiceFactory, -) -from shared.utils.quota_manager import BaseQuotaManager, TokenConfig - - -class AliQuotaManager(BaseQuotaManager): - """Quota-aware Ali DashScope token manager backed by Redis.""" - - SERVICE_PREFIX = "ali" - CURSOR_KEY = "ali:quota:cursor" - user_message = "AI service is busy right now. Please retry shortly." - - @classmethod - def from_settings( - cls, redis_service: Optional[SyncRedisService] = None - ) -> "AliQuotaManager": - """Create an AliQuotaManager from application settings.""" - tokens = cls.parse_tokens_from_settings() - instance = cls(redis_service or SyncRedisServiceFactory.get_service(), tokens) - instance.default_cooldown_seconds = settings.ALI_TOKEN_COOLDOWN_SECONDS - return instance - - @staticmethod - def parse_tokens_from_settings() -> List[TokenConfig]: - default_rpm_limit: int = settings.ALI_TOKEN_RPM_LIMIT - default_daily_limit: int = settings.ALI_TOKEN_DAILY_LIMIT - raw_pool: str = (settings.ALI_API_KEYS or "").strip() - - if raw_pool: - parsed = BaseQuotaManager.parse_token_specs( - raw_pool, - default_rpm_limit=default_rpm_limit, - default_daily_limit=default_daily_limit, - ) - if parsed: - return parsed - - raise ValueError( - "No Ali API keys configured. Set ALI_API_KEYS to a comma/newline-separated list or JSON array." - ) - - -_ali_quota_manager: Optional[AliQuotaManager] = None - - -def get_ali_quota_manager() -> AliQuotaManager: - """Return a singleton AliQuotaManager for worker processes.""" - global _ali_quota_manager - if _ali_quota_manager is None: - _ali_quota_manager = AliQuotaManager.from_settings() - return _ali_quota_manager +__all__ = ["AliQuotaManager", "get_ali_quota_manager"] diff --git a/packages/shared-python/shared/utils/http_clients.py b/packages/shared-python/shared/utils/http_clients.py index 928b3f02c..50693141a 100644 --- a/packages/shared-python/shared/utils/http_clients.py +++ b/packages/shared-python/shared/utils/http_clients.py @@ -1,78 +1,15 @@ -""" -Singleton HTTP clients with connection pooling. -Reduces socket churn under high concurrency by reusing -TCP+TLS connections across requests. -""" +"""Compatibility wrapper for shared.services.http.client_pool.""" -import threading -from typing import Optional - -import httpx - -_DEFAULT_LIMITS = httpx.Limits( - max_keepalive_connections=20, - keepalive_expiry=30, -) - -_DEFAULT_TIMEOUT = httpx.Timeout( - connect=10.0, - read=300.0, - write=60.0, - pool=30.0, +from shared.services.http.client_pool import ( + close_async_client, + close_sync_client, + get_async_client, + get_sync_client, ) -# --- Sync client (gevent-patched threading.Lock) --- -_sync_client: Optional[httpx.Client] = None -_sync_lock = threading.Lock() - - -def get_sync_client() -> httpx.Client: - """Return a shared sync httpx.Client with connection pooling.""" - global _sync_client - if _sync_client is None: - with _sync_lock: - if _sync_client is None: - _sync_client = httpx.Client( - limits=_DEFAULT_LIMITS, - timeout=_DEFAULT_TIMEOUT, - follow_redirects=True, - ) - return _sync_client - - -def close_sync_client() -> None: - """Close the shared sync client. Call on worker shutdown.""" - global _sync_client - with _sync_lock: - if _sync_client is not None: - _sync_client.close() - _sync_client = None - - -# --- Async client --- -_async_client: Optional[httpx.AsyncClient] = None -_async_lock = threading.Lock() - - -def get_async_client() -> httpx.AsyncClient: - """Return a shared async httpx.AsyncClient with connection pooling.""" - global _async_client - if _async_client is None: - with _async_lock: - if _async_client is None: - _async_client = httpx.AsyncClient( - limits=_DEFAULT_LIMITS, - timeout=_DEFAULT_TIMEOUT, - follow_redirects=True, - ) - return _async_client - - -async def close_async_client() -> None: - """Close the shared async client. Call on API shutdown.""" - global _async_client - with _async_lock: - client = _async_client - _async_client = None - if client is not None: - await client.aclose() +__all__ = [ + "close_async_client", + "close_sync_client", + "get_async_client", + "get_sync_client", +] diff --git a/packages/shared-python/shared/utils/iloveapi_quota_manager.py b/packages/shared-python/shared/utils/iloveapi_quota_manager.py index b3b923a41..43549ae60 100644 --- a/packages/shared-python/shared/utils/iloveapi_quota_manager.py +++ b/packages/shared-python/shared/utils/iloveapi_quota_manager.py @@ -1,215 +1,11 @@ -""" -ILoveApiQuotaManager — Redis-backed project pool for iLoveAPI keys. +"""Compatibility wrapper for the iLoveAPI quota manager.""" -Each iLoveAPI public_key+secret_key pair represents a separate project -with its own credit quota (free tier: 2,500 credits/month, officepdf -costs 10 credits/file = ~250 files/month). - -Distributes requests across multiple projects using the shared -BaseQuotaManager, with RPM burst protection, daily file limits, -429 cooldown, and a global in-flight concurrency limiter to prevent -worker pool starvation when iLoveAPI is slow or rate-limited. - -This is an optional enhancement path — all failures fail open to -LibreOffice local conversion. -""" - -from __future__ import annotations - -import json -from typing import List, Optional - -from loguru import logger - -from shared.core.config import settings -from shared.services.redis.redis_sync_service import ( - SyncRedisService, - SyncRedisServiceFactory, +from shared.services.ai.iloveapi_quota_manager import ( + ILoveApiQuotaManager, + get_iloveapi_quota_manager, ) -from shared.utils.quota_manager import BaseQuotaManager, TokenConfig - - -class ILoveApiQuotaManager(BaseQuotaManager): - """Quota-aware iLoveAPI token manager backed by Redis.""" - - SERVICE_PREFIX = "iloveapi" - CURSOR_KEY = "iloveapi:quota:cursor" - user_message = "Document parsing service is busy right now. Please retry shortly." - - INFLIGHT_KEY = "iloveapi:inflight" - INFLIGHT_TTL_SECONDS = 300 # 5-minute safety net; auto-expires stale counts - - _ACQUIRE_INFLIGHT_SCRIPT: str = """ -local key = KEYS[1] -local max_concurrent = tonumber(ARGV[1]) -local ttl_seconds = tonumber(ARGV[2]) - -local current = tonumber(redis.call('GET', key) or '0') -if current >= max_concurrent then - return {0, current} -end - -local new_count = redis.call('INCR', key) --- Set TTL only if not already set (safety net for crash recovery) -if redis.call('TTL', key) < 0 then - redis.call('EXPIRE', key, ttl_seconds) -end - -return {1, new_count} -""" - - _RELEASE_INFLIGHT_SCRIPT: str = """ -local key = KEYS[1] -local current = tonumber(redis.call('GET', key) or '0') -if current <= 0 then - redis.call('DEL', key) - return 0 -end -return redis.call('DECR', key) -""" - - def __init__( - self, - redis_service: SyncRedisService, - tokens: List[TokenConfig], - max_concurrent: int = 5, - ) -> None: - super().__init__(redis_service, tokens) - self.max_concurrent = max_concurrent - - @classmethod - def from_settings( - cls, redis_service: Optional[SyncRedisService] = None - ) -> "ILoveApiQuotaManager": - """Create an ILoveApiQuotaManager from application settings.""" - tokens = cls.parse_tokens_from_settings() - max_concurrent: int = getattr(settings, "ILOVEAPI_MAX_CONCURRENT", 5) - instance = cls( - redis_service or SyncRedisServiceFactory.get_service(), - tokens, - max_concurrent=max_concurrent, - ) - instance.default_cooldown_seconds = 60 - return instance - - @staticmethod - def parse_tokens_from_settings() -> List[TokenConfig]: - default_rpm_limit: int = getattr(settings, "ILOVEAPI_TOKEN_RPM_LIMIT", 25) - default_daily_limit: int = getattr(settings, "ILOVEAPI_TOKEN_DAILY_LIMIT", 250) - raw_pool: str = (getattr(settings, "ILOVEAPI_KEYS", "") or "").strip() - - legacy_pub: str = (getattr(settings, "ILOVEAPI_PUBLIC_KEY", "") or "").strip() - legacy_sec: str = (getattr(settings, "ILOVEAPI_SECRET_KEY", "") or "").strip() - - specs: List[TokenConfig] = [] - if raw_pool: - try: - loaded = json.loads(raw_pool) - if isinstance(loaded, list): - for index, entry in enumerate(loaded): - pub = str(entry.get("public_key") or "").strip() - sec = str(entry.get("secret_key") or "").strip() - if pub and sec: - specs.append( - TokenConfig( - token_id=str( - entry.get("token_id") or f"iloveapi-{index + 1}" - ), - api_key=f"{pub}:{sec}", - rpm_limit=int( - entry.get("rpm_limit") or default_rpm_limit - ), - daily_limit=int( - entry.get("daily_limit") or default_daily_limit - ), - ) - ) - except json.JSONDecodeError: - pass - - if not specs and legacy_pub and legacy_sec: - specs.append( - TokenConfig( - token_id="iloveapi-default", - api_key=f"{legacy_pub}:{legacy_sec}", - rpm_limit=default_rpm_limit, - daily_limit=default_daily_limit, - ) - ) - - if not specs: - raise ValueError( - "No iLoveAPI keys configured. Set ILOVEAPI_KEYS to a JSON array containing public_key and secret_key." - ) - - return specs - - # ------------------------------------------------------------------ - # In-flight concurrency limiter - # ------------------------------------------------------------------ - - def acquire_inflight(self) -> Optional[bool]: - """Try to acquire an in-flight slot. - - Returns: - True: reserved a Redis-backed slot and must release it later. - False: at capacity; caller should fail open to local conversion. - None: Redis failed; caller may proceed, but must not release. - """ - try: - result = self.redis.eval( - self._ACQUIRE_INFLIGHT_SCRIPT, - keys=[self.INFLIGHT_KEY], - args=[self.max_concurrent, self.INFLIGHT_TTL_SECONDS], - ) - if isinstance(result, list) and len(result) >= 2: - acquired = int(result[0]) == 1 - current_count = int(result[1]) - if not acquired: - logger.bind( - service=self.SERVICE_PREFIX, - step="concurrency_exceeded", - current_inflight=current_count, - max_concurrent=self.max_concurrent, - ).warning( - f"iLoveAPI concurrency limit reached ({current_count}/{self.max_concurrent})" - ) - return acquired - except Exception: - logger.opt(exception=True).warning( - "Failed to check iLoveAPI in-flight counter; failing open to allow request" - ) - # Fail open on Redis errors — allow the request through without a lease. - return None - - def release_inflight(self) -> None: - """Release an in-flight slot after conversion completes (success or failure).""" - try: - self.redis.eval( - self._RELEASE_INFLIGHT_SCRIPT, - keys=[self.INFLIGHT_KEY], - args=[], - ) - except Exception: - logger.opt(exception=True).warning( - "Failed to release iLoveAPI in-flight slot" - ) - - def get_inflight_count(self) -> int: - """Return current in-flight count (for observability / testing).""" - try: - raw = self.redis.get(self.INFLIGHT_KEY, 0) - return max(0, int(raw or 0)) - except Exception: - return 0 - - -_iloveapi_quota_manager: Optional[ILoveApiQuotaManager] = None - -def get_iloveapi_quota_manager() -> ILoveApiQuotaManager: - """Return a singleton ILoveApiQuotaManager for worker processes.""" - global _iloveapi_quota_manager - if _iloveapi_quota_manager is None: - _iloveapi_quota_manager = ILoveApiQuotaManager.from_settings() - return _iloveapi_quota_manager +__all__ = [ + "ILoveApiQuotaManager", + "get_iloveapi_quota_manager", +] diff --git a/packages/shared-python/shared/utils/llm_mock.py b/packages/shared-python/shared/utils/llm_mock.py index 8bf155f21..8bf18999e 100644 --- a/packages/shared-python/shared/utils/llm_mock.py +++ b/packages/shared-python/shared/utils/llm_mock.py @@ -1,157 +1,5 @@ -"""Helpers for deterministic mock responses from OpenAI-compatible LLM calls.""" +"""Compatibility wrapper for LLM mock responses.""" -from typing import Any, Dict, List +from shared.services.ai.llm_mock import build_mock_chat_completion_response -from loguru import logger -from openai.types.chat import ChatCompletionMessageParam - - -def build_mock_chat_completion_response( - messages: List[ChatCompletionMessageParam], - model_name: str, -) -> str: - """Return a canned response compatible with the inferred prompt contract.""" - prompt_text = _extract_prompt_text(messages) - task_name = _detect_mock_task(prompt_text) - logger.debug( - "LLM mock enabled, returning canned response: model={}, task={}", - model_name, - task_name, - ) - return _build_mock_response(task_name) - - -def _extract_message_text(content: Any) -> str: - """Flatten chat message content into plain text for mock-task detection.""" - if isinstance(content, str): - return content - - if not isinstance(content, list): - return "" - - parts: List[str] = [] - for item in content: - if isinstance(item, str): - parts.append(item) - continue - - if not isinstance(item, dict): - continue - - item_type = item.get("type") - if item_type == "text" and isinstance(item.get("text"), str): - parts.append(item["text"]) - continue - - if isinstance(item.get("content"), str): - parts.append(item["content"]) - - return "\n".join(parts) - - -def _extract_prompt_text(messages: List[ChatCompletionMessageParam]) -> str: - """Extract all text segments from a chat-completion message list.""" - parts: List[str] = [] - - for message in messages: - if not isinstance(message, dict): - continue - - text_part = _extract_message_text(message.get("content")) - if text_part: - parts.append(text_part) - - return "\n".join(parts) - - -def _detect_mock_task(prompt_text: str) -> str: - """Infer the prompt task so the mock can return a compatible response shape.""" - normalized_prompt = prompt_text.lower() - - if ( - "generate a concise title" in normalized_prompt - and "return only the title" in normalized_prompt - ): - return "fragment-title" - if ( - '"toc_start"' in normalized_prompt - and '"toc_end"' in normalized_prompt - and '"confidence"' in normalized_prompt - ): - return "detect-toc-range" - if ( - "row numbers where headers are located" in normalized_prompt - and '"answer"' in normalized_prompt - ): - return "detect-table-headers" - if '"answer"' in normalized_prompt and '"text" or "image"' in normalized_prompt: - return "judge-image-type" - if ( - '"title"' in normalized_prompt - and '"keywords"' in normalized_prompt - and '"summary"' in normalized_prompt - ): - return "summary-full" - if ( - 'json dictionary format with key "answer"' in normalized_prompt - and "keywords" in normalized_prompt - ): - return "summary-keywords" - if ( - "json array only" in normalized_prompt - and "toc, not body text" in normalized_prompt - ): - return "eval-toc-headings" - if ( - "json array" in normalized_prompt - and "only include rows that you judge to be headings" in normalized_prompt - ): - return "eval-headings" - if "scanned page from an engineering atlas" in normalized_prompt: - return "atlas-page-info" - if "perform ocr operation" in normalized_prompt: - return "ocr-image" - if ( - "you will receive an image" in normalized_prompt - and "line 1: output a short title" in normalized_prompt - ): - return "summary-images" - if ( - "you will receive one or more images and the user's current question" - in normalized_prompt - ): - return "ask-image" - if "summaries of sub-sections from a document section" in normalized_prompt: - return "file-summary" - if ( - "line 1: output a short title" in normalized_prompt - and "line 2 onward" in normalized_prompt - ): - return "summary-titled" - if "extract the main content of the material" in normalized_prompt: - return "summary" - - return "default" - - -def _build_mock_response(task_name: str) -> str: - """Return a canned response compatible with the inferred task contract.""" - response_by_task: Dict[str, str] = { - "fragment-title": "Mock Fragment Title", - "detect-toc-range": '{"toc_start": null, "toc_end": null, "confidence": "low"}', - "detect-table-headers": '{"answer": [0]}', - "judge-image-type": '{"answer": "image"}', - "summary-full": '{"title": "Mock Title", "keywords": "mock", "summary": "Mock summary"}', - "summary-keywords": '{"answer": "mock"}', - "eval-headings": "[]", - "eval-toc-headings": "[]", - "atlas-page-info": "Mock atlas page info", - "ocr-image": "Mock OCR text", - "summary-images": "Mock Image Title\nMock image summary", - "ask-image": "Mock image answer", - "file-summary": "Mock section summary", - "summary-titled": "Mock Title\nMock summary", - "summary": "Mock summary", - "default": "Mock LLM response", - } - return response_by_task.get(task_name, response_by_task["default"]) +__all__ = ["build_mock_chat_completion_response"] diff --git a/packages/shared-python/shared/utils/pinned_outbound_http.py b/packages/shared-python/shared/utils/pinned_outbound_http.py index 808b219c8..1d4e0fd3e 100644 --- a/packages/shared-python/shared/utils/pinned_outbound_http.py +++ b/packages/shared-python/shared/utils/pinned_outbound_http.py @@ -1,313 +1,27 @@ -"""Pinned outbound HTTP helpers.""" - -from __future__ import annotations - -import os -import socket -import tempfile -from collections.abc import Mapping -from dataclasses import dataclass -from typing import Any -from urllib.parse import urlsplit - -import aiohttp -from aiohttp.abc import AbstractResolver -from urllib3 import Retry -from urllib3.connection import HTTPConnection, HTTPSConnection -from urllib3.connectionpool import HTTPConnectionPool, HTTPSConnectionPool -from urllib3.util import Timeout -from urllib3.util.connection import create_connection - -from shared.core.exceptions.domain_exceptions import ValidationException -from shared.utils.url_security import SafePublicHTTPURL - - -@dataclass(frozen=True) -class PinnedOutboundResponse: - """Minimal response metadata for pinned outbound HTTP requests.""" - - status: int - - -@dataclass(frozen=True) -class PinnedDownloadResult: - """Metadata for a pinned HTTP download to a temporary file.""" - - status: int - temp_file_path: str - - -class PinnedIPResolver(AbstractResolver): - """Resolver that always returns the supplied pinned IP address.""" - - def __init__(self, pinned_ip: str) -> None: - self.pinned_ip = pinned_ip - - async def resolve( - self, - host: str, - port: int = 0, - family: int = socket.AF_INET, - ) -> list[dict[str, Any]]: - pinned_family: int = ( - socket.AF_INET6 if ":" in self.pinned_ip else socket.AF_INET - ) - return [ - { - "hostname": host, - "host": self.pinned_ip, - "port": port, - "family": pinned_family, - "proto": 0, - "flags": socket.AI_NUMERICHOST, - } - ] - - async def close(self) -> None: - pass - - -class PinnedHTTPConnection(HTTPConnection): - """HTTP connection that resolves the hostname to a pinned IP address.""" - - def __init__(self, *args: Any, pinned_ip: str, **kwargs: Any) -> None: - self._pinned_ip = pinned_ip - super().__init__(*args, **kwargs) - - def _new_conn(self) -> socket.socket: - return create_connection( - (self._pinned_ip, self.port), - self.timeout, - source_address=self.source_address, - socket_options=self.socket_options, - ) - - -class PinnedHTTPSConnection(HTTPSConnection): - """HTTPS connection that resolves the hostname to a pinned IP address.""" - - def __init__(self, *args: Any, pinned_ip: str, **kwargs: Any) -> None: - self._pinned_ip = pinned_ip - super().__init__(*args, **kwargs) - - def _new_conn(self) -> socket.socket: - return create_connection( - (self._pinned_ip, self.port), - self.timeout, - source_address=self.source_address, - socket_options=self.socket_options, - ) - - -class PinnedHTTPConnectionPool(HTTPConnectionPool): - """HTTP pool that pins DNS resolution to a fixed IP address.""" - - ConnectionCls = PinnedHTTPConnection # pyright: ignore[reportAssignmentType] - - def __init__(self, *args: Any, pinned_ip: str, **kwargs: Any) -> None: - kwargs["pinned_ip"] = pinned_ip - super().__init__(*args, **kwargs) - - -class PinnedHTTPSConnectionPool(HTTPSConnectionPool): - """HTTPS pool that pins DNS resolution to a fixed IP address.""" - - ConnectionCls = PinnedHTTPSConnection # pyright: ignore[reportAssignmentType] - - def __init__(self, *args: Any, pinned_ip: str, **kwargs: Any) -> None: - kwargs["pinned_ip"] = pinned_ip - super().__init__(*args, **kwargs) - - -def _build_host_header(parsed_url: Any) -> str: - hostname = parsed_url.hostname - if not hostname: - return "" - - if ":" in hostname and not hostname.startswith("["): - formatted_host = f"[{hostname}]" - else: - formatted_host = hostname - - if parsed_url.port is not None: - return f"{formatted_host}:{parsed_url.port}" - return formatted_host - - -def _validate_download_url(url: str, field: str) -> SafePublicHTTPURL: - parsed_url = urlsplit(url) - if parsed_url.scheme not in {"http", "https"}: - raise ValidationException( - user_message="Invalid URL", - violations=[ - {"field": field, "description": "URL must use http or https"} - ], - ) - if not parsed_url.hostname: - raise ValidationException( - user_message="Invalid URL", - violations=[{"field": field, "description": "URL must include a hostname"}], - ) - return SafePublicHTTPURL(url) - - -async def send_pinned_outbound_request( - *, - method: str, - url: str, - pinned_ip: str, - timeout_seconds: float, - headers: Mapping[str, str] | None = None, - json_body: Any | None = None, -) -> PinnedOutboundResponse: - """Send an outbound HTTP request through a resolver pinned to a validated IP.""" - validated_url = _validate_download_url(url, "url") - connector = aiohttp.TCPConnector(resolver=PinnedIPResolver(pinned_ip)) - timeout = aiohttp.ClientTimeout(total=timeout_seconds) - - async with aiohttp.ClientSession( - connector=connector, - timeout=timeout, - ) as session: - async with session.request( - method=method, - url=validated_url, - headers=headers, - json=json_body, - allow_redirects=False, - ) as response: - return PinnedOutboundResponse(status=response.status) - - -async def download_pinned_outbound_file_async( - *, - url: str, - pinned_ip: str, - timeout_seconds: float, - user_agent: str, - temp_dir: str | None = None, - field: str = "source_url", -) -> PinnedDownloadResult: - """Download a file with aiohttp while pinning DNS to a validated IP.""" - validated_url = _validate_download_url(url, field) - - temp_file_descriptor, temp_file_path = tempfile.mkstemp(dir=temp_dir) - os.close(temp_file_descriptor) - - connector = aiohttp.TCPConnector(resolver=PinnedIPResolver(pinned_ip)) - timeout = aiohttp.ClientTimeout(total=timeout_seconds) - - try: - async with aiohttp.ClientSession( - connector=connector, - timeout=timeout, - ) as session: - async with session.get( - validated_url, - allow_redirects=False, - headers={"User-Agent": user_agent}, - ) as response: - if not 200 <= response.status < 300: - raise ValidationException( - user_message="Invalid URL", - violations=[ - { - "field": field, - "description": f"URL request failed with status {response.status}", - } - ], - ) - - with open(temp_file_path, "wb") as output_file: - async for chunk in response.content.iter_chunked(65536): - if chunk: - output_file.write(chunk) - - return PinnedDownloadResult( - status=response.status, - temp_file_path=temp_file_path, - ) - except Exception: - if os.path.exists(temp_file_path): - os.remove(temp_file_path) - raise - - -def download_pinned_outbound_file( - *, - url: str, - pinned_ip: str, - timeout_seconds: float, - user_agent: str, - temp_dir: str | None = None, - field: str = "source_url", -) -> PinnedDownloadResult: - """ - Download a file through a resolver pinned to a validated IP. - - Redirects are blocked to prevent redirect-based SSRF. The request connects - to the pre-validated IP while preserving the original Host header. - """ - validated_url = _validate_download_url(url, field) - parsed_url = urlsplit(validated_url) - - temp_file_descriptor, temp_file_path = tempfile.mkstemp(dir=temp_dir) - os.close(temp_file_descriptor) - - try: - connection_pool: HTTPConnectionPool - request_path = parsed_url.path or "/" - if parsed_url.query: - request_path = f"{request_path}?{parsed_url.query}" - - if parsed_url.scheme == "https": - connection_pool = PinnedHTTPSConnectionPool( - parsed_url.hostname, - parsed_url.port or 443, - pinned_ip=pinned_ip, - retries=Retry(total=0, redirect=False), - ) - else: - connection_pool = PinnedHTTPConnectionPool( - parsed_url.hostname, - parsed_url.port or 80, - pinned_ip=pinned_ip, - retries=Retry(total=0, redirect=False), - ) - - response = connection_pool.urlopen( - "GET", - request_path, - timeout=Timeout.from_float(timeout_seconds), - preload_content=False, - redirect=False, - headers={"User-Agent": user_agent, "Host": _build_host_header(parsed_url)}, - ) - try: - if not 200 <= response.status < 300: - raise ValidationException( - user_message="Invalid URL", - violations=[ - { - "field": field, - "description": f"URL request failed with status {response.status}", - } - ], - ) - - with open(temp_file_path, "wb") as output_file: - for chunk in response.stream(65536): - if chunk: - output_file.write(chunk) - finally: - response.release_conn() - response.close() - - return PinnedDownloadResult( - status=response.status, - temp_file_path=temp_file_path, - ) - except Exception: - if os.path.exists(temp_file_path): - os.remove(temp_file_path) - raise +"""Compatibility wrapper for shared.services.http.pinned_outbound.""" + +from shared.services.http.pinned_outbound import ( + PinnedDownloadResult, + PinnedHTTPConnection, + PinnedHTTPConnectionPool, + PinnedHTTPSConnection, + PinnedHTTPSConnectionPool, + PinnedIPResolver, + PinnedOutboundResponse, + download_pinned_outbound_file, + download_pinned_outbound_file_async, + send_pinned_outbound_request, +) + +__all__ = [ + "PinnedDownloadResult", + "PinnedHTTPConnection", + "PinnedHTTPConnectionPool", + "PinnedHTTPSConnection", + "PinnedHTTPSConnectionPool", + "PinnedIPResolver", + "PinnedOutboundResponse", + "download_pinned_outbound_file", + "download_pinned_outbound_file_async", + "send_pinned_outbound_request", +] diff --git a/packages/shared-python/shared/utils/quota_manager.py b/packages/shared-python/shared/utils/quota_manager.py index 735f8fa87..7acf2b8bb 100644 --- a/packages/shared-python/shared/utils/quota_manager.py +++ b/packages/shared-python/shared/utils/quota_manager.py @@ -1,368 +1,5 @@ -""" -BaseQuotaManager — Redis-backed quota management for any token-pool API. +"""Compatibility wrapper for shared quota token-pool services.""" -Design goals: -- Coordinate token usage across workers with Redis -- Enforce per-token RPM and daily quotas before calling the provider -- Keep selection fair with a quota-aware round-robin cursor -- Cool down tokens temporarily when the provider responds with 429 -""" +from shared.services.quota.token_pool import BaseQuotaManager, TokenConfig, TokenLease -from __future__ import annotations - -import json -import time -from dataclasses import dataclass -from typing import Any, Iterable, List, Optional - -from loguru import logger - -from shared.core.exceptions.domain_exceptions import UnavailableException -from shared.services.redis.redis_sync_service import SyncRedisService - - -@dataclass(frozen=True) -class TokenConfig: - """Generic configuration for a single API token.""" - - token_id: str - api_key: str - rpm_limit: int - daily_limit: int - - -@dataclass(frozen=True) -class TokenLease: - """A token reservation returned by acquire_request. - - Callers use ``api_key`` for the Authorization header, and pass - ``token_id`` back to ``mark_rate_limited`` when the provider - signals 429. - """ - - token_id: str - api_key: str - rpm_limit: int - daily_limit: int - - -class BaseQuotaManager: - """Abstract Redis-backed token pool with round-robin and cooldown. - - Subclasses only need to set: - - ``SERVICE_PREFIX`` — Redis key namespace (e.g. ``"ali"`` or ``"mineru"``) - - ``CURSOR_KEY`` — Redis key for the round-robin cursor - - ``default_cooldown_seconds`` — fallback cooldown when no Retry-After header - - ``user_message`` — user-facing error string on exhaustion - """ - - SERVICE_PREFIX: str = "base" - CURSOR_KEY: str = "base:quota:cursor" - - TOKEN_COOLDOWN_TTL_PADDING: int = 60 - MINUTE_WINDOW_SECONDS: int = 60 - DAY_WINDOW_SECONDS: int = 24 * 60 * 60 - - _RESERVE_REQUEST_SCRIPT: str = """ -local minute_limit = tonumber(ARGV[1]) -local day_limit = tonumber(ARGV[2]) -local minute_ttl = tonumber(ARGV[3]) -local day_ttl = tonumber(ARGV[4]) -local now_ts = tonumber(ARGV[5]) - -local cooldown_until = tonumber(redis.call('GET', KEYS[3]) or '0') -if cooldown_until > now_ts then - return {0, 'cooldown', cooldown_until - now_ts} -end - -local minute_count = tonumber(redis.call('GET', KEYS[1]) or '0') -if minute_count >= minute_limit then - local ttl = redis.call('TTL', KEYS[1]) - if ttl < 0 then - ttl = minute_ttl - redis.call('EXPIRE', KEYS[1], minute_ttl) - end - return {0, 'minute', ttl} -end - -local day_count = tonumber(redis.call('GET', KEYS[2]) or '0') -if day_count >= day_limit then - local ttl = redis.call('TTL', KEYS[2]) - if ttl < 0 then - ttl = day_ttl - redis.call('EXPIRE', KEYS[2], day_ttl) - end - return {0, 'day', ttl} -end - -minute_count = redis.call('INCR', KEYS[1]) -if minute_count == 1 then - redis.call('EXPIRE', KEYS[1], minute_ttl) -end - -day_count = redis.call('INCR', KEYS[2]) -if day_count == 1 then - redis.call('EXPIRE', KEYS[2], day_ttl) -end - -local minute_ttl_value = redis.call('TTL', KEYS[1]) -if minute_ttl_value < 0 then - minute_ttl_value = minute_ttl -end - -local day_ttl_value = redis.call('TTL', KEYS[2]) -if day_ttl_value < 0 then - day_ttl_value = day_ttl -end - -return {1, minute_count, day_count, minute_ttl_value, day_ttl_value} -""" - - # ---- overridable defaults ---- - default_cooldown_seconds: int = 60 - user_message: str = "Service is busy right now. Please retry shortly." - - def __init__( - self, redis_service: SyncRedisService, token_configs: Iterable[TokenConfig] - ) -> None: - self.redis = redis_service - self.tokens: List[TokenConfig] = list(token_configs) - if not self.tokens: - raise ValueError(f"{type(self).__name__} requires at least one token") - - # ------------------------------------------------------------------ - # Public API - # ------------------------------------------------------------------ - - def acquire_request( - self, - operation: str, - preferred_token_id: Optional[str] = None, - ) -> TokenLease: - """Reserve request capacity on one token before calling the provider.""" - now_ts = int(time.time()) - candidate_tokens = self._ordered_tokens(preferred_token_id) - earliest_retry_after: Optional[int] = None - retry_period = "minute" - - for token in candidate_tokens: - reserved, retry_after, quota_period = self._reserve_token_capacity( - token, now_ts - ) - if reserved: - logger.debug( - f"Reserved {self.SERVICE_PREFIX} token {token.token_id} " - f"for operation={operation}" - ) - return TokenLease( - token_id=token.token_id, - api_key=token.api_key, - rpm_limit=token.rpm_limit, - daily_limit=token.daily_limit, - ) - retry_after = max(1, retry_after) - if earliest_retry_after is None or retry_after < earliest_retry_after: - earliest_retry_after = retry_after - retry_period = quota_period - - logger.bind( - service=self.SERVICE_PREFIX, - step="quota_exhausted", - operation=operation, - preferred_token_id=preferred_token_id, - retry_after=earliest_retry_after or self.default_cooldown_seconds, - retry_period=retry_period, - token_count=len(self.tokens), - ).warning(f"{self.SERVICE_PREFIX} quota exhausted") - raise UnavailableException( - internal_message=( - f"{self.SERVICE_PREFIX} quota exhausted for operation={operation}; " - f"no token has available capacity right now" - ), - retry_after=earliest_retry_after or self.default_cooldown_seconds, - limit=max(token.rpm_limit for token in self.tokens), - period=retry_period, - user_message=self.user_message, - ) - - def mark_rate_limited( - self, token_id: str, retry_after: Optional[int] = None - ) -> None: - """Temporarily cool a token when the provider signals rate limiting.""" - retry_after_seconds = max(1, retry_after or self.default_cooldown_seconds) - cooldown_until = int(time.time()) + retry_after_seconds - cooldown_key = self._cooldown_key(token_id) - self.redis.set( - cooldown_key, - cooldown_until, - ttl=retry_after_seconds + self.TOKEN_COOLDOWN_TTL_PADDING, - ) - logger.warning( - f"{self.SERVICE_PREFIX} token {token_id} cooling down for " - f"{retry_after_seconds}s (until {cooldown_until})" - ) - - # ------------------------------------------------------------------ - # Token parsing utility (shared across subclasses) - # ------------------------------------------------------------------ - - @staticmethod - def parse_token_specs( - raw_value: str, - *, - default_rpm_limit: int, - default_daily_limit: int, - ) -> List[TokenConfig]: - """Parse token configs from JSON, newline, or comma-separated env values.""" - specs: List[TokenConfig] = [] - - def build_token(entry: Any, index: int) -> TokenConfig: - if isinstance(entry, dict): - api_key = str(entry.get("api_key") or entry.get("key") or "").strip() - if not api_key: - raise ValueError("Token entry missing api_key") - token_id = str( - entry.get("token_id") or entry.get("id") or f"token-{index + 1}" - ) - rpm_limit = int(entry.get("rpm_limit") or default_rpm_limit) - daily_limit = int(entry.get("daily_limit") or default_daily_limit) - return TokenConfig( - token_id=token_id, - api_key=api_key, - rpm_limit=rpm_limit, - daily_limit=daily_limit, - ) - - entry_text = str(entry).strip() - if not entry_text: - raise ValueError("Empty token entry") - - token_id = f"token-{index + 1}" - api_key = entry_text - if "=" in entry_text: - maybe_id, maybe_key = entry_text.split("=", 1) - if maybe_key.strip(): - token_id = maybe_id.strip() or token_id - api_key = maybe_key.strip() - if not api_key: - raise ValueError("Token entry missing api_key") - - return TokenConfig( - token_id=token_id, - api_key=api_key, - rpm_limit=default_rpm_limit, - daily_limit=default_daily_limit, - ) - - try: - loaded = json.loads(raw_value) - except json.JSONDecodeError: - loaded = None - - if isinstance(loaded, list): - for index, entry in enumerate(loaded): - specs.append(build_token(entry, index)) - return specs - - for index, line in enumerate(raw_value.replace(",", "\n").splitlines()): - line = line.strip() - if not line: - continue - specs.append(build_token(line, index)) - - return specs - - # ------------------------------------------------------------------ - # Internal helpers - # ------------------------------------------------------------------ - - def _ordered_tokens(self, preferred_token_id: Optional[str]) -> List[TokenConfig]: - if preferred_token_id: - preferred = [ - token for token in self.tokens if token.token_id == preferred_token_id - ] - remaining = [ - token - for token in self._ordered_tokens(None) - if token.token_id != preferred_token_id - ] - return preferred + remaining - - raw_cursor = self.redis.get(self.CURSOR_KEY, 0) - try: - cursor = int(raw_cursor or 0) - except (TypeError, ValueError): - cursor = 0 - - ordered = self.tokens[cursor:] + self.tokens[:cursor] - next_cursor = (cursor + 1) % len(self.tokens) - self.redis.set(self.CURSOR_KEY, next_cursor, ttl=self.DAY_WINDOW_SECONDS) - return ordered - - def _reserve_token_capacity( - self, token: TokenConfig, now_ts: int - ) -> tuple[bool, int, str]: - minute_key = self._minute_key(token.token_id, now_ts) - day_key = self._day_key(token.token_id, now_ts) - cooldown_key = self._cooldown_key(token.token_id) - - result = self.redis.eval( - self._RESERVE_REQUEST_SCRIPT, - keys=[minute_key, day_key, cooldown_key], - args=[ - token.rpm_limit, - token.daily_limit, - self._seconds_until_next_minute(now_ts), - self._seconds_until_next_day(now_ts), - now_ts, - ], - ) - - if isinstance(result, list) and result and int(result[0]) == 1: - return True, 0, "minute" - - if isinstance(result, list) and len(result) >= 3: - period = ( - str(result[1]) - if result[1] in {"minute", "day", "cooldown"} - else "minute" - ) - if period == "cooldown": - period = "minute" - return False, int(result[2]), period - - return False, self.default_cooldown_seconds, "minute" - - @staticmethod - def _seconds_until_next_minute(now_ts: int) -> int: - return max( - 1, - BaseQuotaManager.MINUTE_WINDOW_SECONDS - - (now_ts % BaseQuotaManager.MINUTE_WINDOW_SECONDS), - ) - - @staticmethod - def _seconds_until_next_day(now_ts: int) -> int: - return max( - 1, - BaseQuotaManager.DAY_WINDOW_SECONDS - - (now_ts % BaseQuotaManager.DAY_WINDOW_SECONDS), - ) - - def _minute_key(self, token_id: str, now_ts: int) -> str: - minute_bucket = now_ts // self.MINUTE_WINDOW_SECONDS - slot_tag = self._token_slot_tag(token_id) - return f"{self.SERVICE_PREFIX}:quota:{{{slot_tag}}}:minute:{minute_bucket}" - - def _day_key(self, token_id: str, now_ts: int) -> str: - day_bucket = now_ts // self.DAY_WINDOW_SECONDS - slot_tag = self._token_slot_tag(token_id) - return f"{self.SERVICE_PREFIX}:quota:{{{slot_tag}}}:day:{day_bucket}" - - def _cooldown_key(self, token_id: str) -> str: - slot_tag = self._token_slot_tag(token_id) - return f"{self.SERVICE_PREFIX}:quota:{{{slot_tag}}}:cooldown" - - @staticmethod - def _token_slot_tag(token_id: str) -> str: - return str(token_id).replace("{", "_").replace("}", "_") +__all__ = ["BaseQuotaManager", "TokenConfig", "TokenLease"] diff --git a/packages/shared-python/shared/utils/redis_key_builder.py b/packages/shared-python/shared/utils/redis_key_builder.py index 62d3d5b28..720293354 100644 --- a/packages/shared-python/shared/utils/redis_key_builder.py +++ b/packages/shared-python/shared/utils/redis_key_builder.py @@ -1,287 +1,9 @@ -"""Helpers for standardized Redis key naming.""" +"""Compatibility wrapper for Redis key helpers.""" -from enum import Enum -from typing import Any, Dict +from shared.services.redis.key_builder import ( + RedisKeyBuilder, + RedisKeyType, + redis_key_builder, +) - -class RedisKeyType(Enum): - """Enumeration of Redis key types.""" - - USER = "user" - TASK = "task" - CONVERSATION = "conversation" - KNOWLEDGE_BASE = "kb" - SESSION = "session" - CACHE = "cache" - QUEUE = "queue" - COUNTER = "counter" - SET = "set" - HASH = "hash" - LIST = "list" - - -class RedisKeyBuilder: - """Redis key builder.""" - - def __init__(self, prefix: str = "knowhere-api"): - self.prefix = prefix - - def build_key(self, key_type: RedisKeyType, *parts: str, **kwargs) -> str: - """ - Build a Redis key. - - Args: - key_type: Key type. - *parts: Positional key segments. - **kwargs: Additional named segments. - - Returns: - Full Redis key. - """ - key_parts = [self.prefix, key_type.value] - key_parts.extend(parts) - - # Append any extra key-value segments. - if kwargs: - for key, value in sorted(kwargs.items()): - key_parts.append(f"{key}:{value}") - - return ":".join(key_parts) - - # ==================== User Keys ==================== - - def user_config(self, username: str) -> str: - """User config key.""" - return self.build_key(RedisKeyType.USER, username, "config") - - def user_session(self, username: str) -> str: - """User session key.""" - return self.build_key(RedisKeyType.USER, username, "session") - - def user_permissions(self, username: str) -> str: - """User permissions key.""" - return self.build_key(RedisKeyType.USER, username, "permissions") - - def user_activity(self, username: str) -> str: - """User activity key.""" - return self.build_key(RedisKeyType.USER, username, "activity") - - # ==================== Task Keys ==================== - - def task_status(self, task_id: str) -> str: - """Task status key.""" - return self.build_key(RedisKeyType.TASK, task_id, "status") - - def task_result(self, task_id: str) -> str: - """Task result key.""" - return self.build_key(RedisKeyType.TASK, task_id, "result") - - def task_metadata(self, task_id: str) -> str: - """Task metadata key.""" - return self.build_key(RedisKeyType.TASK, task_id, "metadata") - - def task_progress(self, task_id: str) -> str: - """Task progress key.""" - return self.build_key(RedisKeyType.TASK, task_id, "progress") - - def task_info(self, task_id: str) -> str: - """Task info key.""" - return self.build_key(RedisKeyType.TASK, task_id, "info") - - def task_queue(self, queue_name: str) -> str: - """Task queue key.""" - return self.build_key(RedisKeyType.QUEUE, queue_name) - - # ==================== Conversation Keys ==================== - - def conversation_state(self, conversation_id: str) -> str: - """Conversation state key.""" - return self.build_key(RedisKeyType.CONVERSATION, conversation_id, "state") - - def conversation_history(self, conversation_id: str) -> str: - """Conversation history key.""" - return self.build_key(RedisKeyType.CONVERSATION, conversation_id, "history") - - def conversation_context(self, conversation_id: str) -> str: - """Conversation context key.""" - return self.build_key(RedisKeyType.CONVERSATION, conversation_id, "context") - - # ==================== Knowledge Base Keys ==================== - - def kb_status(self, user_id: str) -> str: - """Knowledge-base status key.""" - return self.build_key(RedisKeyType.KNOWLEDGE_BASE, user_id, "status") - - def kb_vectors(self, user_id: str) -> str: - """Knowledge-base vectors key.""" - return self.build_key(RedisKeyType.KNOWLEDGE_BASE, user_id, "vectors") - - def kb_metadata(self, user_id: str) -> str: - """Knowledge-base metadata key.""" - return self.build_key(RedisKeyType.KNOWLEDGE_BASE, user_id, "metadata") - - def kb_index(self, user_id: str) -> str: - """Knowledge-base index key.""" - return self.build_key(RedisKeyType.KNOWLEDGE_BASE, user_id, "index") - - # ==================== Session Keys ==================== - - def session_data(self, session_id: str) -> str: - """Session data key.""" - return self.build_key(RedisKeyType.SESSION, session_id, "data") - - def session_expiry(self, session_id: str) -> str: - """Session expiry key.""" - return self.build_key(RedisKeyType.SESSION, session_id, "expiry") - - # ==================== Cache Keys ==================== - - def cache_data(self, cache_key: str) -> str: - """Cache data key.""" - return self.build_key(RedisKeyType.CACHE, cache_key) - - def cache_metadata(self, cache_key: str) -> str: - """Cache metadata key.""" - return self.build_key(RedisKeyType.CACHE, cache_key, "metadata") - - # ==================== Counter Keys ==================== - - def counter_user_requests(self, username: str) -> str: - """User-request counter key.""" - return self.build_key(RedisKeyType.COUNTER, "user_requests", username) - - def counter_api_calls(self, api_name: str) -> str: - """API-call counter key.""" - return self.build_key(RedisKeyType.COUNTER, "api_calls", api_name) - - def counter_task_completed(self, task_type: str) -> str: - """Task-completed counter key.""" - return self.build_key(RedisKeyType.COUNTER, "task_completed", task_type) - - def rate_limit_api(self, user_id: str, api_name: str) -> str: - """API rate-limit key.""" - return self.build_key(RedisKeyType.COUNTER, "rate_limit", api_name, user_id) - - # ==================== Set Keys ==================== - - def set_active_users(self) -> str: - """Active-users set key.""" - return self.build_key(RedisKeyType.SET, "active_users") - - def set_online_users(self) -> str: - """Online-users set key.""" - return self.build_key(RedisKeyType.SET, "online_users") - - def set_processing_tasks(self) -> str: - """Processing-tasks set key.""" - return self.build_key(RedisKeyType.SET, "processing_tasks") - - # ==================== Hash Keys ==================== - - def hash_user_stats(self, username: str) -> str: - """User-stats hash key.""" - return self.build_key(RedisKeyType.HASH, "user_stats", username) - - def hash_system_stats(self) -> str: - """System-stats hash key.""" - return self.build_key(RedisKeyType.HASH, "system_stats") - - def hash_task_stats(self, task_id: str) -> str: - """Task-stats hash key.""" - return self.build_key(RedisKeyType.HASH, "task_stats", task_id) - - # ==================== List Keys ==================== - - def list_recent_activities(self, username: str) -> str: - """Recent-activities list key.""" - return self.build_key(RedisKeyType.LIST, "recent_activities", username) - - def list_error_logs(self) -> str: - """Error-logs list key.""" - return self.build_key(RedisKeyType.LIST, "error_logs") - - def list_audit_logs(self) -> str: - """Audit-logs list key.""" - return self.build_key(RedisKeyType.LIST, "audit_logs") - - # ==================== Lock Keys ==================== - - def lock_job_processing(self, job_id: str) -> str: - """Distributed lock key for exclusive job processing.""" - return f"lock:job_processing:{job_id}" - - # ==================== Utility Methods ==================== - - def parse_key(self, key: str) -> Dict[str, Any]: - """ - Parse a Redis key into its components. - - Args: - key: Redis key. - - Returns: - Parsed key information. - """ - if not key.startswith(self.prefix): - raise ValueError(f"Key '{key}' is not a valid Redis key format") - - parts = key[len(self.prefix) :].lstrip(":").split(":") - - if len(parts) < 2: - raise ValueError(f"Key '{key}' has an invalid format") - - result = { - "prefix": self.prefix, - "type": parts[0], - "parts": parts[1:], - "original_key": key, - } - - return result - - def is_key_type(self, key: str, key_type: RedisKeyType) -> bool: - """ - Check whether a key belongs to the requested type. - - Args: - key: Redis key. - key_type: Key type. - - Returns: - Whether the key matches the requested type. - """ - try: - parsed = self.parse_key(key) - return parsed["type"] == key_type.value - except ValueError: - return False - - def get_key_ttl(self, key_type: RedisKeyType) -> int: - """ - Get the default TTL for a key type. - - Args: - key_type: Key type. - - Returns: - TTL in seconds. - """ - ttl_mapping = { - RedisKeyType.USER: 86400, # 1 day (user_config cache). - RedisKeyType.TASK: 86400, # 1 day. - RedisKeyType.CONVERSATION: 3600 * 2, # 2 hours. - RedisKeyType.KNOWLEDGE_BASE: 86400 * 30, # 30 days. - RedisKeyType.SESSION: 3600, # 1 hour. - RedisKeyType.CACHE: 3600, # 1 hour. - RedisKeyType.QUEUE: 86400, # 1 day. - RedisKeyType.COUNTER: 86400, # 1 day. - RedisKeyType.SET: 86400, # 1 day. - RedisKeyType.HASH: 86400, # 1 day. - RedisKeyType.LIST: 86400, # 1 day. - } - - return ttl_mapping.get(key_type, 3600) # Default to 1 hour. - - -# Global key-builder instance. -redis_key_builder = RedisKeyBuilder() +__all__ = ["RedisKeyBuilder", "RedisKeyType", "redis_key_builder"] diff --git a/packages/shared-python/shared/utils/redis_retry.py b/packages/shared-python/shared/utils/redis_retry.py index e1915605a..f9b838c02 100644 --- a/packages/shared-python/shared/utils/redis_retry.py +++ b/packages/shared-python/shared/utils/redis_retry.py @@ -1,146 +1,5 @@ -"""Retry helpers for Redis operations.""" +"""Compatibility wrapper for Redis retry helpers.""" -import asyncio -import random -from typing import Any, Callable +from shared.services.redis.retry_policy import RedisHealthChecker, RedisRetry -from loguru import logger - -from shared.core.exceptions.redis_exceptions import ( - RedisConnectionError, - RedisOperationError, - RedisTimeoutError, -) - - -class RedisRetry: - """Retry helper for Redis operations.""" - - @staticmethod - async def with_retry( - operation: Callable, - max_retries: int = 3, - base_delay: float = 1.0, - max_delay: float = 60.0, - exponential_base: float = 2.0, - jitter: bool = True, - ) -> Any: - """ - Execute a Redis operation with retry support. - - Args: - operation: Operation callback to execute. - max_retries: Maximum retry count. - base_delay: Base delay in seconds. - max_delay: Maximum delay in seconds. - exponential_base: Exponential backoff base. - jitter: Whether to add random jitter. - - Returns: - Operation result. - - Raises: - RedisOperationError: Raised when all retries are exhausted. - """ - last_exception = None - - for attempt in range(max_retries + 1): - try: - return await operation() - except (RedisConnectionError, RedisOperationError, RedisTimeoutError) as e: - last_exception = e - - if attempt == max_retries: - logger.error( - f"Redis operation failed after {max_retries} retries: {e}" - ) - raise RedisOperationError( - internal_message=f"Redis operation failed after {max_retries} retries: {str(e)}", - original_exception=e, - ) - - # Compute the retry delay. - delay = min(base_delay * (exponential_base**attempt), max_delay) - - # Add random jitter when enabled. - if jitter: - delay *= 0.5 + random.random() * 0.5 - - logger.warning( - f"Redis operation failed, retrying in {delay:.2f}s (attempt {attempt + 1}/{max_retries + 1}): {e}" - ) - await asyncio.sleep(delay) - - # Should not reach here - raise RedisOperationError( - internal_message=f"Redis operation failed: {str(last_exception)}", - original_exception=last_exception, - ) - - @staticmethod - async def with_circuit_breaker( - operation: Callable, - failure_threshold: int = 5, - recovery_timeout: float = 60.0, - expected_exception: type = Exception, - ) -> Any: - """ - Execute a Redis operation with a circuit-breaker placeholder. - - Args: - operation: Operation callback to execute. - failure_threshold: Failure threshold. - recovery_timeout: Recovery timeout. - expected_exception: Expected exception type. - - Returns: - Operation result. - """ - # Circuit-breaker logic can be implemented here later. - # For now, execute the operation directly. - return await operation() - - -class RedisHealthChecker: - """Health checker for Redis connections.""" - - def __init__(self, redis_client): - self.redis_client = redis_client - self._is_healthy = True - self._last_check = 0 - self._check_interval = 30 # Check every 30 seconds. - - async def is_healthy(self) -> bool: - """Check whether Redis is healthy.""" - import time - - current_time = time.time() - - # Return the cached result if the last check is still fresh. - if current_time - self._last_check < self._check_interval: - return self._is_healthy - - try: - # Use PING to validate the connection. - await self.redis_client.ping() - self._is_healthy = True - self._last_check = current_time - return True - except Exception as e: - logger.error(f"Redis health check failed: {e}") - self._is_healthy = False - self._last_check = current_time - return False - - async def wait_for_healthy(self, timeout: float = 30.0) -> bool: - """Wait for Redis to become healthy again.""" - import time - - start_time = time.time() - - while time.time() - start_time < timeout: - if await self.is_healthy(): - return True - await asyncio.sleep(1) - - return False +__all__ = ["RedisHealthChecker", "RedisRetry"] diff --git a/packages/shared-python/shared/utils/url_file_type.py b/packages/shared-python/shared/utils/url_file_type.py index 7314d7e23..290299c0f 100644 --- a/packages/shared-python/shared/utils/url_file_type.py +++ b/packages/shared-python/shared/utils/url_file_type.py @@ -1,205 +1,19 @@ -""" -Resolve file extension from a URL. - -Tries the URL path first, then falls back to a HEAD request to read Content-Type. -Provides both async (for API) and sync (for worker) variants. -""" - -import os -from urllib.parse import urljoin, urlparse - -from loguru import logger - -from shared.core.config import settings -from shared.core.exceptions.domain_exceptions import ValidationException -from shared.utils.url_security import ( - HTTPURLValidationResult, - SafePublicHTTPURL, - validate_http_url_and_resolve_ip, +"""Compatibility wrapper for shared.services.http.url_file_type.""" + +from shared.services.http.url_file_type import ( + CONTENT_TYPE_TO_EXTENSION, + MAX_SAFE_REDIRECTS, + REDIRECT_STATUS_CODES, + URL_VALIDATION_DESCRIPTIONS, + resolve_file_extension_async, + resolve_file_extension_sync, ) -# Content-Type to file extension mapping -CONTENT_TYPE_TO_EXTENSION: dict[str, str] = { - "application/pdf": ".pdf", - "application/msword": ".doc", - "application/vnd.openxmlformats-officedocument.wordprocessingml.document": ".docx", - "application/vnd.ms-excel": ".xls", - "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": ".xlsx", - "application/vnd.ms-powerpoint": ".ppt", - "application/vnd.openxmlformats-officedocument.presentationml.presentation": ".pptx", - "text/csv": ".csv", - "text/plain": ".txt", - "text/markdown": ".md", - "image/png": ".png", - "image/jpeg": ".jpg", - "image/gif": ".gif", - "image/bmp": ".bmp", - "image/tiff": ".tiff", - "image/svg+xml": ".svg", -} - -REDIRECT_STATUS_CODES: set[int] = {301, 302, 303, 307, 308} -MAX_SAFE_REDIRECTS: int = 5 -URL_VALIDATION_DESCRIPTIONS: dict[str, str] = { - "unsupported_scheme": "URL must use http or https", - "missing_hostname": "URL must include a hostname", - "hostname_resolution_failed": "URL hostname could not be resolved", - "invalid_resolved_address": "URL resolved to an invalid IP", - "hostname_not_allowed": "URL host is not allowed", -} - - -def _validate_source_url(url: str, field: str) -> SafePublicHTTPURL: - validation = validate_http_url_and_resolve_ip(url) - if not validation.is_valid: - _raise_url_validation_error(field, validation) - return SafePublicHTTPURL(url) - - -def _validate_redirect_url( - url: str, - redirect_url: str, - field: str, -) -> SafePublicHTTPURL: - return _validate_source_url(urljoin(url, redirect_url), field=field) - - -def _raise_url_validation_error( - field: str, - validation: HTTPURLValidationResult, -) -> None: - description = URL_VALIDATION_DESCRIPTIONS.get( - str(validation.failure_reason), - "URL is invalid", - ) - - raise ValidationException( - user_message="Invalid URL", - violations=[{"field": field, "description": description}], - internal_message=validation.error_message, - ) - - -def _extension_from_path(url: str) -> str | None: - """Extract a recognised file extension from the URL path.""" - path = urlparse(url).path - ext = os.path.splitext(path)[1].lower() - if ext and ext in settings.get_supported_extensions(): - return ext - return None - - -def _extension_from_content_type(content_type: str | None) -> str | None: - """Map a Content-Type header value to a supported file extension.""" - if not content_type: - return None - # Strip parameters like "; charset=utf-8" - mime = content_type.split(";")[0].strip().lower() - ext = CONTENT_TYPE_TO_EXTENSION.get(mime) - if ext and ext in settings.get_supported_extensions(): - return ext - return None - - -async def resolve_file_extension_async(url: str) -> str | None: - """ - Resolve file extension from a URL (async version for API layer). - - 1. Try extracting extension from URL path. - 2. If that fails, send a HEAD request and read Content-Type. - 3. Return None if neither method produces a supported extension. - """ - safe_url = _validate_source_url(url, field="source_url") - - ext = _extension_from_path(safe_url) - if ext: - return ext - - try: - from shared.utils.http_clients import get_async_client - - client = get_async_client() - response = None - request_url: SafePublicHTTPURL = safe_url - for _ in range(MAX_SAFE_REDIRECTS + 1): - response = await client.head(request_url, follow_redirects=False) - if response.status_code not in REDIRECT_STATUS_CODES: - break - - location = response.headers.get("location") - if not location: - break - request_url = _validate_redirect_url( - request_url, - location, - field="source_url", - ) - - if response is None: - return None - content_type = response.headers.get("content-type") - ext = _extension_from_content_type(content_type) - if ext: - logger.info( - f"Resolved file extension from Content-Type header: {ext}" - ) - return ext - except ValidationException: - raise - except Exception as exc: - logger.warning( - f"HEAD request failed for URL file type detection: {exc}" - ) - - return None - - -def resolve_file_extension_sync(url: str) -> str | None: - """ - Resolve file extension from a URL (sync version for worker layer). - - Same logic as async variant but uses the shared sync httpx client. - """ - safe_url = _validate_source_url(url, field="source_url") - - ext = _extension_from_path(safe_url) - if ext: - return ext - - try: - from shared.utils.http_clients import get_sync_client - - client = get_sync_client() - response = None - request_url: SafePublicHTTPURL = safe_url - for _ in range(MAX_SAFE_REDIRECTS + 1): - response = client.head(request_url, follow_redirects=False) - if response.status_code not in REDIRECT_STATUS_CODES: - break - - location = response.headers.get("location") - if not location: - break - request_url = _validate_redirect_url( - request_url, - location, - field="source_url", - ) - - if response is None: - return None - content_type = response.headers.get("content-type") - ext = _extension_from_content_type(content_type) - if ext: - logger.info( - f"Resolved file extension from Content-Type header: {ext}" - ) - return ext - except ValidationException: - raise - except Exception as exc: - logger.warning( - f"HEAD request failed for URL file type detection: {exc}" - ) - - return None +__all__ = [ + "CONTENT_TYPE_TO_EXTENSION", + "MAX_SAFE_REDIRECTS", + "REDIRECT_STATUS_CODES", + "URL_VALIDATION_DESCRIPTIONS", + "resolve_file_extension_async", + "resolve_file_extension_sync", +] diff --git a/packages/shared-python/shared/utils/url_security.py b/packages/shared-python/shared/utils/url_security.py index bc54c1107..ca19d4dc0 100644 --- a/packages/shared-python/shared/utils/url_security.py +++ b/packages/shared-python/shared/utils/url_security.py @@ -1,243 +1,23 @@ -import asyncio -import ipaddress -import socket -from dataclasses import dataclass -from typing import Literal, TypeAlias, cast -from urllib.parse import urlparse - -SocketAddress: TypeAlias = tuple[object, ...] -AddressInfo: TypeAlias = tuple[int, int, int, str, SocketAddress] -AllowedIPAddress: TypeAlias = ipaddress.IPv4Address | ipaddress.IPv6Address -URLValidationFailureReason: TypeAlias = Literal[ - "unsupported_scheme", - "missing_hostname", - "hostname_resolution_failed", - "invalid_resolved_address", - "hostname_not_allowed", - "validation_failed", -] -HTTP_SCHEMES: frozenset[str] = frozenset({"http", "https"}) -DEVELOPMENT_ENVIRONMENTS: frozenset[str] = frozenset({"dev", "development", "local"}) -BLOCKED_HOSTNAMES: frozenset[str] = frozenset( - {"localhost", "ip6-localhost", "ip6-loopback"} +"""Compatibility wrapper for shared.services.http.url_security.""" + +from shared.services.http.url_security import ( + BLOCKED_HOSTNAMES, + DEVELOPMENT_ENVIRONMENTS, + HTTP_SCHEMES, + HTTPURLValidationResult, + SafePublicHTTPURL, + URLValidationFailureReason, + validate_http_url_and_resolve_ip, + validate_http_url_and_resolve_ip_async, ) - -class SafePublicHTTPURL(str): - """A URL string that has passed public HTTP SSRF validation.""" - - -@dataclass(frozen=True) -class HTTPURLValidationResult: - """Validated HTTP URL and its pinned IP address.""" - - is_valid: bool - url: str - hostname: str | None = None - validated_ip: str | None = None - error_message: str | None = None - failure_reason: URLValidationFailureReason | None = None - - -def validate_http_url_and_resolve_ip( - url: str, -) -> HTTPURLValidationResult: - """Validate an HTTP URL policy and return the pinned resolved IP address.""" - try: - can_use_private_hosts = _can_use_private_hosts() - hostname_result = _validate_url_hostname(url) - if not hostname_result.is_valid: - return hostname_result - - hostname = cast(str, hostname_result.hostname) - if not can_use_private_hosts and _is_blocked_hostname(hostname): - return _build_url_validation_failure( - url, - f"Hostname {hostname} failed validation", - "hostname_not_allowed", - ) - - try: - address_infos = cast(list[AddressInfo], socket.getaddrinfo(hostname, None)) - except socket.gaierror as exc: - return _build_url_validation_failure( - url, - f"Unable to resolve hostname {hostname}: {exc}", - "hostname_resolution_failed", - ) - - return _validate_resolved_addresses( - url, - hostname, - address_infos, - can_use_private_hosts=can_use_private_hosts, - ) - except Exception as exc: - return _build_url_validation_failure( - url, - f"URL validation failed: {exc}", - "validation_failed", - ) - - -async def validate_http_url_and_resolve_ip_async( - url: str, -) -> HTTPURLValidationResult: - """Validate an HTTP URL policy asynchronously and return the pinned IP address.""" - try: - can_use_private_hosts = _can_use_private_hosts() - hostname_result = _validate_url_hostname(url) - if not hostname_result.is_valid: - return hostname_result - - hostname = cast(str, hostname_result.hostname) - if not can_use_private_hosts and _is_blocked_hostname(hostname): - return _build_url_validation_failure( - url, - f"Hostname {hostname} failed validation", - "hostname_not_allowed", - ) - - try: - loop = asyncio.get_running_loop() - address_infos = cast( - list[AddressInfo], - await loop.getaddrinfo(hostname, None), - ) - except socket.gaierror as exc: - return _build_url_validation_failure( - url, - f"Unable to resolve hostname {hostname}: {exc}", - "hostname_resolution_failed", - ) - - return _validate_resolved_addresses( - url, - hostname, - address_infos, - can_use_private_hosts=can_use_private_hosts, - ) - except Exception as exc: - return _build_url_validation_failure( - url, - f"URL validation failed: {exc}", - "validation_failed", - ) - - -def _can_use_private_hosts() -> bool: - from shared.core.config import app_config - - return app_config.ENVIRONMENT.lower() in DEVELOPMENT_ENVIRONMENTS - - -def _extract_resolved_addresses(address_infos: list[AddressInfo]) -> list[str]: - resolved_addresses: list[str] = [] - seen_addresses: set[str] = set() - for family, _, _, _, socket_address in address_infos: - if family not in (socket.AF_INET, socket.AF_INET6): - continue - if not socket_address: - continue - address = socket_address[0] - if not isinstance(address, str): - continue - if address and address not in seen_addresses: - resolved_addresses.append(address) - seen_addresses.add(address) - return resolved_addresses - - -def _is_blocked_hostname(hostname: str) -> bool: - normalized_hostname = hostname.rstrip(".").lower() - return normalized_hostname in BLOCKED_HOSTNAMES or normalized_hostname.endswith( - (".localhost", ".local") - ) - - -def _is_public_ip_address(ip_address: AllowedIPAddress) -> bool: - return ip_address.is_global and not _is_blocked_ip_address(ip_address) - - -def _is_blocked_ip_address(ip_address: AllowedIPAddress) -> bool: - return ( - ip_address.is_private - or ip_address.is_loopback - or ip_address.is_link_local - or ip_address.is_multicast - or ip_address.is_reserved - or ip_address.is_unspecified - ) - - -def _validate_url_hostname(url: str) -> HTTPURLValidationResult: - parsed_url = urlparse(url) - if parsed_url.scheme not in HTTP_SCHEMES: - return _build_url_validation_failure( - url, - f"Unsupported URL scheme: {parsed_url.scheme}", - "unsupported_scheme", - ) - - hostname = parsed_url.hostname - if not hostname: - return _build_url_validation_failure( - url, - "URL must include a hostname", - "missing_hostname", - ) - - return HTTPURLValidationResult(is_valid=True, url=url, hostname=hostname) - - -def _validate_resolved_addresses( - url: str, - hostname: str, - address_infos: list[AddressInfo], - *, - can_use_private_hosts: bool, -) -> HTTPURLValidationResult: - resolved_addresses = _extract_resolved_addresses(address_infos) - if not resolved_addresses: - return _build_url_validation_failure( - url, - f"Unable to resolve hostname {hostname}", - "hostname_resolution_failed", - ) - - for address in resolved_addresses: - try: - ip_address = ipaddress.ip_address(address) - except ValueError: - return _build_url_validation_failure( - url, - f"Hostname {hostname} resolved to invalid IP {address}", - "invalid_resolved_address", - ) - - if not can_use_private_hosts and not _is_public_ip_address(ip_address): - return _build_url_validation_failure( - url, - f"Hostname {hostname} failed validation", - "hostname_not_allowed", - ) - - return HTTPURLValidationResult( - is_valid=True, - url=url, - hostname=hostname, - validated_ip=resolved_addresses[0], - ) - - -def _build_url_validation_failure( - url: str, - error_message: str, - failure_reason: URLValidationFailureReason, -) -> HTTPURLValidationResult: - return HTTPURLValidationResult( - is_valid=False, - url=url, - error_message=error_message, - failure_reason=failure_reason, - ) +__all__ = [ + "BLOCKED_HOSTNAMES", + "DEVELOPMENT_ENVIRONMENTS", + "HTTP_SCHEMES", + "HTTPURLValidationResult", + "SafePublicHTTPURL", + "URLValidationFailureReason", + "validate_http_url_and_resolve_ip", + "validate_http_url_and_resolve_ip_async", +] From 8a98eb942a60e91fff0a2af5c47afee6d2dba740 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 09:42:01 +0000 Subject: [PATCH 56/78] refactor: remove internal compatibility wrappers --- AGENTS.md | 2 +- .../services/document_parser/pdf_parser.py | 11 ++------ docs/adr/0002-use-typed-workflow-outcomes.md | 10 +++---- .../utils/OpenAICompatibleClientSync.py | 8 ------ .../shared/utils/ali_quota_manager.py | 5 ---- .../shared/utils/http_clients.py | 15 ----------- .../shared/utils/iloveapi_quota_manager.py | 11 -------- .../shared-python/shared/utils/llm_mock.py | 5 ---- .../shared/utils/pinned_outbound_http.py | 27 ------------------- .../shared/utils/quota_manager.py | 5 ---- .../shared/utils/redis_key_builder.py | 9 ------- .../shared-python/shared/utils/redis_retry.py | 5 ---- .../shared/utils/url_file_type.py | 19 ------------- .../shared/utils/url_security.py | 23 ---------------- 14 files changed, 8 insertions(+), 147 deletions(-) delete mode 100644 packages/shared-python/shared/utils/OpenAICompatibleClientSync.py delete mode 100644 packages/shared-python/shared/utils/ali_quota_manager.py delete mode 100644 packages/shared-python/shared/utils/http_clients.py delete mode 100644 packages/shared-python/shared/utils/iloveapi_quota_manager.py delete mode 100644 packages/shared-python/shared/utils/llm_mock.py delete mode 100644 packages/shared-python/shared/utils/pinned_outbound_http.py delete mode 100644 packages/shared-python/shared/utils/quota_manager.py delete mode 100644 packages/shared-python/shared/utils/redis_key_builder.py delete mode 100644 packages/shared-python/shared/utils/redis_retry.py delete mode 100644 packages/shared-python/shared/utils/url_file_type.py delete mode 100644 packages/shared-python/shared/utils/url_security.py diff --git a/AGENTS.md b/AGENTS.md index a52642f81..25c00e802 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -53,7 +53,7 @@ knowhereapi-main/ │ ├── services/http/ # Public URL validation and outbound HTTP │ ├── services/redis/ # Redis state, key language, and retry policy │ ├── services/quota/ # Shared token-pool quota primitives -│ └── utils/ # Compatibility and generic text/chunk utilities +│ └── utils/ # Generic text, chunk, and API helpers └── deploy/ # Docker Compose & deployment scripts ``` diff --git a/apps/worker/app/services/document_parser/pdf_parser.py b/apps/worker/app/services/document_parser/pdf_parser.py index ffc693980..fec266090 100755 --- a/apps/worker/app/services/document_parser/pdf_parser.py +++ b/apps/worker/app/services/document_parser/pdf_parser.py @@ -208,13 +208,6 @@ def _fast_path_worker(queue, pdf_path, output_dir, image_dir): ) -def upload_and_parse( - pdf_url: str, filename: str, output_dir: str, s3_key: str | None = None -) -> None: - """Compatibility wrapper for the extracted MinerU workflow module.""" - parse_via_full(pdf_url, filename, output_dir, s3_key=s3_key) - - def parse_pdfs( pdf_path, filename, @@ -259,14 +252,14 @@ def parse_pdfs( # ) # else: # with stage_timer("pdf.extract.standard", filename=filename): - # upload_and_parse(pdf_path, filename, output_dir, s3_key=s3_key) + # parse_via_full(pdf_path, filename, output_dir, s3_key=s3_key) # _inject_page_markers(output_dir) logger.info( f"🛡️ Conservative mode: forcing MinerU (standard) for {filename} [route={route}]" ) with stage_timer("pdf.extract.standard", filename=filename): - upload_and_parse(pdf_path, filename, output_dir, s3_key=s3_key) + parse_via_full(pdf_path, filename, output_dir, s3_key=s3_key) # Inject page markers from MinerU layout.json _inject_page_markers(output_dir) diff --git a/docs/adr/0002-use-typed-workflow-outcomes.md b/docs/adr/0002-use-typed-workflow-outcomes.md index 4e7b8ec50..fce8dddd3 100644 --- a/docs/adr/0002-use-typed-workflow-outcomes.md +++ b/docs/adr/0002-use-typed-workflow-outcomes.md @@ -17,12 +17,12 @@ than a yes/no answer. Current examples are `JobTransitionOutcome`, `WorkloadEstimate`, `ParseOutput`, `ParseArtifact`, `GeneratedResultPackage`, and `PostCommitEffectPlan`. -Boolean and tuple facades may remain where they are public compatibility -interfaces, but the implementation should concentrate reason-bearing behavior -behind typed outcomes. +Boolean and tuple facades may remain only where there is a real public contract. +The implementation should concentrate reason-bearing behavior behind typed +outcomes. ## Consequences Contract tests can assert stable behavior and failure reasons without reaching -into private helper call order. Compatibility wrappers should be kept small and -named as wrappers, not used as the primary implementation seam. +into private helper call order. Avoid internal-only compatibility facades; if +there is no external consumer, update callers to the real module boundary. diff --git a/packages/shared-python/shared/utils/OpenAICompatibleClientSync.py b/packages/shared-python/shared/utils/OpenAICompatibleClientSync.py deleted file mode 100644 index c9b31d6b8..000000000 --- a/packages/shared-python/shared/utils/OpenAICompatibleClientSync.py +++ /dev/null @@ -1,8 +0,0 @@ -"""Compatibility wrapper for the OpenAI-compatible sync client.""" - -from shared.services.ai.openai_compatible_client_sync import ( - OpenAICompatibleClientSync, - get_openai_client, -) - -__all__ = ["OpenAICompatibleClientSync", "get_openai_client"] diff --git a/packages/shared-python/shared/utils/ali_quota_manager.py b/packages/shared-python/shared/utils/ali_quota_manager.py deleted file mode 100644 index 304a36e8e..000000000 --- a/packages/shared-python/shared/utils/ali_quota_manager.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Compatibility wrapper for the Ali quota manager.""" - -from shared.services.ai.ali_quota_manager import AliQuotaManager, get_ali_quota_manager - -__all__ = ["AliQuotaManager", "get_ali_quota_manager"] diff --git a/packages/shared-python/shared/utils/http_clients.py b/packages/shared-python/shared/utils/http_clients.py deleted file mode 100644 index 50693141a..000000000 --- a/packages/shared-python/shared/utils/http_clients.py +++ /dev/null @@ -1,15 +0,0 @@ -"""Compatibility wrapper for shared.services.http.client_pool.""" - -from shared.services.http.client_pool import ( - close_async_client, - close_sync_client, - get_async_client, - get_sync_client, -) - -__all__ = [ - "close_async_client", - "close_sync_client", - "get_async_client", - "get_sync_client", -] diff --git a/packages/shared-python/shared/utils/iloveapi_quota_manager.py b/packages/shared-python/shared/utils/iloveapi_quota_manager.py deleted file mode 100644 index 43549ae60..000000000 --- a/packages/shared-python/shared/utils/iloveapi_quota_manager.py +++ /dev/null @@ -1,11 +0,0 @@ -"""Compatibility wrapper for the iLoveAPI quota manager.""" - -from shared.services.ai.iloveapi_quota_manager import ( - ILoveApiQuotaManager, - get_iloveapi_quota_manager, -) - -__all__ = [ - "ILoveApiQuotaManager", - "get_iloveapi_quota_manager", -] diff --git a/packages/shared-python/shared/utils/llm_mock.py b/packages/shared-python/shared/utils/llm_mock.py deleted file mode 100644 index 8bf18999e..000000000 --- a/packages/shared-python/shared/utils/llm_mock.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Compatibility wrapper for LLM mock responses.""" - -from shared.services.ai.llm_mock import build_mock_chat_completion_response - -__all__ = ["build_mock_chat_completion_response"] diff --git a/packages/shared-python/shared/utils/pinned_outbound_http.py b/packages/shared-python/shared/utils/pinned_outbound_http.py deleted file mode 100644 index 1d4e0fd3e..000000000 --- a/packages/shared-python/shared/utils/pinned_outbound_http.py +++ /dev/null @@ -1,27 +0,0 @@ -"""Compatibility wrapper for shared.services.http.pinned_outbound.""" - -from shared.services.http.pinned_outbound import ( - PinnedDownloadResult, - PinnedHTTPConnection, - PinnedHTTPConnectionPool, - PinnedHTTPSConnection, - PinnedHTTPSConnectionPool, - PinnedIPResolver, - PinnedOutboundResponse, - download_pinned_outbound_file, - download_pinned_outbound_file_async, - send_pinned_outbound_request, -) - -__all__ = [ - "PinnedDownloadResult", - "PinnedHTTPConnection", - "PinnedHTTPConnectionPool", - "PinnedHTTPSConnection", - "PinnedHTTPSConnectionPool", - "PinnedIPResolver", - "PinnedOutboundResponse", - "download_pinned_outbound_file", - "download_pinned_outbound_file_async", - "send_pinned_outbound_request", -] diff --git a/packages/shared-python/shared/utils/quota_manager.py b/packages/shared-python/shared/utils/quota_manager.py deleted file mode 100644 index 7acf2b8bb..000000000 --- a/packages/shared-python/shared/utils/quota_manager.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Compatibility wrapper for shared quota token-pool services.""" - -from shared.services.quota.token_pool import BaseQuotaManager, TokenConfig, TokenLease - -__all__ = ["BaseQuotaManager", "TokenConfig", "TokenLease"] diff --git a/packages/shared-python/shared/utils/redis_key_builder.py b/packages/shared-python/shared/utils/redis_key_builder.py deleted file mode 100644 index 720293354..000000000 --- a/packages/shared-python/shared/utils/redis_key_builder.py +++ /dev/null @@ -1,9 +0,0 @@ -"""Compatibility wrapper for Redis key helpers.""" - -from shared.services.redis.key_builder import ( - RedisKeyBuilder, - RedisKeyType, - redis_key_builder, -) - -__all__ = ["RedisKeyBuilder", "RedisKeyType", "redis_key_builder"] diff --git a/packages/shared-python/shared/utils/redis_retry.py b/packages/shared-python/shared/utils/redis_retry.py deleted file mode 100644 index f9b838c02..000000000 --- a/packages/shared-python/shared/utils/redis_retry.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Compatibility wrapper for Redis retry helpers.""" - -from shared.services.redis.retry_policy import RedisHealthChecker, RedisRetry - -__all__ = ["RedisHealthChecker", "RedisRetry"] diff --git a/packages/shared-python/shared/utils/url_file_type.py b/packages/shared-python/shared/utils/url_file_type.py deleted file mode 100644 index 290299c0f..000000000 --- a/packages/shared-python/shared/utils/url_file_type.py +++ /dev/null @@ -1,19 +0,0 @@ -"""Compatibility wrapper for shared.services.http.url_file_type.""" - -from shared.services.http.url_file_type import ( - CONTENT_TYPE_TO_EXTENSION, - MAX_SAFE_REDIRECTS, - REDIRECT_STATUS_CODES, - URL_VALIDATION_DESCRIPTIONS, - resolve_file_extension_async, - resolve_file_extension_sync, -) - -__all__ = [ - "CONTENT_TYPE_TO_EXTENSION", - "MAX_SAFE_REDIRECTS", - "REDIRECT_STATUS_CODES", - "URL_VALIDATION_DESCRIPTIONS", - "resolve_file_extension_async", - "resolve_file_extension_sync", -] diff --git a/packages/shared-python/shared/utils/url_security.py b/packages/shared-python/shared/utils/url_security.py deleted file mode 100644 index ca19d4dc0..000000000 --- a/packages/shared-python/shared/utils/url_security.py +++ /dev/null @@ -1,23 +0,0 @@ -"""Compatibility wrapper for shared.services.http.url_security.""" - -from shared.services.http.url_security import ( - BLOCKED_HOSTNAMES, - DEVELOPMENT_ENVIRONMENTS, - HTTP_SCHEMES, - HTTPURLValidationResult, - SafePublicHTTPURL, - URLValidationFailureReason, - validate_http_url_and_resolve_ip, - validate_http_url_and_resolve_ip_async, -) - -__all__ = [ - "BLOCKED_HOSTNAMES", - "DEVELOPMENT_ENVIRONMENTS", - "HTTP_SCHEMES", - "HTTPURLValidationResult", - "SafePublicHTTPURL", - "URLValidationFailureReason", - "validate_http_url_and_resolve_ip", - "validate_http_url_and_resolve_ip_async", -] From 95a8fde355ca0a1dca810bebba0819d7da0b52d8 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 09:55:42 +0000 Subject: [PATCH 57/78] refactor: group demo services by feature --- CONTEXT.md | 3 +-- apps/api/app/api/v1/routes/demo.py | 16 ++++++++-------- apps/api/app/services/demo/__init__.py | 1 + .../document_service.py} | 4 ++-- .../source_catalog.py} | 4 ++-- .../source_materializer.py} | 2 +- .../source_projection.py} | 0 .../contract/test_demo_documents_contract.py | 16 ++++++++-------- 8 files changed, 23 insertions(+), 23 deletions(-) create mode 100644 apps/api/app/services/demo/__init__.py rename apps/api/app/services/{demo_document_service.py => demo/document_service.py} (95%) rename apps/api/app/services/{demo_source_catalog.py => demo/source_catalog.py} (98%) rename apps/api/app/services/{demo_source_materializer.py => demo/source_materializer.py} (99%) rename apps/api/app/services/{demo_source_projection.py => demo/source_projection.py} (100%) diff --git a/CONTEXT.md b/CONTEXT.md index 9b69831a5..b5701ef67 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -293,8 +293,7 @@ exceptions. ### Demo Source Materialization - `app/api/v1/routes/demo.py` -- `app/services/demo_document_service.py` -- `app/services/demo_source_materializer.py` +- `app/services/demo/*` ### Billing Workflow diff --git a/apps/api/app/api/v1/routes/demo.py b/apps/api/app/api/v1/routes/demo.py index 330a5167c..42508c5e1 100644 --- a/apps/api/app/api/v1/routes/demo.py +++ b/apps/api/app/api/v1/routes/demo.py @@ -4,7 +4,7 @@ from typing import Any -from app.services.demo_document_service import DemoDocumentService +from app.services.demo.document_service import DemoDocumentService from app.services.rate_limit.dependencies import CurrentUser, with_current_user from fastapi import APIRouter, Depends, Query from fastapi.responses import FileResponse @@ -16,7 +16,7 @@ router = APIRouter(tags=["Demo Documents"]) -_demo_document_service = DemoDocumentService() +_demo_service = DemoDocumentService() class DemoMaterializeRequest(BaseModel): @@ -33,7 +33,7 @@ class DemoMaterializeRequest(BaseModel): @router.get("/catalog") async def get_demo_catalog() -> dict[str, Any]: """Return API-owned canonical demo source metadata and curated Q/A.""" - return _demo_document_service.get_catalog() + return _demo_service.get_catalog() @router.get("/sources/{demo_source_id}/chunks") @@ -43,7 +43,7 @@ async def list_demo_source_chunks( page_size: int = Query(50, ge=1, le=200, description="Items per page"), ) -> dict[str, Any]: """Return paginated canonical chunks for a demo source.""" - response = _demo_document_service.list_chunks( + response = _demo_service.list_chunks( demo_source_id=demo_source_id, page=page, page_size=page_size, @@ -59,7 +59,7 @@ async def get_demo_source_chunk( demo_chunk_id: str, ) -> dict[str, Any]: """Return one canonical demo chunk for citation focusing.""" - response = _demo_document_service.get_chunk( + response = _demo_service.get_chunk( demo_source_id=demo_source_id, demo_chunk_id=demo_chunk_id, ) @@ -75,7 +75,7 @@ async def get_demo_source_chunk( @router.get("/sources/{demo_source_id}/original") async def get_demo_source_original(demo_source_id: str) -> FileResponse: """Return the canonical original file for preview.""" - file_path = _demo_document_service.get_original_file_path( + file_path = _demo_service.get_original_file_path( demo_source_id=demo_source_id, ) if file_path is None: @@ -95,7 +95,7 @@ async def get_demo_source_asset( asset_path: str, ) -> FileResponse: """Return a canonical parsed media or table asset for preview.""" - file_path = _demo_document_service.get_asset_file_path( + file_path = _demo_service.get_asset_file_path( demo_source_id=demo_source_id, asset_path=asset_path, ) @@ -118,7 +118,7 @@ async def materialize_demo_sources( """Copy canonical demo sources into the authenticated user's namespace.""" namespace = (payload.namespace or "default").strip() or "default" try: - materialized_sources = await _demo_document_service.materialize_sources( + materialized_sources = await _demo_service.materialize_sources( db, user_id=current_user.user_id, namespace=namespace, diff --git a/apps/api/app/services/demo/__init__.py b/apps/api/app/services/demo/__init__.py new file mode 100644 index 000000000..a25f9979d --- /dev/null +++ b/apps/api/app/services/demo/__init__.py @@ -0,0 +1 @@ +"""Demo source catalog and materialization services.""" diff --git a/apps/api/app/services/demo_document_service.py b/apps/api/app/services/demo/document_service.py similarity index 95% rename from apps/api/app/services/demo_document_service.py rename to apps/api/app/services/demo/document_service.py index 99d735f75..f0a63bd1c 100644 --- a/apps/api/app/services/demo_document_service.py +++ b/apps/api/app/services/demo/document_service.py @@ -5,8 +5,8 @@ from pathlib import Path from typing import Any -from app.services.demo_source_catalog import DemoSourceCatalog -from app.services.demo_source_materializer import ( +from app.services.demo.source_catalog import DemoSourceCatalog +from app.services.demo.source_materializer import ( DemoSourceMaterializer, MaterializedDemoSource, ) diff --git a/apps/api/app/services/demo_source_catalog.py b/apps/api/app/services/demo/source_catalog.py similarity index 98% rename from apps/api/app/services/demo_source_catalog.py rename to apps/api/app/services/demo/source_catalog.py index ec6328af7..1b0add285 100644 --- a/apps/api/app/services/demo_source_catalog.py +++ b/apps/api/app/services/demo/source_catalog.py @@ -9,7 +9,7 @@ from pathlib import Path from typing import Any -from app.services.demo_source_projection import DemoSourceProjection +from app.services.demo.source_projection import DemoSourceProjection @dataclass(frozen=True) @@ -39,7 +39,7 @@ class DemoSourceDefinition: examples: tuple[DemoExampleDefinition, ...] -_DATA_ROOT = Path(__file__).resolve().parents[1] / "data" / "demo_documents" +_DATA_ROOT = Path(__file__).resolve().parents[2] / "data" / "demo_documents" _ASSET_DIRECTORY_NAMES = frozenset({"images", "tables"}) _DEMO_SOURCE_DEFINITIONS: tuple[DemoSourceDefinition, ...] = ( DemoSourceDefinition( diff --git a/apps/api/app/services/demo_source_materializer.py b/apps/api/app/services/demo/source_materializer.py similarity index 99% rename from apps/api/app/services/demo_source_materializer.py rename to apps/api/app/services/demo/source_materializer.py index 6b418ff3c..df9e2f0c8 100644 --- a/apps/api/app/services/demo_source_materializer.py +++ b/apps/api/app/services/demo/source_materializer.py @@ -10,7 +10,7 @@ from pathlib import Path from uuid import uuid4 -from app.services.demo_source_catalog import DemoSourceCatalog, DemoSourceDefinition +from app.services.demo.source_catalog import DemoSourceCatalog, DemoSourceDefinition from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession diff --git a/apps/api/app/services/demo_source_projection.py b/apps/api/app/services/demo/source_projection.py similarity index 100% rename from apps/api/app/services/demo_source_projection.py rename to apps/api/app/services/demo/source_projection.py diff --git a/apps/api/tests/contract/test_demo_documents_contract.py b/apps/api/tests/contract/test_demo_documents_contract.py index 6975a378f..db850c4c9 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_source_materializer as demo_source_materializer + import app.services.demo.source_materializer as source_materializer_module monkeypatch.setattr( - demo_source_materializer, + source_materializer_module, "get_result_storage", lambda: fake_result_storage, ) @@ -281,10 +281,10 @@ async def test_should_materialize_each_normalized_demo_source_once_per_request( fake_result_storage = FakeResultStorage() async with developer_api_client_factory() as api_client: - import app.services.demo_source_materializer as demo_source_materializer + import app.services.demo.source_materializer as source_materializer_module monkeypatch.setattr( - demo_source_materializer, + source_materializer_module, "get_result_storage", lambda: fake_result_storage, ) @@ -341,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_source_materializer as demo_source_materializer + import app.services.demo.source_materializer as source_materializer_module monkeypatch.setattr( - demo_source_materializer, + source_materializer_module, "get_result_storage", lambda: fake_result_storage, ) @@ -437,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_source_materializer as demo_source_materializer + import app.services.demo.source_materializer as source_materializer_module monkeypatch.setattr( - demo_source_materializer, + source_materializer_module, "get_result_storage", lambda: fake_result_storage, ) From 16dda96f41ffee231eeadc881aa14b9f85d1c520 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 10:08:44 +0000 Subject: [PATCH 58/78] refactor: remove empty app package markers --- apps/api/app/__init__.py | 5 +++++ apps/api/app/api/__init__.py | 0 apps/api/app/api/v1/__init__.py | 0 apps/api/app/api/v1/routes/__init__.py | 0 apps/api/app/core/response/__init__.py | 0 apps/api/app/services/demo/__init__.py | 1 - apps/api/app/services/s3_events/__init__.py | 1 - apps/api/tests/__init__.py | 2 +- apps/api/tests/support/__init__.py | 2 +- apps/worker/app/__init__.py | 5 +++++ apps/worker/app/core/tasks/__init__.py | 1 - apps/worker/app/services/document_ingestion/__init__.py | 1 - .../app/services/document_parser/orchestration/__init__.py | 1 - apps/worker/app/services/webhook/__init__.py | 1 - apps/worker/app/services/workload/__init__.py | 1 - apps/worker/tests/support/__init__.py | 0 16 files changed, 12 insertions(+), 9 deletions(-) delete mode 100644 apps/api/app/api/__init__.py delete mode 100644 apps/api/app/api/v1/__init__.py delete mode 100644 apps/api/app/api/v1/routes/__init__.py delete mode 100644 apps/api/app/core/response/__init__.py delete mode 100644 apps/api/app/services/demo/__init__.py delete mode 100644 apps/api/app/services/s3_events/__init__.py delete mode 100644 apps/worker/app/core/tasks/__init__.py delete mode 100644 apps/worker/app/services/document_ingestion/__init__.py delete mode 100644 apps/worker/app/services/document_parser/orchestration/__init__.py delete mode 100644 apps/worker/app/services/webhook/__init__.py delete mode 100644 apps/worker/app/services/workload/__init__.py delete mode 100644 apps/worker/tests/support/__init__.py diff --git a/apps/api/app/__init__.py b/apps/api/app/__init__.py index e69de29bb..75f316e90 100644 --- a/apps/api/app/__init__.py +++ b/apps/api/app/__init__.py @@ -0,0 +1,5 @@ +"""API application package boundary. + +This prevents namespace-package merging with the worker's separate `app` +package when both service paths are on `PYTHONPATH`. +""" diff --git a/apps/api/app/api/__init__.py b/apps/api/app/api/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/apps/api/app/api/v1/__init__.py b/apps/api/app/api/v1/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/apps/api/app/api/v1/routes/__init__.py b/apps/api/app/api/v1/routes/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/apps/api/app/core/response/__init__.py b/apps/api/app/core/response/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/apps/api/app/services/demo/__init__.py b/apps/api/app/services/demo/__init__.py deleted file mode 100644 index a25f9979d..000000000 --- a/apps/api/app/services/demo/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Demo source catalog and materialization services.""" diff --git a/apps/api/app/services/s3_events/__init__.py b/apps/api/app/services/s3_events/__init__.py deleted file mode 100644 index b68b3a02b..000000000 --- a/apps/api/app/services/s3_events/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""S3-compatible storage event services.""" diff --git a/apps/api/tests/__init__.py b/apps/api/tests/__init__.py index 8b1378917..e8b0948aa 100644 --- a/apps/api/tests/__init__.py +++ b/apps/api/tests/__init__.py @@ -1 +1 @@ - +"""API test package for tests.support imports.""" diff --git a/apps/api/tests/support/__init__.py b/apps/api/tests/support/__init__.py index 8b1378917..535994872 100644 --- a/apps/api/tests/support/__init__.py +++ b/apps/api/tests/support/__init__.py @@ -1 +1 @@ - +"""Shared support helpers for API tests.""" diff --git a/apps/worker/app/__init__.py b/apps/worker/app/__init__.py index e69de29bb..d6bf73bb1 100644 --- a/apps/worker/app/__init__.py +++ b/apps/worker/app/__init__.py @@ -0,0 +1,5 @@ +"""Worker application package boundary. + +This prevents namespace-package merging with the API's separate `app` package +when both service paths are on `PYTHONPATH`. +""" diff --git a/apps/worker/app/core/tasks/__init__.py b/apps/worker/app/core/tasks/__init__.py deleted file mode 100644 index ba382c975..000000000 --- a/apps/worker/app/core/tasks/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Celery task package for worker-side knowledge-base jobs.""" diff --git a/apps/worker/app/services/document_ingestion/__init__.py b/apps/worker/app/services/document_ingestion/__init__.py deleted file mode 100644 index 191bf4665..000000000 --- a/apps/worker/app/services/document_ingestion/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Worker-side Document Ingestion modules.""" diff --git a/apps/worker/app/services/document_parser/orchestration/__init__.py b/apps/worker/app/services/document_parser/orchestration/__init__.py deleted file mode 100644 index 6f12d8497..000000000 --- a/apps/worker/app/services/document_parser/orchestration/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Document parser orchestration modules.""" diff --git a/apps/worker/app/services/webhook/__init__.py b/apps/worker/app/services/webhook/__init__.py deleted file mode 100644 index 3040aede5..000000000 --- a/apps/worker/app/services/webhook/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Worker webhook services.""" diff --git a/apps/worker/app/services/workload/__init__.py b/apps/worker/app/services/workload/__init__.py deleted file mode 100644 index 3ad3226c0..000000000 --- a/apps/worker/app/services/workload/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Worker workload adapters.""" diff --git a/apps/worker/tests/support/__init__.py b/apps/worker/tests/support/__init__.py deleted file mode 100644 index e69de29bb..000000000 From 78061171c36ed5ff2de259068b5183c4262b4e0f Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 10:21:04 +0000 Subject: [PATCH 59/78] refactor: remove init-only modules --- apps/api/app/__init__.py | 5 ----- apps/api/app/api/__init__.py | 0 apps/api/app/api/v1/__init__.py | 0 apps/api/app/api/v1/routes/__init__.py | 0 apps/api/app/core/response/__init__.py | 0 apps/api/app/services/demo/__init__.py | 1 + apps/api/app/services/redis/__init__.py | 22 ------------------- apps/api/app/services/s3_events/__init__.py | 1 + apps/api/tests/__init__.py | 2 +- apps/api/tests/support/__init__.py | 2 +- apps/worker/app/__init__.py | 5 ----- apps/worker/app/core/tasks/__init__.py | 1 + apps/worker/app/services/billing/__init__.py | 7 ------ .../services/document_ingestion/__init__.py | 1 + .../document_parser/orchestration/__init__.py | 1 + apps/worker/app/services/workload/__init__.py | 1 + apps/worker/tests/support/__init__.py | 0 17 files changed, 8 insertions(+), 41 deletions(-) create mode 100644 apps/api/app/api/__init__.py create mode 100644 apps/api/app/api/v1/__init__.py create mode 100644 apps/api/app/api/v1/routes/__init__.py create mode 100644 apps/api/app/core/response/__init__.py create mode 100644 apps/api/app/services/demo/__init__.py delete mode 100644 apps/api/app/services/redis/__init__.py create mode 100644 apps/api/app/services/s3_events/__init__.py create mode 100644 apps/worker/app/core/tasks/__init__.py delete mode 100644 apps/worker/app/services/billing/__init__.py create mode 100644 apps/worker/app/services/document_ingestion/__init__.py create mode 100644 apps/worker/app/services/document_parser/orchestration/__init__.py create mode 100644 apps/worker/app/services/workload/__init__.py create mode 100644 apps/worker/tests/support/__init__.py diff --git a/apps/api/app/__init__.py b/apps/api/app/__init__.py index 75f316e90..e69de29bb 100644 --- a/apps/api/app/__init__.py +++ b/apps/api/app/__init__.py @@ -1,5 +0,0 @@ -"""API application package boundary. - -This prevents namespace-package merging with the worker's separate `app` -package when both service paths are on `PYTHONPATH`. -""" diff --git a/apps/api/app/api/__init__.py b/apps/api/app/api/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/apps/api/app/api/v1/__init__.py b/apps/api/app/api/v1/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/apps/api/app/api/v1/routes/__init__.py b/apps/api/app/api/v1/routes/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/apps/api/app/core/response/__init__.py b/apps/api/app/core/response/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/apps/api/app/services/demo/__init__.py b/apps/api/app/services/demo/__init__.py new file mode 100644 index 000000000..a25f9979d --- /dev/null +++ b/apps/api/app/services/demo/__init__.py @@ -0,0 +1 @@ +"""Demo source catalog and materialization services.""" diff --git a/apps/api/app/services/redis/__init__.py b/apps/api/app/services/redis/__init__.py deleted file mode 100644 index a7c271e34..000000000 --- a/apps/api/app/services/redis/__init__.py +++ /dev/null @@ -1,22 +0,0 @@ -"""API-specific Redis service exports.""" - -# Re-export shared Redis services used by the API runtime. -from shared.services.redis import ( - JobInfoRedisService, - RedisService, - RedisServiceFactory, - UserRedisService, -) -from shared.services.redis.job_metadata_service import JobMetadataService -from shared.services.redis.rate_limit_service import RateLimitService -from shared.services.redis.task_redis_service import TaskRedisService - -__all__ = [ - "JobMetadataService", - "TaskRedisService", - "RateLimitService", - "RedisService", # Re-exported from shared. - "RedisServiceFactory", # Re-exported from shared. - "UserRedisService", # Re-exported from shared for both Worker and API. - "JobInfoRedisService", # Re-exported from shared for both Worker and API. -] 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/tests/__init__.py b/apps/api/tests/__init__.py index e8b0948aa..8b1378917 100644 --- a/apps/api/tests/__init__.py +++ b/apps/api/tests/__init__.py @@ -1 +1 @@ -"""API test package for tests.support imports.""" + diff --git a/apps/api/tests/support/__init__.py b/apps/api/tests/support/__init__.py index 535994872..8b1378917 100644 --- a/apps/api/tests/support/__init__.py +++ b/apps/api/tests/support/__init__.py @@ -1 +1 @@ -"""Shared support helpers for API tests.""" + diff --git a/apps/worker/app/__init__.py b/apps/worker/app/__init__.py index d6bf73bb1..e69de29bb 100644 --- a/apps/worker/app/__init__.py +++ b/apps/worker/app/__init__.py @@ -1,5 +0,0 @@ -"""Worker application package boundary. - -This prevents namespace-package merging with the API's separate `app` package -when both service paths are on `PYTHONPATH`. -""" diff --git a/apps/worker/app/core/tasks/__init__.py b/apps/worker/app/core/tasks/__init__.py new file mode 100644 index 000000000..ba382c975 --- /dev/null +++ b/apps/worker/app/core/tasks/__init__.py @@ -0,0 +1 @@ +"""Celery task package for worker-side knowledge-base jobs.""" diff --git a/apps/worker/app/services/billing/__init__.py b/apps/worker/app/services/billing/__init__.py deleted file mode 100644 index 50a959c99..000000000 --- a/apps/worker/app/services/billing/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -""" -Billing services for the worker. -""" - -from shared.services.billing.credits_sync_service import SyncCreditsService - -__all__ = ["SyncCreditsService"] 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/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/workload/__init__.py b/apps/worker/app/services/workload/__init__.py new file mode 100644 index 000000000..3ad3226c0 --- /dev/null +++ b/apps/worker/app/services/workload/__init__.py @@ -0,0 +1 @@ +"""Worker workload adapters.""" diff --git a/apps/worker/tests/support/__init__.py b/apps/worker/tests/support/__init__.py new file mode 100644 index 000000000..e69de29bb From bb334ffd0a1bc4d064a464e59cf927410e48198e Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 10:34:58 +0000 Subject: [PATCH 60/78] refactor: remove unused image client bootstrap --- apps/api/app/core/image_cli.py | 20 -------------------- apps/api/main.py | 4 ---- 2 files changed, 24 deletions(-) delete mode 100644 apps/api/app/core/image_cli.py diff --git a/apps/api/app/core/image_cli.py b/apps/api/app/core/image_cli.py deleted file mode 100644 index 522ca99bb..000000000 --- a/apps/api/app/core/image_cli.py +++ /dev/null @@ -1,20 +0,0 @@ -from typing import Optional - -import httpx - - -class ImageCli: - """Image-processing client.""" - - http_client: Optional[httpx.AsyncClient] = None - - -http_client: Optional[httpx.AsyncClient] = None - - -def get_http_client() -> httpx.AsyncClient: - if http_client is None: - raise RuntimeError( - "HTTP client has not been initialized. Is it in the lifespan manager?" - ) - return http_client diff --git a/apps/api/main.py b/apps/api/main.py index 4b0dfde44..8caac693f 100644 --- a/apps/api/main.py +++ b/apps/api/main.py @@ -1,6 +1,5 @@ import os from pathlib import Path -import httpx import uvicorn from fastapi import FastAPI from starlette.routing import Route @@ -18,7 +17,6 @@ from contextlib import asynccontextmanager from app.api.api_router import api_router from app.core.middleware import setup_cors, LoggingMiddleware -from app.core.image_cli import ImageCli from app.core.exception_handlers import setup_exception_handlers from app.mcp import create_retrieval_mcp_server from app.services.rate_limit.rule_loader import load_rules @@ -59,8 +57,6 @@ async def lifespan(app: FastAPI): await redis_pool_manager.init_pool() logger.info("Redis connection pool created.") - ImageCli.http_client = httpx.AsyncClient(timeout=30.0, follow_redirects=True) - # Initialize rate limiter rules from DB. # Changes now require a pod restart to take effect. from app.services.rate_limit.config import RateLimitConfig From 4249d6cf66af72faff9f9a4f9ade0c737d537d51 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 10:52:08 +0000 Subject: [PATCH 61/78] refactor: move current user auth out of core --- apps/api/app/api/dependencies/__init__.py | 1 + apps/api/app/api/dependencies/auth.py | 23 +++ apps/api/app/core/dependencies.py | 160 ------------------ apps/api/app/mcp/retrieval_server.py | 9 +- .../current_user_authentication_service.py | 96 +++++++++++ .../dashboard_jwt_authentication_service.py | 91 ++++++++++ .../app/services/rate_limit/dependencies.py | 2 +- .../contract/test_job_creation_contract.py | 6 +- 8 files changed, 221 insertions(+), 167 deletions(-) create mode 100644 apps/api/app/api/dependencies/__init__.py create mode 100644 apps/api/app/api/dependencies/auth.py delete mode 100644 apps/api/app/core/dependencies.py create mode 100644 apps/api/app/services/auth/current_user_authentication_service.py create mode 100644 apps/api/app/services/auth/dashboard_jwt_authentication_service.py diff --git a/apps/api/app/api/dependencies/__init__.py b/apps/api/app/api/dependencies/__init__.py new file mode 100644 index 000000000..e0170a914 --- /dev/null +++ b/apps/api/app/api/dependencies/__init__.py @@ -0,0 +1 @@ +"""FastAPI dependency adapters.""" diff --git a/apps/api/app/api/dependencies/auth.py b/apps/api/app/api/dependencies/auth.py new file mode 100644 index 000000000..4ebec70cc --- /dev/null +++ b/apps/api/app/api/dependencies/auth.py @@ -0,0 +1,23 @@ +"""FastAPI authentication dependency adapters.""" + +from app.services.auth.current_user_authentication_service import ( + get_current_user_authentication_service, +) +from fastapi import Depends, Header +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.database import get_db + + +async def get_current_user_id( + authorization: str | None = Header( + default=None, + description="Bearer OR internal signature auth", + ), + db: AsyncSession = Depends(get_db), +) -> str: + """Authenticate the caller and return the current user ID.""" + return await get_current_user_authentication_service().authenticate_authorization_header( + db, + authorization, + ) diff --git a/apps/api/app/core/dependencies.py b/apps/api/app/core/dependencies.py deleted file mode 100644 index 153177cd4..000000000 --- a/apps/api/app/core/dependencies.py +++ /dev/null @@ -1,160 +0,0 @@ -import threading -from datetime import timedelta -from typing import Any - -import jwt -from app.services.auth.api_key_authentication_service import ( - APIKeyAuthenticationService, -) -from fastapi import Depends, Header, Request -from jwt import PyJWKClient -from loguru import logger -from sqlalchemy import 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 ( - AuthException, -) -from shared.models.database.user import User -from shared.utils.api_keys import is_api_key_token - -# Standard JWKS endpoint path (fixed, following OpenID Connect convention) -JWKS_ENDPOINT_PATH = "/api/auth/jwks" - -# Cache settings: 1 hour in seconds -JWKS_CACHE_TTL_SECONDS = 60 * 60 # 3600 seconds - -# Cached PyJWKClient instance -_jwks_client: PyJWKClient | None = None -_jwks_client_lock = threading.Lock() -_api_key_authentication_service = APIKeyAuthenticationService() - - -def _get_jwks_client() -> PyJWKClient: - """ - Get or create a cached PyJWKClient instance. - - The JWKS endpoint is constructed from INTERNAL_DASHBOARD_ENDPOINT + fixed path. - PyJWKClient caches the JWKS response for 1 hour. - """ - global _jwks_client - - if _jwks_client is None: - with _jwks_client_lock: - if _jwks_client is None: - jwks_url = f"{settings.INTERNAL_DASHBOARD_ENDPOINT}{JWKS_ENDPOINT_PATH}" - _jwks_client = PyJWKClient( - jwks_url, - cache_jwk_set=True, - lifespan=JWKS_CACHE_TTL_SECONDS, - timeout=30, - ) - logger.info(f"Initialized JWKS client with endpoint: {jwks_url}") - - return _jwks_client - - -def _get_verification_key(token: str) -> Any: - """ - Get the verification key for the JWT from the JWKS endpoint. - - Uses PyJWKClient's built-in cache with 1 hour TTL. - """ - try: - jwks_client = _get_jwks_client() - signing_key = jwks_client.get_signing_key_from_jwt(token) - return signing_key.key - except jwt.PyJWKClientError as e: - logger.error(f"Failed to fetch JWKS: {e}") - raise AuthException( - internal_message=f"Failed to fetch verification key from JWKS endpoint: {e}" - ) - except jwt.PyJWKSetError as e: - logger.error(f"Invalid JWKS format: {e}") - raise AuthException(internal_message=f"Invalid JWKS format: {e}") - - -def decode_jwt_token(token: str) -> str: - """ - Decode and validate JWT token using JWKS. - - Expected JWT claims: - - id/sub: User ID - - exp: Expiration - """ - try: - key = _get_verification_key(token) - - # Decode and Verify - payload = jwt.decode( - token, - key, - algorithms=["HS256", "RS256", "EdDSA"], - leeway=timedelta(seconds=30), - options={"verify_aud": False}, - ) - - # Extract user_id from 'id' claim - user_id = payload.get("id") - - if not user_id: - raise AuthException(user_message="Token missing 'id' claim") - - return user_id - - except jwt.ExpiredSignatureError: - raise AuthException(user_message="Token has expired") - except jwt.InvalidTokenError as e: - logger.warning(f"Invalid JWT token: {e}") - raise AuthException(user_message="Invalid token") - - -async def _ensure_authenticated_user_exists( - db: AsyncSession, - user_id: str, -) -> None: - result = await db.execute(select(User.id).where(User.id == user_id).limit(1)) - if result.scalar_one_or_none() is not None: - return - - raise AuthException( - user_message="Invalid authentication credentials", - internal_message=( - "Authenticated user id is not present in the user table: " - f"user_id={user_id}" - ), - ) - - -async def get_current_user_id( - request: Request, - authorization: str | None = Header( - default=None, description="Bearer OR internal signature auth" - ), - db: AsyncSession = Depends(get_db), -) -> str: - """Authenticate the caller and return user_id.""" - if not authorization: - raise AuthException( - user_message="Authentication required. Provide Authorization header." - ) - - # Parse Authorization header - scheme, _, token = authorization.partition(" ") - if scheme.lower() != "bearer" or not token: - raise AuthException(user_message="Invalid Authorization header format") - - # Mode 1: API Key verification (for external clients) - if is_api_key_token(token): - user_id = await _api_key_authentication_service.validate_api_key(db, token) - if user_id: - return user_id - - raise AuthException(user_message="Invalid API Key") - - # Mode 2: JWT verification (for Dashboard/Internal) - user_id = decode_jwt_token(token) - await _ensure_authenticated_user_exists(db, user_id) - return user_id diff --git a/apps/api/app/mcp/retrieval_server.py b/apps/api/app/mcp/retrieval_server.py index 2121cdb4e..2b0cbcbcc 100644 --- a/apps/api/app/mcp/retrieval_server.py +++ b/apps/api/app/mcp/retrieval_server.py @@ -2,7 +2,9 @@ from typing import Annotated, Any, AsyncContextManager, Callable -from app.core.dependencies import get_current_user_id +from app.services.auth.current_user_authentication_service import ( + get_current_user_authentication_service, +) from mcp.server.fastmcp import Context, FastMCP from mcp.server.transport_security import TransportSecuritySettings from pydantic import Field @@ -81,10 +83,9 @@ async def resolve_mcp_user_id(*, ctx: Context | None, db: AsyncSession) -> str: request = get_mcp_request(ctx) headers = getattr(request, "headers", {}) or {} authorization = get_header(headers, "authorization") - return await get_current_user_id( - request=request, + return await get_current_user_authentication_service().authenticate_authorization_header( + db, authorization=authorization, - db=db, ) diff --git a/apps/api/app/services/auth/current_user_authentication_service.py b/apps/api/app/services/auth/current_user_authentication_service.py new file mode 100644 index 000000000..a32c02068 --- /dev/null +++ b/apps/api/app/services/auth/current_user_authentication_service.py @@ -0,0 +1,96 @@ +"""Current-user authentication workflow.""" + +from __future__ import annotations + +from app.services.auth.api_key_authentication_service import ( + APIKeyAuthenticationService, +) +from app.services.auth.dashboard_jwt_authentication_service import ( + DashboardJWTAuthenticationService, + get_dashboard_jwt_authentication_service, +) +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.exceptions.domain_exceptions import AuthException +from shared.models.database.user import User +from shared.utils.api_keys import is_api_key_token + + +class CurrentUserAuthenticationService: + """Authenticate API keys or Dashboard JWTs and return the current user ID.""" + + def __init__( + self, + *, + api_key_authentication_service: APIKeyAuthenticationService | None = None, + dashboard_jwt_authentication_service: ( + DashboardJWTAuthenticationService | None + ) = None, + ) -> None: + self._api_key_authentication_service = ( + api_key_authentication_service or APIKeyAuthenticationService() + ) + self._dashboard_jwt_authentication_service = ( + dashboard_jwt_authentication_service + or get_dashboard_jwt_authentication_service() + ) + + async def authenticate_authorization_header( + self, + session: AsyncSession, + authorization: str | None, + ) -> str: + """Authenticate an Authorization header and return the owning user ID.""" + token = self._extract_bearer_token(authorization) + + if is_api_key_token(token): + user_id = await self._api_key_authentication_service.validate_api_key( + session, + token, + ) + if user_id: + return user_id + raise AuthException(user_message="Invalid API Key") + + user_id = self._dashboard_jwt_authentication_service.decode_user_id(token) + await self._ensure_authenticated_user_exists(session, user_id) + return user_id + + @staticmethod + def _extract_bearer_token(authorization: str | None) -> str: + if not authorization: + raise AuthException( + user_message="Authentication required. Provide Authorization header." + ) + + scheme, _, token = authorization.partition(" ") + if scheme.lower() != "bearer" or not token: + raise AuthException(user_message="Invalid Authorization header format") + + return token + + @staticmethod + async def _ensure_authenticated_user_exists( + session: AsyncSession, + user_id: str, + ) -> None: + result = await session.execute(select(User.id).where(User.id == user_id).limit(1)) + if result.scalar_one_or_none() is not None: + return + + raise AuthException( + user_message="Invalid authentication credentials", + internal_message=( + "Authenticated user id is not present in the user table: " + f"user_id={user_id}" + ), + ) + + +_current_user_authentication_service = CurrentUserAuthenticationService() + + +def get_current_user_authentication_service() -> CurrentUserAuthenticationService: + """Return the process-wide current-user authentication service.""" + return _current_user_authentication_service diff --git a/apps/api/app/services/auth/dashboard_jwt_authentication_service.py b/apps/api/app/services/auth/dashboard_jwt_authentication_service.py new file mode 100644 index 000000000..07cc6ef4c --- /dev/null +++ b/apps/api/app/services/auth/dashboard_jwt_authentication_service.py @@ -0,0 +1,91 @@ +"""Dashboard JWT authentication workflow.""" + +from __future__ import annotations + +import threading +from datetime import timedelta +from typing import Any + +import jwt +from jwt import PyJWKClient +from loguru import logger + +from shared.core.config import settings +from shared.core.exceptions.domain_exceptions import AuthException + +JWKS_ENDPOINT_PATH = "/api/auth/jwks" +JWKS_CACHE_TTL_SECONDS = 60 * 60 + + +class DashboardJWTAuthenticationService: + """Validate Dashboard-issued JWTs through the configured JWKS endpoint.""" + + def __init__(self) -> None: + self._jwks_client: PyJWKClient | None = None + self._jwks_client_lock = threading.Lock() + + def decode_user_id(self, token: str) -> str: + """Decode and validate a JWT, returning its authenticated user ID.""" + try: + key = self._get_verification_key(token) + payload: dict[str, Any] = jwt.decode( + token, + key, + algorithms=["HS256", "RS256", "EdDSA"], + leeway=timedelta(seconds=30), + options={"verify_aud": False}, + ) + user_id = payload.get("id") + if not isinstance(user_id, str) or not user_id: + raise AuthException(user_message="Token missing 'id' claim") + return user_id + except jwt.ExpiredSignatureError: + raise AuthException(user_message="Token has expired") + except jwt.InvalidTokenError as exc: + logger.warning(f"Invalid JWT token: {exc}") + raise AuthException(user_message="Invalid token") + + def _get_verification_key(self, token: str) -> Any: + """Resolve the JWT verification key from the Dashboard JWKS endpoint.""" + try: + jwks_client = self._get_jwks_client() + signing_key = jwks_client.get_signing_key_from_jwt(token) + return signing_key.key + except jwt.PyJWKClientError as exc: + logger.error(f"Failed to fetch JWKS: {exc}") + raise AuthException( + internal_message=( + "Failed to fetch verification key from JWKS endpoint: " + f"{exc}" + ) + ) + except jwt.PyJWKSetError as exc: + logger.error(f"Invalid JWKS format: {exc}") + raise AuthException(internal_message=f"Invalid JWKS format: {exc}") + + def _get_jwks_client(self) -> PyJWKClient: + """Return a cached JWKS client for Dashboard token verification.""" + if self._jwks_client is None: + with self._jwks_client_lock: + if self._jwks_client is None: + jwks_url = ( + f"{settings.INTERNAL_DASHBOARD_ENDPOINT}" + f"{JWKS_ENDPOINT_PATH}" + ) + self._jwks_client = PyJWKClient( + jwks_url, + cache_jwk_set=True, + lifespan=JWKS_CACHE_TTL_SECONDS, + timeout=30, + ) + logger.info(f"Initialized JWKS client with endpoint: {jwks_url}") + + return self._jwks_client + + +_dashboard_jwt_authentication_service = DashboardJWTAuthenticationService() + + +def get_dashboard_jwt_authentication_service() -> DashboardJWTAuthenticationService: + """Return the process-wide Dashboard JWT authentication service.""" + return _dashboard_jwt_authentication_service diff --git a/apps/api/app/services/rate_limit/dependencies.py b/apps/api/app/services/rate_limit/dependencies.py index 888f9a570..c76785dc2 100644 --- a/apps/api/app/services/rate_limit/dependencies.py +++ b/apps/api/app/services/rate_limit/dependencies.py @@ -2,7 +2,7 @@ from typing import AsyncGenerator -from app.core.dependencies import get_current_user_id +from app.api.dependencies.auth import get_current_user_id from app.services.rate_limit.data_structures import CurrentUser from app.services.rate_limit.job_admission_service import JobAdmissionService from fastapi import Depends, Request diff --git a/apps/api/tests/contract/test_job_creation_contract.py b/apps/api/tests/contract/test_job_creation_contract.py index b56f482f6..d39a53296 100644 --- a/apps/api/tests/contract/test_job_creation_contract.py +++ b/apps/api/tests/contract/test_job_creation_contract.py @@ -429,10 +429,12 @@ async def test_should_reject_authenticated_user_id_missing_from_user_table( } async with api_client_factory() as api_client: - from app.core import dependencies as auth_dependencies + from app.services.auth.dashboard_jwt_authentication_service import ( + get_dashboard_jwt_authentication_service, + ) monkeypatch.setattr( - auth_dependencies, + get_dashboard_jwt_authentication_service(), "_get_verification_key", lambda _token: jwt_secret, ) From 74e77fea59d16ea7c8ed3a2833e481808b3617d3 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 11:37:15 +0000 Subject: [PATCH 62/78] refactor: isolate job admission request adapters --- .../api/app/api/dependencies/job_admission.py | 71 +++++++++++++++++++ apps/api/app/api/v1/routes/api_key.py | 4 +- apps/api/app/api/v1/routes/billing.py | 3 +- apps/api/app/api/v1/routes/demo.py | 3 +- apps/api/app/api/v1/routes/documents.py | 3 +- apps/api/app/api/v1/routes/guest.py | 2 +- apps/api/app/api/v1/routes/jobs.py | 8 +-- apps/api/app/api/v1/routes/retrieval.py | 3 +- apps/api/app/api/v1/routes/webhook.py | 3 +- apps/api/app/api/v1/routes/webhook_secrets.py | 3 +- .../services/document_ingestion/service.py | 9 ++- .../services/rate_limit/data_structures.py | 9 +++ .../app/services/rate_limit/dependencies.py | 48 ------------- .../job_admission_route_policy_service.py | 52 ++++++-------- .../rate_limit/job_admission_service.py | 19 +++-- 15 files changed, 134 insertions(+), 106 deletions(-) create mode 100644 apps/api/app/api/dependencies/job_admission.py delete mode 100644 apps/api/app/services/rate_limit/dependencies.py diff --git a/apps/api/app/api/dependencies/job_admission.py b/apps/api/app/api/dependencies/job_admission.py new file mode 100644 index 000000000..768ccc872 --- /dev/null +++ b/apps/api/app/api/dependencies/job_admission.py @@ -0,0 +1,71 @@ +"""FastAPI adapters for the Job Admission workflow.""" + +from typing import AsyncGenerator + +from app.api.dependencies.auth import get_current_user_id +from app.services.rate_limit.data_structures import ( + CurrentUser, + RouteAdmissionContext, +) +from app.services.rate_limit.job_admission_service import JobAdmissionService +from fastapi import Depends, Request + +_job_admission_service = JobAdmissionService() + + +def get_route_admission_context(request: Request) -> RouteAdmissionContext: + """Extract route facts needed by Job Admission from a FastAPI request.""" + return RouteAdmissionContext( + method=request.method, + path=_get_route_path(request), + limit_identifier=_get_route_limit_identifier(request), + ) + + +async def with_current_user( + route_context: RouteAdmissionContext = Depends(get_route_admission_context), + user_id: str = Depends(get_current_user_id), +) -> AsyncGenerator[CurrentUser, None]: + current_user = await _job_admission_service.resolve_current_user( + route_context=route_context, + user_id=user_id, + ) + yield current_user + + +async def require_billing_limits( + current_user: CurrentUser = Depends(with_current_user), +) -> AsyncGenerator[CurrentUser, None]: + await _job_admission_service.enforce_billing_limits(current_user=current_user) + yield current_user + + +async def require_route_system_limit( + route_context: RouteAdmissionContext = Depends(get_route_admission_context), +) -> None: + await _job_admission_service.enforce_route_system_limit( + route_context=route_context, + ) + + +def _get_route_path(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(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 _get_route_path(request) diff --git a/apps/api/app/api/v1/routes/api_key.py b/apps/api/app/api/v1/routes/api_key.py index 2dbcefce4..85a3f6836 100644 --- a/apps/api/app/api/v1/routes/api_key.py +++ b/apps/api/app/api/v1/routes/api_key.py @@ -3,10 +3,10 @@ """ from app.services.auth.api_key_management_service import APIKeyManagementService -from app.services.rate_limit.dependencies import ( - CurrentUser, +from app.api.dependencies.job_admission import ( with_current_user, ) +from app.services.rate_limit.data_structures import CurrentUser from fastapi import APIRouter, Depends from sqlalchemy.ext.asyncio import AsyncSession diff --git a/apps/api/app/api/v1/routes/billing.py b/apps/api/app/api/v1/routes/billing.py index 34e1334f7..45388fd9e 100644 --- a/apps/api/app/api/v1/routes/billing.py +++ b/apps/api/app/api/v1/routes/billing.py @@ -6,7 +6,8 @@ BillingWorkflowService, ParseUsageResponse, ) -from app.services.rate_limit.dependencies import CurrentUser, with_current_user +from app.api.dependencies.job_admission import with_current_user +from app.services.rate_limit.data_structures import CurrentUser from fastapi import APIRouter, Depends, Query, Request from sqlalchemy.ext.asyncio import AsyncSession diff --git a/apps/api/app/api/v1/routes/demo.py b/apps/api/app/api/v1/routes/demo.py index 42508c5e1..320d9b232 100644 --- a/apps/api/app/api/v1/routes/demo.py +++ b/apps/api/app/api/v1/routes/demo.py @@ -5,7 +5,8 @@ from typing import Any from app.services.demo.document_service import DemoDocumentService -from app.services.rate_limit.dependencies import CurrentUser, with_current_user +from app.api.dependencies.job_admission import with_current_user +from app.services.rate_limit.data_structures import CurrentUser from fastapi import APIRouter, Depends, Query from fastapi.responses import FileResponse from pydantic import BaseModel, Field diff --git a/apps/api/app/api/v1/routes/documents.py b/apps/api/app/api/v1/routes/documents.py index 6d21a73cc..853599cad 100644 --- a/apps/api/app/api/v1/routes/documents.py +++ b/apps/api/app/api/v1/routes/documents.py @@ -5,7 +5,8 @@ from typing import Literal from app.services.document_service import DocumentService -from app.services.rate_limit.dependencies import CurrentUser, with_current_user +from app.api.dependencies.job_admission import with_current_user +from app.services.rate_limit.data_structures import CurrentUser from fastapi import APIRouter, Depends, Query from sqlalchemy.ext.asyncio import AsyncSession diff --git a/apps/api/app/api/v1/routes/guest.py b/apps/api/app/api/v1/routes/guest.py index 42935f7e1..700b5e5a1 100644 --- a/apps/api/app/api/v1/routes/guest.py +++ b/apps/api/app/api/v1/routes/guest.py @@ -1,7 +1,7 @@ """Guest registration routes.""" from app.services.guest.guest_registration_service import GuestRegistrationService -from app.services.rate_limit.dependencies import require_route_system_limit +from app.api.dependencies.job_admission import require_route_system_limit from fastapi import APIRouter, Depends from sqlalchemy.ext.asyncio import AsyncSession diff --git a/apps/api/app/api/v1/routes/jobs.py b/apps/api/app/api/v1/routes/jobs.py index e65c7c6c4..d6828eb31 100644 --- a/apps/api/app/api/v1/routes/jobs.py +++ b/apps/api/app/api/v1/routes/jobs.py @@ -12,12 +12,12 @@ get_job_result_for_user, list_jobs_for_user, ) -from app.services.rate_limit.dependencies import ( - CurrentUser, +from app.api.dependencies.job_admission import ( require_billing_limits, with_current_user, ) -from fastapi import APIRouter, Depends, Query, Request +from app.services.rate_limit.data_structures import CurrentUser +from fastapi import APIRouter, Depends, Query from sqlalchemy.ext.asyncio import AsyncSession from shared.core.database import get_db @@ -40,7 +40,6 @@ @router.post("/", include_in_schema=False) async def create_job( # pyright: ignore[reportGeneralTypeIssues] payload: JobCreate, - http_request: Request, current_user: CurrentUser = Depends(require_billing_limits), db: AsyncSession = Depends(get_db), ): @@ -51,7 +50,6 @@ async def create_job( # pyright: ignore[reportGeneralTypeIssues] db, payload=payload, current_user=current_user, - request=http_request, ) diff --git a/apps/api/app/api/v1/routes/retrieval.py b/apps/api/app/api/v1/routes/retrieval.py index 1ef332e4b..9e0df0a9b 100644 --- a/apps/api/app/api/v1/routes/retrieval.py +++ b/apps/api/app/api/v1/routes/retrieval.py @@ -4,7 +4,8 @@ from typing import Literal -from app.services.rate_limit.dependencies import CurrentUser, with_current_user +from app.api.dependencies.job_admission import with_current_user +from app.services.rate_limit.data_structures import CurrentUser from fastapi import APIRouter, Depends from pydantic import BaseModel, Field, field_validator from sqlalchemy.ext.asyncio import AsyncSession diff --git a/apps/api/app/api/v1/routes/webhook.py b/apps/api/app/api/v1/routes/webhook.py index 4e155cd11..053079119 100644 --- a/apps/api/app/api/v1/routes/webhook.py +++ b/apps/api/app/api/v1/routes/webhook.py @@ -9,7 +9,8 @@ from app.repositories.job_repository import JobRepository from app.repositories.webhook_repository import WebhookRepository -from app.services.rate_limit.dependencies import CurrentUser, with_current_user +from app.api.dependencies.job_admission import with_current_user +from app.services.rate_limit.data_structures import CurrentUser from fastapi import APIRouter, Depends, Query from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession diff --git a/apps/api/app/api/v1/routes/webhook_secrets.py b/apps/api/app/api/v1/routes/webhook_secrets.py index c0f63bf44..1f3a28708 100644 --- a/apps/api/app/api/v1/routes/webhook_secrets.py +++ b/apps/api/app/api/v1/routes/webhook_secrets.py @@ -6,7 +6,8 @@ from typing import List, Optional -from app.services.rate_limit.dependencies import CurrentUser, with_current_user +from app.api.dependencies.job_admission import with_current_user +from app.services.rate_limit.data_structures import CurrentUser from fastapi import APIRouter, Depends from pydantic import BaseModel, Field from sqlalchemy.ext.asyncio import AsyncSession diff --git a/apps/api/app/services/document_ingestion/service.py b/apps/api/app/services/document_ingestion/service.py index 03f5da8fe..7313c07f8 100644 --- a/apps/api/app/services/document_ingestion/service.py +++ b/apps/api/app/services/document_ingestion/service.py @@ -17,8 +17,7 @@ resolve_effective_document_scope, ) from app.services.rate_limit.data_structures import CurrentUser -from app.services.rate_limit.dependencies import enforce_job_creation_capacity -from fastapi import Request +from app.services.rate_limit.job_admission_service import JobAdmissionService from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession @@ -47,11 +46,13 @@ def __init__( *, creation_service: DocumentIngestionCreationService | None = None, confirmation_service: DocumentIngestionConfirmationService | None = None, + job_admission_service: JobAdmissionService | None = None, ) -> None: self._creation_service = creation_service or DocumentIngestionCreationService() self._confirmation_service = ( confirmation_service or DocumentIngestionConfirmationService() ) + self._job_admission_service = job_admission_service or JobAdmissionService() async def create_job( self, @@ -59,7 +60,6 @@ async def create_job( *, payload: JobCreate, current_user: CurrentUser, - request: Request, ) -> JobResponse: try: job_id = f"job_{uuid.uuid4().hex[:12]}" @@ -70,8 +70,7 @@ async def create_job( current_user=current_user, ) - await enforce_job_creation_capacity( - request=request, + await self._job_admission_service.enforce_job_creation_capacity( db=db, current_user=current_user, ) diff --git a/apps/api/app/services/rate_limit/data_structures.py b/apps/api/app/services/rate_limit/data_structures.py index 0467c79aa..af206917b 100644 --- a/apps/api/app/services/rate_limit/data_structures.py +++ b/apps/api/app/services/rate_limit/data_structures.py @@ -15,6 +15,15 @@ class CurrentUser: user_tier: str +@dataclass(frozen=True) +class RouteAdmissionContext: + """HTTP route facts needed by the Job Admission workflow.""" + + method: str + path: str + limit_identifier: str + + @dataclass(frozen=True) class TierLimits: """ diff --git a/apps/api/app/services/rate_limit/dependencies.py b/apps/api/app/services/rate_limit/dependencies.py deleted file mode 100644 index c76785dc2..000000000 --- a/apps/api/app/services/rate_limit/dependencies.py +++ /dev/null @@ -1,48 +0,0 @@ -"""FastAPI adapters for the Job Admission module.""" - -from typing import AsyncGenerator - -from app.api.dependencies.auth import get_current_user_id -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 sqlalchemy.ext.asyncio import AsyncSession -from shared.core.database import get_db -_job_admission_service = JobAdmissionService() - -async def with_current_user( - request: Request, - user_id: str = Depends(get_current_user_id), -) -> AsyncGenerator[CurrentUser, None]: - current_user = await _job_admission_service.resolve_current_user( - request=request, - user_id=user_id, - ) - yield current_user - - -async def require_billing_limits( - request: Request, - current_user: CurrentUser = Depends(with_current_user), - _db: AsyncSession = Depends(get_db), -) -> AsyncGenerator[CurrentUser, None]: - 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: - await _job_admission_service.enforce_route_system_limit(request=request) - - -async def enforce_job_creation_capacity( - request: Request, - db: AsyncSession, - current_user: CurrentUser, -) -> None: - del request - await _job_admission_service.enforce_job_creation_capacity( - db=db, - current_user=current_user, - ) 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 index d455f9927..ce89931f6 100644 --- 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 @@ -3,9 +3,9 @@ from fnmatch import fnmatch from app.services.rate_limit.config import RateLimitConfig +from app.services.rate_limit.data_structures import RouteAdmissionContext 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, @@ -36,12 +36,15 @@ class JobAdmissionRoutePolicyService: async def enforce_user_system_limit( self, *, - request: Request, + route_context: RouteAdmissionContext, config: RateLimitConfig, user_id: str, ) -> None: - route_path = self._get_route_path(request) - rule = find_system_rule(request.method, route_path, config.system_rules) + rule = find_system_rule( + route_context.method, + route_context.path, + config.system_rules, + ) limiter = RateLimiter(config) await limiter.check_system_limit( identifier=user_id, @@ -50,14 +53,21 @@ async def enforce_user_system_limit( period=rule.period, ) - async def enforce_route_system_limit(self, *, request: Request) -> None: + async def enforce_route_system_limit( + self, + *, + route_context: RouteAdmissionContext, + ) -> 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) + route_identifier = route_context.limit_identifier + rule = find_system_rule( + route_context.method, + route_context.path, + config.system_rules, + ) limiter = RateLimiter(config) try: @@ -81,14 +91,13 @@ async def enforce_route_system_limit(self, *, request: Request) -> None: def enforce_guest_api_key_scope( self, *, - request: Request, + route_context: RouteAdmissionContext, 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): + if self._is_guest_api_key_route_allowed(route_context.path): return raise PermissionDeniedException( @@ -96,27 +105,6 @@ def enforce_guest_api_key_scope( 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 "/" 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 554b8a36d..130a880c3 100644 --- a/apps/api/app/services/rate_limit/job_admission_service.py +++ b/apps/api/app/services/rate_limit/job_admission_service.py @@ -1,7 +1,7 @@ from __future__ import annotations from app.services.rate_limit.config import RateLimitConfig -from app.services.rate_limit.data_structures import CurrentUser +from app.services.rate_limit.data_structures import CurrentUser, RouteAdmissionContext from app.services.rate_limit.job_admission_capacity_service import ( JobAdmissionCapacityService, ) @@ -9,7 +9,6 @@ JobAdmissionRoutePolicyService, ) from app.services.rate_limit.tier_service import TierService -from fastapi import Request from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession @@ -32,12 +31,12 @@ def __init__( async def resolve_current_user( self, *, - request: Request, + route_context: RouteAdmissionContext, user_id: str, ) -> CurrentUser: user_tier = await TierService.get_tier(user_id) self._route_policy_service.enforce_guest_api_key_scope( - request=request, + route_context=route_context, user_tier=user_tier, ) current_user = CurrentUser(user_id=user_id, user_tier=user_tier) @@ -49,7 +48,7 @@ async def resolve_current_user( try: await self._route_policy_service.enforce_user_system_limit( - request=request, + route_context=route_context, config=config, user_id=user_id, ) @@ -72,8 +71,14 @@ async def enforce_billing_limits( ) -> None: await self._capacity_service.enforce_billing_limits(current_user=current_user) - async def enforce_route_system_limit(self, *, request: Request) -> None: - await self._route_policy_service.enforce_route_system_limit(request=request) + async def enforce_route_system_limit( + self, + *, + route_context: RouteAdmissionContext, + ) -> None: + await self._route_policy_service.enforce_route_system_limit( + route_context=route_context, + ) async def enforce_job_creation_capacity( self, From e585654749976c89b01030ec3369f1a6820188ed Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 11:39:41 +0000 Subject: [PATCH 63/78] refactor: return qstash callback outcomes from workflow --- .../api/app/api/v1/routes/qstash_callbacks.py | 13 ++++++++-- .../app/services/qstash_callback_service.py | 25 +++++++++++++------ 2 files changed, 28 insertions(+), 10 deletions(-) diff --git a/apps/api/app/api/v1/routes/qstash_callbacks.py b/apps/api/app/api/v1/routes/qstash_callbacks.py index 495cb22ea..4034ec773 100644 --- a/apps/api/app/api/v1/routes/qstash_callbacks.py +++ b/apps/api/app/api/v1/routes/qstash_callbacks.py @@ -3,6 +3,7 @@ from __future__ import annotations from app.services import qstash_callback_service +from app.services.qstash_callback_service import QStashCallbackOutcome from fastapi import APIRouter, Request, Response router = APIRouter(tags=["QStash Callbacks"]) @@ -25,7 +26,9 @@ async def handle_qstash_callback(request: Request) -> Response: ): return Response(status_code=401, content="Invalid signature") - return qstash_callback_service.handle_qstash_success_callback(raw_body) + return _to_response( + qstash_callback_service.handle_qstash_success_callback(raw_body) + ) @router.post("/qstash/failure") @@ -45,4 +48,10 @@ async def handle_qstash_failure(request: Request) -> Response: ): return Response(status_code=401, content="Invalid signature") - return qstash_callback_service.handle_qstash_failure_callback(raw_body) + return _to_response( + qstash_callback_service.handle_qstash_failure_callback(raw_body) + ) + + +def _to_response(outcome: QStashCallbackOutcome) -> Response: + return Response(status_code=outcome.status_code, content=outcome.content) diff --git a/apps/api/app/services/qstash_callback_service.py b/apps/api/app/services/qstash_callback_service.py index bc7e9518c..dfd16f49b 100644 --- a/apps/api/app/services/qstash_callback_service.py +++ b/apps/api/app/services/qstash_callback_service.py @@ -1,11 +1,11 @@ from __future__ import annotations import json +from dataclasses import dataclass 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 @@ -15,6 +15,12 @@ from shared.models.database.webhook_log import WebhookLog +@dataclass(frozen=True) +class QStashCallbackOutcome: + status_code: int + content: str + + 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: @@ -52,13 +58,13 @@ def verify_qstash_signature(raw_body: bytes, signature: str, url: str) -> bool: return False -def handle_qstash_success_callback(raw_body: bytes) -> Response: +def handle_qstash_success_callback(raw_body: bytes) -> QStashCallbackOutcome: 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)") + return QStashCallbackOutcome(status_code=200, content="OK (no event_id)") retried = data.get("retried", 0) logger.info( @@ -74,13 +80,13 @@ def handle_qstash_success_callback(raw_body: bytes) -> Response: ) -def handle_qstash_failure_callback(raw_body: bytes) -> Response: +def handle_qstash_failure_callback(raw_body: bytes) -> QStashCallbackOutcome: 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)") + return QStashCallbackOutcome(status_code=200, content="OK (no event_id)") retried = data.get("retried", 0) max_retries = data.get("maxRetries", 0) @@ -179,7 +185,7 @@ def process_qstash_callback( event_id: str, callback_status: str, log_label: str, -) -> Response: +) -> QStashCallbackOutcome: response_status_code = get_response_status_code(data.get("status")) response_body = data.get("body", "") qstash_message_id = data.get("sourceMessageId") @@ -202,7 +208,10 @@ def process_qstash_callback( 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)") + return QStashCallbackOutcome( + 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) @@ -231,4 +240,4 @@ def process_qstash_callback( db.add(log) db.commit() - return Response(status_code=200, content="OK") + return QStashCallbackOutcome(status_code=200, content="OK") From 74d1a2ec0f7e20c5e9e11467b63891956cde89ff Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 11:42:30 +0000 Subject: [PATCH 64/78] refactor: group document lifecycle service --- apps/api/app/api/v1/routes/documents.py | 2 +- apps/api/app/services/documents/__init__.py | 1 + .../{document_service.py => documents/lifecycle_service.py} | 4 +--- 3 files changed, 3 insertions(+), 4 deletions(-) create mode 100644 apps/api/app/services/documents/__init__.py rename apps/api/app/services/{document_service.py => documents/lifecycle_service.py} (99%) diff --git a/apps/api/app/api/v1/routes/documents.py b/apps/api/app/api/v1/routes/documents.py index 853599cad..b13a7cb4d 100644 --- a/apps/api/app/api/v1/routes/documents.py +++ b/apps/api/app/api/v1/routes/documents.py @@ -4,8 +4,8 @@ from typing import Literal -from app.services.document_service import DocumentService from app.api.dependencies.job_admission import with_current_user +from app.services.documents.lifecycle_service import DocumentService from app.services.rate_limit.data_structures import CurrentUser from fastapi import APIRouter, Depends, Query from sqlalchemy.ext.asyncio import AsyncSession diff --git a/apps/api/app/services/documents/__init__.py b/apps/api/app/services/documents/__init__.py new file mode 100644 index 000000000..53fc5491c --- /dev/null +++ b/apps/api/app/services/documents/__init__.py @@ -0,0 +1 @@ +"""Document lifecycle workflow package.""" diff --git a/apps/api/app/services/document_service.py b/apps/api/app/services/documents/lifecycle_service.py similarity index 99% rename from apps/api/app/services/document_service.py rename to apps/api/app/services/documents/lifecycle_service.py index 13dddd3e8..93b9fb300 100644 --- a/apps/api/app/services/document_service.py +++ b/apps/api/app/services/documents/lifecycle_service.py @@ -1,6 +1,4 @@ -""" -Application service for document lifecycle routes. -""" +"""Application workflow for document lifecycle routes.""" from __future__ import annotations From afc7f5ad9d9b3597646cbf75767ba95b6c547887 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 11:46:59 +0000 Subject: [PATCH 65/78] refactor: group shared job lifecycle services --- apps/worker/app/core/tasks/base_task.py | 2 +- .../app/services/document_ingestion/processing_run.py | 2 +- apps/worker/app/services/workload/url_upload_service.py | 2 +- .../shared/services/jobs/lifecycle/__init__.py | 1 + .../lifecycle/failure_finalizer.py} | 4 ++-- .../lifecycle/post_commit_effects.py} | 4 ++-- .../lifecycle/publication.py} | 0 .../lifecycle/result_writer.py} | 0 .../{job_lifecycle_sync.py => jobs/lifecycle/service.py} | 6 +++--- .../lifecycle/success_finalizer.py} | 8 ++++---- .../lifecycle/webhook_outbox.py} | 0 packages/shared-python/shared/testing/contract_runtime.py | 2 +- 12 files changed, 16 insertions(+), 15 deletions(-) create mode 100644 packages/shared-python/shared/services/jobs/lifecycle/__init__.py rename packages/shared-python/shared/services/{job_failure_sync.py => jobs/lifecycle/failure_finalizer.py} (95%) rename packages/shared-python/shared/services/{job_post_commit_effects_sync.py => jobs/lifecycle/post_commit_effects.py} (91%) rename packages/shared-python/shared/services/{job_publication_sync.py => jobs/lifecycle/publication.py} (100%) rename packages/shared-python/shared/services/{job_result_sync.py => jobs/lifecycle/result_writer.py} (100%) rename packages/shared-python/shared/services/{job_lifecycle_sync.py => jobs/lifecycle/service.py} (96%) rename packages/shared-python/shared/services/{job_success_sync.py => jobs/lifecycle/success_finalizer.py} (91%) rename packages/shared-python/shared/services/{job_webhook_outbox_sync.py => jobs/lifecycle/webhook_outbox.py} (100%) diff --git a/apps/worker/app/core/tasks/base_task.py b/apps/worker/app/core/tasks/base_task.py index 4a9f65543..dee61c98f 100644 --- a/apps/worker/app/core/tasks/base_task.py +++ b/apps/worker/app/core/tasks/base_task.py @@ -44,7 +44,7 @@ def on_failure(self, exc, task_id, args, kwargs, einfo): # Finalize failure directly to the database. if job_id: try: - from shared.services.job_lifecycle_sync import ( + from shared.services.jobs.lifecycle.service import ( get_sync_job_lifecycle_service, ) diff --git a/apps/worker/app/services/document_ingestion/processing_run.py b/apps/worker/app/services/document_ingestion/processing_run.py index 88e63ef22..d629eef69 100644 --- a/apps/worker/app/services/document_ingestion/processing_run.py +++ b/apps/worker/app/services/document_ingestion/processing_run.py @@ -35,7 +35,7 @@ from loguru import logger from shared.models.schemas.job_metadata import JobMetadataHelper -from shared.services.job_lifecycle_sync import get_sync_job_lifecycle_service +from shared.services.jobs.lifecycle.service import get_sync_job_lifecycle_service from shared.services.redis.distributed_lock import RedisJobLock from shared.services.redis.redis_sync_service import ( SyncRedisServiceFactory, diff --git a/apps/worker/app/services/workload/url_upload_service.py b/apps/worker/app/services/workload/url_upload_service.py index cad1e2744..27cf81616 100644 --- a/apps/worker/app/services/workload/url_upload_service.py +++ b/apps/worker/app/services/workload/url_upload_service.py @@ -13,7 +13,7 @@ upload_temp_file_to_source_storage, verify_source_upload, ) -from shared.services.job_lifecycle_sync import get_sync_job_lifecycle_service +from shared.services.jobs.lifecycle.service import get_sync_job_lifecycle_service from shared.services.redis.redis_sync_service import SyncRedisServiceFactory diff --git a/packages/shared-python/shared/services/jobs/lifecycle/__init__.py b/packages/shared-python/shared/services/jobs/lifecycle/__init__.py new file mode 100644 index 000000000..93da44f12 --- /dev/null +++ b/packages/shared-python/shared/services/jobs/lifecycle/__init__.py @@ -0,0 +1 @@ +"""Synchronous job lifecycle finalization workflow.""" diff --git a/packages/shared-python/shared/services/job_failure_sync.py b/packages/shared-python/shared/services/jobs/lifecycle/failure_finalizer.py similarity index 95% rename from packages/shared-python/shared/services/job_failure_sync.py rename to packages/shared-python/shared/services/jobs/lifecycle/failure_finalizer.py index ad3a313dd..41c2f4e64 100644 --- a/packages/shared-python/shared/services/job_failure_sync.py +++ b/packages/shared-python/shared/services/jobs/lifecycle/failure_finalizer.py @@ -11,8 +11,8 @@ from shared.core.state_machine.service_sync import SyncStateMachineService from shared.models.database.job import Job from shared.services.billing.credits_sync_service import SyncCreditsService -from shared.services.job_post_commit_effects_sync import PostCommitEffectPlan -from shared.services.job_webhook_outbox_sync import SyncJobWebhookOutbox +from shared.services.jobs.lifecycle.post_commit_effects import PostCommitEffectPlan +from shared.services.jobs.lifecycle.webhook_outbox import SyncJobWebhookOutbox from shared.utils.error_details import normalize_error_details diff --git a/packages/shared-python/shared/services/job_post_commit_effects_sync.py b/packages/shared-python/shared/services/jobs/lifecycle/post_commit_effects.py similarity index 91% rename from packages/shared-python/shared/services/job_post_commit_effects_sync.py rename to packages/shared-python/shared/services/jobs/lifecycle/post_commit_effects.py index 35546ed71..0575d5d42 100644 --- a/packages/shared-python/shared/services/job_post_commit_effects_sync.py +++ b/packages/shared-python/shared/services/jobs/lifecycle/post_commit_effects.py @@ -3,8 +3,8 @@ from dataclasses import dataclass, field from typing import Any -from shared.services.job_publication_sync import SyncJobPublicationFinalizer -from shared.services.job_webhook_outbox_sync import SyncJobWebhookOutbox +from shared.services.jobs.lifecycle.publication import SyncJobPublicationFinalizer +from shared.services.jobs.lifecycle.webhook_outbox import SyncJobWebhookOutbox @dataclass(frozen=True) diff --git a/packages/shared-python/shared/services/job_publication_sync.py b/packages/shared-python/shared/services/jobs/lifecycle/publication.py similarity index 100% rename from packages/shared-python/shared/services/job_publication_sync.py rename to packages/shared-python/shared/services/jobs/lifecycle/publication.py diff --git a/packages/shared-python/shared/services/job_result_sync.py b/packages/shared-python/shared/services/jobs/lifecycle/result_writer.py similarity index 100% rename from packages/shared-python/shared/services/job_result_sync.py rename to packages/shared-python/shared/services/jobs/lifecycle/result_writer.py diff --git a/packages/shared-python/shared/services/job_lifecycle_sync.py b/packages/shared-python/shared/services/jobs/lifecycle/service.py similarity index 96% rename from packages/shared-python/shared/services/job_lifecycle_sync.py rename to packages/shared-python/shared/services/jobs/lifecycle/service.py index f160f41af..e0d82d832 100644 --- a/packages/shared-python/shared/services/job_lifecycle_sync.py +++ b/packages/shared-python/shared/services/jobs/lifecycle/service.py @@ -17,12 +17,12 @@ from sqlalchemy.orm import Session from shared.core.database_sync import get_sync_db_context -from shared.services.job_failure_sync import SyncJobFailureFinalizer -from shared.services.job_post_commit_effects_sync import ( +from shared.services.jobs.lifecycle.failure_finalizer import SyncJobFailureFinalizer +from shared.services.jobs.lifecycle.post_commit_effects import ( PostCommitEffectPlan, SyncJobPostCommitEffectRunner, ) -from shared.services.job_success_sync import SyncJobSuccessFinalizer +from shared.services.jobs.lifecycle.success_finalizer import SyncJobSuccessFinalizer from shared.services.redis.redis_sync_service import ( SyncRedisServiceFactory, ) diff --git a/packages/shared-python/shared/services/job_success_sync.py b/packages/shared-python/shared/services/jobs/lifecycle/success_finalizer.py similarity index 91% rename from packages/shared-python/shared/services/job_success_sync.py rename to packages/shared-python/shared/services/jobs/lifecycle/success_finalizer.py index 1f7b8b20b..c325edefc 100644 --- a/packages/shared-python/shared/services/job_success_sync.py +++ b/packages/shared-python/shared/services/jobs/lifecycle/success_finalizer.py @@ -7,10 +7,10 @@ from sqlalchemy.orm import Session from shared.core.state_machine.service_sync import SyncStateMachineService -from shared.services.job_post_commit_effects_sync import PostCommitEffectPlan -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 +from shared.services.jobs.lifecycle.post_commit_effects import PostCommitEffectPlan +from shared.services.jobs.lifecycle.publication import SyncJobPublicationFinalizer +from shared.services.jobs.lifecycle.result_writer import SyncJobResultWriter +from shared.services.jobs.lifecycle.webhook_outbox import SyncJobWebhookOutbox @dataclass(frozen=True) diff --git a/packages/shared-python/shared/services/job_webhook_outbox_sync.py b/packages/shared-python/shared/services/jobs/lifecycle/webhook_outbox.py similarity index 100% rename from packages/shared-python/shared/services/job_webhook_outbox_sync.py rename to packages/shared-python/shared/services/jobs/lifecycle/webhook_outbox.py diff --git a/packages/shared-python/shared/testing/contract_runtime.py b/packages/shared-python/shared/testing/contract_runtime.py index 53c139510..2865226a3 100644 --- a/packages/shared-python/shared/testing/contract_runtime.py +++ b/packages/shared-python/shared/testing/contract_runtime.py @@ -65,7 +65,7 @@ "shared.core.state_machine.service_sync", "shared.services.billing.credits_sync_service", "shared.services.billing.work_billing_service", - "shared.services.job_lifecycle_sync", + "shared.services.jobs.lifecycle.service", "shared.services.retrieval.app_service", "shared.services.retrieval.publication_service", "shared.services.webhook", From b8ee5097cfd8f907922ff06851e9b1c163141296 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 11:48:31 +0000 Subject: [PATCH 66/78] docs: refresh architecture context map --- CONTEXT.md | 32 +++++++++++++++++++++++++++----- 1 file changed, 27 insertions(+), 5 deletions(-) diff --git a/CONTEXT.md b/CONTEXT.md index b5701ef67..ac3b6fa43 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -122,7 +122,8 @@ quota. ### Job Admission Route Policy The route-aware part of Job Admission that enforces guest API key scope and -system limits. +system limits from plain route-admission context built by HTTP dependency +adapters. ### Job Admission Capacity @@ -231,8 +232,10 @@ managers such as Ali, iLoveAPI, and MinerU. ### HTTP Adapters `apps/api/app/api/v1/routes/*` +`apps/api/app/api/dependencies/*` -These modules translate HTTP requests into application workflow calls. +These modules translate HTTP requests and dependency context into application +workflow calls. ### Application Workflows @@ -254,7 +257,8 @@ These modules own database reads and writes for API-side workflows. These modules own the lower-level implementations for publication, retrieval, state machines, storage, Redis-backed metadata, billing primitives, and core -exceptions. +exceptions. Shared Job lifecycle finalization lives under +`packages/shared-python/shared/services/jobs/lifecycle/*`. ## apps/api Workflow Ownership @@ -276,13 +280,15 @@ exceptions. ### Job Admission +- `app/api/dependencies/auth.py` +- `app/api/dependencies/job_admission.py` +- `app/services/auth/*` - `app/services/rate_limit/*` -- `app/core/dependencies.py` ### Document Lifecycle - `app/api/v1/routes/documents.py` -- `app/services/document_service.py` +- `app/services/documents/lifecycle_service.py` - `app/repositories/document_repository.py` ### Retrieval @@ -330,6 +336,22 @@ headers, acknowledges malformed or unsafe events, and triggers upload handoff. - `app/api/v1/routes/qstash_callbacks.py` - `app/services/qstash_callback_service.py` +The route owns QStash HTTP signature verification and HTTP response projection. +The workflow owns callback parsing, event status resolution, and webhook log +side effects. + +## Shared Workflow Ownership + +### Job Lifecycle Finalization + +- `shared/services/jobs/lifecycle/service.py` +- `shared/services/jobs/lifecycle/success_finalizer.py` +- `shared/services/jobs/lifecycle/failure_finalizer.py` +- `shared/services/jobs/lifecycle/result_writer.py` +- `shared/services/jobs/lifecycle/publication.py` +- `shared/services/jobs/lifecycle/post_commit_effects.py` +- `shared/services/jobs/lifecycle/webhook_outbox.py` + ## apps/worker Workflow Ownership ### Worker Document Parsing From 189b71cd322e4e778a46e8cf41a88171348c6dda Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 12:02:41 +0000 Subject: [PATCH 67/78] refactor: group webhook callback services --- CONTEXT.md | 4 +- .../api/app/api/v1/routes/qstash_callbacks.py | 4 +- apps/api/app/services/webhook/__init__.py | 1 + .../{ => webhook}/qstash_callback_service.py | 0 apps/api/app/services/webhook_service.py | 182 ------------------ .../contract/test_qstash_callback_contract.py | 24 ++- 6 files changed, 23 insertions(+), 192 deletions(-) create mode 100644 apps/api/app/services/webhook/__init__.py rename apps/api/app/services/{ => webhook}/qstash_callback_service.py (100%) delete mode 100644 apps/api/app/services/webhook_service.py diff --git a/CONTEXT.md b/CONTEXT.md index ac3b6fa43..90ee0fa1b 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -318,7 +318,7 @@ exceptions. Shared Job lifecycle finalization lives under - `app/api/v1/routes/webhook.py` - `app/api/v1/routes/webhook_secrets.py` -- `app/services/webhook_service.py` +- `app/services/webhook/*` - `app/repositories/webhook_repository.py` ### Internal Storage Events @@ -334,7 +334,7 @@ headers, acknowledges malformed or unsafe events, and triggers upload handoff. ### Async Callbacks - `app/api/v1/routes/qstash_callbacks.py` -- `app/services/qstash_callback_service.py` +- `app/services/webhook/qstash_callback_service.py` The route owns QStash HTTP signature verification and HTTP response projection. The workflow owns callback parsing, event status resolution, and webhook log diff --git a/apps/api/app/api/v1/routes/qstash_callbacks.py b/apps/api/app/api/v1/routes/qstash_callbacks.py index 4034ec773..041146adf 100644 --- a/apps/api/app/api/v1/routes/qstash_callbacks.py +++ b/apps/api/app/api/v1/routes/qstash_callbacks.py @@ -2,8 +2,8 @@ from __future__ import annotations -from app.services import qstash_callback_service -from app.services.qstash_callback_service import QStashCallbackOutcome +from app.services.webhook import qstash_callback_service +from app.services.webhook.qstash_callback_service import QStashCallbackOutcome from fastapi import APIRouter, Request, Response router = APIRouter(tags=["QStash Callbacks"]) diff --git a/apps/api/app/services/webhook/__init__.py b/apps/api/app/services/webhook/__init__.py new file mode 100644 index 000000000..c2528a3fc --- /dev/null +++ b/apps/api/app/services/webhook/__init__.py @@ -0,0 +1 @@ +"""Webhook delivery workflows.""" diff --git a/apps/api/app/services/qstash_callback_service.py b/apps/api/app/services/webhook/qstash_callback_service.py similarity index 100% rename from apps/api/app/services/qstash_callback_service.py rename to apps/api/app/services/webhook/qstash_callback_service.py diff --git a/apps/api/app/services/webhook_service.py b/apps/api/app/services/webhook_service.py deleted file mode 100644 index 25b1f1b21..000000000 --- a/apps/api/app/services/webhook_service.py +++ /dev/null @@ -1,182 +0,0 @@ -""" -Consolidated Webhook Service - -Single responsibility: create WebhookEvents and publish them for async delivery. -""" - -from datetime import datetime, timezone -from typing import Any, Dict, Optional - -from loguru import logger -from sqlalchemy.ext.asyncio import AsyncSession - -from shared.core.exceptions.webhook_exceptions import WebhookConfigException -from shared.core.response import build_standard_error_response -from shared.models.database.webhook import WebhookEvent, WebhookEventStatus - - -class WebhookService: - """ - Webhook Service - Transactional Outbox Pattern. - - Creates WebhookEvent records and publishes them via QStash for async delivery. - """ - - async def create_job_completion_event( - self, db: AsyncSession, job_id: str, webhook_url: str - ) -> WebhookEvent: - """ - Create webhook event for job completion. - - Args: - db: Database session (in active transaction) - job_id: Job ID - webhook_url: Webhook URL - - Returns: - Created WebhookEvent - """ - # Build minimal payload - dispatcher will enrich at delivery time - payload: Dict[str, Any] = { - "event": "job.completed", - "job_id": job_id, - "status": "completed", - "completed_at": datetime.now(timezone.utc).isoformat(), - # NOTE: result_url and result added by dispatcher at delivery - } - - # Create event - event = await self._create_event( - db=db, job_id=job_id, target_url=webhook_url, payload=payload - ) - - logger.info( - f"Job completion webhook event created: event_id={event.id}, job_id={job_id}" - ) - return event - - async def create_job_failure_event( - self, - db: AsyncSession, - job_id: str, - error_message: str, - error_type: Optional[str] = None, - error_code: str = "UNKNOWN", - error_details: Optional[Dict[str, Any]] = None, - webhook_url: Optional[str] = None, - ) -> WebhookEvent: - """ - Create webhook event for job failure. - - Args: - db: Database session (in active transaction) - job_id: Job ID - error_message: Error message - error_type: Error type (optional) - error_code: Error code - error_details: Structured error details (optional) - webhook_url: Webhook URL - - Returns: - Created WebhookEvent - """ - # Build payload with standardized error format - payload: Dict[str, Any] = { - "event": "job.failed", - "job_id": job_id, - "status": "failed", - "failed_at": datetime.now(timezone.utc).isoformat(), - "error": build_standard_error_response( - code=error_code, - message=error_message, - request_id=job_id, - details=error_details, - ), - } - - if webhook_url is None: - raise WebhookConfigException( - internal_message="Missing webhook target_url", - user_message="Webhook URL is required.", - ) - - # Create event - event = await self._create_event( - db=db, job_id=job_id, target_url=webhook_url, payload=payload - ) - - logger.info( - f"Job failure webhook event created: event_id={event.id}, job_id={job_id}" - ) - return event - - async def publish_to_queue(self, event_id: str) -> bool: - """Publish a webhook event for async delivery via QStash after commit.""" - return await self._publish_via_qstash(event_id) - - async def _publish_via_qstash(self, event_id: str) -> bool: - """Publish via QStash for managed delivery and retry.""" - try: - from shared.services.webhook.qstash_publisher import ( - get_qstash_webhook_publisher, - ) - - publisher = get_qstash_webhook_publisher() - message_id = publisher.publish_event(event_id) - if message_id: - logger.info( - f"Webhook published via QStash: event_id={event_id}, message_id={message_id}" - ) - return True - logger.warning( - f"QStash publish returned no message_id: event_id={event_id}" - ) - return False - except Exception as exc: - logger.error(f"QStash publish failed: event_id={event_id}, error={exc}") - return False - - async def _create_event( - self, db: AsyncSession, job_id: str, target_url: str, payload: Dict[str, Any] - ) -> WebhookEvent: - """ - Create a WebhookEvent record. - """ - # Validate - if not target_url: - raise WebhookConfigException( - internal_message="Missing webhook target_url", - user_message="Webhook URL is required.", - ) - - if not target_url.startswith(("http://", "https://")): - raise WebhookConfigException( - internal_message=f"Invalid webhook scheme: {target_url}", - user_message="Webhook URL must start with http:// or https://.", - ) - - # Create event - event = WebhookEvent( - job_id=job_id, - target_url=target_url, - payload=payload, - status=WebhookEventStatus.PENDING, - attempts=0, - ) - - db.add(event) - await db.flush() # Get ID, but don't commit (caller controls transaction) - - return event - - -# Singleton -_webhook_service: Optional[WebhookService] = None - - -def get_webhook_service() -> WebhookService: - """Get singleton WebhookService instance.""" - global _webhook_service - if _webhook_service is None: - _webhook_service = WebhookService() - return _webhook_service diff --git a/apps/api/tests/contract/test_qstash_callback_contract.py b/apps/api/tests/contract/test_qstash_callback_contract.py index 48c78cee3..c9e0c8f40 100644 --- a/apps/api/tests/contract/test_qstash_callback_contract.py +++ b/apps/api/tests/contract/test_qstash_callback_contract.py @@ -48,7 +48,9 @@ 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.services.qstash_callback_service") + qstash_module = importlib.import_module( + "app.services.webhook.qstash_callback_service" + ) monkeypatch.setattr(qstash_module, "verify_qstash_signature", lambda *args: False) response = await api_client.post( "/api/v1/webhooks/qstash/callback", @@ -70,7 +72,9 @@ 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.services.qstash_callback_service") + qstash_module = importlib.import_module( + "app.services.webhook.qstash_callback_service" + ) monkeypatch.setattr(qstash_module, "verify_qstash_signature", lambda *args: True) response = await api_client.post( "/api/v1/webhooks/qstash/callback", @@ -133,7 +137,9 @@ 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.services.qstash_callback_service") + qstash_module = importlib.import_module( + "app.services.webhook.qstash_callback_service" + ) monkeypatch.setattr(qstash_module, "verify_qstash_signature", lambda *args: True) response = await api_client.post( "/api/v1/webhooks/qstash/callback", @@ -197,7 +203,9 @@ 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.services.qstash_callback_service") + qstash_module = importlib.import_module( + "app.services.webhook.qstash_callback_service" + ) monkeypatch.setattr(qstash_module, "verify_qstash_signature", lambda *args: True) response = await api_client.post( "/api/v1/webhooks/qstash/callback", @@ -257,7 +265,9 @@ 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.services.qstash_callback_service") + qstash_module = importlib.import_module( + "app.services.webhook.qstash_callback_service" + ) monkeypatch.setattr(qstash_module, "verify_qstash_signature", lambda *args: True) response = await api_client.post( "/api/v1/webhooks/qstash/failure", @@ -318,7 +328,9 @@ 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.services.qstash_callback_service") + qstash_module = importlib.import_module( + "app.services.webhook.qstash_callback_service" + ) monkeypatch.setattr(qstash_module, "verify_qstash_signature", lambda *args: True) response = await api_client.post( "/api/v1/webhooks/qstash/callback", From 410734b8b8d93b86803e4acbf1c07305c88f4663 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 12:22:29 +0000 Subject: [PATCH 68/78] refactor: structure worker document parser modules --- AGENTS.md | 54 +++++++++---------- CONTEXT.md | 22 +++++--- .../tools/probe_sample_pages.py | 2 +- .../document_agent/tools/probe_vlm_inspect.py | 2 +- .../document_ingestion/page_estimator.py | 4 +- .../document_ingestion/processing_run.py | 4 +- .../{ => assets}/image_compressor.py | 0 .../{ => assets}/inline_asset.py | 2 +- .../{ => conversion}/legacy_converter.py | 2 +- .../atlas/classifier.py} | 2 +- .../atlas/parser.py} | 12 ++--- .../docx/asset_accumulator.py} | 2 +- .../docx/asset_store.py} | 0 .../docx/block_stream.py} | 2 +- .../{doc_parser.py => formats/docx/parser.py} | 36 ++++++------- .../docx/table_html.py} | 0 .../{toc_docx.py => formats/docx/toc.py} | 4 +- .../excel/structure_parser.py} | 0 .../excel/table_parser.py} | 16 +++--- .../fragment/parser.py} | 2 +- .../html/parser.py} | 2 +- .../image/parser.py} | 12 ++--- .../markdown/deferred_summary.py} | 10 ++-- .../markdown/deferred_task.py} | 0 .../markdown/image_asset.py} | 10 ++-- .../markdown/parse_state.py} | 6 +-- .../markdown/parser.py} | 30 +++++------ .../markdown/table_asset.py} | 12 ++--- .../{pdf_parser.py => formats/pdf/parser.py} | 10 ++-- .../pdf/pptx_rendering.py} | 2 +- .../{ => formats/pdf}/pymupdf_subprocess.py | 0 .../pdf/rendered_transform.py} | 6 +-- .../pptx/parser.py} | 10 ++-- .../{txt_parser.py => formats/text/parser.py} | 0 .../orchestration/format_adapters.py | 22 ++++---- .../orchestration/parse_session.py | 6 +-- .../orchestration/postprocess.py | 4 +- .../{ => profiling}/doc_profile_model.py | 0 .../{ => profiling}/doc_profile_pdf.py | 4 +- .../{ => profiling}/doc_profiler.py | 6 +-- .../mineru/client.py} | 2 +- .../mineru/pdf_service.py} | 8 +-- .../mineru/quota_manager.py} | 0 .../mineru/task_polling.py} | 4 +- .../{ => structure}/heading_candidates.py | 6 +-- .../{ => structure}/heading_hierarchy.py | 2 +- .../{ => structure}/heading_llm_executor.py | 6 +-- .../{ => structure}/heading_tree.py | 0 .../{ => structure}/layout_parser.py | 14 ++--- .../{ => structure}/metadata_extractor.py | 0 .../{ => structure}/toc_hierarchy.py | 8 +-- .../{ => structure}/toc_parser.py | 8 +-- .../{ => support}/identifiers.py | 0 .../{ => support}/internal_parse_name.py | 0 .../{ => support}/parser_log_utils.py | 0 .../{ => support}/parser_rows.py | 0 .../{ => support}/path_helpers.py | 0 .../{ => support}/stage_profiler.py | 0 .../{ => support}/text_helpers.py | 0 .../{ => tables}/dataframe_helpers.py | 2 +- .../{ => tables}/dataframe_html_renderer.py | 0 .../{ => tables}/table_asset_writer.py | 2 +- .../{ => tables}/table_frame_parser.py | 6 +-- .../{ => tables}/table_parser.py | 6 +-- .../{ => tables}/table_text_parser.py | 0 65 files changed, 200 insertions(+), 194 deletions(-) rename apps/worker/app/services/document_parser/{ => assets}/image_compressor.py (100%) rename apps/worker/app/services/document_parser/{ => assets}/inline_asset.py (93%) rename apps/worker/app/services/document_parser/{ => conversion}/legacy_converter.py (97%) rename apps/worker/app/services/document_parser/{atlas_classifier.py => formats/atlas/classifier.py} (98%) rename apps/worker/app/services/document_parser/{atlas_parser.py => formats/atlas/parser.py} (96%) rename apps/worker/app/services/document_parser/{docx_asset_accumulator.py => formats/docx/asset_accumulator.py} (96%) rename apps/worker/app/services/document_parser/{docx_asset_store.py => formats/docx/asset_store.py} (100%) rename apps/worker/app/services/document_parser/{docx_block_stream.py => formats/docx/block_stream.py} (99%) rename apps/worker/app/services/document_parser/{doc_parser.py => formats/docx/parser.py} (94%) rename apps/worker/app/services/document_parser/{docx_table_html.py => formats/docx/table_html.py} (100%) rename apps/worker/app/services/document_parser/{toc_docx.py => formats/docx/toc.py} (98%) rename apps/worker/app/services/document_parser/{excel_structure_parser.py => formats/excel/structure_parser.py} (100%) rename apps/worker/app/services/document_parser/{excel_table_parser.py => formats/excel/table_parser.py} (93%) rename apps/worker/app/services/document_parser/{fragment_parser.py => formats/fragment/parser.py} (98%) rename apps/worker/app/services/document_parser/{html_parser.py => formats/html/parser.py} (99%) rename apps/worker/app/services/document_parser/{image_parser.py => formats/image/parser.py} (95%) rename apps/worker/app/services/document_parser/{markdown_deferred_summary.py => formats/markdown/deferred_summary.py} (96%) rename apps/worker/app/services/document_parser/{markdown_deferred_task.py => formats/markdown/deferred_task.py} (100%) rename apps/worker/app/services/document_parser/{markdown_image_asset.py => formats/markdown/image_asset.py} (94%) rename apps/worker/app/services/document_parser/{markdown_parse_state.py => formats/markdown/parse_state.py} (96%) rename apps/worker/app/services/document_parser/{md_parser.py => formats/markdown/parser.py} (93%) rename apps/worker/app/services/document_parser/{markdown_table_asset.py => formats/markdown/table_asset.py} (85%) rename apps/worker/app/services/document_parser/{pdf_parser.py => formats/pdf/parser.py} (95%) rename apps/worker/app/services/document_parser/{pptx_pdf_rendering.py => formats/pdf/pptx_rendering.py} (95%) rename apps/worker/app/services/document_parser/{ => formats/pdf}/pymupdf_subprocess.py (100%) rename apps/worker/app/services/document_parser/{rendered_pdf_transform.py => formats/pdf/rendered_transform.py} (91%) rename apps/worker/app/services/document_parser/{pptx_parser.py => formats/pptx/parser.py} (98%) rename apps/worker/app/services/document_parser/{txt_parser.py => formats/text/parser.py} (100%) rename apps/worker/app/services/document_parser/{ => profiling}/doc_profile_model.py (100%) rename apps/worker/app/services/document_parser/{ => profiling}/doc_profile_pdf.py (99%) rename apps/worker/app/services/document_parser/{ => profiling}/doc_profiler.py (78%) rename apps/worker/app/services/document_parser/{mineru_client.py => providers/mineru/client.py} (96%) rename apps/worker/app/services/document_parser/{mineru_pdf_service.py => providers/mineru/pdf_service.py} (97%) rename apps/worker/app/services/document_parser/{mineru_quota_manager.py => providers/mineru/quota_manager.py} (100%) rename apps/worker/app/services/document_parser/{mineru_task_polling.py => providers/mineru/task_polling.py} (98%) rename apps/worker/app/services/document_parser/{ => structure}/heading_candidates.py (98%) rename apps/worker/app/services/document_parser/{ => structure}/heading_hierarchy.py (93%) rename apps/worker/app/services/document_parser/{ => structure}/heading_llm_executor.py (98%) rename apps/worker/app/services/document_parser/{ => structure}/heading_tree.py (100%) rename apps/worker/app/services/document_parser/{ => structure}/layout_parser.py (98%) rename apps/worker/app/services/document_parser/{ => structure}/metadata_extractor.py (100%) rename apps/worker/app/services/document_parser/{ => structure}/toc_hierarchy.py (93%) rename apps/worker/app/services/document_parser/{ => structure}/toc_parser.py (98%) rename apps/worker/app/services/document_parser/{ => support}/identifiers.py (100%) rename apps/worker/app/services/document_parser/{ => support}/internal_parse_name.py (100%) rename apps/worker/app/services/document_parser/{ => support}/parser_log_utils.py (100%) rename apps/worker/app/services/document_parser/{ => support}/parser_rows.py (100%) rename apps/worker/app/services/document_parser/{ => support}/path_helpers.py (100%) rename apps/worker/app/services/document_parser/{ => support}/stage_profiler.py (100%) rename apps/worker/app/services/document_parser/{ => support}/text_helpers.py (100%) rename apps/worker/app/services/document_parser/{ => tables}/dataframe_helpers.py (98%) rename apps/worker/app/services/document_parser/{ => tables}/dataframe_html_renderer.py (100%) rename apps/worker/app/services/document_parser/{ => tables}/table_asset_writer.py (95%) rename apps/worker/app/services/document_parser/{ => tables}/table_frame_parser.py (98%) rename apps/worker/app/services/document_parser/{ => tables}/table_parser.py (86%) rename apps/worker/app/services/document_parser/{ => tables}/table_text_parser.py (100%) diff --git a/AGENTS.md b/AGENTS.md index 25c00e802..4881c4720 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -74,13 +74,13 @@ flowchart TB subgraph PARSE["② Document Parsing (Worker)"] Queue --> Router["parse_service.checkerboard_inject_parse"] - Router --> Profiler["doc_profiler.profile_document"] - Profiler --> PDF["pdf_parser → MinerU"] - Profiler --> DOCX["doc_parser.parse_docx"] - Profiler --> PPTX["pptx_parser → iLoveAPI → PDF"] - Profiler --> XLSX["table_parser.parse_xlsx"] - Profiler --> MD["md_parser.parse_md"] - Profiler --> IMG["image_parser.parse_image"] + Router --> Profiler["profiling.doc_profiler.profile_document"] + Profiler --> PDF["formats.pdf.parser → MinerU"] + Profiler --> DOCX["formats.docx.parser.parse_docx"] + Profiler --> PPTX["formats.pptx.parser → iLoveAPI → PDF"] + Profiler --> XLSX["formats.excel.table_parser.parse_xlsx"] + Profiler --> MD["formats.markdown.parser.parse_md"] + Profiler --> IMG["formats.image.parser.parse_image"] PDF --> DF["pd.DataFrame (ALL_DF_COLS)"] DOCX --> DF PPTX --> DF @@ -126,7 +126,7 @@ This is the legacy tuple-compatible entry for all file types. Internally, `checkerboard_parse_output()` and parser adapters use the typed `ParseOutput` contract. The parser flow: -1. **Profiles** the document via `doc_profiler.profile_document()` to detect +1. **Profiles** the document via `profiling.doc_profiler.profile_document()` to detect file type, page count, and special categories (e.g. `atlas`). 2. **Routes** to the appropriate parser based on file extension. 3. **Post-processes**: cleans up unreferenced images, compresses PNG→JPG. @@ -137,22 +137,22 @@ contract. The parser flow: | Extension | Parser Module | Strategy | |:---|:---|:---| -| `.pdf` | `pdf_parser.parse_pdfs` | MinerU API → `md_parser` → `layout_parser.pred_titles` | -| `.docx` | `doc_parser.parse_docx` + `convert_doc2dics` | OXML iteration → heading detection → hierarchical tree | -| `.doc` | `legacy_converter.doc_to_docx` → `.docx` pipeline | LibreOffice headless conversion first | -| `.pptx` | `pptx_parser.parse_pptx` | iLoveAPI PPTX→PDF → MinerU pipeline | -| `.xlsx` | `table_parser.parse_xlsx` | Sheet-by-sheet HTML table extraction | -| `.xls` | `legacy_converter.xls_to_xlsx` → `.xlsx` pipeline | LibreOffice conversion first | -| `.md` | `md_parser.parse_md` | Markdown heading parsing + LLM summaries | -| `.txt` | `txt_parser.parse_texts` → `md_parser` | Read lines then route to MD parser | -| `.png/.jpg` | `image_parser.parse_image` | VLM image description + OCR | -| `.fragment` | `fragment_parser.parse_fragment` | Raw text fragment ingestion | - -### Heading Detection: `layout_parser.pred_titles()` +| `.pdf` | `formats.pdf.parser.parse_pdfs` | MinerU API → Markdown parser → `structure.layout_parser.pred_titles` | +| `.docx` | `formats.docx.parser.parse_docx` + `convert_doc2dics` | OXML iteration → heading detection → hierarchical tree | +| `.doc` | `conversion.legacy_converter.doc_to_docx` → `.docx` pipeline | LibreOffice headless conversion first | +| `.pptx` | `formats.pptx.parser.parse_pptx` | iLoveAPI PPTX→PDF → MinerU pipeline | +| `.xlsx` | `formats.excel.table_parser.parse_xlsx` | Sheet-by-sheet HTML table extraction | +| `.xls` | `conversion.legacy_converter.xls_to_xlsx` → `.xlsx` pipeline | LibreOffice conversion first | +| `.md` | `formats.markdown.parser.parse_md` | Markdown heading parsing + LLM summaries | +| `.txt` | `formats.text.parser.parse_texts` → Markdown parser | Read lines then route to MD parser | +| `.png/.jpg` | `formats.image.parser.parse_image` | VLM image description + OCR | +| `.fragment` | `formats.fragment.parser.parse_fragment` | Raw text fragment ingestion | + +### Heading Detection: `structure.layout_parser.pred_titles()` The core hierarchical recognition module. Determines heading levels using: -1. **TOC-first**: If a DOCX TOC exists (`toc_parser.build_docx_toc_hierarchies`), +1. **TOC-first**: If a DOCX TOC exists (`structure.toc_parser.build_docx_toc_hierarchies`), use it as ground truth for heading levels. 2. **Regex patterns**: Match numbered headings like `1.2.3`, `第X章`, `（一）`. 3. **LLM smart parse**: When `smart_title_parse=True`, send candidate headings @@ -161,7 +161,7 @@ The core hierarchical recognition module. Determines heading levels using: 4. **Font clustering (PDF)**: K-means on span heights from MinerU `layout.json` to group headings into 5 discrete tiers. -### DOCX Parsing Deep Dive: `doc_parser.py` +### DOCX Parsing Deep Dive: `formats/docx/parser.py` ```mermaid flowchart LR @@ -194,11 +194,11 @@ Key logic in `parse_docx()`: ```mermaid flowchart LR - PDF[pdf_parser] --> MinerU[MinerU Cloud API] + PDF[formats.pdf.parser] --> MinerU[MinerU Cloud API] MinerU --> MDFile[Markdown + layout.json] - MDFile --> MDParser[md_parser.parse_md] + MDFile --> MDParser[formats.markdown.parser.parse_md] MDParser --> EvalHeadings[eval_md_headings + layout.json] - EvalHeadings --> PredTitles[layout_parser.pred_titles] + EvalHeadings --> PredTitles[structure.layout_parser.pred_titles] PredTitles --> Chunks[Hierarchical Chunks] ``` @@ -210,7 +210,7 @@ flowchart LR | Heading hierarchy recognition | `HIERARCHY_LLM_MODEL` | Falls back to `NORMOL_MODEL` | | Image description (VLM) | `IMAGE_MODEL` | `qwen3.5-flash` | | Image OCR / Q&A | `IMAGE_MODEL_MAX` | `qwen3.5-flash` | -| Atlas classification | VLM via `atlas_classifier` | `IMAGE_MODEL` | +| Atlas classification | VLM via `formats.atlas.classifier` | `IMAGE_MODEL` | --- @@ -428,7 +428,7 @@ Used by agentic retrieval for 2-level section browsing. Structure: ``` The `HIERARCHY` field is a nested dict representing the full heading tree -discovered by `layout_parser.pred_titles()`. Each key is a heading title; +discovered by `structure.layout_parser.pred_titles()`. Each key is a heading title; its value is a dict of child headings (empty `{}` for leaf nodes). ### Intermediate DataFrame (`ALL_DF_COLS`) diff --git a/CONTEXT.md b/CONTEXT.md index 90ee0fa1b..9ac741eab 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -362,20 +362,26 @@ side effects. - `app/services/document_parser/orchestration/route_parse.py` - `app/services/document_parser/orchestration/format_router.py` - `app/services/document_parser/orchestration/format_adapters.py` +- `app/services/document_parser/formats/*` +- `app/services/document_parser/providers/*` +- `app/services/document_parser/structure/*` +- `app/services/document_parser/tables/*` +- `app/services/document_parser/assets/*` +- `app/services/document_parser/support/*` ### 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` +- `app/services/document_parser/formats/pdf/rendered_transform.py` +- `app/services/document_parser/formats/pdf/pptx_rendering.py` +- `app/services/document_parser/formats/pdf/parser.py` +- `app/services/document_parser/formats/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` +- `app/services/document_parser/structure/heading_hierarchy.py` +- `app/services/document_parser/structure/layout_parser.py` +- `app/services/document_parser/formats/markdown/parser.py` +- `app/services/document_parser/formats/docx/parser.py` ## Invariants diff --git a/apps/worker/app/services/document_agent/tools/probe_sample_pages.py b/apps/worker/app/services/document_agent/tools/probe_sample_pages.py index 1ac8018e8..b5b7e1122 100644 --- a/apps/worker/app/services/document_agent/tools/probe_sample_pages.py +++ b/apps/worker/app/services/document_agent/tools/probe_sample_pages.py @@ -6,7 +6,7 @@ import statistics from typing import Any, Literal -from app.services.document_parser.pymupdf_subprocess import run_in_child_process, worker +from app.services.document_parser.formats.pdf.pymupdf_subprocess import run_in_child_process, worker SampleStrategy = Literal["stratified", "uniform", "key_pages"] diff --git a/apps/worker/app/services/document_agent/tools/probe_vlm_inspect.py b/apps/worker/app/services/document_agent/tools/probe_vlm_inspect.py index f8949826d..b073236a5 100644 --- a/apps/worker/app/services/document_agent/tools/probe_vlm_inspect.py +++ b/apps/worker/app/services/document_agent/tools/probe_vlm_inspect.py @@ -7,7 +7,7 @@ import tempfile from typing import Any -from app.services.document_parser.pymupdf_subprocess import run_in_child_process, worker +from app.services.document_parser.formats.pdf.pymupdf_subprocess import run_in_child_process, worker from loguru import logger from openai.types.chat import ( ChatCompletionContentPartImageParam, diff --git a/apps/worker/app/services/document_ingestion/page_estimator.py b/apps/worker/app/services/document_ingestion/page_estimator.py index bdc97a4d6..f2ed89cdb 100644 --- a/apps/worker/app/services/document_ingestion/page_estimator.py +++ b/apps/worker/app/services/document_ingestion/page_estimator.py @@ -136,7 +136,7 @@ def _count_docx(cls, file_path: str) -> int: @classmethod def _count_doc(cls, file_path: str) -> int: """Estimate pages for DOC by converting it to DOCX first.""" - from app.services.document_parser.legacy_converter import doc_to_docx + from app.services.document_parser.conversion.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) @@ -159,7 +159,7 @@ def _count_xlsx(cls, file_path: str) -> int: @classmethod def _count_xls(cls, file_path: str) -> int: """Estimate pages for XLS by converting it to XLSX first.""" - from app.services.document_parser.legacy_converter import xls_to_xlsx + from app.services.document_parser.conversion.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) diff --git a/apps/worker/app/services/document_ingestion/processing_run.py b/apps/worker/app/services/document_ingestion/processing_run.py index d629eef69..d3c31ca0a 100644 --- a/apps/worker/app/services/document_ingestion/processing_run.py +++ b/apps/worker/app/services/document_ingestion/processing_run.py @@ -31,7 +31,7 @@ create_task_workspace, download_s3_file_to_temp, ) -from app.services.document_parser.stage_profiler import stage_timer +from app.services.document_parser.support.stage_profiler import stage_timer from loguru import logger from shared.models.schemas.job_metadata import JobMetadataHelper @@ -111,7 +111,7 @@ def _run_parse_job( 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 ( + from app.services.document_parser.support.internal_parse_name import ( prepare_internal_parse_input, ) from app.services.document_parser import parse_service diff --git a/apps/worker/app/services/document_parser/image_compressor.py b/apps/worker/app/services/document_parser/assets/image_compressor.py similarity index 100% rename from apps/worker/app/services/document_parser/image_compressor.py rename to apps/worker/app/services/document_parser/assets/image_compressor.py diff --git a/apps/worker/app/services/document_parser/inline_asset.py b/apps/worker/app/services/document_parser/assets/inline_asset.py similarity index 93% rename from apps/worker/app/services/document_parser/inline_asset.py rename to apps/worker/app/services/document_parser/assets/inline_asset.py index 4cb16b531..9e792ee1d 100644 --- a/apps/worker/app/services/document_parser/inline_asset.py +++ b/apps/worker/app/services/document_parser/assets/inline_asset.py @@ -1,6 +1,6 @@ from __future__ import annotations -from app.services.document_parser.parser_rows import ParsedRow +from app.services.document_parser.support.parser_rows import ParsedRow def build_image_asset_row( diff --git a/apps/worker/app/services/document_parser/legacy_converter.py b/apps/worker/app/services/document_parser/conversion/legacy_converter.py similarity index 97% rename from apps/worker/app/services/document_parser/legacy_converter.py rename to apps/worker/app/services/document_parser/conversion/legacy_converter.py index f4b18fdc2..a602d67b8 100644 --- a/apps/worker/app/services/document_parser/legacy_converter.py +++ b/apps/worker/app/services/document_parser/conversion/legacy_converter.py @@ -4,7 +4,7 @@ import tempfile from pathlib import Path -from app.services.document_parser.parser_log_utils import truncate_log_value +from app.services.document_parser.support.parser_log_utils import truncate_log_value from shared.core.exceptions.domain_exceptions import LibreOfficeServiceException diff --git a/apps/worker/app/services/document_parser/atlas_classifier.py b/apps/worker/app/services/document_parser/formats/atlas/classifier.py similarity index 98% rename from apps/worker/app/services/document_parser/atlas_classifier.py rename to apps/worker/app/services/document_parser/formats/atlas/classifier.py index ed58ad49f..87c43e862 100644 --- a/apps/worker/app/services/document_parser/atlas_classifier.py +++ b/apps/worker/app/services/document_parser/formats/atlas/classifier.py @@ -17,7 +17,7 @@ import tempfile from typing import Optional -from app.services.document_parser.pymupdf_subprocess import run_in_child_process, worker +from app.services.document_parser.formats.pdf.pymupdf_subprocess import run_in_child_process, worker from loguru import logger from openai.types.chat import ( ChatCompletionContentPartImageParam, diff --git a/apps/worker/app/services/document_parser/atlas_parser.py b/apps/worker/app/services/document_parser/formats/atlas/parser.py similarity index 96% rename from apps/worker/app/services/document_parser/atlas_parser.py rename to apps/worker/app/services/document_parser/formats/atlas/parser.py index fef806add..6c5815c29 100644 --- a/apps/worker/app/services/document_parser/atlas_parser.py +++ b/apps/worker/app/services/document_parser/formats/atlas/parser.py @@ -19,11 +19,11 @@ from concurrent.futures import ThreadPoolExecutor, as_completed 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 app.services.document_parser.tables.dataframe_helpers import process_dup_paths_df +from app.services.document_parser.support.identifiers import gen_str_codes, get_str_time +from app.services.document_parser.support.parser_rows import ParsedRow, ParsedRowsBuilder +from app.services.document_parser.formats.pdf.pymupdf_subprocess import run_in_child_process, worker +from app.services.document_parser.structure.toc_parser import detect_tocs_in_texts from loguru import logger from shared.core.config import settings @@ -117,7 +117,7 @@ def _vlm_extract_page_info(output_dir: str, img_name: str) -> str: Returns extracted info string, or empty string on failure. """ - from app.services.document_parser.image_parser import _get_vision_client, ask_image + from app.services.document_parser.formats.image.parser import _get_vision_client, ask_image img_path = os.path.join(output_dir, "images", img_name) if not os.path.exists(img_path): diff --git a/apps/worker/app/services/document_parser/docx_asset_accumulator.py b/apps/worker/app/services/document_parser/formats/docx/asset_accumulator.py similarity index 96% rename from apps/worker/app/services/document_parser/docx_asset_accumulator.py rename to apps/worker/app/services/document_parser/formats/docx/asset_accumulator.py index ab0d67509..f69526e95 100644 --- a/apps/worker/app/services/document_parser/docx_asset_accumulator.py +++ b/apps/worker/app/services/document_parser/formats/docx/asset_accumulator.py @@ -4,7 +4,7 @@ from dataclasses import dataclass, field from typing import Any -from app.services.document_parser.docx_asset_store import DocxAssetStore +from app.services.document_parser.formats.docx.asset_store import DocxAssetStore ImageHandler = Callable[ diff --git a/apps/worker/app/services/document_parser/docx_asset_store.py b/apps/worker/app/services/document_parser/formats/docx/asset_store.py similarity index 100% rename from apps/worker/app/services/document_parser/docx_asset_store.py rename to apps/worker/app/services/document_parser/formats/docx/asset_store.py diff --git a/apps/worker/app/services/document_parser/docx_block_stream.py b/apps/worker/app/services/document_parser/formats/docx/block_stream.py similarity index 99% rename from apps/worker/app/services/document_parser/docx_block_stream.py rename to apps/worker/app/services/document_parser/formats/docx/block_stream.py index 9832e0e54..dafe28aba 100644 --- a/apps/worker/app/services/document_parser/docx_block_stream.py +++ b/apps/worker/app/services/document_parser/formats/docx/block_stream.py @@ -4,7 +4,7 @@ import io import zipfile -from app.services.document_parser.toc_docx import detect_doc_tocs, detect_sdt_toc +from app.services.document_parser.formats.docx.toc 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 diff --git a/apps/worker/app/services/document_parser/doc_parser.py b/apps/worker/app/services/document_parser/formats/docx/parser.py similarity index 94% rename from apps/worker/app/services/document_parser/doc_parser.py rename to apps/worker/app/services/document_parser/formats/docx/parser.py index eadcf238f..2f3a92551 100755 --- a/apps/worker/app/services/document_parser/doc_parser.py +++ b/apps/worker/app/services/document_parser/formats/docx/parser.py @@ -3,34 +3,34 @@ import os 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 ( +from app.services.document_parser.tables.dataframe_helpers import process_dup_paths_df +from app.services.document_parser.formats.docx.asset_accumulator import DocxAssetAccumulator +from app.services.document_parser.formats.docx.asset_store import DocxAssetStore +from app.services.document_parser.formats.docx.block_stream import iter_block_items +from app.services.document_parser.formats.docx.table_html import table2html +from app.services.document_parser.support.identifiers import gen_str_codes, get_str_time +from app.services.document_parser.assets.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 ( +from app.services.document_parser.support.parser_rows import ParsedRow, ParsedRowsBuilder +from app.services.document_parser.support.path_helpers import ( find_matches_parsing, process_path_texts, remove_spaces, ) -from app.services.document_parser.heading_hierarchy import ( +from app.services.document_parser.structure.heading_hierarchy import ( HeadingHierarchyInput, predict_heading_hierarchy, ) -from app.services.document_parser.image_parser import ( +from app.services.document_parser.formats.image.parser import ( _get_vision_client, ask_image, perceptual_hash, ) -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 app.services.document_parser.tables.table_text_parser import sanitize_table_name_from_header +from app.services.document_parser.formats.docx.toc import build_docx_toc_hierarchies +from app.services.document_parser.formats.text.parser import postprocess_leaf_dics from docx.text.paragraph import Paragraph from loguru import logger @@ -72,7 +72,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.document_parser.text_helpers import truncate_text + from app.services.document_parser.support.text_helpers import truncate_text try: content_list = headings_stack[-1].get("content", []) @@ -141,7 +141,7 @@ def handle_image( llm_title = None llm_summary = None if smart_summary: - from app.services.document_parser.txt_parser import split_title_summary + from app.services.document_parser.formats.text.parser import split_title_summary # 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. @@ -220,7 +220,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.document_parser.text_helpers import truncate_text + from app.services.document_parser.support.text_helpers import truncate_text first_row_text = "" first_col_text = "" @@ -392,7 +392,7 @@ def handle_table( llm_summary = None tb_keywords = "" if summary_table: - from app.services.document_parser.txt_parser import ( + from app.services.document_parser.formats.text.parser import ( extract_title_keywords_summary, ) diff --git a/apps/worker/app/services/document_parser/docx_table_html.py b/apps/worker/app/services/document_parser/formats/docx/table_html.py similarity index 100% rename from apps/worker/app/services/document_parser/docx_table_html.py rename to apps/worker/app/services/document_parser/formats/docx/table_html.py diff --git a/apps/worker/app/services/document_parser/toc_docx.py b/apps/worker/app/services/document_parser/formats/docx/toc.py similarity index 98% rename from apps/worker/app/services/document_parser/toc_docx.py rename to apps/worker/app/services/document_parser/formats/docx/toc.py index 86d612add..deb77ecfd 100644 --- a/apps/worker/app/services/document_parser/toc_docx.py +++ b/apps/worker/app/services/document_parser/formats/docx/toc.py @@ -2,11 +2,11 @@ import re -from app.services.document_parser.heading_candidates import ( +from app.services.document_parser.structure.heading_candidates import ( judge_by_conditions, remove_by_conditions, ) -from app.services.document_parser.toc_hierarchy import build_toc_hierarchy_payload +from app.services.document_parser.structure.toc_hierarchy import build_toc_hierarchy_payload import lxml.etree as etree TOC_TITLE_KEYWORDS = {"目录", "目次", "contents", "table of contents"} diff --git a/apps/worker/app/services/document_parser/excel_structure_parser.py b/apps/worker/app/services/document_parser/formats/excel/structure_parser.py similarity index 100% rename from apps/worker/app/services/document_parser/excel_structure_parser.py rename to apps/worker/app/services/document_parser/formats/excel/structure_parser.py diff --git a/apps/worker/app/services/document_parser/excel_table_parser.py b/apps/worker/app/services/document_parser/formats/excel/table_parser.py similarity index 93% rename from apps/worker/app/services/document_parser/excel_table_parser.py rename to apps/worker/app/services/document_parser/formats/excel/table_parser.py index 589aba7a3..3f5499ea3 100644 --- a/apps/worker/app/services/document_parser/excel_table_parser.py +++ b/apps/worker/app/services/document_parser/formats/excel/table_parser.py @@ -7,16 +7,16 @@ 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 ( +from app.services.document_parser.tables.dataframe_helpers import process_dup_paths_df +from app.services.document_parser.formats.excel.structure_parser import parse_excel_structure +from app.services.document_parser.support.identifiers import gen_str_codes, get_str_time +from app.services.document_parser.support.parser_rows import ParsedRow, ParsedRowsBuilder +from app.services.document_parser.support.path_helpers import remove_spaces +from app.services.document_parser.tables.table_asset_writer import ( TableAssetInput, write_table_asset, ) -from app.services.document_parser.table_frame_parser import ( +from app.services.document_parser.tables.table_frame_parser import ( parse_headers, parse_tb_contents, parse_tb_keywords, @@ -265,7 +265,7 @@ def _summarize_excel_table( mechanical_keywords = parse_tb_keywords(table_frame) if llm_parameters["summary_table"]: - from app.services.document_parser.txt_parser import ( + from app.services.document_parser.formats.text.parser import ( extract_title_keywords_summary, ) diff --git a/apps/worker/app/services/document_parser/fragment_parser.py b/apps/worker/app/services/document_parser/formats/fragment/parser.py similarity index 98% rename from apps/worker/app/services/document_parser/fragment_parser.py rename to apps/worker/app/services/document_parser/formats/fragment/parser.py index 420d6b084..6ab680cc4 100644 --- a/apps/worker/app/services/document_parser/fragment_parser.py +++ b/apps/worker/app/services/document_parser/formats/fragment/parser.py @@ -8,7 +8,7 @@ import os from typing import Any, Optional -from app.services.document_parser.md_parser import parse_md +from app.services.document_parser.formats.markdown.parser import parse_md from loguru import logger from openai.types.chat import ChatCompletionMessageParam diff --git a/apps/worker/app/services/document_parser/html_parser.py b/apps/worker/app/services/document_parser/formats/html/parser.py similarity index 99% rename from apps/worker/app/services/document_parser/html_parser.py rename to apps/worker/app/services/document_parser/formats/html/parser.py index f412067d6..bc39731b2 100644 --- a/apps/worker/app/services/document_parser/html_parser.py +++ b/apps/worker/app/services/document_parser/formats/html/parser.py @@ -407,7 +407,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.document_parser.text_helpers import truncate_text + from app.services.document_parser.support.text_helpers import truncate_text soup = BeautifulSoup(html_str, "html.parser") table = soup.find("table") diff --git a/apps/worker/app/services/document_parser/image_parser.py b/apps/worker/app/services/document_parser/formats/image/parser.py similarity index 95% rename from apps/worker/app/services/document_parser/image_parser.py rename to apps/worker/app/services/document_parser/formats/image/parser.py index 726fb6a99..4562a2135 100755 --- a/apps/worker/app/services/document_parser/image_parser.py +++ b/apps/worker/app/services/document_parser/formats/image/parser.py @@ -8,9 +8,9 @@ from pathlib import Path 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.tables.dataframe_helpers import process_dup_paths_df +from app.services.document_parser.support.identifiers import gen_str_codes, get_str_time +from app.services.document_parser.support.parser_rows import ParsedRow, ParsedRowsBuilder from loguru import logger from PIL import Image @@ -196,7 +196,7 @@ def detect_summary_img_md(line, last_context, kb_dir, mode=False): try: llm_resp = ask_image(client, kb_dir, paths_=[ip]) if llm_resp: - from app.services.document_parser.txt_parser import ( + from app.services.document_parser.formats.text.parser import ( split_title_summary, ) @@ -296,7 +296,7 @@ def parse_image( size_cut=False, ) if llm_resp: - from app.services.document_parser.txt_parser import split_title_summary + from app.services.document_parser.formats.text.parser import split_title_summary img_title, image_summary = split_title_summary(llm_resp) else: @@ -305,7 +305,7 @@ def parse_image( else: # For non-text images, split title from summary-images response if base_llm_paras["summary_image"] and image_content != filename: - from app.services.document_parser.txt_parser import split_title_summary + from app.services.document_parser.formats.text.parser import split_title_summary img_title, image_summary = split_title_summary(image_content) else: diff --git a/apps/worker/app/services/document_parser/markdown_deferred_summary.py b/apps/worker/app/services/document_parser/formats/markdown/deferred_summary.py similarity index 96% rename from apps/worker/app/services/document_parser/markdown_deferred_summary.py rename to apps/worker/app/services/document_parser/formats/markdown/deferred_summary.py index 49327122c..394ed06cc 100644 --- a/apps/worker/app/services/document_parser/markdown_deferred_summary.py +++ b/apps/worker/app/services/document_parser/formats/markdown/deferred_summary.py @@ -6,16 +6,16 @@ from typing import Literal, TypeGuard import gevent -from app.services.document_parser.markdown_deferred_task import ( +from app.services.document_parser.formats.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_text_parser import sanitize_table_name_from_header -from app.services.document_parser.txt_parser import ( +from app.services.document_parser.formats.image.parser import _get_vision_client, ask_image +from app.services.document_parser.support.stage_profiler import stage_timer +from app.services.document_parser.tables.table_text_parser import sanitize_table_name_from_header +from app.services.document_parser.formats.text.parser import ( extract_title_keywords_summary, split_title_summary, ) diff --git a/apps/worker/app/services/document_parser/markdown_deferred_task.py b/apps/worker/app/services/document_parser/formats/markdown/deferred_task.py similarity index 100% rename from apps/worker/app/services/document_parser/markdown_deferred_task.py rename to apps/worker/app/services/document_parser/formats/markdown/deferred_task.py diff --git a/apps/worker/app/services/document_parser/markdown_image_asset.py b/apps/worker/app/services/document_parser/formats/markdown/image_asset.py similarity index 94% rename from apps/worker/app/services/document_parser/markdown_image_asset.py rename to apps/worker/app/services/document_parser/formats/markdown/image_asset.py index 6672bc178..88a684e56 100644 --- a/apps/worker/app/services/document_parser/markdown_image_asset.py +++ b/apps/worker/app/services/document_parser/formats/markdown/image_asset.py @@ -5,14 +5,14 @@ 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 ( +from app.services.document_parser.support.identifiers import gen_str_codes +from app.services.document_parser.formats.image.parser import perceptual_hash +from app.services.document_parser.assets.inline_asset import build_image_asset_row +from app.services.document_parser.formats.markdown.deferred_task import ( ImageDeferredSummaryTask, MarkdownDeferredSummaryTask, ) -from app.services.document_parser.markdown_parse_state import ParserRowValues +from app.services.document_parser.formats.markdown.parse_state import ParserRowValues from loguru import logger from shared.utils.chunk_refs import build_chunk_ref diff --git a/apps/worker/app/services/document_parser/markdown_parse_state.py b/apps/worker/app/services/document_parser/formats/markdown/parse_state.py similarity index 96% rename from apps/worker/app/services/document_parser/markdown_parse_state.py rename to apps/worker/app/services/document_parser/formats/markdown/parse_state.py index 51f97b954..b7c6da508 100644 --- a/apps/worker/app/services/document_parser/markdown_parse_state.py +++ b/apps/worker/app/services/document_parser/formats/markdown/parse_state.py @@ -7,12 +7,12 @@ 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 ( +from app.services.document_parser.tables.dataframe_helpers import process_dup_paths_df +from app.services.document_parser.formats.markdown.deferred_task import ( MarkdownDeferredSummaryTask, TextDeferredSummaryTask, ) -from app.services.document_parser.parser_rows import ParsedRow, ParsedRowsBuilder +from app.services.document_parser.support.parser_rows import ParsedRow, ParsedRowsBuilder ParserRowValues = list[str | int] diff --git a/apps/worker/app/services/document_parser/md_parser.py b/apps/worker/app/services/document_parser/formats/markdown/parser.py similarity index 93% rename from apps/worker/app/services/document_parser/md_parser.py rename to apps/worker/app/services/document_parser/formats/markdown/parser.py index 7aea710fa..a33be14e9 100755 --- a/apps/worker/app/services/document_parser/md_parser.py +++ b/apps/worker/app/services/document_parser/formats/markdown/parser.py @@ -4,42 +4,42 @@ import re import shutil -from app.services.document_parser.identifiers import gen_str_codes, get_str_time -from app.services.document_parser.markdown_deferred_summary import ( +from app.services.document_parser.support.identifiers import gen_str_codes, get_str_time +from app.services.document_parser.formats.markdown.deferred_summary import ( MarkdownDeferredSummaryInput, apply_markdown_deferred_summaries, ) -from app.services.document_parser.markdown_image_asset import ( +from app.services.document_parser.formats.markdown.image_asset import ( MarkdownImageAssetRequest, build_markdown_image_name, build_markdown_image_asset, ) -from app.services.document_parser.markdown_parse_state import MarkdownParseState -from app.services.document_parser.markdown_table_asset import ( +from app.services.document_parser.formats.markdown.parse_state import MarkdownParseState +from app.services.document_parser.formats.markdown.table_asset import ( MarkdownTableAssetRequest, build_markdown_table_asset, ) -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 ( +from app.services.document_parser.support.parser_rows import ParsedRow +from app.services.document_parser.support.path_helpers import find_matches_parsing +from app.services.document_parser.formats.html.parser import ( merge_html_tables, ) -from app.services.document_parser.image_parser import ( +from app.services.document_parser.formats.image.parser import ( MD_IMAGE_PATTERN, detect_summary_img_md, ) -from app.services.document_parser.heading_hierarchy import ( +from app.services.document_parser.structure.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_text_parser import ( +from app.services.document_parser.structure.heading_candidates import md_heading_match +from app.services.document_parser.support.stage_profiler import stage_timer +from app.services.document_parser.tables.table_text_parser import ( extract_tables_by_forms, identify_tables, ) -from app.services.document_parser.toc_parser import detect_tocs_in_texts -from app.services.document_parser.txt_parser import extract_title_keywords_summary +from app.services.document_parser.structure.toc_parser import detect_tocs_in_texts +from app.services.document_parser.formats.text.parser import extract_title_keywords_summary from loguru import logger from shared.core.config import settings diff --git a/apps/worker/app/services/document_parser/markdown_table_asset.py b/apps/worker/app/services/document_parser/formats/markdown/table_asset.py similarity index 85% rename from apps/worker/app/services/document_parser/markdown_table_asset.py rename to apps/worker/app/services/document_parser/formats/markdown/table_asset.py index c0985b7a9..19f795f4f 100644 --- a/apps/worker/app/services/document_parser/markdown_table_asset.py +++ b/apps/worker/app/services/document_parser/formats/markdown/table_asset.py @@ -4,15 +4,15 @@ 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 ( +from app.services.document_parser.formats.html.parser import first_cols_rows_html +from app.services.document_parser.support.identifiers import gen_str_codes +from app.services.document_parser.assets.inline_asset import build_table_asset_row +from app.services.document_parser.formats.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 app.services.document_parser.formats.markdown.parse_state import ParserRowValues +from app.services.document_parser.tables.table_text_parser import sanitize_table_name_from_header from shared.utils.chunk_refs import build_chunk_ref from app.services.common.file_utils import path_handle diff --git a/apps/worker/app/services/document_parser/pdf_parser.py b/apps/worker/app/services/document_parser/formats/pdf/parser.py similarity index 95% rename from apps/worker/app/services/document_parser/pdf_parser.py rename to apps/worker/app/services/document_parser/formats/pdf/parser.py index fec266090..6a8dfaf4a 100755 --- a/apps/worker/app/services/document_parser/pdf_parser.py +++ b/apps/worker/app/services/document_parser/formats/pdf/parser.py @@ -3,10 +3,10 @@ import os import re -from app.services.document_parser.md_parser import parse_md -from app.services.document_parser.mineru_pdf_service import parse_via_full -from app.services.document_parser.pymupdf_subprocess import worker -from app.services.document_parser.stage_profiler import stage_timer +from app.services.document_parser.formats.markdown.parser import parse_md +from app.services.document_parser.providers.mineru.pdf_service import parse_via_full +from app.services.document_parser.formats.pdf.pymupdf_subprocess import worker +from app.services.document_parser.support.stage_profiler import stage_timer from loguru import logger @@ -223,7 +223,7 @@ def parse_pdfs( # ── Atlas routing: bypass MinerU entirely, use PyMuPDF for per-page chunking ── if profile and profile.doc_category == "atlas": logger.info(f"📐 Atlas detected, bypassing MinerU for {filename}") - from app.services.document_parser.atlas_parser import parse_atlas + from app.services.document_parser.formats.atlas.parser import parse_atlas return parse_atlas( pdf_path, output_dir, base_llm_paras, relative_root, profile=profile diff --git a/apps/worker/app/services/document_parser/pptx_pdf_rendering.py b/apps/worker/app/services/document_parser/formats/pdf/pptx_rendering.py similarity index 95% rename from apps/worker/app/services/document_parser/pptx_pdf_rendering.py rename to apps/worker/app/services/document_parser/formats/pdf/pptx_rendering.py index 6e0a881a9..de5373ead 100644 --- a/apps/worker/app/services/document_parser/pptx_pdf_rendering.py +++ b/apps/worker/app/services/document_parser/formats/pdf/pptx_rendering.py @@ -1,7 +1,7 @@ import os import tempfile -from app.services.document_parser.pymupdf_subprocess import run_in_child_process, worker +from app.services.document_parser.formats.pdf.pymupdf_subprocess import run_in_child_process, worker from loguru import logger diff --git a/apps/worker/app/services/document_parser/pymupdf_subprocess.py b/apps/worker/app/services/document_parser/formats/pdf/pymupdf_subprocess.py similarity index 100% rename from apps/worker/app/services/document_parser/pymupdf_subprocess.py rename to apps/worker/app/services/document_parser/formats/pdf/pymupdf_subprocess.py diff --git a/apps/worker/app/services/document_parser/rendered_pdf_transform.py b/apps/worker/app/services/document_parser/formats/pdf/rendered_transform.py similarity index 91% rename from apps/worker/app/services/document_parser/rendered_pdf_transform.py rename to apps/worker/app/services/document_parser/formats/pdf/rendered_transform.py index cab7ca822..aa5f9adda 100644 --- a/apps/worker/app/services/document_parser/rendered_pdf_transform.py +++ b/apps/worker/app/services/document_parser/formats/pdf/rendered_transform.py @@ -4,11 +4,11 @@ import pandas as pd -from app.services.document_parser.mineru_pdf_service import ( +from app.services.document_parser.providers.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 app.services.document_parser.formats.pdf.parser import parse_pdfs +from app.services.document_parser.formats.pdf.pptx_rendering import render_pdf_to_image_pdf from loguru import logger from shared.core.config import settings diff --git a/apps/worker/app/services/document_parser/pptx_parser.py b/apps/worker/app/services/document_parser/formats/pptx/parser.py similarity index 98% rename from apps/worker/app/services/document_parser/pptx_parser.py rename to apps/worker/app/services/document_parser/formats/pptx/parser.py index 8329fb18a..fb8a93acc 100755 --- a/apps/worker/app/services/document_parser/pptx_parser.py +++ b/apps/worker/app/services/document_parser/formats/pptx/parser.py @@ -6,13 +6,13 @@ import jwt import requests -from app.services.document_parser.path_helpers import find_images -from app.services.document_parser.legacy_converter import ( +from app.services.document_parser.support.path_helpers import find_images +from app.services.document_parser.conversion.legacy_converter import ( _convert_with_libreoffice, ) -from app.services.document_parser.md_parser import parse_md -from app.services.document_parser.parser_log_utils import truncate_log_value -from app.services.document_parser.rendered_pdf_transform import ( +from app.services.document_parser.formats.markdown.parser import parse_md +from app.services.document_parser.support.parser_log_utils import truncate_log_value +from app.services.document_parser.formats.pdf.rendered_transform import ( build_rendered_pdf_s3_key, parse_cached_rendered_pdf, parse_rendered_pdf_bytes, diff --git a/apps/worker/app/services/document_parser/txt_parser.py b/apps/worker/app/services/document_parser/formats/text/parser.py similarity index 100% rename from apps/worker/app/services/document_parser/txt_parser.py rename to apps/worker/app/services/document_parser/formats/text/parser.py 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 e9b5a8846..55d0e4e74 100644 --- a/apps/worker/app/services/document_parser/orchestration/format_adapters.py +++ b/apps/worker/app/services/document_parser/orchestration/format_adapters.py @@ -22,7 +22,7 @@ class FragmentParseAdapter: document_format: object def parse(self, session: ParseSession) -> ParseOutput: - from app.services.document_parser.fragment_parser import parse_fragment + from app.services.document_parser.formats.fragment.parser import parse_fragment full_output_dir, _relative_root, parsed_df = parse_fragment( session.fragment_content, @@ -39,8 +39,8 @@ class TextParseAdapter: document_format: object def parse(self, session: ParseSession) -> ParseOutput: - from app.services.document_parser.md_parser import parse_md - from app.services.document_parser.txt_parser import parse_texts + from app.services.document_parser.formats.markdown.parser import parse_md + from app.services.document_parser.formats.text.parser import parse_texts text_lines = parse_texts(file_path=session.file_full_path, baseurl=session.base_url) parsed_df = parse_md( @@ -58,7 +58,7 @@ class ImageParseAdapter: document_format: object def parse(self, session: ParseSession) -> ParseOutput: - from app.services.document_parser.image_parser import parse_image + from app.services.document_parser.formats.image.parser import parse_image parsed_df = parse_image( session.file_full_path, @@ -76,7 +76,7 @@ class PdfParseAdapter: document_format: object def parse(self, session: ParseSession) -> ParseOutput: - from app.services.document_parser.pdf_parser import parse_pdfs + from app.services.document_parser.formats.pdf.parser import parse_pdfs parsed_df = parse_pdfs( session.file_full_path, @@ -95,7 +95,7 @@ class DocParseAdapter: document_format: object def parse(self, session: ParseSession) -> ParseOutput: - from app.services.document_parser.legacy_converter import doc_to_docx + from app.services.document_parser.conversion.legacy_converter import doc_to_docx converted_docx_path, _ = doc_to_docx( session.file_full_path, @@ -117,7 +117,7 @@ class XlsParseAdapter: document_format: object def parse(self, session: ParseSession) -> ParseOutput: - from app.services.document_parser.legacy_converter import xls_to_xlsx + from app.services.document_parser.conversion.legacy_converter import xls_to_xlsx converted_xlsx_path, _ = xls_to_xlsx( session.file_full_path, @@ -139,7 +139,7 @@ class PptxParseAdapter: document_format: object def parse(self, session: ParseSession) -> ParseOutput: - from app.services.document_parser.pptx_parser import parse_pptx + from app.services.document_parser.formats.pptx.parser import parse_pptx parsed_df = parse_pptx( session.file_full_path, @@ -159,7 +159,7 @@ class MarkdownParseAdapter: document_format: object def parse(self, session: ParseSession) -> ParseOutput: - from app.services.document_parser.md_parser import parse_md + from app.services.document_parser.formats.markdown.parser import parse_md parsed_df = parse_md( session.full_output_dir, @@ -183,7 +183,7 @@ def _parse_docx_path( docx_path: str, session: ParseSession, ) -> ParseOutput: - from app.services.document_parser.doc_parser import convert_doc2dics, parse_docx + from app.services.document_parser.formats.docx.parser import convert_doc2dics, parse_docx parsed_structure, dataframe_list = parse_docx( docx_path, @@ -207,7 +207,7 @@ def _parse_xlsx_path( xlsx_path: str, session: ParseSession, ) -> ParseOutput: - from app.services.document_parser.excel_table_parser import parse_xlsx + from app.services.document_parser.formats.excel.table_parser import parse_xlsx parsed_df = parse_xlsx( xlsx_path, 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 9750a331d..14670b8f8 100644 --- a/apps/worker/app/services/document_parser/orchestration/parse_session.py +++ b/apps/worker/app/services/document_parser/orchestration/parse_session.py @@ -4,10 +4,10 @@ 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.formats.atlas.classifier import classify_atlas_with_vlm +from app.services.document_parser.profiling.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 app.services.document_parser.support.stage_profiler import stage_timer from loguru import logger from shared.core.config import settings diff --git a/apps/worker/app/services/document_parser/orchestration/postprocess.py b/apps/worker/app/services/document_parser/orchestration/postprocess.py index 9cebc7c12..0fbb9f5c8 100644 --- a/apps/worker/app/services/document_parser/orchestration/postprocess.py +++ b/apps/worker/app/services/document_parser/orchestration/postprocess.py @@ -4,11 +4,11 @@ import re import pandas as pd -from app.services.document_parser.image_compressor import ( +from app.services.document_parser.assets.image_compressor import ( apply_rename_map_to_dataframe, compress_output_images, ) -from app.services.document_parser.stage_profiler import stage_timer +from app.services.document_parser.support.stage_profiler import stage_timer from loguru import logger diff --git a/apps/worker/app/services/document_parser/doc_profile_model.py b/apps/worker/app/services/document_parser/profiling/doc_profile_model.py similarity index 100% rename from apps/worker/app/services/document_parser/doc_profile_model.py rename to apps/worker/app/services/document_parser/profiling/doc_profile_model.py diff --git a/apps/worker/app/services/document_parser/doc_profile_pdf.py b/apps/worker/app/services/document_parser/profiling/doc_profile_pdf.py similarity index 99% rename from apps/worker/app/services/document_parser/doc_profile_pdf.py rename to apps/worker/app/services/document_parser/profiling/doc_profile_pdf.py index 4f9c447e6..32a13840a 100644 --- a/apps/worker/app/services/document_parser/doc_profile_pdf.py +++ b/apps/worker/app/services/document_parser/profiling/doc_profile_pdf.py @@ -4,11 +4,11 @@ import math from typing import Any -from app.services.document_parser.doc_profile_model import ( +from app.services.document_parser.profiling.doc_profile_model import ( DocProfile, publish_profile_result, ) -from app.services.document_parser.pymupdf_subprocess import run_in_child_process, worker +from app.services.document_parser.formats.pdf.pymupdf_subprocess import run_in_child_process, worker from loguru import logger # Thresholds diff --git a/apps/worker/app/services/document_parser/doc_profiler.py b/apps/worker/app/services/document_parser/profiling/doc_profiler.py similarity index 78% rename from apps/worker/app/services/document_parser/doc_profiler.py rename to apps/worker/app/services/document_parser/profiling/doc_profiler.py index bf728ccf3..184f13642 100644 --- a/apps/worker/app/services/document_parser/doc_profiler.py +++ b/apps/worker/app/services/document_parser/profiling/doc_profiler.py @@ -5,14 +5,14 @@ DocProfile, driving routing decisions and type annotations. Usage: - from app.services.document_parser.doc_profiler import profile_document + from app.services.document_parser.profiling.doc_profiler import profile_document profile = profile_document("/path/to/file.pdf") """ import os -from app.services.document_parser.doc_profile_model import DocProfile -from app.services.document_parser.doc_profile_pdf import profile_pdf +from app.services.document_parser.profiling.doc_profile_model import DocProfile +from app.services.document_parser.profiling.doc_profile_pdf import profile_pdf def profile_document(file_path: str, filename: str = "") -> DocProfile: diff --git a/apps/worker/app/services/document_parser/mineru_client.py b/apps/worker/app/services/document_parser/providers/mineru/client.py similarity index 96% rename from apps/worker/app/services/document_parser/mineru_client.py rename to apps/worker/app/services/document_parser/providers/mineru/client.py index 1cf5e3b97..b86f77531 100644 --- a/apps/worker/app/services/document_parser/mineru_client.py +++ b/apps/worker/app/services/document_parser/providers/mineru/client.py @@ -3,7 +3,7 @@ from typing import Any, Optional import requests -from app.services.document_parser.mineru_quota_manager import get_mineru_quota_manager +from app.services.document_parser.providers.mineru.quota_manager import get_mineru_quota_manager from loguru import logger from requests.adapters import HTTPAdapter from urllib3.util.retry import Retry diff --git a/apps/worker/app/services/document_parser/mineru_pdf_service.py b/apps/worker/app/services/document_parser/providers/mineru/pdf_service.py similarity index 97% rename from apps/worker/app/services/document_parser/mineru_pdf_service.py rename to apps/worker/app/services/document_parser/providers/mineru/pdf_service.py index 8be565518..4a3b6ec50 100644 --- a/apps/worker/app/services/document_parser/mineru_pdf_service.py +++ b/apps/worker/app/services/document_parser/providers/mineru/pdf_service.py @@ -2,18 +2,18 @@ from typing import Optional import requests -from app.services.document_parser.mineru_client import ( +from app.services.document_parser.providers.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 ( +from app.services.document_parser.providers.mineru.quota_manager import get_mineru_quota_manager +from app.services.document_parser.providers.mineru.task_polling import ( get_batch_status, poll_mineru_task, ) -from app.services.document_parser.parser_log_utils import truncate_log_value +from app.services.document_parser.support.parser_log_utils import truncate_log_value from shared.core.config import settings from shared.core.constants import APIConstants diff --git a/apps/worker/app/services/document_parser/mineru_quota_manager.py b/apps/worker/app/services/document_parser/providers/mineru/quota_manager.py similarity index 100% rename from apps/worker/app/services/document_parser/mineru_quota_manager.py rename to apps/worker/app/services/document_parser/providers/mineru/quota_manager.py diff --git a/apps/worker/app/services/document_parser/mineru_task_polling.py b/apps/worker/app/services/document_parser/providers/mineru/task_polling.py similarity index 98% rename from apps/worker/app/services/document_parser/mineru_task_polling.py rename to apps/worker/app/services/document_parser/providers/mineru/task_polling.py index c09c1ea86..689ac69ab 100644 --- a/apps/worker/app/services/document_parser/mineru_task_polling.py +++ b/apps/worker/app/services/document_parser/providers/mineru/task_polling.py @@ -5,13 +5,13 @@ from typing import Any, Optional import requests -from app.services.document_parser.mineru_client import ( +from app.services.document_parser.providers.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.providers.mineru.quota_manager import get_mineru_quota_manager from loguru import logger from shared.core.config import settings diff --git a/apps/worker/app/services/document_parser/heading_candidates.py b/apps/worker/app/services/document_parser/structure/heading_candidates.py similarity index 98% rename from apps/worker/app/services/document_parser/heading_candidates.py rename to apps/worker/app/services/document_parser/structure/heading_candidates.py index ccf293e16..f9d060349 100644 --- a/apps/worker/app/services/document_parser/heading_candidates.py +++ b/apps/worker/app/services/document_parser/structure/heading_candidates.py @@ -10,7 +10,7 @@ from loguru import logger from pandas import Index -from app.services.document_parser.text_helpers import count_cn_en +from app.services.document_parser.support.text_helpers import count_cn_en HEADING_COLUMNS = Index(["id", "heading", "level", "reason"]) @@ -145,7 +145,7 @@ def filter_markdown_headings( meta_ctx = None if layout_json_path: try: - from app.services.document_parser.metadata_extractor import MetadataContext + from app.services.document_parser.structure.metadata_extractor import MetadataContext meta_ctx = MetadataContext(md_lines, layout_json_path) except Exception as exc: @@ -258,7 +258,7 @@ def _is_non_heading_markdown_line(line: str) -> bool: def _estimate_markdown_heading_level(line: str, meta_ctx: Any | None): - from app.services.document_parser.metadata_extractor import detect_and_strip_md_bold + from app.services.document_parser.structure.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) diff --git a/apps/worker/app/services/document_parser/heading_hierarchy.py b/apps/worker/app/services/document_parser/structure/heading_hierarchy.py similarity index 93% rename from apps/worker/app/services/document_parser/heading_hierarchy.py rename to apps/worker/app/services/document_parser/structure/heading_hierarchy.py index e53075b62..20a9b170e 100644 --- a/apps/worker/app/services/document_parser/heading_hierarchy.py +++ b/apps/worker/app/services/document_parser/structure/heading_hierarchy.py @@ -5,7 +5,7 @@ import pandas as pd -from app.services.document_parser.layout_parser import pred_titles +from app.services.document_parser.structure.layout_parser import pred_titles @dataclass(frozen=True) diff --git a/apps/worker/app/services/document_parser/heading_llm_executor.py b/apps/worker/app/services/document_parser/structure/heading_llm_executor.py similarity index 98% rename from apps/worker/app/services/document_parser/heading_llm_executor.py rename to apps/worker/app/services/document_parser/structure/heading_llm_executor.py index 2f72789e6..6f314416b 100644 --- a/apps/worker/app/services/document_parser/heading_llm_executor.py +++ b/apps/worker/app/services/document_parser/structure/heading_llm_executor.py @@ -8,9 +8,9 @@ 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 app.services.document_parser.structure.metadata_extractor import clean_md_text_for_llm +from app.services.document_parser.support.stage_profiler import stage_timer +from app.services.document_parser.support.text_helpers import count_cn_en, truncate_text_by_tokens from loguru import logger from shared.core.exceptions.domain_exceptions import WorkerHandlingException diff --git a/apps/worker/app/services/document_parser/heading_tree.py b/apps/worker/app/services/document_parser/structure/heading_tree.py similarity index 100% rename from apps/worker/app/services/document_parser/heading_tree.py rename to apps/worker/app/services/document_parser/structure/heading_tree.py diff --git a/apps/worker/app/services/document_parser/layout_parser.py b/apps/worker/app/services/document_parser/structure/layout_parser.py similarity index 98% rename from apps/worker/app/services/document_parser/layout_parser.py rename to apps/worker/app/services/document_parser/structure/layout_parser.py index ec7775904..47f0650af 100755 --- a/apps/worker/app/services/document_parser/layout_parser.py +++ b/apps/worker/app/services/document_parser/structure/layout_parser.py @@ -3,27 +3,27 @@ import gevent import pandas as pd -from app.services.document_parser.heading_candidates import ( +from app.services.document_parser.structure.heading_candidates import ( filter_document_headings, filter_markdown_headings, postprocess_headings, ) -from app.services.document_parser.heading_llm_executor import ( +from app.services.document_parser.structure.heading_llm_executor import ( build_level_mapping, execute_level_mapping, execute_llm_heading_hierarchy, ) -from app.services.document_parser.heading_tree import ( +from app.services.document_parser.structure.heading_tree import ( build_tree_from_dataframe as build_heading_tree_from_dataframe, ) -from app.services.document_parser.heading_tree import ( +from app.services.document_parser.structure.heading_tree import ( remove_isolated_nodes as remove_isolated_heading_nodes, ) -from app.services.document_parser.heading_tree import ( +from app.services.document_parser.structure.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_text_parser import df2md +from app.services.document_parser.support.stage_profiler import stage_timer +from app.services.document_parser.tables.table_text_parser import df2md from gevent.pool import Pool as GeventPool from loguru import logger diff --git a/apps/worker/app/services/document_parser/metadata_extractor.py b/apps/worker/app/services/document_parser/structure/metadata_extractor.py similarity index 100% rename from apps/worker/app/services/document_parser/metadata_extractor.py rename to apps/worker/app/services/document_parser/structure/metadata_extractor.py diff --git a/apps/worker/app/services/document_parser/toc_hierarchy.py b/apps/worker/app/services/document_parser/structure/toc_hierarchy.py similarity index 93% rename from apps/worker/app/services/document_parser/toc_hierarchy.py rename to apps/worker/app/services/document_parser/structure/toc_hierarchy.py index b88420e32..439223b38 100644 --- a/apps/worker/app/services/document_parser/toc_hierarchy.py +++ b/apps/worker/app/services/document_parser/structure/toc_hierarchy.py @@ -1,10 +1,10 @@ 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_text_parser import df2md -from app.services.document_parser.text_helpers import normalize_md +from app.services.document_parser.structure.layout_parser import hiearchy_llm +from app.services.document_parser.support.stage_profiler import stage_timer +from app.services.document_parser.tables.table_text_parser import df2md +from app.services.document_parser.support.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/structure/toc_parser.py similarity index 98% rename from apps/worker/app/services/document_parser/toc_parser.py rename to apps/worker/app/services/document_parser/structure/toc_parser.py index 5c52ef5c4..9380884f3 100644 --- a/apps/worker/app/services/document_parser/toc_parser.py +++ b/apps/worker/app/services/document_parser/structure/toc_parser.py @@ -12,10 +12,10 @@ 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.stage_profiler import stage_timer -from app.services.document_parser.table_text_parser import df2md +from app.services.document_parser.structure.toc_hierarchy import eval_toc_levels +from app.services.document_parser.support.text_helpers import normalize_md, truncate_text_by_tokens +from app.services.document_parser.support.stage_profiler import stage_timer +from app.services.document_parser.tables.table_text_parser import df2md from gevent.pool import Pool as GeventPool from loguru import logger diff --git a/apps/worker/app/services/document_parser/identifiers.py b/apps/worker/app/services/document_parser/support/identifiers.py similarity index 100% rename from apps/worker/app/services/document_parser/identifiers.py rename to apps/worker/app/services/document_parser/support/identifiers.py diff --git a/apps/worker/app/services/document_parser/internal_parse_name.py b/apps/worker/app/services/document_parser/support/internal_parse_name.py similarity index 100% rename from apps/worker/app/services/document_parser/internal_parse_name.py rename to apps/worker/app/services/document_parser/support/internal_parse_name.py diff --git a/apps/worker/app/services/document_parser/parser_log_utils.py b/apps/worker/app/services/document_parser/support/parser_log_utils.py similarity index 100% rename from apps/worker/app/services/document_parser/parser_log_utils.py rename to apps/worker/app/services/document_parser/support/parser_log_utils.py diff --git a/apps/worker/app/services/document_parser/parser_rows.py b/apps/worker/app/services/document_parser/support/parser_rows.py similarity index 100% rename from apps/worker/app/services/document_parser/parser_rows.py rename to apps/worker/app/services/document_parser/support/parser_rows.py diff --git a/apps/worker/app/services/document_parser/path_helpers.py b/apps/worker/app/services/document_parser/support/path_helpers.py similarity index 100% rename from apps/worker/app/services/document_parser/path_helpers.py rename to apps/worker/app/services/document_parser/support/path_helpers.py diff --git a/apps/worker/app/services/document_parser/stage_profiler.py b/apps/worker/app/services/document_parser/support/stage_profiler.py similarity index 100% rename from apps/worker/app/services/document_parser/stage_profiler.py rename to apps/worker/app/services/document_parser/support/stage_profiler.py diff --git a/apps/worker/app/services/document_parser/text_helpers.py b/apps/worker/app/services/document_parser/support/text_helpers.py similarity index 100% rename from apps/worker/app/services/document_parser/text_helpers.py rename to apps/worker/app/services/document_parser/support/text_helpers.py diff --git a/apps/worker/app/services/document_parser/dataframe_helpers.py b/apps/worker/app/services/document_parser/tables/dataframe_helpers.py similarity index 98% rename from apps/worker/app/services/document_parser/dataframe_helpers.py rename to apps/worker/app/services/document_parser/tables/dataframe_helpers.py index 2173286b3..cef9b6446 100644 --- a/apps/worker/app/services/document_parser/dataframe_helpers.py +++ b/apps/worker/app/services/document_parser/tables/dataframe_helpers.py @@ -5,7 +5,7 @@ import pandas as pd -from app.services.document_parser.identifiers import gen_str_codes +from app.services.document_parser.support.identifiers import gen_str_codes def flatten_list(nested_list: list[Any]) -> list[Any]: diff --git a/apps/worker/app/services/document_parser/dataframe_html_renderer.py b/apps/worker/app/services/document_parser/tables/dataframe_html_renderer.py similarity index 100% rename from apps/worker/app/services/document_parser/dataframe_html_renderer.py rename to apps/worker/app/services/document_parser/tables/dataframe_html_renderer.py diff --git a/apps/worker/app/services/document_parser/table_asset_writer.py b/apps/worker/app/services/document_parser/tables/table_asset_writer.py similarity index 95% rename from apps/worker/app/services/document_parser/table_asset_writer.py rename to apps/worker/app/services/document_parser/tables/table_asset_writer.py index 91599f14d..679ee028e 100644 --- a/apps/worker/app/services/document_parser/table_asset_writer.py +++ b/apps/worker/app/services/document_parser/tables/table_asset_writer.py @@ -3,7 +3,7 @@ import os from dataclasses import dataclass -from app.services.document_parser.parser_rows import ParsedRow +from app.services.document_parser.support.parser_rows import ParsedRow @dataclass(frozen=True) diff --git a/apps/worker/app/services/document_parser/table_frame_parser.py b/apps/worker/app/services/document_parser/tables/table_frame_parser.py similarity index 98% rename from apps/worker/app/services/document_parser/table_frame_parser.py rename to apps/worker/app/services/document_parser/tables/table_frame_parser.py index a71cabf65..76d7965dd 100644 --- a/apps/worker/app/services/document_parser/table_frame_parser.py +++ b/apps/worker/app/services/document_parser/tables/table_frame_parser.py @@ -8,9 +8,9 @@ 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 app.services.document_parser.tables.dataframe_html_renderer import df2html +from app.services.document_parser.support.identifiers import gen_str_codes +from app.services.document_parser.support.path_helpers import flatten_dic2paths from loguru import logger from shared.services.ai.prompt_service import build_prompt diff --git a/apps/worker/app/services/document_parser/table_parser.py b/apps/worker/app/services/document_parser/tables/table_parser.py similarity index 86% rename from apps/worker/app/services/document_parser/table_parser.py rename to apps/worker/app/services/document_parser/tables/table_parser.py index 2c189c0c7..7cdafe6f3 100755 --- a/apps/worker/app/services/document_parser/table_parser.py +++ b/apps/worker/app/services/document_parser/tables/table_parser.py @@ -1,7 +1,7 @@ from __future__ import annotations import pandas as pd -from app.services.document_parser.table_frame_parser import ( +from app.services.document_parser.tables.table_frame_parser import ( format_tb_scope as format_tb_scope, multiindex_to_tree as multiindex_to_tree, parse_headers as parse_headers, @@ -12,7 +12,7 @@ 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 ( +from app.services.document_parser.tables.table_text_parser import ( clean_html_tb as clean_html_tb, df2md as df2md, extract_tables_by_forms as extract_tables_by_forms, @@ -32,7 +32,7 @@ def parse_xlsx( use_precision_mode: bool = True, include_hidden_sheets: bool = False, ) -> pd.DataFrame: - from app.services.document_parser.excel_table_parser import ( + from app.services.document_parser.formats.excel.table_parser import ( parse_xlsx as parse_excel_xlsx, ) diff --git a/apps/worker/app/services/document_parser/table_text_parser.py b/apps/worker/app/services/document_parser/tables/table_text_parser.py similarity index 100% rename from apps/worker/app/services/document_parser/table_text_parser.py rename to apps/worker/app/services/document_parser/tables/table_text_parser.py From 607b7ea9bb1e50fc0fbf5545a97eafacbefbbebf Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 12:34:56 +0000 Subject: [PATCH 69/78] refactor: group retrieval internals by ownership --- AGENTS.md | 13 +++++++++++-- CONTEXT.md | 9 ++++++++- .../app/services/documents/lifecycle_service.py | 2 +- apps/api/main.py | 2 +- apps/api/tests/contract/test_retrieval_contract.py | 6 +++--- .../services/connect_builder/summary_builder.py | 2 +- .../shared/services/retrieval/__init__.py | 6 +++--- .../retrieval/agentic/discovery_selection.py | 2 +- .../services/retrieval/agentic/discovery_tools.py | 4 ++-- .../shared/services/retrieval/agentic/evidence.py | 4 ++-- .../services/retrieval/agentic/section_tree.py | 2 +- .../retrieval/agentic/selection_hydration.py | 4 ++-- .../shared/services/retrieval/app_service.py | 4 ++-- .../services/retrieval/execution/__init__.py | 1 + .../retrieval/{ => execution}/legacy_route.py | 14 +++++++------- .../{execution_plan.py => execution/plan.py} | 8 ++++---- .../retrieval/{ => execution}/query_request.py | 2 +- .../{ => execution}/reference_resolver.py | 6 +++--- .../{ => execution}/response_projection.py | 4 ++-- .../retrieval/{ => execution}/route_types.py | 0 .../{execution_routes.py => execution/routes.py} | 12 ++++++------ .../shared/services/retrieval/graph/__init__.py | 1 + .../{graph_keywords.py => graph/keywords.py} | 0 .../query_service.py} | 2 +- .../{graph_service.py => graph/service.py} | 2 +- .../services/retrieval/hydration/__init__.py | 1 + .../services/retrieval/{ => hydration}/assets.py | 2 +- .../connected.py} | 2 +- .../{path_hydration.py => hydration/path.py} | 4 ++-- .../reference.py} | 2 +- .../retrieval/{ => hydration}/result_assembly.py | 4 ++-- .../retrieval/{ => hydration}/row_utils.py | 2 +- .../services/retrieval/publication_service.py | 4 ++-- .../shared/services/retrieval/search/__init__.py | 1 + .../services/retrieval/{ => search}/channels.py | 4 ++-- .../retrieval/{ => search}/lexical_ranker.py | 0 .../retrieval/{ => search}/lexical_text.py | 0 .../services/retrieval/{ => search}/ranking.py | 4 ++-- .../retrieval/{ => search}/scoped_corpus.py | 2 +- .../services/retrieval/{ => search}/scoring.py | 0 .../retrieval/{ => search}/section_filters.py | 0 .../shared/services/retrieval/stats/__init__.py | 1 + .../{hit_stats_recorder.py => stats/recorder.py} | 2 +- .../{hit_stats_service.py => stats/service.py} | 0 .../services/retrieval/workflow/run_request.py | 2 +- 45 files changed, 85 insertions(+), 64 deletions(-) create mode 100644 packages/shared-python/shared/services/retrieval/execution/__init__.py rename packages/shared-python/shared/services/retrieval/{ => execution}/legacy_route.py (95%) rename packages/shared-python/shared/services/retrieval/{execution_plan.py => execution/plan.py} (96%) rename packages/shared-python/shared/services/retrieval/{ => execution}/query_request.py (97%) rename packages/shared-python/shared/services/retrieval/{ => execution}/reference_resolver.py (92%) rename packages/shared-python/shared/services/retrieval/{ => execution}/response_projection.py (92%) rename packages/shared-python/shared/services/retrieval/{ => execution}/route_types.py (100%) rename packages/shared-python/shared/services/retrieval/{execution_routes.py => execution/routes.py} (90%) create mode 100644 packages/shared-python/shared/services/retrieval/graph/__init__.py rename packages/shared-python/shared/services/retrieval/{graph_keywords.py => graph/keywords.py} (100%) rename packages/shared-python/shared/services/retrieval/{graph_query_service.py => graph/query_service.py} (98%) rename packages/shared-python/shared/services/retrieval/{graph_service.py => graph/service.py} (99%) create mode 100644 packages/shared-python/shared/services/retrieval/hydration/__init__.py rename packages/shared-python/shared/services/retrieval/{ => hydration}/assets.py (96%) rename packages/shared-python/shared/services/retrieval/{connected_hydration.py => hydration/connected.py} (98%) rename packages/shared-python/shared/services/retrieval/{path_hydration.py => hydration/path.py} (98%) rename packages/shared-python/shared/services/retrieval/{reference_hydration.py => hydration/reference.py} (98%) rename packages/shared-python/shared/services/retrieval/{ => hydration}/result_assembly.py (95%) rename packages/shared-python/shared/services/retrieval/{ => hydration}/row_utils.py (96%) create mode 100644 packages/shared-python/shared/services/retrieval/search/__init__.py rename packages/shared-python/shared/services/retrieval/{ => search}/channels.py (98%) rename packages/shared-python/shared/services/retrieval/{ => search}/lexical_ranker.py (100%) rename packages/shared-python/shared/services/retrieval/{ => search}/lexical_text.py (100%) rename packages/shared-python/shared/services/retrieval/{ => search}/ranking.py (97%) rename packages/shared-python/shared/services/retrieval/{ => search}/scoped_corpus.py (97%) rename packages/shared-python/shared/services/retrieval/{ => search}/scoring.py (100%) rename packages/shared-python/shared/services/retrieval/{ => search}/section_filters.py (100%) create mode 100644 packages/shared-python/shared/services/retrieval/stats/__init__.py rename packages/shared-python/shared/services/retrieval/{hit_stats_recorder.py => stats/recorder.py} (96%) rename packages/shared-python/shared/services/retrieval/{hit_stats_service.py => stats/service.py} (100%) diff --git a/AGENTS.md b/AGENTS.md index 4881c4720..f6c4c0e6f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -533,6 +533,15 @@ debug CSVs (`preds_*.csv`) are saved alongside for troubleshooting. `shared/services/retrieval/app_service.py` → `run_retrieval_query()` +Core retrieval internals are grouped by ownership: + +- `execution/`: request shaping, route selection, legacy route execution, and public response projection. +- `search/`: lexical channels, scoring, section filters, and candidate ranking. +- `hydration/`: row/path/reference hydration, inline assets, and result assembly. +- `graph/`: document graph publication/query support. +- `stats/`: retrieval hit recording. +- `workflow/` and `agentic/`: agentic planning and navigation flows. + ### Two Retrieval Modes The system supports two modes, controlled globally by `RETRIEVAL_AGENTIC_ENABLED` and locally via the per-request `use_agentic` toggle. @@ -549,7 +558,7 @@ flowchart LR T --> RRF RRF --> Graph[Legacy Graph Routing] Graph --> Rank[Dual-priority ranking] - Rank --> Assemble[assemble_retrieval_results] + Rank --> Assemble[hydration.result_assembly] ``` **Channel weights** (default): path=1.0, content=2.0, term=1.5 @@ -584,7 +593,7 @@ Unlike legacy retrieval which relied on static `hydrate_mode` tags, hydration is ### Result Assembly -`assemble_retrieval_results()`: +`hydration.result_assembly.assemble_retrieval_results()`: 1. Filters by `exclude_document_ids` and `exclude_sections` 2. Filters by `allowed_chunk_types` (data_type parameter) diff --git a/CONTEXT.md b/CONTEXT.md index 9ac741eab..dd37639ec 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -294,7 +294,14 @@ exceptions. Shared Job lifecycle finalization lives under ### Retrieval - `app/api/v1/routes/retrieval.py` -- shared retrieval modules in `packages/shared-python/shared/services/retrieval/*` +- `packages/shared-python/shared/services/retrieval/app_service.py` +- `packages/shared-python/shared/services/retrieval/execution/*` +- `packages/shared-python/shared/services/retrieval/search/*` +- `packages/shared-python/shared/services/retrieval/hydration/*` +- `packages/shared-python/shared/services/retrieval/graph/*` +- `packages/shared-python/shared/services/retrieval/stats/*` +- `packages/shared-python/shared/services/retrieval/workflow/*` +- `packages/shared-python/shared/services/retrieval/agentic/*` ### Demo Source Materialization diff --git a/apps/api/app/services/documents/lifecycle_service.py b/apps/api/app/services/documents/lifecycle_service.py index 93b9fb300..9986f29ad 100644 --- a/apps/api/app/services/documents/lifecycle_service.py +++ b/apps/api/app/services/documents/lifecycle_service.py @@ -14,7 +14,7 @@ from shared.services.retrieval.cache_service import ( invalidate_retrieval_cache_namespaces, ) -from shared.services.retrieval.graph_service import DocumentGraphService, GraphScope +from shared.services.retrieval.graph.service import DocumentGraphService, GraphScope def document_payload(document) -> dict[str, Any]: diff --git a/apps/api/main.py b/apps/api/main.py index 8caac693f..22b058c67 100644 --- a/apps/api/main.py +++ b/apps/api/main.py @@ -79,7 +79,7 @@ async def lifespan(app: FastAPI): yield try: - from shared.services.retrieval.hit_stats_recorder import ( + from shared.services.retrieval.stats.recorder import ( drain_retrieval_hit_stats_updates, ) diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index 489e2c365..2ee929d60 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -451,15 +451,15 @@ async def fake_graph_routing(*_args: object, **_kwargs: object) -> list[dict[str return [] monkeypatch.setattr( - "shared.services.retrieval.legacy_route.path_channel", + "shared.services.retrieval.execution.legacy_route.path_channel", fake_path_channel, ) monkeypatch.setattr( - "shared.services.retrieval.legacy_route.content_channel", + "shared.services.retrieval.execution.legacy_route.content_channel", fake_content_channel, ) monkeypatch.setattr( - "shared.services.retrieval.legacy_route.list_graph_routed_chunks", + "shared.services.retrieval.execution.legacy_route.list_graph_routed_chunks", fake_graph_routing, ) diff --git a/apps/worker/app/services/connect_builder/summary_builder.py b/apps/worker/app/services/connect_builder/summary_builder.py index 44f472076..c52348112 100644 --- a/apps/worker/app/services/connect_builder/summary_builder.py +++ b/apps/worker/app/services/connect_builder/summary_builder.py @@ -347,7 +347,7 @@ def build_section_summary_lookup(file_dir: str) -> Dict[str, str]: Returns: Dict mapping section_path → summary string (empty dict on any error). """ - from shared.services.retrieval.lexical_text import section_path_from_chunk_path + from shared.services.retrieval.search.lexical_text import section_path_from_chunk_path doc_nav = _load_doc_nav(file_dir) if doc_nav is None: diff --git a/packages/shared-python/shared/services/retrieval/__init__.py b/packages/shared-python/shared/services/retrieval/__init__.py index 4c1a56a95..259fa6489 100644 --- a/packages/shared-python/shared/services/retrieval/__init__.py +++ b/packages/shared-python/shared/services/retrieval/__init__.py @@ -6,9 +6,9 @@ invalidate_retrieval_cache_namespaces, set_cached_retrieval_query_result, ) -from .graph_query_service import GraphQueryService -from .graph_service import DocumentGraphService, GraphScope -from .hit_stats_service import record_retrieval_hits +from .graph.query_service import GraphQueryService +from .graph.service import DocumentGraphService, GraphScope +from .stats.service import record_retrieval_hits from .llm_adapter import create_retrieval_llm_fn, create_retrieval_planner_fn __all__ = [ diff --git a/packages/shared-python/shared/services/retrieval/agentic/discovery_selection.py b/packages/shared-python/shared/services/retrieval/agentic/discovery_selection.py index d54df4e5f..e47d9e4e3 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/discovery_selection.py +++ b/packages/shared-python/shared/services/retrieval/agentic/discovery_selection.py @@ -17,7 +17,7 @@ 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.search.lexical_text import normalize_section_path from shared.services.retrieval.llm_adapter import LLMFn diff --git a/packages/shared-python/shared/services/retrieval/agentic/discovery_tools.py b/packages/shared-python/shared/services/retrieval/agentic/discovery_tools.py index 44ac0f038..8589bba00 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/discovery_tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/discovery_tools.py @@ -22,9 +22,9 @@ 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.search.channels import content_channel, path_channel, term_channel from shared.services.retrieval.llm_adapter import LLMFn -from shared.services.retrieval.scoring import ( +from shared.services.retrieval.search.scoring import ( merge_channels_rrf, merge_same_section_rows, normalize_row_scores, diff --git a/packages/shared-python/shared/services/retrieval/agentic/evidence.py b/packages/shared-python/shared/services/retrieval/agentic/evidence.py index 3d54fb47b..2f97f557c 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/evidence.py +++ b/packages/shared-python/shared/services/retrieval/agentic/evidence.py @@ -9,8 +9,8 @@ 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 build_retrieval_asset_url_map -from shared.services.retrieval.hit_stats_service import compute_importance_score +from shared.services.retrieval.hydration.assets import build_retrieval_asset_url_map +from shared.services.retrieval.stats.service import compute_importance_score from shared.utils.token_estimate import estimate_tokens 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 53507f35d..3a285262b 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/section_tree.py +++ b/packages/shared-python/shared/services/retrieval/agentic/section_tree.py @@ -7,7 +7,7 @@ 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.services.retrieval.search.lexical_text import normalize_section_path, split_section_path async def load_child_sections( diff --git a/packages/shared-python/shared/services/retrieval/agentic/selection_hydration.py b/packages/shared-python/shared/services/retrieval/agentic/selection_hydration.py index 2fa86c690..66d92ec2f 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/selection_hydration.py +++ b/packages/shared-python/shared/services/retrieval/agentic/selection_hydration.py @@ -6,8 +6,8 @@ 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 +from shared.services.retrieval.hydration.connected import hydrate_connected_target_rows +from shared.services.retrieval.hydration.path import hydrate_paths_to_rows async def hydrate_path_selections_into_node( diff --git a/packages/shared-python/shared/services/retrieval/app_service.py b/packages/shared-python/shared/services/retrieval/app_service.py index 0644c478e..a5e3a55df 100644 --- a/packages/shared-python/shared/services/retrieval/app_service.py +++ b/packages/shared-python/shared/services/retrieval/app_service.py @@ -4,10 +4,10 @@ from sqlalchemy.ext.asyncio import AsyncSession -from shared.services.retrieval.execution_plan import ( +from shared.services.retrieval.execution.plan import ( run_retrieval_query as execute_retrieval_query, ) -from shared.services.retrieval.scoring import merge_channels_rrf +from shared.services.retrieval.search.scoring import merge_channels_rrf __all__ = ["merge_channels_rrf", "run_retrieval_query"] diff --git a/packages/shared-python/shared/services/retrieval/execution/__init__.py b/packages/shared-python/shared/services/retrieval/execution/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/execution/__init__.py @@ -0,0 +1 @@ + diff --git a/packages/shared-python/shared/services/retrieval/legacy_route.py b/packages/shared-python/shared/services/retrieval/execution/legacy_route.py similarity index 95% rename from packages/shared-python/shared/services/retrieval/legacy_route.py rename to packages/shared-python/shared/services/retrieval/execution/legacy_route.py index e8642840b..6b1976f53 100644 --- a/packages/shared-python/shared/services/retrieval/legacy_route.py +++ b/packages/shared-python/shared/services/retrieval/execution/legacy_route.py @@ -6,16 +6,16 @@ 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 ( +from shared.services.retrieval.search.channels import content_channel, path_channel, term_channel +from shared.services.retrieval.graph.query_service import GraphQueryService +from shared.services.retrieval.search.ranking import rank_retrieval_candidates +from shared.services.retrieval.execution.response_projection import attach_citation +from shared.services.retrieval.hydration.result_assembly import assemble_retrieval_results +from shared.services.retrieval.execution.route_types import ( RetrievalRouteContext, RetrievalRouteOutcome, ) -from shared.services.retrieval.scoring import ( +from shared.services.retrieval.search.scoring import ( get_row_path, merge_channels_rrf, merge_same_section_rows, diff --git a/packages/shared-python/shared/services/retrieval/execution_plan.py b/packages/shared-python/shared/services/retrieval/execution/plan.py similarity index 96% rename from packages/shared-python/shared/services/retrieval/execution_plan.py rename to packages/shared-python/shared/services/retrieval/execution/plan.py index d94cc8f94..89fc85048 100644 --- a/packages/shared-python/shared/services/retrieval/execution_plan.py +++ b/packages/shared-python/shared/services/retrieval/execution/plan.py @@ -10,14 +10,14 @@ get_cached_retrieval_query_result, set_cached_retrieval_query_result, ) -from shared.services.retrieval.execution_routes import run_retrieval_route -from shared.services.retrieval.hit_stats_recorder import ( +from shared.services.retrieval.execution.routes import run_retrieval_route +from shared.services.retrieval.stats.recorder import ( schedule_retrieval_hit_stats_update, ) -from shared.services.retrieval.response_projection import ( +from shared.services.retrieval.execution.response_projection import ( project_public_retrieval_response, ) -from shared.services.retrieval.query_request import RetrievalQuery +from shared.services.retrieval.execution.query_request import RetrievalQuery async def run_retrieval_query( diff --git a/packages/shared-python/shared/services/retrieval/query_request.py b/packages/shared-python/shared/services/retrieval/execution/query_request.py similarity index 97% rename from packages/shared-python/shared/services/retrieval/query_request.py rename to packages/shared-python/shared/services/retrieval/execution/query_request.py index 253b2e7b2..a233df233 100644 --- a/packages/shared-python/shared/services/retrieval/query_request.py +++ b/packages/shared-python/shared/services/retrieval/execution/query_request.py @@ -5,7 +5,7 @@ from sqlalchemy.ext.asyncio import AsyncSession -from shared.services.retrieval.route_types import RetrievalRouteContext +from shared.services.retrieval.execution.route_types import RetrievalRouteContext from shared.services.retrieval.settings import ( INTERNAL_RECALL_K_MULTIPLIER, resolve_allowed_chunk_types, diff --git a/packages/shared-python/shared/services/retrieval/reference_resolver.py b/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py similarity index 92% rename from packages/shared-python/shared/services/retrieval/reference_resolver.py rename to packages/shared-python/shared/services/retrieval/execution/reference_resolver.py index 7845ef177..c44095f50 100644 --- a/packages/shared-python/shared/services/retrieval/reference_resolver.py +++ b/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py @@ -5,11 +5,11 @@ from sqlalchemy.ext.asyncio import AsyncSession -from shared.services.retrieval.reference_hydration import hydrate_referenced_chunk_rows -from shared.services.retrieval.response_projection import ( +from shared.services.retrieval.hydration.reference import hydrate_referenced_chunk_rows +from shared.services.retrieval.execution.response_projection import ( enrich_referenced_chunks_with_asset_urls, ) -from shared.services.retrieval.row_utils import build_reference_lookup_key +from shared.services.retrieval.hydration.row_utils import build_reference_lookup_key @dataclass(frozen=True) diff --git a/packages/shared-python/shared/services/retrieval/response_projection.py b/packages/shared-python/shared/services/retrieval/execution/response_projection.py similarity index 92% rename from packages/shared-python/shared/services/retrieval/response_projection.py rename to packages/shared-python/shared/services/retrieval/execution/response_projection.py index f12638a5c..14e84bf4b 100644 --- a/packages/shared-python/shared/services/retrieval/response_projection.py +++ b/packages/shared-python/shared/services/retrieval/execution/response_projection.py @@ -2,8 +2,8 @@ from typing import Any -from shared.services.retrieval.assets import enrich_rows_with_retrieval_asset_urls -from shared.services.retrieval.row_utils import ( +from shared.services.retrieval.hydration.assets import enrich_rows_with_retrieval_asset_urls +from shared.services.retrieval.hydration.row_utils import ( PUBLIC_RESULT_FIELDS, PUBLIC_SOURCE_FIELDS, ) diff --git a/packages/shared-python/shared/services/retrieval/route_types.py b/packages/shared-python/shared/services/retrieval/execution/route_types.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/route_types.py rename to packages/shared-python/shared/services/retrieval/execution/route_types.py diff --git a/packages/shared-python/shared/services/retrieval/execution_routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py similarity index 90% rename from packages/shared-python/shared/services/retrieval/execution_routes.py rename to packages/shared-python/shared/services/retrieval/execution/routes.py index 2cd4f5ea5..d7f13538a 100644 --- a/packages/shared-python/shared/services/retrieval/execution_routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -4,17 +4,17 @@ from loguru import logger -from shared.services.retrieval.legacy_route import run_legacy_retrieval_route -from shared.services.retrieval.reference_resolver import resolve_workflow_references -from shared.services.retrieval.result_assembly import assemble_retrieval_results -from shared.services.retrieval.response_projection import ( +from shared.services.retrieval.execution.legacy_route import run_legacy_retrieval_route +from shared.services.retrieval.execution.reference_resolver import resolve_workflow_references +from shared.services.retrieval.hydration.result_assembly import assemble_retrieval_results +from shared.services.retrieval.execution.response_projection import ( attach_citation, ) -from shared.services.retrieval.route_types import ( +from shared.services.retrieval.execution.route_types import ( RetrievalRouteContext, RetrievalRouteOutcome, ) -from shared.services.retrieval.scoped_corpus import ( +from shared.services.retrieval.search.scoped_corpus import ( count_scoped_chunks, load_all_scoped_chunks, ) diff --git a/packages/shared-python/shared/services/retrieval/graph/__init__.py b/packages/shared-python/shared/services/retrieval/graph/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/graph/__init__.py @@ -0,0 +1 @@ + diff --git a/packages/shared-python/shared/services/retrieval/graph_keywords.py b/packages/shared-python/shared/services/retrieval/graph/keywords.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/graph_keywords.py rename to packages/shared-python/shared/services/retrieval/graph/keywords.py diff --git a/packages/shared-python/shared/services/retrieval/graph_query_service.py b/packages/shared-python/shared/services/retrieval/graph/query_service.py similarity index 98% rename from packages/shared-python/shared/services/retrieval/graph_query_service.py rename to packages/shared-python/shared/services/retrieval/graph/query_service.py index de912a0ca..819d99469 100644 --- a/packages/shared-python/shared/services/retrieval/graph_query_service.py +++ b/packages/shared-python/shared/services/retrieval/graph/query_service.py @@ -8,7 +8,7 @@ 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 +from shared.services.retrieval.search.section_filters import is_excluded_section _SECTION_EXCLUSION_PAGE_MULTIPLIER = 2 diff --git a/packages/shared-python/shared/services/retrieval/graph_service.py b/packages/shared-python/shared/services/retrieval/graph/service.py similarity index 99% rename from packages/shared-python/shared/services/retrieval/graph_service.py rename to packages/shared-python/shared/services/retrieval/graph/service.py index 0d96a76e3..0d1cd1456 100644 --- a/packages/shared-python/shared/services/retrieval/graph_service.py +++ b/packages/shared-python/shared/services/retrieval/graph/service.py @@ -13,7 +13,7 @@ GraphEdge, GraphNode, ) -from shared.services.retrieval.graph_keywords import ( +from shared.services.retrieval.graph.keywords import ( KEYWORD_SCORE_WEIGHT, MIN_KEYWORD_OVERLAP, MIN_SCORE_THRESHOLD, diff --git a/packages/shared-python/shared/services/retrieval/hydration/__init__.py b/packages/shared-python/shared/services/retrieval/hydration/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/hydration/__init__.py @@ -0,0 +1 @@ + diff --git a/packages/shared-python/shared/services/retrieval/assets.py b/packages/shared-python/shared/services/retrieval/hydration/assets.py similarity index 96% rename from packages/shared-python/shared/services/retrieval/assets.py rename to packages/shared-python/shared/services/retrieval/hydration/assets.py index 035825430..99d1bb4a2 100644 --- a/packages/shared-python/shared/services/retrieval/assets.py +++ b/packages/shared-python/shared/services/retrieval/hydration/assets.py @@ -4,7 +4,7 @@ from loguru import logger -from shared.services.retrieval.row_utils import MEDIA_CHUNK_TYPES, normalize_chunk_type +from shared.services.retrieval.hydration.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/hydration/connected.py similarity index 98% rename from packages/shared-python/shared/services/retrieval/connected_hydration.py rename to packages/shared-python/shared/services/retrieval/hydration/connected.py index c5c2a181b..f7c299b77 100644 --- a/packages/shared-python/shared/services/retrieval/connected_hydration.py +++ b/packages/shared-python/shared/services/retrieval/hydration/connected.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.row_utils import ( +from shared.services.retrieval.hydration.row_utils import ( filter_excluded_rows, iter_connected_target_ids, normalize_chunk_type, diff --git a/packages/shared-python/shared/services/retrieval/path_hydration.py b/packages/shared-python/shared/services/retrieval/hydration/path.py similarity index 98% rename from packages/shared-python/shared/services/retrieval/path_hydration.py rename to packages/shared-python/shared/services/retrieval/hydration/path.py index d973f8a8b..de2ec15a2 100644 --- a/packages/shared-python/shared/services/retrieval/path_hydration.py +++ b/packages/shared-python/shared/services/retrieval/hydration/path.py @@ -8,8 +8,8 @@ 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 +from shared.services.retrieval.search.lexical_text import normalize_section_path +from shared.services.retrieval.search.scoring import get_row_path async def hydrate_paths_to_rows( diff --git a/packages/shared-python/shared/services/retrieval/reference_hydration.py b/packages/shared-python/shared/services/retrieval/hydration/reference.py similarity index 98% rename from packages/shared-python/shared/services/retrieval/reference_hydration.py rename to packages/shared-python/shared/services/retrieval/hydration/reference.py index 299024e2e..5c6a0556e 100644 --- a/packages/shared-python/shared/services/retrieval/reference_hydration.py +++ b/packages/shared-python/shared/services/retrieval/hydration/reference.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.row_utils import ( +from shared.services.retrieval.hydration.row_utils import ( ReferenceLookupKey, build_reference_lookup_key, ) diff --git a/packages/shared-python/shared/services/retrieval/result_assembly.py b/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py similarity index 95% rename from packages/shared-python/shared/services/retrieval/result_assembly.py rename to packages/shared-python/shared/services/retrieval/hydration/result_assembly.py index 1ef948078..17a26a43a 100644 --- a/packages/shared-python/shared/services/retrieval/result_assembly.py +++ b/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py @@ -4,8 +4,8 @@ from sqlalchemy.ext.asyncio import AsyncSession -from shared.services.retrieval.connected_hydration import hydrate_connected_target_rows -from shared.services.retrieval.row_utils import ( +from shared.services.retrieval.hydration.connected import hydrate_connected_target_rows +from shared.services.retrieval.hydration.row_utils import ( clean_content, filter_excluded_rows, iter_connected_target_ids, diff --git a/packages/shared-python/shared/services/retrieval/row_utils.py b/packages/shared-python/shared/services/retrieval/hydration/row_utils.py similarity index 96% rename from packages/shared-python/shared/services/retrieval/row_utils.py rename to packages/shared-python/shared/services/retrieval/hydration/row_utils.py index 0784830b0..d163836ca 100644 --- a/packages/shared-python/shared/services/retrieval/row_utils.py +++ b/packages/shared-python/shared/services/retrieval/hydration/row_utils.py @@ -3,7 +3,7 @@ import re from typing import Any -from shared.services.retrieval.section_filters import is_excluded_section +from shared.services.retrieval.search.section_filters import is_excluded_section MEDIA_CHUNK_TYPES = {'image', 'table'} PUBLIC_RESULT_FIELDS = { diff --git a/packages/shared-python/shared/services/retrieval/publication_service.py b/packages/shared-python/shared/services/retrieval/publication_service.py index 80aaa4ae0..40ae68045 100644 --- a/packages/shared-python/shared/services/retrieval/publication_service.py +++ b/packages/shared-python/shared/services/retrieval/publication_service.py @@ -19,8 +19,8 @@ from shared.models.database.document import Document, DocumentChunk, DocumentSection from shared.models.database.job import Job from shared.models.database.job_result import JobResult -from shared.services.retrieval.graph_service import DocumentGraphService, GraphScope -from shared.services.retrieval.lexical_text import ( +from shared.services.retrieval.graph.service import DocumentGraphService, GraphScope +from shared.services.retrieval.search.lexical_text import ( build_content_lexical_text, build_content_search_text, build_path_lexical_text, diff --git a/packages/shared-python/shared/services/retrieval/search/__init__.py b/packages/shared-python/shared/services/retrieval/search/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/search/__init__.py @@ -0,0 +1 @@ + diff --git a/packages/shared-python/shared/services/retrieval/channels.py b/packages/shared-python/shared/services/retrieval/search/channels.py similarity index 98% rename from packages/shared-python/shared/services/retrieval/channels.py rename to packages/shared-python/shared/services/retrieval/search/channels.py index 35966ed23..bd881769e 100644 --- a/packages/shared-python/shared/services/retrieval/channels.py +++ b/packages/shared-python/shared/services/retrieval/search/channels.py @@ -11,11 +11,11 @@ from sqlalchemy import text from sqlalchemy.ext.asyncio import AsyncSession -from shared.services.retrieval.lexical_ranker import ( +from shared.services.retrieval.search.lexical_ranker import ( rank_rows_by_bm25, tokenize_query_for_ranker, ) -from shared.services.retrieval.section_filters import is_excluded_section +from shared.services.retrieval.search.section_filters import is_excluded_section _SCOPED_CORPUS_CTE = """ diff --git a/packages/shared-python/shared/services/retrieval/lexical_ranker.py b/packages/shared-python/shared/services/retrieval/search/lexical_ranker.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/lexical_ranker.py rename to packages/shared-python/shared/services/retrieval/search/lexical_ranker.py diff --git a/packages/shared-python/shared/services/retrieval/lexical_text.py b/packages/shared-python/shared/services/retrieval/search/lexical_text.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/lexical_text.py rename to packages/shared-python/shared/services/retrieval/search/lexical_text.py diff --git a/packages/shared-python/shared/services/retrieval/ranking.py b/packages/shared-python/shared/services/retrieval/search/ranking.py similarity index 97% rename from packages/shared-python/shared/services/retrieval/ranking.py rename to packages/shared-python/shared/services/retrieval/search/ranking.py index bd1cfd9fd..392aa1f13 100644 --- a/packages/shared-python/shared/services/retrieval/ranking.py +++ b/packages/shared-python/shared/services/retrieval/search/ranking.py @@ -8,8 +8,8 @@ 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 +from shared.services.retrieval.stats.service import compute_importance_score +from shared.services.retrieval.search.scoring import get_row_path def get_candidate_key(row: dict[str, Any]) -> str: diff --git a/packages/shared-python/shared/services/retrieval/scoped_corpus.py b/packages/shared-python/shared/services/retrieval/search/scoped_corpus.py similarity index 97% rename from packages/shared-python/shared/services/retrieval/scoped_corpus.py rename to packages/shared-python/shared/services/retrieval/search/scoped_corpus.py index 012ba7129..3f7e4b3d3 100644 --- a/packages/shared-python/shared/services/retrieval/scoped_corpus.py +++ b/packages/shared-python/shared/services/retrieval/search/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.section_filters import is_excluded_section +from shared.services.retrieval.search.section_filters import is_excluded_section async def count_scoped_chunks( diff --git a/packages/shared-python/shared/services/retrieval/scoring.py b/packages/shared-python/shared/services/retrieval/search/scoring.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/scoring.py rename to packages/shared-python/shared/services/retrieval/search/scoring.py diff --git a/packages/shared-python/shared/services/retrieval/section_filters.py b/packages/shared-python/shared/services/retrieval/search/section_filters.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/section_filters.py rename to packages/shared-python/shared/services/retrieval/search/section_filters.py diff --git a/packages/shared-python/shared/services/retrieval/stats/__init__.py b/packages/shared-python/shared/services/retrieval/stats/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/stats/__init__.py @@ -0,0 +1 @@ + diff --git a/packages/shared-python/shared/services/retrieval/hit_stats_recorder.py b/packages/shared-python/shared/services/retrieval/stats/recorder.py similarity index 96% rename from packages/shared-python/shared/services/retrieval/hit_stats_recorder.py rename to packages/shared-python/shared/services/retrieval/stats/recorder.py index b101292fa..05b018678 100644 --- a/packages/shared-python/shared/services/retrieval/hit_stats_recorder.py +++ b/packages/shared-python/shared/services/retrieval/stats/recorder.py @@ -5,7 +5,7 @@ from loguru import logger -from shared.services.retrieval.hit_stats_service import record_retrieval_hits +from shared.services.retrieval.stats.service import record_retrieval_hits _pending_retrieval_hit_stat_tasks: set[asyncio.Task[None]] = set() diff --git a/packages/shared-python/shared/services/retrieval/hit_stats_service.py b/packages/shared-python/shared/services/retrieval/stats/service.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/hit_stats_service.py rename to packages/shared-python/shared/services/retrieval/stats/service.py diff --git a/packages/shared-python/shared/services/retrieval/workflow/run_request.py b/packages/shared-python/shared/services/retrieval/workflow/run_request.py index 4eff26d91..05ee21cba 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/run_request.py +++ b/packages/shared-python/shared/services/retrieval/workflow/run_request.py @@ -2,7 +2,7 @@ from dataclasses import dataclass -from shared.services.retrieval.route_types import RetrievalRouteContext +from shared.services.retrieval.execution.route_types import RetrievalRouteContext from shared.services.retrieval.settings import INTERNAL_RECALL_K_MULTIPLIER from shared.services.retrieval.workflow.types import PlannedStep From 9d1d635d2f2a2419a5de9299b774ad97fb7eabd3 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 12:46:42 +0000 Subject: [PATCH 70/78] refactor: structure agentic retrieval modules --- AGENTS.md | 6 +++++- CONTEXT.md | 5 ++++- apps/api/tests/contract/test_retrieval_contract.py | 2 +- .../services/retrieval/agentic/core/__init__.py | 1 + .../retrieval/agentic/{ => core}/budget.py | 0 .../retrieval/agentic/{ => core}/runtime.py | 4 ++-- .../services/retrieval/agentic/{ => core}/trace.py | 2 +- .../services/retrieval/agentic/{ => core}/types.py | 2 +- .../retrieval/agentic/discovery/__init__.py | 1 + .../{discovery_phase.py => discovery/phase.py} | 6 +++--- .../selection.py} | 6 +++--- .../{discovery_tools.py => discovery/tools.py} | 6 +++--- .../retrieval/agentic/evidence/__init__.py | 1 + .../agentic/{evidence.py => evidence/builder.py} | 6 +++--- .../{evidence_renderer.py => evidence/renderer.py} | 2 +- .../retrieval/agentic/navigation/__init__.py | 1 + .../{asset_tools.py => navigation/assets.py} | 0 .../document.py} | 10 +++++----- .../agentic/{ => navigation}/knowledge_map.py | 0 .../agentic/{ => navigation}/section_counts.py | 0 .../{ => navigation}/section_prompt_projection.py | 0 .../agentic/{ => navigation}/section_tree.py | 2 +- .../{ => navigation}/selection_hydration.py | 4 ++-- .../{navigation_tools.py => navigation/tools.py} | 12 ++++++------ .../services/retrieval/agentic/orchestrator.py | 14 +++++++------- .../shared/services/retrieval/agentic/policy.py | 4 ++-- .../shared/services/retrieval/agentic/tools.py | 12 +++++------- .../services/retrieval/workflow/orchestrator.py | 2 +- .../services/retrieval/workflow/plan_service.py | 2 +- .../shared/services/retrieval/workflow/planner.py | 2 +- .../services/retrieval/workflow/step_runner.py | 4 ++-- .../services/retrieval/workflow/synthesizer.py | 2 +- .../shared/services/retrieval/workflow/wallet.py | 2 +- 33 files changed, 66 insertions(+), 57 deletions(-) create mode 100644 packages/shared-python/shared/services/retrieval/agentic/core/__init__.py rename packages/shared-python/shared/services/retrieval/agentic/{ => core}/budget.py (100%) rename packages/shared-python/shared/services/retrieval/agentic/{ => core}/runtime.py (96%) rename packages/shared-python/shared/services/retrieval/agentic/{ => core}/trace.py (98%) rename packages/shared-python/shared/services/retrieval/agentic/{ => core}/types.py (99%) create mode 100644 packages/shared-python/shared/services/retrieval/agentic/discovery/__init__.py rename packages/shared-python/shared/services/retrieval/agentic/{discovery_phase.py => discovery/phase.py} (97%) rename packages/shared-python/shared/services/retrieval/agentic/{discovery_selection.py => discovery/selection.py} (96%) rename packages/shared-python/shared/services/retrieval/agentic/{discovery_tools.py => discovery/tools.py} (97%) create mode 100644 packages/shared-python/shared/services/retrieval/agentic/evidence/__init__.py rename packages/shared-python/shared/services/retrieval/agentic/{evidence.py => evidence/builder.py} (98%) rename packages/shared-python/shared/services/retrieval/agentic/{evidence_renderer.py => evidence/renderer.py} (99%) create mode 100644 packages/shared-python/shared/services/retrieval/agentic/navigation/__init__.py rename packages/shared-python/shared/services/retrieval/agentic/{asset_tools.py => navigation/assets.py} (100%) rename packages/shared-python/shared/services/retrieval/agentic/{document_navigation.py => navigation/document.py} (97%) rename packages/shared-python/shared/services/retrieval/agentic/{ => navigation}/knowledge_map.py (100%) rename packages/shared-python/shared/services/retrieval/agentic/{ => navigation}/section_counts.py (100%) rename packages/shared-python/shared/services/retrieval/agentic/{ => navigation}/section_prompt_projection.py (100%) rename packages/shared-python/shared/services/retrieval/agentic/{ => navigation}/section_tree.py (98%) rename packages/shared-python/shared/services/retrieval/agentic/{ => navigation}/selection_hydration.py (95%) rename packages/shared-python/shared/services/retrieval/agentic/{navigation_tools.py => navigation/tools.py} (92%) diff --git a/AGENTS.md b/AGENTS.md index f6c4c0e6f..d2fb2443d 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -540,7 +540,11 @@ Core retrieval internals are grouped by ownership: - `hydration/`: row/path/reference hydration, inline assets, and result assembly. - `graph/`: document graph publication/query support. - `stats/`: retrieval hit recording. -- `workflow/` and `agentic/`: agentic planning and navigation flows. +- `workflow/`: query planning, step execution, synthesis, and wallet state. +- `agentic/core/`: agentic run types, token budgets, runtime config, and traces. +- `agentic/discovery/`: bottom discovery and document selection. +- `agentic/navigation/`: section-tree navigation, selection hydration, and asset tools. +- `agentic/evidence/`: evidence tree rendering and budget trimming. ### Two Retrieval Modes diff --git a/CONTEXT.md b/CONTEXT.md index dd37639ec..0e5d8244b 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -301,7 +301,10 @@ exceptions. Shared Job lifecycle finalization lives under - `packages/shared-python/shared/services/retrieval/graph/*` - `packages/shared-python/shared/services/retrieval/stats/*` - `packages/shared-python/shared/services/retrieval/workflow/*` -- `packages/shared-python/shared/services/retrieval/agentic/*` +- `packages/shared-python/shared/services/retrieval/agentic/core/*` +- `packages/shared-python/shared/services/retrieval/agentic/discovery/*` +- `packages/shared-python/shared/services/retrieval/agentic/navigation/*` +- `packages/shared-python/shared/services/retrieval/agentic/evidence/*` ### Demo Source Materialization diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index 2ee929d60..898a476ae 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -10,7 +10,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from tests.support.contract_database import ContractDatabase -from shared.services.retrieval.agentic.types import AgenticResult +from shared.services.retrieval.agentic.core.types import AgenticResult from shared.services.retrieval.workflow.run_request import WorkflowRunRequest from shared.services.retrieval.workflow.types import PlannedStep, QueryPlan, WorkflowResult diff --git a/packages/shared-python/shared/services/retrieval/agentic/core/__init__.py b/packages/shared-python/shared/services/retrieval/agentic/core/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/core/__init__.py @@ -0,0 +1 @@ + diff --git a/packages/shared-python/shared/services/retrieval/agentic/budget.py b/packages/shared-python/shared/services/retrieval/agentic/core/budget.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/agentic/budget.py rename to packages/shared-python/shared/services/retrieval/agentic/core/budget.py diff --git a/packages/shared-python/shared/services/retrieval/agentic/runtime.py b/packages/shared-python/shared/services/retrieval/agentic/core/runtime.py similarity index 96% rename from packages/shared-python/shared/services/retrieval/agentic/runtime.py rename to packages/shared-python/shared/services/retrieval/agentic/core/runtime.py index d658cd562..7f71cf076 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/runtime.py +++ b/packages/shared-python/shared/services/retrieval/agentic/core/runtime.py @@ -9,8 +9,8 @@ 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.agentic.core.budget import BudgetExceeded, BudgetPoolName +from shared.services.retrieval.agentic.core.types import AgentRunConfig, AgentState from shared.services.retrieval.llm_adapter import LLMFn, current_llm_usage from shared.utils.token_estimate import estimate_tokens diff --git a/packages/shared-python/shared/services/retrieval/agentic/trace.py b/packages/shared-python/shared/services/retrieval/agentic/core/trace.py similarity index 98% rename from packages/shared-python/shared/services/retrieval/agentic/trace.py rename to packages/shared-python/shared/services/retrieval/agentic/core/trace.py index 87df0c94c..08b98e642 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/trace.py +++ b/packages/shared-python/shared/services/retrieval/agentic/core/trace.py @@ -15,7 +15,7 @@ from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession -from shared.services.retrieval.agentic.types import AgentRunConfig, ToolResult +from shared.services.retrieval.agentic.core.types import AgentRunConfig, ToolResult def _now_utc() -> datetime: diff --git a/packages/shared-python/shared/services/retrieval/agentic/types.py b/packages/shared-python/shared/services/retrieval/agentic/core/types.py similarity index 99% rename from packages/shared-python/shared/services/retrieval/agentic/types.py rename to packages/shared-python/shared/services/retrieval/agentic/core/types.py index 92c8b7f2f..79981b9f6 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/types.py +++ b/packages/shared-python/shared/services/retrieval/agentic/core/types.py @@ -10,7 +10,7 @@ from dataclasses import dataclass, field from typing import Any -from shared.services.retrieval.agentic.budget import BudgetLedger +from shared.services.retrieval.agentic.core.budget import BudgetLedger @dataclass diff --git a/packages/shared-python/shared/services/retrieval/agentic/discovery/__init__.py b/packages/shared-python/shared/services/retrieval/agentic/discovery/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/discovery/__init__.py @@ -0,0 +1 @@ + diff --git a/packages/shared-python/shared/services/retrieval/agentic/discovery_phase.py b/packages/shared-python/shared/services/retrieval/agentic/discovery/phase.py similarity index 97% rename from packages/shared-python/shared/services/retrieval/agentic/discovery_phase.py rename to packages/shared-python/shared/services/retrieval/agentic/discovery/phase.py index f244cadfa..5a95061ed 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/discovery_phase.py +++ b/packages/shared-python/shared/services/retrieval/agentic/discovery/phase.py @@ -9,9 +9,9 @@ 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.agentic.core.budget import BudgetExceeded +from shared.services.retrieval.agentic.core.trace import TraceRecorder +from shared.services.retrieval.agentic.core.types import AgentState, CandidateDoc, ToolResult from shared.services.retrieval.llm_adapter import LLMFn diff --git a/packages/shared-python/shared/services/retrieval/agentic/discovery_selection.py b/packages/shared-python/shared/services/retrieval/agentic/discovery/selection.py similarity index 96% rename from packages/shared-python/shared/services/retrieval/agentic/discovery_selection.py rename to packages/shared-python/shared/services/retrieval/agentic/discovery/selection.py index e47d9e4e3..55103b94b 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/discovery_selection.py +++ b/packages/shared-python/shared/services/retrieval/agentic/discovery/selection.py @@ -7,16 +7,16 @@ from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession -from shared.services.retrieval.agentic.budget import BudgetExceeded +from shared.services.retrieval.agentic.core.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 ( +from shared.services.retrieval.agentic.navigation.selection_hydration import ( hydrate_path_selections_into_node, ) -from shared.services.retrieval.agentic.types import DocTreeNode +from shared.services.retrieval.agentic.core.types import DocTreeNode from shared.services.retrieval.search.lexical_text import normalize_section_path from shared.services.retrieval.llm_adapter import LLMFn diff --git a/packages/shared-python/shared/services/retrieval/agentic/discovery_tools.py b/packages/shared-python/shared/services/retrieval/agentic/discovery/tools.py similarity index 97% rename from packages/shared-python/shared/services/retrieval/agentic/discovery_tools.py rename to packages/shared-python/shared/services/retrieval/agentic/discovery/tools.py index 8589bba00..e574d1b94 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/discovery_tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/discovery/tools.py @@ -14,14 +14,14 @@ 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.core.budget import BudgetExceeded +from shared.services.retrieval.agentic.navigation.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.agentic.core.types import ToolResult from shared.services.retrieval.search.channels import content_channel, path_channel, term_channel from shared.services.retrieval.llm_adapter import LLMFn from shared.services.retrieval.search.scoring import ( diff --git a/packages/shared-python/shared/services/retrieval/agentic/evidence/__init__.py b/packages/shared-python/shared/services/retrieval/agentic/evidence/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/evidence/__init__.py @@ -0,0 +1 @@ + diff --git a/packages/shared-python/shared/services/retrieval/agentic/evidence.py b/packages/shared-python/shared/services/retrieval/agentic/evidence/builder.py similarity index 98% rename from packages/shared-python/shared/services/retrieval/agentic/evidence.py rename to packages/shared-python/shared/services/retrieval/agentic/evidence/builder.py index 2f97f557c..6da147f42 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/evidence.py +++ b/packages/shared-python/shared/services/retrieval/agentic/evidence/builder.py @@ -7,8 +7,8 @@ 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.agentic.core.budget import BudgetLedger +from shared.services.retrieval.agentic.core.types import DocTreeNode from shared.services.retrieval.hydration.assets import build_retrieval_asset_url_map from shared.services.retrieval.stats.service import compute_importance_score from shared.utils.token_estimate import estimate_tokens @@ -165,7 +165,7 @@ async def render_evidence( ) -> str: del db - from shared.services.retrieval.agentic.evidence_renderer import render_unified_doc_tree + 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)) diff --git a/packages/shared-python/shared/services/retrieval/agentic/evidence_renderer.py b/packages/shared-python/shared/services/retrieval/agentic/evidence/renderer.py similarity index 99% rename from packages/shared-python/shared/services/retrieval/agentic/evidence_renderer.py rename to packages/shared-python/shared/services/retrieval/agentic/evidence/renderer.py index 22326d484..016e245d4 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/evidence_renderer.py +++ b/packages/shared-python/shared/services/retrieval/agentic/evidence/renderer.py @@ -3,7 +3,7 @@ from typing import Any, cast -from shared.services.retrieval.agentic.types import DocTreeNode +from shared.services.retrieval.agentic.core.types import DocTreeNode def render_unified_doc_tree( diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/__init__.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/__init__.py @@ -0,0 +1 @@ + diff --git a/packages/shared-python/shared/services/retrieval/agentic/asset_tools.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/assets.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/agentic/asset_tools.py rename to packages/shared-python/shared/services/retrieval/agentic/navigation/assets.py diff --git a/packages/shared-python/shared/services/retrieval/agentic/document_navigation.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/document.py similarity index 97% rename from packages/shared-python/shared/services/retrieval/agentic/document_navigation.py rename to packages/shared-python/shared/services/retrieval/agentic/navigation/document.py index 14b3d5f63..6720bd28c 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/document_navigation.py +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/document.py @@ -7,11 +7,11 @@ 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 ( +from shared.services.retrieval.agentic.core.budget import BudgetExceeded +from shared.services.retrieval.agentic.evidence.builder import reconcile_deferred_assets +from shared.services.retrieval.agentic.core.runtime import AgentLlmBudget +from shared.services.retrieval.agentic.core.trace import TraceRecorder +from shared.services.retrieval.agentic.core.types import ( AgentRunConfig, AgentState, CandidateDoc, diff --git a/packages/shared-python/shared/services/retrieval/agentic/knowledge_map.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/knowledge_map.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/agentic/knowledge_map.py rename to packages/shared-python/shared/services/retrieval/agentic/navigation/knowledge_map.py diff --git a/packages/shared-python/shared/services/retrieval/agentic/section_counts.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/section_counts.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/agentic/section_counts.py rename to packages/shared-python/shared/services/retrieval/agentic/navigation/section_counts.py diff --git a/packages/shared-python/shared/services/retrieval/agentic/section_prompt_projection.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/section_prompt_projection.py similarity index 100% rename from packages/shared-python/shared/services/retrieval/agentic/section_prompt_projection.py rename to packages/shared-python/shared/services/retrieval/agentic/navigation/section_prompt_projection.py diff --git a/packages/shared-python/shared/services/retrieval/agentic/section_tree.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/section_tree.py similarity index 98% rename from packages/shared-python/shared/services/retrieval/agentic/section_tree.py rename to packages/shared-python/shared/services/retrieval/agentic/navigation/section_tree.py index 3a285262b..7125fdf4f 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/section_tree.py +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/section_tree.py @@ -6,7 +6,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from shared.models.database.document import DocumentSection -from shared.services.retrieval.agentic.section_counts import attach_section_counts +from shared.services.retrieval.agentic.navigation.section_counts import attach_section_counts from shared.services.retrieval.search.lexical_text import normalize_section_path, split_section_path diff --git a/packages/shared-python/shared/services/retrieval/agentic/selection_hydration.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/selection_hydration.py similarity index 95% rename from packages/shared-python/shared/services/retrieval/agentic/selection_hydration.py rename to packages/shared-python/shared/services/retrieval/agentic/navigation/selection_hydration.py index 66d92ec2f..e75e32759 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/selection_hydration.py +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/selection_hydration.py @@ -4,8 +4,8 @@ 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.agentic.core.types import DocTreeNode +from shared.services.retrieval.agentic.navigation import assets as asset_tools from shared.services.retrieval.hydration.connected import hydrate_connected_target_rows from shared.services.retrieval.hydration.path import hydrate_paths_to_rows diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py similarity index 92% rename from packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py rename to packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py index 98f8fd424..79753a2f3 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation_tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py @@ -11,22 +11,22 @@ from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession -from shared.services.retrieval.agentic.asset_tools import ( +from shared.services.retrieval.agentic.navigation.assets import ( build_asset_tools_block, count_assets_under_scope, ) -from shared.services.retrieval.agentic.budget import BudgetExceeded +from shared.services.retrieval.agentic.core.budget import BudgetExceeded from shared.services.retrieval.agentic.prompts import ( ACTION_PROMPT, format_budget_block, parse_action_response, ) -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 ( +from shared.services.retrieval.agentic.navigation.section_prompt_projection import format_items_for_llm +from shared.services.retrieval.agentic.navigation.section_tree import load_child_sections +from shared.services.retrieval.agentic.navigation.selection_hydration import ( hydrate_path_selections_into_node, ) -from shared.services.retrieval.agentic.types import DocTreeNode +from shared.services.retrieval.agentic.core.types import DocTreeNode from shared.services.retrieval.llm_adapter import LLMFn diff --git a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py index 3dfa91a80..4ffdda5d8 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py +++ b/packages/shared-python/shared/services/retrieval/agentic/orchestrator.py @@ -23,27 +23,27 @@ from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession -from shared.services.retrieval.agentic.budget import BudgetExceeded, BudgetLedger -from shared.services.retrieval.agentic.discovery_phase import ( +from shared.services.retrieval.agentic.core.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 ( +from shared.services.retrieval.agentic.navigation.document import DocumentNavigationRunner +from shared.services.retrieval.agentic.evidence.builder import ( build_asset_url_map as _build_asset_url_map, collect_media_chunks_all as _collect_media_chunks_all, 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 ( +from shared.services.retrieval.agentic.core.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 ( +from shared.services.retrieval.agentic.core.trace import TraceRecorder +from shared.services.retrieval.agentic.core.types import ( AgentRunConfig, AgentState, AgenticResult, diff --git a/packages/shared-python/shared/services/retrieval/agentic/policy.py b/packages/shared-python/shared/services/retrieval/agentic/policy.py index 25323580a..6517fc861 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/policy.py +++ b/packages/shared-python/shared/services/retrieval/agentic/policy.py @@ -18,8 +18,8 @@ from loguru import logger -from shared.services.retrieval.agentic.budget import BudgetExceeded -from shared.services.retrieval.agentic.types import AgentRunConfig, AgentState +from shared.services.retrieval.agentic.core.budget import BudgetExceeded +from shared.services.retrieval.agentic.core.types import AgentRunConfig, AgentState from shared.services.retrieval.llm_adapter import LLMFn from shared.utils.token_estimate import estimate_tokens diff --git a/packages/shared-python/shared/services/retrieval/agentic/tools.py b/packages/shared-python/shared/services/retrieval/agentic/tools.py index 4fde17d1d..bce60a71b 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/tools.py @@ -9,13 +9,11 @@ from sqlalchemy.ext.asyncio import AsyncSession -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.agentic.core.types import DocTreeNode, ToolResult +from shared.services.retrieval.agentic.discovery import selection as discovery_selection +from shared.services.retrieval.agentic.discovery import tools as discovery_tools +from shared.services.retrieval.agentic.navigation import assets as asset_tools +from shared.services.retrieval.agentic.navigation import tools as navigation_tools from shared.services.retrieval.llm_adapter import LLMFn diff --git a/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py b/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py index e830ee5fe..3fec8dbda 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py +++ b/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py @@ -10,7 +10,7 @@ from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession -from shared.services.retrieval.agentic.budget import BudgetLedger +from shared.services.retrieval.agentic.core.budget import BudgetLedger from shared.services.retrieval.agentic.orchestrator import _load_budget_inventory from shared.services.retrieval.llm_adapter import ( create_retrieval_llm_fn, diff --git a/packages/shared-python/shared/services/retrieval/workflow/plan_service.py b/packages/shared-python/shared/services/retrieval/workflow/plan_service.py index a5e3a0d43..11e7897f3 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/plan_service.py +++ b/packages/shared-python/shared/services/retrieval/workflow/plan_service.py @@ -3,7 +3,7 @@ from loguru import logger -from shared.services.retrieval.agentic.budget import BudgetLedger +from shared.services.retrieval.agentic.core.budget import BudgetLedger from shared.services.retrieval.cache_service import ( get_cached_workflow_plan, set_cached_workflow_plan, diff --git a/packages/shared-python/shared/services/retrieval/workflow/planner.py b/packages/shared-python/shared/services/retrieval/workflow/planner.py index e34713db6..7f4937b80 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/planner.py +++ b/packages/shared-python/shared/services/retrieval/workflow/planner.py @@ -8,7 +8,7 @@ from loguru import logger -from shared.services.retrieval.agentic.budget import BudgetExceeded, BudgetLedger +from shared.services.retrieval.agentic.core.budget import BudgetExceeded, BudgetLedger from shared.services.retrieval.llm_adapter import LLMFn, current_llm_usage from shared.services.retrieval.workflow.types import FinalStrategy, OutputRole, PlannedStep, QueryPlan, StepKind from shared.utils.token_estimate import estimate_tokens diff --git a/packages/shared-python/shared/services/retrieval/workflow/step_runner.py b/packages/shared-python/shared/services/retrieval/workflow/step_runner.py index 772b334d3..18a2e7059 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/step_runner.py +++ b/packages/shared-python/shared/services/retrieval/workflow/step_runner.py @@ -8,9 +8,9 @@ from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession -from shared.services.retrieval.agentic.budget import BudgetLedger +from shared.services.retrieval.agentic.core.budget import BudgetLedger from shared.services.retrieval.agentic.orchestrator import RetrievalAgent -from shared.services.retrieval.agentic.types import AgenticResult +from shared.services.retrieval.agentic.core.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.run_request import WorkflowStepRequest diff --git a/packages/shared-python/shared/services/retrieval/workflow/synthesizer.py b/packages/shared-python/shared/services/retrieval/workflow/synthesizer.py index 89d4ec031..a81bb517d 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/synthesizer.py +++ b/packages/shared-python/shared/services/retrieval/workflow/synthesizer.py @@ -3,7 +3,7 @@ import re -from shared.services.retrieval.agentic.budget import BudgetLedger +from shared.services.retrieval.agentic.core.budget import BudgetLedger from shared.services.retrieval.llm_adapter import LLMFn, current_llm_usage from shared.services.retrieval.workflow.types import PlannedStep, QueryPlan, StepResult from shared.utils.token_estimate import estimate_tokens diff --git a/packages/shared-python/shared/services/retrieval/workflow/wallet.py b/packages/shared-python/shared/services/retrieval/workflow/wallet.py index 0a540a120..624f03f7c 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/wallet.py +++ b/packages/shared-python/shared/services/retrieval/workflow/wallet.py @@ -5,7 +5,7 @@ import os from dataclasses import dataclass, field -from shared.services.retrieval.agentic.budget import BudgetLedger +from shared.services.retrieval.agentic.core.budget import BudgetLedger from shared.services.retrieval.workflow.types import PlannedStep, QueryPlan From d1cc76f65aa1d548c22dd5003f35d07285c02825 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 13:04:00 +0000 Subject: [PATCH 71/78] refactor: split current user dependency from job admission --- CONTEXT.md | 6 +++ apps/api/app/api/dependencies/current_user.py | 25 ++++++++++ .../api/app/api/dependencies/job_admission.py | 48 ++----------------- .../app/api/dependencies/route_admission.py | 36 ++++++++++++++ apps/api/app/api/v1/routes/api_key.py | 4 +- apps/api/app/api/v1/routes/billing.py | 2 +- apps/api/app/api/v1/routes/demo.py | 2 +- apps/api/app/api/v1/routes/documents.py | 2 +- apps/api/app/api/v1/routes/jobs.py | 6 +-- apps/api/app/api/v1/routes/retrieval.py | 2 +- apps/api/app/api/v1/routes/webhook.py | 2 +- apps/api/app/api/v1/routes/webhook_secrets.py | 2 +- 12 files changed, 79 insertions(+), 58 deletions(-) create mode 100644 apps/api/app/api/dependencies/current_user.py create mode 100644 apps/api/app/api/dependencies/route_admission.py diff --git a/CONTEXT.md b/CONTEXT.md index 0e5d8244b..43f7a0e7d 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -281,10 +281,16 @@ exceptions. Shared Job lifecycle finalization lives under ### Job Admission - `app/api/dependencies/auth.py` +- `app/api/dependencies/current_user.py` +- `app/api/dependencies/route_admission.py` - `app/api/dependencies/job_admission.py` - `app/services/auth/*` - `app/services/rate_limit/*` +`auth.py`, `current_user.py`, and `route_admission.py` are HTTP dependency +adapters. `job_admission.py` owns only the route-level billing and system-limit +admission dependencies. + ### Document Lifecycle - `app/api/v1/routes/documents.py` diff --git a/apps/api/app/api/dependencies/current_user.py b/apps/api/app/api/dependencies/current_user.py new file mode 100644 index 000000000..4c5593c1a --- /dev/null +++ b/apps/api/app/api/dependencies/current_user.py @@ -0,0 +1,25 @@ +"""FastAPI current-user dependency adapters.""" + +from typing import AsyncGenerator + +from app.api.dependencies.auth import get_current_user_id +from app.api.dependencies.route_admission import get_route_admission_context +from app.services.rate_limit.data_structures import ( + CurrentUser, + RouteAdmissionContext, +) +from app.services.rate_limit.job_admission_service import JobAdmissionService +from fastapi import Depends + +_job_admission_service = JobAdmissionService() + + +async def with_current_user( + route_context: RouteAdmissionContext = Depends(get_route_admission_context), + user_id: str = Depends(get_current_user_id), +) -> AsyncGenerator[CurrentUser, None]: + current_user = await _job_admission_service.resolve_current_user( + route_context=route_context, + user_id=user_id, + ) + yield current_user diff --git a/apps/api/app/api/dependencies/job_admission.py b/apps/api/app/api/dependencies/job_admission.py index 768ccc872..ab87c1dc0 100644 --- a/apps/api/app/api/dependencies/job_admission.py +++ b/apps/api/app/api/dependencies/job_admission.py @@ -2,37 +2,18 @@ from typing import AsyncGenerator -from app.api.dependencies.auth import get_current_user_id +from app.api.dependencies.current_user import with_current_user +from app.api.dependencies.route_admission import get_route_admission_context from app.services.rate_limit.data_structures import ( CurrentUser, RouteAdmissionContext, ) from app.services.rate_limit.job_admission_service import JobAdmissionService -from fastapi import Depends, Request +from fastapi import Depends _job_admission_service = JobAdmissionService() -def get_route_admission_context(request: Request) -> RouteAdmissionContext: - """Extract route facts needed by Job Admission from a FastAPI request.""" - return RouteAdmissionContext( - method=request.method, - path=_get_route_path(request), - limit_identifier=_get_route_limit_identifier(request), - ) - - -async def with_current_user( - route_context: RouteAdmissionContext = Depends(get_route_admission_context), - user_id: str = Depends(get_current_user_id), -) -> AsyncGenerator[CurrentUser, None]: - current_user = await _job_admission_service.resolve_current_user( - route_context=route_context, - user_id=user_id, - ) - yield current_user - - async def require_billing_limits( current_user: CurrentUser = Depends(with_current_user), ) -> AsyncGenerator[CurrentUser, None]: @@ -46,26 +27,3 @@ async def require_route_system_limit( await _job_admission_service.enforce_route_system_limit( route_context=route_context, ) - - -def _get_route_path(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(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 _get_route_path(request) diff --git a/apps/api/app/api/dependencies/route_admission.py b/apps/api/app/api/dependencies/route_admission.py new file mode 100644 index 000000000..08457ba09 --- /dev/null +++ b/apps/api/app/api/dependencies/route_admission.py @@ -0,0 +1,36 @@ +"""FastAPI route-fact adapters for the Job Admission workflow.""" + +from app.services.rate_limit.data_structures import RouteAdmissionContext +from fastapi import Request + + +def get_route_admission_context(request: Request) -> RouteAdmissionContext: + """Extract route facts needed by Job Admission from a FastAPI request.""" + return RouteAdmissionContext( + method=request.method, + path=_get_route_path(request), + limit_identifier=_get_route_limit_identifier(request), + ) + + +def _get_route_path(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(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 _get_route_path(request) diff --git a/apps/api/app/api/v1/routes/api_key.py b/apps/api/app/api/v1/routes/api_key.py index 85a3f6836..e42fb3dcf 100644 --- a/apps/api/app/api/v1/routes/api_key.py +++ b/apps/api/app/api/v1/routes/api_key.py @@ -3,9 +3,7 @@ """ from app.services.auth.api_key_management_service import APIKeyManagementService -from app.api.dependencies.job_admission import ( - with_current_user, -) +from app.api.dependencies.current_user import with_current_user from app.services.rate_limit.data_structures import CurrentUser from fastapi import APIRouter, Depends from sqlalchemy.ext.asyncio import AsyncSession diff --git a/apps/api/app/api/v1/routes/billing.py b/apps/api/app/api/v1/routes/billing.py index 45388fd9e..94379c33c 100644 --- a/apps/api/app/api/v1/routes/billing.py +++ b/apps/api/app/api/v1/routes/billing.py @@ -6,7 +6,7 @@ BillingWorkflowService, ParseUsageResponse, ) -from app.api.dependencies.job_admission import with_current_user +from app.api.dependencies.current_user import with_current_user from app.services.rate_limit.data_structures import CurrentUser from fastapi import APIRouter, Depends, Query, Request from sqlalchemy.ext.asyncio import AsyncSession diff --git a/apps/api/app/api/v1/routes/demo.py b/apps/api/app/api/v1/routes/demo.py index 320d9b232..c8559efc7 100644 --- a/apps/api/app/api/v1/routes/demo.py +++ b/apps/api/app/api/v1/routes/demo.py @@ -5,7 +5,7 @@ from typing import Any from app.services.demo.document_service import DemoDocumentService -from app.api.dependencies.job_admission import with_current_user +from app.api.dependencies.current_user import with_current_user from app.services.rate_limit.data_structures import CurrentUser from fastapi import APIRouter, Depends, Query from fastapi.responses import FileResponse diff --git a/apps/api/app/api/v1/routes/documents.py b/apps/api/app/api/v1/routes/documents.py index b13a7cb4d..1cbcf7ad9 100644 --- a/apps/api/app/api/v1/routes/documents.py +++ b/apps/api/app/api/v1/routes/documents.py @@ -4,7 +4,7 @@ from typing import Literal -from app.api.dependencies.job_admission import with_current_user +from app.api.dependencies.current_user import with_current_user from app.services.documents.lifecycle_service import DocumentService from app.services.rate_limit.data_structures import CurrentUser from fastapi import APIRouter, Depends, Query diff --git a/apps/api/app/api/v1/routes/jobs.py b/apps/api/app/api/v1/routes/jobs.py index d6828eb31..ce9c486f7 100644 --- a/apps/api/app/api/v1/routes/jobs.py +++ b/apps/api/app/api/v1/routes/jobs.py @@ -12,10 +12,8 @@ get_job_result_for_user, list_jobs_for_user, ) -from app.api.dependencies.job_admission import ( - require_billing_limits, - with_current_user, -) +from app.api.dependencies.current_user import with_current_user +from app.api.dependencies.job_admission import require_billing_limits from app.services.rate_limit.data_structures import CurrentUser from fastapi import APIRouter, Depends, Query from sqlalchemy.ext.asyncio import AsyncSession diff --git a/apps/api/app/api/v1/routes/retrieval.py b/apps/api/app/api/v1/routes/retrieval.py index 9e0df0a9b..61037c4b3 100644 --- a/apps/api/app/api/v1/routes/retrieval.py +++ b/apps/api/app/api/v1/routes/retrieval.py @@ -4,7 +4,7 @@ from typing import Literal -from app.api.dependencies.job_admission import with_current_user +from app.api.dependencies.current_user import with_current_user from app.services.rate_limit.data_structures import CurrentUser from fastapi import APIRouter, Depends from pydantic import BaseModel, Field, field_validator diff --git a/apps/api/app/api/v1/routes/webhook.py b/apps/api/app/api/v1/routes/webhook.py index 053079119..19eb84f30 100644 --- a/apps/api/app/api/v1/routes/webhook.py +++ b/apps/api/app/api/v1/routes/webhook.py @@ -9,7 +9,7 @@ from app.repositories.job_repository import JobRepository from app.repositories.webhook_repository import WebhookRepository -from app.api.dependencies.job_admission import with_current_user +from app.api.dependencies.current_user import with_current_user from app.services.rate_limit.data_structures import CurrentUser from fastapi import APIRouter, Depends, Query from loguru import logger diff --git a/apps/api/app/api/v1/routes/webhook_secrets.py b/apps/api/app/api/v1/routes/webhook_secrets.py index 1f3a28708..0fb4f9bf2 100644 --- a/apps/api/app/api/v1/routes/webhook_secrets.py +++ b/apps/api/app/api/v1/routes/webhook_secrets.py @@ -6,7 +6,7 @@ from typing import List, Optional -from app.api.dependencies.job_admission import with_current_user +from app.api.dependencies.current_user import with_current_user from app.services.rate_limit.data_structures import CurrentUser from fastapi import APIRouter, Depends from pydantic import BaseModel, Field From cf8e0296a56dc59b59febdb7778ac2b4f6b97156 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 13:50:25 +0000 Subject: [PATCH 72/78] refactor: rename document ingestion worker dispatch --- AGENTS.md | 18 +- CONTEXT.md | 7 + apps/api/app/mcp/retrieval_server.py | 8 +- .../document_ingestion/creation_service.py | 15 +- .../document_ingestion/handoff_service.py | 69 ++++--- .../document_ingestion/worker_dispatcher.py | 82 ++++++++ apps/api/app/services/knowledge/__init__.py | 7 - .../app/services/knowledge/kb_orchestrator.py | 187 ------------------ apps/api/main.py | 9 +- .../tests/contract/test_documents_contract.py | 2 +- .../contract/test_job_creation_contract.py | 44 ++--- .../tests/contract/test_s3_event_contract.py | 101 +++++++--- .../tests/migrations/test_schema_contract.py | 2 +- apps/api/tests/support/contract_database.py | 2 +- apps/worker/app/core/tasks/__init__.py | 2 +- apps/worker/app/core/tasks/base_task.py | 10 +- ...b_tasks.py => document_ingestion_tasks.py} | 92 +++++++-- apps/worker/app/core/worker_bootstrap.py | 6 +- apps/worker/app/services/__init__.py | 2 - .../document_ingestion/job_state_gate.py | 4 +- .../document_ingestion/processing_run.py | 4 +- .../services/document_ingestion/workspace.py | 5 +- apps/worker/tests/contract/README.md | 4 +- .../contract/test_parse_task_contract.py | 48 ++--- .../contract/test_url_upload_contract.py | 76 +++++-- .../test_worker_bootstrap_contract.py | 47 ++++- .../worker/tests/support/contract_database.py | 2 +- .../shared-python/shared/core/celery_app.py | 11 +- .../shared/core/celery_router.py | 34 +++- .../shared-python/shared/core/config/base.py | 2 +- .../shared/core/config/celery.py | 14 +- .../shared/core/tasks/__init__.py | 4 +- .../shared/models/database/job.py | 2 +- .../shared/services/redis/distributed_lock.py | 2 +- 34 files changed, 527 insertions(+), 397 deletions(-) create mode 100644 apps/api/app/services/document_ingestion/worker_dispatcher.py delete mode 100644 apps/api/app/services/knowledge/__init__.py delete mode 100644 apps/api/app/services/knowledge/kb_orchestrator.py rename apps/worker/app/core/tasks/{kb_tasks.py => document_ingestion_tasks.py} (52%) diff --git a/AGENTS.md b/AGENTS.md index d2fb2443d..7017e6e82 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -33,7 +33,7 @@ knowhereapi-main/ │ ├── api/ # FastAPI REST API (port 5005) │ │ ├── app/ │ │ │ ├── api/v1/routes/ # Endpoint handlers -│ │ │ ├── services/ # Business logic (auth, knowledge, billing) +│ │ │ ├── services/ # Business logic (auth, ingestion, billing) │ │ │ └── repositories/ # Data access layer │ │ └── main.py # Entrypoint, runs migrations on start │ ├── worker/ # Celery worker for async document processing @@ -214,17 +214,17 @@ flowchart LR --- -## Persisted Knowledge Base Schema (On-Disk Output) +## Persisted Document Corpus Schema (On-Disk Output) -After parsing and chunk conversion, results are persisted to `~/.knowhere/{kb_name}/`. +After parsing and chunk conversion, results are persisted to `~/.knowhere/{corpus_name}/`. This on-disk structure is the **authoritative persisted format** — the intermediate DataFrame is an internal detail. Below is the complete schema. -### KB-Level Directory Layout +### Corpus-Level Directory Layout ```text -~/.knowhere/{kb_name}/ -├── knowledge_graph.json # KB-wide graph: file metadata + cross-doc edges +~/.knowhere/{corpus_name}/ +├── knowledge_graph.json # corpus-wide graph: file metadata + cross-doc edges ├── chunk_stats.json # Per-chunk retrieval hit analytics {chunk_id → stats} ├── {source_file_name}/ # One directory per ingested document │ ├── chunks.json # All parsed chunks for this document @@ -239,12 +239,12 @@ DataFrame is an internal detail. Below is the complete schema. │ └── toc_hierarchies.json # Debug: extracted TOC structure (DOCX only) ``` -### `knowledge_graph.json` — KB-Wide Graph +### `knowledge_graph.json` — Corpus-Wide Graph ```json { "version": "2.0", - "kb_id": "test_kb", + "corpus_id": "test-corpus", "stats": { "total_files": 3, "total_chunks": 364, "total_cross_file_edges": 0 }, "files": { "AI_Security_Report.docx": { @@ -605,7 +605,7 @@ Unlike legacy retrieval which relied on static `hydrate_mode` tags, hydration is 4. Cleans asset path references from content 5. Attaches citation: `{document_id, chunk_id, source_file_name, section_path}` -### Small KB Optimization +### Small Corpus Optimization When `total_chunks <= top_k`, skips the full pipeline and returns all chunks directly (router: `small_kb_all`). diff --git a/CONTEXT.md b/CONTEXT.md index 43f7a0e7d..24d47b7bf 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -64,6 +64,11 @@ 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. +### Document Ingestion Worker Dispatch + +The API-side handoff that advances an uploaded file Job to pending state and +enqueues the worker parse task with the user-aware Celery queue policy. + ### Worker Document Parsing The worker-side workflow that turns a source file into parsed DataFrame rows, @@ -268,7 +273,9 @@ exceptions. Shared Job lifecycle finalization lives under - `app/services/document_ingestion/service.py` - `app/services/document_ingestion/creation_service.py` - `app/services/document_ingestion/confirmation_service.py` +- `app/services/document_ingestion/handoff_service.py` - `app/services/document_ingestion/scope_service.py` +- `app/services/document_ingestion/worker_dispatcher.py` - `app/repositories/job_repository.py` ### Job Read diff --git a/apps/api/app/mcp/retrieval_server.py b/apps/api/app/mcp/retrieval_server.py index 2b0cbcbcc..b3f33a0c0 100644 --- a/apps/api/app/mcp/retrieval_server.py +++ b/apps/api/app/mcp/retrieval_server.py @@ -97,7 +97,7 @@ def create_retrieval_mcp_server( server = FastMCP( "knowhere-retrieval", instructions=( - "Use this server to search knowledge. " + "Use this server to search published documents. " "If you need information before answering, try searching with this tool." ), streamable_http_path=streamable_http_path, @@ -106,10 +106,10 @@ def create_retrieval_mcp_server( ) @server.tool( - name="kb.query", - description="Search for information and return relevant knowledge snippets.", + name="retrieval.query", + description="Search published documents and return relevant snippets.", ) - async def kb_query( + async def query_documents( query: Annotated[str, Field(description="What you want to search for.")], top_k: Annotated[ int, Field(description="Maximum number of results to return.") diff --git a/apps/api/app/services/document_ingestion/creation_service.py b/apps/api/app/services/document_ingestion/creation_service.py index acecb1974..9ca6d2ed6 100644 --- a/apps/api/app/services/document_ingestion/creation_service.py +++ b/apps/api/app/services/document_ingestion/creation_service.py @@ -31,7 +31,10 @@ from shared.services.storage.file_upload_service import FileUploadService from shared.services.http.url_file_type import resolve_file_extension_async -_JOB_TYPE_KB_MANAGEMENT = "kb_management" +_DOCUMENT_INGESTION_JOB_TYPE = "document_ingestion" +_URL_UPLOAD_TASK_NAME = ( + "app.core.tasks.document_ingestion_tasks.upload_url_file_task" +) JobMetadata = dict[str, object] UploadHeaders = dict[str, str] @@ -95,7 +98,7 @@ async def _create_waiting_job( db=db, job_id=job_id, user_id=user_id, - job_type=_JOB_TYPE_KB_MANAGEMENT, + job_type=_DOCUMENT_INGESTION_JOB_TYPE, source_type=source_type, file_path=None, webhook_url=webhook_url, @@ -134,7 +137,7 @@ async def _cache_job_creation_state( "s3_key": s3_key, "user_id": user_id, "webhook_enabled": webhook_enabled, - "job_type": _JOB_TYPE_KB_MANAGEMENT, + "job_type": _DOCUMENT_INGESTION_JOB_TYPE, "source_type": source_type, "created_at": datetime.now(timezone.utc).isoformat(), } @@ -313,10 +316,8 @@ 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 = celery_app.signature(_URL_UPLOAD_TASK_NAME) upload_url_file_task.apply_async( args=[job_id, source_url, user_id], - kwargs={"job_type": _JOB_TYPE_KB_MANAGEMENT}, + kwargs={"job_type": _DOCUMENT_INGESTION_JOB_TYPE}, ) diff --git a/apps/api/app/services/document_ingestion/handoff_service.py b/apps/api/app/services/document_ingestion/handoff_service.py index 68b1cd94a..c8de180bd 100644 --- a/apps/api/app/services/document_ingestion/handoff_service.py +++ b/apps/api/app/services/document_ingestion/handoff_service.py @@ -1,8 +1,10 @@ from __future__ import annotations -from typing import Any +from typing import Protocol -from app.services.knowledge.kb_orchestrator import KBOrchestrator +from app.services.document_ingestion.worker_dispatcher import ( + DocumentIngestionWorkerDispatcher, +) from loguru import logger from sqlalchemy.ext.asyncio import AsyncSession @@ -14,7 +16,19 @@ from shared.core.state_machine.service import AsyncStateMachineService from shared.core.state_machine.states import JobStatus -_JOB_TYPE_KB_MANAGEMENT = "kb_management" +_DOCUMENT_PARSE_JOB_TYPE = "document_ingestion" +_LEGACY_DOCUMENT_PARSE_JOB_TYPE = "kb_management" +_SUPPORTED_DOCUMENT_PARSE_JOB_TYPES = frozenset( + { + _DOCUMENT_PARSE_JOB_TYPE, + _LEGACY_DOCUMENT_PARSE_JOB_TYPE, + } +) + + +class _UploadedFileJob(Protocol): + job_id: str + job_type: str class DocumentIngestionHandoffService: @@ -24,19 +38,34 @@ def __init__( self, *, state_machine: AsyncStateMachineService | None = None, - orchestrator: KBOrchestrator | None = None, + worker_dispatcher: DocumentIngestionWorkerDispatcher | None = None, ) -> None: self._state_machine = state_machine or AsyncStateMachineService() - self._orchestrator = orchestrator or KBOrchestrator() + self._worker_dispatcher = ( + worker_dispatcher or DocumentIngestionWorkerDispatcher() + ) async def start_uploaded_file_workflow( self, db: AsyncSession, *, - job: Any, + job: _UploadedFileJob, user_id: str, trigger: str, ) -> None: + if job.job_type not in _SUPPORTED_DOCUMENT_PARSE_JOB_TYPES: + raise ValidationException( + user_message="Unsupported job type", + violations=[ + { + "field": "job_type", + "description": ( + f"Job type '{job.job_type}' is not supported" + ), + } + ], + ) + outcome = await self._state_machine.transition_outcome( db, job.job_id, @@ -55,33 +84,21 @@ async def start_uploaded_file_workflow( f"Could not advance uploaded job {job.job_id} to pending: " f"{outcome.reason}" ), - retry_after=settings.KB_TASK_RETRY_COUNTDOWN, + retry_after=settings.DOCUMENT_INGESTION_TASK_RETRY_COUNTDOWN, user_message="Job state is still settling. Retrying shortly.", ) - 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, + await self._worker_dispatcher.start_uploaded_file_parse( 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: + async def mark_upload_expired( + self, + db: AsyncSession, + *, + job: _UploadedFileJob, + ) -> None: outcome = await self._state_machine.mark_failed_outcome( db, job.job_id, diff --git a/apps/api/app/services/document_ingestion/worker_dispatcher.py b/apps/api/app/services/document_ingestion/worker_dispatcher.py new file mode 100644 index 000000000..e340c8647 --- /dev/null +++ b/apps/api/app/services/document_ingestion/worker_dispatcher.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +from celery import signature +from celery.canvas import Signature +from loguru import logger + +from shared.core.celery_router import CeleryTaskRouter, task_router +from shared.core.exceptions.domain_exceptions import WorkerHandlingException + +_DOCUMENT_PARSE_JOB_TYPE = "document_ingestion" +_DOCUMENT_PARSE_TASK_NAME = "app.core.tasks.document_ingestion_tasks.parse_task" + + +class DocumentIngestionWorkerDispatcher: + """Dispatch uploaded Document Ingestion Jobs to worker parsing.""" + + def __init__( + self, + *, + celery_task_router: CeleryTaskRouter | None = None, + ) -> None: + self._task_router = celery_task_router or task_router + + async def start_uploaded_file_parse(self, *, job_id: str, user_id: str) -> str: + task_signature = self._build_uploaded_file_parse_signature( + job_id=job_id, + user_id=user_id, + ) + result = task_signature.apply_async() + if result is None or result.id is None: + raise WorkerHandlingException( + internal_message=( + "Failed to start Document Ingestion worker parse: " + "missing Celery task id" + ) + ) + + task_id = str(result.id) + signature_options = task_signature.options or {} + queue_name = signature_options.get("queue") + logger.info( + "Document Ingestion worker parse started: " + f"job_id={job_id}, task_id={task_id}, queue={queue_name}" + ) + return task_id + + def _build_uploaded_file_parse_signature( + self, + *, + job_id: str, + user_id: str, + ) -> Signature: + queue_name = self._task_router.get_queue_for_job( + _DOCUMENT_PARSE_JOB_TYPE, + user_id, + ) + task_kwargs: dict[str, str] = { + "user_id": user_id, + "job_type": _DOCUMENT_PARSE_JOB_TYPE, + } + task_signature = signature( + _DOCUMENT_PARSE_TASK_NAME, + args=[job_id], + kwargs=task_kwargs, + ) + if task_signature is None: + raise WorkerHandlingException( + internal_message=( + "Failed to build Document Ingestion worker parse: " + "missing Celery signature" + ) + ) + + configured_signature = task_signature.set(queue=queue_name) + if configured_signature is None: + raise WorkerHandlingException( + internal_message=( + "Failed to configure Document Ingestion worker parse queue" + ) + ) + + return configured_signature diff --git a/apps/api/app/services/knowledge/__init__.py b/apps/api/app/services/knowledge/__init__.py deleted file mode 100644 index f95391f1e..000000000 --- a/apps/api/app/services/knowledge/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -""" -Knowledge-base services. -""" - -from .kb_orchestrator import KBOrchestrator - -__all__ = ["KBOrchestrator"] diff --git a/apps/api/app/services/knowledge/kb_orchestrator.py b/apps/api/app/services/knowledge/kb_orchestrator.py deleted file mode 100644 index b71b9af20..000000000 --- a/apps/api/app/services/knowledge/kb_orchestrator.py +++ /dev/null @@ -1,187 +0,0 @@ -""" -Knowledge-base workflow orchestration. -""" - -from typing import Optional - -from loguru import logger - -from shared.core.celery_router import task_router - -# Tasks now run in the Worker service and are referenced by task name. -from shared.core.exceptions.domain_exceptions import ( - KnowhereException, - WorkerHandlingException, -) - - -class KBOrchestrator: - """Coordinate knowledge-base processing jobs.""" - - def __init__(self): - self.task_router = task_router - - async def start_workflow( - self, - db, - job_id: str, - source_type: str, - file_path: Optional[str] = None, - file_url: Optional[str] = None, - user_id: Optional[str] = None, - ) -> str: - """ - Start the knowledge-base workflow. - - Args: - db: Database session. - job_id: Job identifier. - source_type: Source type. - file_path: Uploaded file path when direct upload is used. - file_url: Source URL when URL ingestion is used. - user_id: User identifier. - - Returns: - str: Celery task identifier. - """ - try: - # When the source is a URL, recover file_url from job metadata if needed. - if source_type == "url" and not file_url: - from app.repositories.job_repository import JobRepository - - from shared.models.schemas.job_metadata import JobMetadataHelper - from shared.services.redis import RedisServiceFactory - - job_repo = JobRepository() - redis_service = RedisServiceFactory.get_service() - job_metadata = await job_repo.get_job_metadata( - db, job_id, redis_service - ) - file_url = JobMetadataHelper.get_field(job_metadata, "file_url") - - # Resolve the queue name for this job. - effective_user_id = user_id or "" - queue_name = self.task_router.get_queue_for_job( - "kb_management", effective_user_id - ) - task_kwargs = { - "user_id": effective_user_id, - "job_type": "kb_management", - } - - # Start the single worker task. Upload is already complete via S3. - # The task handles parsing, vectorization, ZIP generation, S3 upload, - # and result publication. Webhook and email delivery stay in the API. - from celery import signature - - task_signature = signature( - "app.core.tasks.kb_tasks.parse_task", - args=[job_id], - kwargs=task_kwargs, - ) - if task_signature is None: - raise WorkerHandlingException( - internal_message=( - "Failed to build knowledge-base workflow: missing Celery signature" - ) - ) - task_signature = task_signature.set(queue=queue_name) - if task_signature is None: - raise WorkerHandlingException( - internal_message=( - "Failed to configure knowledge-base workflow queue" - ) - ) - - # Enqueue the task. - result = task_signature.apply_async() - if result is None or result.id is None: - raise WorkerHandlingException( - internal_message=( - "Failed to start knowledge-base workflow: missing Celery task id" - ) - ) - - logger.info( - f"Knowledge-base workflow started: job_id={job_id}, task_id={result.id}, queue={queue_name}" - ) - - return result.id - - except KnowhereException: - raise - except Exception as e: - logger.error(f"Failed to start knowledge-base workflow: {e}") - raise WorkerHandlingException( - internal_message=f"Failed to start knowledge-base workflow: {str(e)}", - original_exception=e, - ) - - def create_workflow_chain( - self, job_id: str, user_id: str, queue_name: Optional[str] = None - ): - """ - Create the workflow signature for tests or manual execution. - - Args: - job_id: Job identifier. - user_id: User identifier. - queue_name: Optional explicit queue name. - - Returns: - signature: Celery task signature. - """ - if not queue_name: - queue_name = self.task_router.get_queue_for_job("kb_management", user_id) - task_kwargs = { - "user_id": user_id, - "job_type": "kb_management", - } - - from celery import signature - - # Return the single-task signature. Parsing, vectorization, ZIP generation, - # S3 upload, and publication happen in the worker. Webhook and email - # delivery remain in the API service. - task_signature = signature( - "app.core.tasks.kb_tasks.parse_task", - args=[job_id], - kwargs=task_kwargs, - ) - if task_signature is None: - raise WorkerHandlingException( - internal_message=( - "Failed to build knowledge-base workflow: missing Celery signature" - ) - ) - configured_signature = task_signature.set(queue=queue_name) - if configured_signature is None: - raise WorkerHandlingException( - internal_message="Failed to configure knowledge-base workflow queue" - ) - return configured_signature - - def cancel_workflow(self, workflow_id: str) -> bool: - """ - Cancel a workflow. - - Args: - workflow_id: Workflow identifier. - - Returns: - bool: Whether cancellation succeeded. - """ - try: - from shared.core.celery_app import get_celery_app - - celery_app = get_celery_app() - - result = celery_app.AsyncResult(workflow_id) - result.revoke(terminate=True) - - logger.info(f"Workflow cancelled: {workflow_id}") - return True - - except Exception as e: - logger.error(f"Failed to cancel workflow: {e}") - return False diff --git a/apps/api/main.py b/apps/api/main.py index 22b058c67..db75099bf 100644 --- a/apps/api/main.py +++ b/apps/api/main.py @@ -71,7 +71,7 @@ async def lifespan(app: FastAPI): mcp_server = getattr(app.state, "retrieval_mcp_server", None) mcp_session_manager = getattr(mcp_server, "session_manager", None) - logger.info("knowledge library API service started!") + logger.info("Document API service started!") if mcp_session_manager is not None: async with mcp_session_manager.run(): yield @@ -94,7 +94,7 @@ async def lifespan(app: FastAPI): except Exception as e: logger.error(f"async HTTP client close failed: {e}") - logger.info("knowledge library API service stopped!") + logger.info("Document API service stopped!") await safe_dispose_engine(engine) logger.info("database engine connection pool disposed.") logger.info("service stopped.") @@ -135,7 +135,7 @@ def create_app() -> FastAPI: @app.get("/", tags=["Root"]) async def read_root(): - return {"message": f"Welcome to {app.title} - Knowledge Base API Service!"} + return {"message": f"Welcome to {app.title} - Document API Service!"} @app.api_route("/health", methods=["GET", "HEAD"], tags=["Health"]) async def health_check(): @@ -169,11 +169,10 @@ async def health_check(): return app -# Worker settings removed as DsTasks.py was deleted app = create_app() if __name__ == "__main__": - logger.info("Knowledge Base API service starting...") + logger.info("Document API service starting...") port = 5005 reload = False # Enable hot reload host = "0.0.0.0" diff --git a/apps/api/tests/contract/test_documents_contract.py b/apps/api/tests/contract/test_documents_contract.py index bb1f8dcad..d0b401839 100644 --- a/apps/api/tests/contract/test_documents_contract.py +++ b/apps/api/tests/contract/test_documents_contract.py @@ -310,7 +310,7 @@ async def _insert_document_revision_with_chunks( ) VALUES ( :job_id, :user_id, - 'kb_management', + 'document_ingestion', 'done', 'url', FALSE, diff --git a/apps/api/tests/contract/test_job_creation_contract.py b/apps/api/tests/contract/test_job_creation_contract.py index d39a53296..95bb57cb2 100644 --- a/apps/api/tests/contract/test_job_creation_contract.py +++ b/apps/api/tests/contract/test_job_creation_contract.py @@ -163,7 +163,7 @@ async def _insert_active_job( { "job_id": job_id, "user_id": user_id, - "job_type": "kb_management", + "job_type": "document_ingestion", "status": status, "source_type": "file", "webhook_enabled": False, @@ -244,7 +244,7 @@ async def test_should_create_a_waiting_file_job_for_an_authenticated_developer( original_request = cast(dict[str, object], job_metadata["original_request"]) assert job_row["user_id"] == "local-dev-user" - assert job_row["job_type"] == "kb_management" + assert job_row["job_type"] == "document_ingestion" assert job_row["status"] == "waiting-file" assert job_row["source_type"] == "file" assert job_row["s3_key"] == f"uploads/{job_id}.pdf" @@ -272,7 +272,7 @@ async def test_should_create_a_waiting_file_job_for_an_authenticated_developer( assert cached_metadata["source_file_name"] == payload["file_name"] assert cached_job_info["job_id"] == job_id assert cached_job_info["user_id"] == "local-dev-user" - assert cached_job_info["job_type"] == "kb_management" + assert cached_job_info["job_type"] == "document_ingestion" assert cached_job_info["source_type"] == "file" assert cached_job_info["s3_key"] == f"uploads/{job_id}.pdf" assert cached_job_info["webhook_enabled"] is False @@ -675,7 +675,7 @@ def resolve_public_address( original_request = cast(dict[str, object], job_metadata["original_request"]) assert job_row["user_id"] == "local-dev-user" - assert job_row["job_type"] == "kb_management" + assert job_row["job_type"] == "document_ingestion" assert job_row["status"] == "waiting-file" assert job_row["source_type"] == "url" assert job_row["s3_key"] == f"uploads/{job_id}.pdf" @@ -705,15 +705,15 @@ def resolve_public_address( assert cached_metadata["source_url"] == payload["source_url"] assert cached_job_info["job_id"] == job_id assert cached_job_info["user_id"] == "local-dev-user" - assert cached_job_info["job_type"] == "kb_management" + assert cached_job_info["job_type"] == "document_ingestion" assert cached_job_info["source_type"] == "url" assert cached_job_info["s3_key"] == f"uploads/{job_id}.pdf" assert cached_job_info["webhook_enabled"] is False assert scheduled_tasks == [ { - "task_name": "app.core.tasks.kb_tasks.upload_url_file_task", + "task_name": "app.core.tasks.document_ingestion_tasks.upload_url_file_task", "args": [job_id, payload["source_url"], "local-dev-user"], - "kwargs": {"job_type": "kb_management"}, + "kwargs": {"job_type": "document_ingestion"}, } ] @@ -842,9 +842,9 @@ def resolve_private_address( assert job_metadata["source_url"] == source_url assert scheduled_tasks == [ { - "task_name": "app.core.tasks.kb_tasks.upload_url_file_task", + "task_name": "app.core.tasks.document_ingestion_tasks.upload_url_file_task", "args": [job_id, source_url, "local-dev-user"], - "kwargs": {"job_type": "kb_management"}, + "kwargs": {"job_type": "document_ingestion"}, } ] @@ -982,28 +982,22 @@ async def _fake_verify_s3_file_exists( assert bucket is None return {"exists": True, "s3_key": s3_key} - async def _fake_start_workflow( + async def _fake_start_uploaded_file_parse( self: object, - db: object, + *, job_id: str, - source_type: str, - file_path: str | None, - file_url: str | None, user_id: str, - ) -> None: + ) -> str: started_workflows.append( { "job_id": job_id, - "source_type": source_type, "user_id": user_id, - "file_path": file_path, - "file_url": file_url, - "db_bound": db is not None, } ) + return "contract-task-id" async with developer_api_client_factory() as api_client: - import app.services.knowledge.kb_orchestrator as kb_orchestrator_module + import app.services.document_ingestion.worker_dispatcher as dispatcher_module import shared.services.storage.file_upload_service as file_upload_service_module monkeypatch.setattr( @@ -1012,9 +1006,9 @@ async def _fake_start_workflow( _fake_verify_s3_file_exists, ) monkeypatch.setattr( - kb_orchestrator_module.KBOrchestrator, - "start_workflow", - _fake_start_workflow, + dispatcher_module.DocumentIngestionWorkerDispatcher, + "start_uploaded_file_parse", + _fake_start_uploaded_file_parse, ) create_response = await api_client.post("/api/v1/jobs", json=payload) @@ -1036,11 +1030,7 @@ async def _fake_start_workflow( assert started_workflows == [ { "job_id": job_id, - "source_type": "file", "user_id": "local-dev-user", - "file_path": None, - "file_url": None, - "db_bound": True, } ] diff --git a/apps/api/tests/contract/test_s3_event_contract.py b/apps/api/tests/contract/test_s3_event_contract.py index 65c1c6378..b1348c0cc 100644 --- a/apps/api/tests/contract/test_s3_event_contract.py +++ b/apps/api/tests/contract/test_s3_event_contract.py @@ -31,7 +31,9 @@ def _build_s3_event_payload(job_id: str) -> dict[str, object]: } -async def _insert_waiting_file_job() -> tuple[str, str]: +async def _insert_waiting_file_job( + *, job_type: str = "document_ingestion" +) -> tuple[str, str]: user_id = f"contract-s3-user-{uuid4().hex[:12]}" job_id = f"job_{uuid4().hex[:12]}" @@ -40,7 +42,7 @@ async def _insert_waiting_file_job() -> tuple[str, str]: job_id=job_id, user_id=user_id, status="waiting-file", - job_type="kb_management", + job_type=job_type, source_type="file", s3_key=f"uploads/{job_id}.pdf", ) @@ -67,29 +69,24 @@ async def test_should_accept_a_direct_upload_complete_event_advance_the_waiting_ api_client_factory: Callable[[], AbstractAsyncContextManager[AsyncClient]], monkeypatch: MonkeyPatch, ) -> None: - workflow_calls: list[dict[str, str | None]] = [] + workflow_calls: list[dict[str, str]] = [] user_id: str = "" job_id: str = "" - class FakeKBOrchestrator: - async def start_workflow( + class FakeDocumentIngestionWorkerDispatcher: + async def start_uploaded_file_parse( self, - db, + *, job_id: str, - source_type: str, - file_path: str | None, - file_url: str | None, user_id: str, - ) -> None: + ) -> str: workflow_calls.append( { "job_id": job_id, - "source_type": source_type, - "file_path": file_path, - "file_url": file_url, "user_id": user_id, } ) + return "contract-task-id" async with api_client_factory() as api_client: user_id, job_id = await _insert_waiting_file_job() @@ -98,8 +95,62 @@ async def start_workflow( ) monkeypatch.setattr( handoff_service, - "KBOrchestrator", - FakeKBOrchestrator, + "DocumentIngestionWorkerDispatcher", + FakeDocumentIngestionWorkerDispatcher, + ) + response = await api_client.post( + "/api/v1/internal/s3-events", + json=_build_s3_event_payload(job_id), + ) + + assert response.status_code == 200 + assert response.json() == {"message": "Event handled successfully"} + + job_row = await ContractDatabase.fetch_job(job_id) + + assert job_row is not None + assert job_row["status"] == "pending" + assert workflow_calls == [ + { + "job_id": job_id, + "user_id": user_id, + } + ] + + +@pytest.mark.asyncio +async def test_should_accept_a_pre_rename_waiting_file_job_type_during_upload_handoff( + api_client_factory: Callable[[], AbstractAsyncContextManager[AsyncClient]], + monkeypatch: MonkeyPatch, +) -> None: + workflow_calls: list[dict[str, str]] = [] + user_id: str = "" + job_id: str = "" + + class FakeDocumentIngestionWorkerDispatcher: + async def start_uploaded_file_parse( + self, + *, + job_id: str, + user_id: str, + ) -> str: + workflow_calls.append( + { + "job_id": job_id, + "user_id": user_id, + } + ) + return "contract-task-id" + + async with api_client_factory() as api_client: + user_id, job_id = await _insert_waiting_file_job(job_type="kb_management") + handoff_service = importlib.import_module( + "app.services.document_ingestion.handoff_service" + ) + monkeypatch.setattr( + handoff_service, + "DocumentIngestionWorkerDispatcher", + FakeDocumentIngestionWorkerDispatcher, ) response = await api_client.post( "/api/v1/internal/s3-events", @@ -116,9 +167,6 @@ async def start_workflow( assert workflow_calls == [ { "job_id": job_id, - "source_type": "file", - "file_path": None, - "file_url": None, "user_id": user_id, } ] @@ -131,17 +179,14 @@ async def test_should_accept_an_sns_wrapped_upload_complete_event_and_advance_a_ ) -> None: job_id: str = "" - class FakeKBOrchestrator: - async def start_workflow( + class FakeDocumentIngestionWorkerDispatcher: + async def start_uploaded_file_parse( self, - db, + *, job_id: str, - source_type: str, - file_path: str | None, - file_url: str | None, user_id: str, - ) -> None: - return None + ) -> str: + return "contract-task-id" async with api_client_factory() as api_client: _, job_id = await _insert_waiting_file_job() @@ -150,8 +195,8 @@ async def start_workflow( ) monkeypatch.setattr( handoff_service, - "KBOrchestrator", - FakeKBOrchestrator, + "DocumentIngestionWorkerDispatcher", + FakeDocumentIngestionWorkerDispatcher, ) response = await api_client.post( "/api/v1/internal/s3-events", diff --git a/apps/api/tests/migrations/test_schema_contract.py b/apps/api/tests/migrations/test_schema_contract.py index 7cbf27eb7..6d1a9db6a 100644 --- a/apps/api/tests/migrations/test_schema_contract.py +++ b/apps/api/tests/migrations/test_schema_contract.py @@ -88,7 +88,7 @@ def _insert_job( { "job_id": job_id, "user_id": user_id, - "job_type": "kb_management", + "job_type": "document_ingestion", "status": status, "source_type": "file", "webhook_enabled": False, diff --git a/apps/api/tests/support/contract_database.py b/apps/api/tests/support/contract_database.py index 664fff00a..2944cf14a 100644 --- a/apps/api/tests/support/contract_database.py +++ b/apps/api/tests/support/contract_database.py @@ -356,7 +356,7 @@ async def insert_job( *, job_id: str, user_id: str, - job_type: str = "kb_management", + job_type: str = "document_ingestion", status: str = "pending", source_type: str = "file", file_path: str | None = None, diff --git a/apps/worker/app/core/tasks/__init__.py b/apps/worker/app/core/tasks/__init__.py index ba382c975..35cdd963b 100644 --- a/apps/worker/app/core/tasks/__init__.py +++ b/apps/worker/app/core/tasks/__init__.py @@ -1 +1 @@ -"""Celery task package for worker-side knowledge-base jobs.""" +"""Celery task package for worker-side Document Ingestion jobs.""" diff --git a/apps/worker/app/core/tasks/base_task.py b/apps/worker/app/core/tasks/base_task.py index dee61c98f..c3474884f 100644 --- a/apps/worker/app/core/tasks/base_task.py +++ b/apps/worker/app/core/tasks/base_task.py @@ -1,5 +1,5 @@ """ -Base Celery task class for KB worker tasks. +Base Celery task class for worker-side Document Ingestion tasks. Provides centralized exception handling with direct DB writes for failure finalization. """ @@ -13,15 +13,15 @@ from shared.core.logging import LogEvent -class KBBaseTask(Task): - """Knowledge Base base task class - provides centralized exception handling.""" +class DocumentIngestionBaseTask(Task): + """Base task class for worker-side Document Ingestion error handling.""" def on_success(self, retval, task_id, args, kwargs): """Task success callback.""" logger.bind( event=LogEvent.WORKER_TASK_COMPLETE.value, task_id=task_id, - ).info("KB task completed successfully") + ).info("Document Ingestion task completed successfully") def on_failure(self, exc, task_id, args, kwargs, einfo): """Task failure callback — finalize failure directly to the database.""" @@ -84,4 +84,4 @@ def on_retry(self, exc, task_id, args, kwargs, einfo): task_id=task_id, job_id=job_id, retry_count=self.request.retries, - ).warning(f"KB task retrying: {exc}") + ).warning(f"Document Ingestion task retrying: {exc}") diff --git a/apps/worker/app/core/tasks/kb_tasks.py b/apps/worker/app/core/tasks/document_ingestion_tasks.py similarity index 52% rename from apps/worker/app/core/tasks/kb_tasks.py rename to apps/worker/app/core/tasks/document_ingestion_tasks.py index 13e9d67f7..1683a05f0 100644 --- a/apps/worker/app/core/tasks/kb_tasks.py +++ b/apps/worker/app/core/tasks/document_ingestion_tasks.py @@ -1,12 +1,11 @@ """ -Knowledge Base Management Celery Tasks +Document Ingestion Celery Tasks Sync implementation for gevent worker pool. All I/O operations use sync services that yield cooperatively under gevent. """ -# Base task class -from app.core.tasks.base_task import KBBaseTask +from app.core.tasks.base_task import DocumentIngestionBaseTask 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 @@ -25,16 +24,19 @@ celery_app = get_celery_app() +_DOCUMENT_INGESTION_JOB_TYPE = "document_ingestion" +_LEGACY_DOCUMENT_INGESTION_JOB_TYPE = "kb_management" + @celery_app.task( bind=True, - base=KBBaseTask, - name="app.core.tasks.kb_tasks.upload_url_file_task", + base=DocumentIngestionBaseTask, + name="app.core.tasks.document_ingestion_tasks.upload_url_file_task", ignore_result=True, autoretry_for=RETRYABLE_EXCEPTIONS, retry_kwargs={ - "countdown": settings.KB_TASK_RETRY_COUNTDOWN, - "max_retries": settings.KB_TASK_MAX_RETRIES, + "countdown": settings.DOCUMENT_INGESTION_TASK_RETRY_COUNTDOWN, + "max_retries": settings.DOCUMENT_INGESTION_TASK_MAX_RETRIES, }, ) def upload_url_file_task( @@ -45,6 +47,38 @@ def upload_url_file_task( job_type: str | None = None, ): """Download file from URL and upload to S3.""" + return _run_upload_url_file_task(self, job_id, source_url, user_id, job_type) + + +@celery_app.task( + bind=True, + base=DocumentIngestionBaseTask, + name="app.core.tasks.kb_tasks.upload_url_file_task", + ignore_result=True, + autoretry_for=RETRYABLE_EXCEPTIONS, + retry_kwargs={ + "countdown": settings.DOCUMENT_INGESTION_TASK_RETRY_COUNTDOWN, + "max_retries": settings.DOCUMENT_INGESTION_TASK_MAX_RETRIES, + }, +) +def legacy_upload_url_file_task( + self, + job_id: str, + source_url: str, + user_id: str | None = None, + job_type: str | None = _LEGACY_DOCUMENT_INGESTION_JOB_TYPE, +): + """Drain pre-rename URL-upload messages from legacy broker queues.""" + return _run_upload_url_file_task(self, job_id, source_url, user_id, job_type) + + +def _run_upload_url_file_task( + self, + job_id: str, + source_url: str, + user_id: str | None = None, + job_type: str | None = None, +): with log_context(task_id=self.request.id): if not job_id: raise WorkerHandlingException( @@ -69,19 +103,53 @@ def _upload_url_file( @celery_app.task( bind=True, - base=KBBaseTask, - name="app.core.tasks.kb_tasks.parse_task", + base=DocumentIngestionBaseTask, + name="app.core.tasks.document_ingestion_tasks.parse_task", ignore_result=True, autoretry_for=RETRYABLE_EXCEPTIONS, retry_kwargs={ - "countdown": settings.KB_TASK_RETRY_COUNTDOWN, - "max_retries": settings.KB_TASK_MAX_RETRIES, + "countdown": settings.DOCUMENT_INGESTION_TASK_RETRY_COUNTDOWN, + "max_retries": settings.DOCUMENT_INGESTION_TASK_MAX_RETRIES, }, ) def parse_task( - self, job_id: str, user_id: str | None = None, job_type: str = "kb_management" + self, + job_id: str, + user_id: str | None = None, + job_type: str = _DOCUMENT_INGESTION_JOB_TYPE, ): """Parse and vectorize task (file already uploaded to S3).""" + return _run_parse_task(self, job_id, user_id, job_type) + + +@celery_app.task( + bind=True, + base=DocumentIngestionBaseTask, + name="app.core.tasks.kb_tasks.parse_task", + ignore_result=True, + autoretry_for=RETRYABLE_EXCEPTIONS, + retry_kwargs={ + "countdown": settings.DOCUMENT_INGESTION_TASK_RETRY_COUNTDOWN, + "max_retries": settings.DOCUMENT_INGESTION_TASK_MAX_RETRIES, + }, +) +def legacy_parse_task( + self, + job_id: str, + user_id: str | None = None, + job_type: str = _LEGACY_DOCUMENT_INGESTION_JOB_TYPE, +): + """Drain pre-rename parse messages from legacy broker queues.""" + return _run_parse_task(self, job_id, user_id, job_type) + + +def _run_parse_task( + self, + job_id: str, + user_id: str | None = None, + job_type: str = _DOCUMENT_INGESTION_JOB_TYPE, +): + del job_type with log_context(task_id=self.request.id): if not job_id: raise WorkerHandlingException( diff --git a/apps/worker/app/core/worker_bootstrap.py b/apps/worker/app/core/worker_bootstrap.py index 2689bb71c..44a40c60c 100644 --- a/apps/worker/app/core/worker_bootstrap.py +++ b/apps/worker/app/core/worker_bootstrap.py @@ -15,7 +15,7 @@ def _register_task_modules() -> None: """Import task modules for Celery side-effect registration.""" - import app.core.tasks.kb_tasks # noqa: F401 + import app.core.tasks.document_ingestion_tasks # noqa: F401 import app.core.tasks.stale_job_sweeper # noqa: F401 import app.core.tasks.webhook_tasks # noqa: F401 @@ -114,7 +114,9 @@ def run_worker() -> None: f"--loglevel={log_level}", f"--hostname={node_name}", "-Q", - "kb_high,kb_medium,kb_low,ai_high_priority,default", + "document_ingestion_high,document_ingestion_medium," + "document_ingestion_low,kb_high,kb_medium,kb_low," + "ai_high_priority,default", "--without-gossip", "--without-mingle", ] diff --git a/apps/worker/app/services/__init__.py b/apps/worker/app/services/__init__.py index 62d28a4ab..8f5b2f6ca 100644 --- a/apps/worker/app/services/__init__.py +++ b/apps/worker/app/services/__init__.py @@ -1,5 +1,3 @@ """Worker service package.""" -# KBOrchestrator was removed and now lives only in the API service. - __all__ = [] diff --git a/apps/worker/app/services/document_ingestion/job_state_gate.py b/apps/worker/app/services/document_ingestion/job_state_gate.py index c8b963105..4c9b232f3 100644 --- a/apps/worker/app/services/document_ingestion/job_state_gate.py +++ b/apps/worker/app/services/document_ingestion/job_state_gate.py @@ -52,7 +52,7 @@ def mark_job_running(job_id: str, redis_service: Any) -> bool: f"Parse task started before upload transition completed for job {job_id}; " f"current_state={current_state}" ), - retry_after=settings.KB_TASK_RETRY_COUNTDOWN, + retry_after=settings.DOCUMENT_INGESTION_TASK_RETRY_COUNTDOWN, user_message="Job is not ready for processing yet. Retrying shortly.", ) @@ -85,7 +85,7 @@ def mark_job_running(job_id: str, redis_service: Any) -> bool: f"Failed to transition job {job_id} from pending to running; " f"reason={outcome.reason}" ), - retry_after=settings.KB_TASK_RETRY_COUNTDOWN, + retry_after=settings.DOCUMENT_INGESTION_TASK_RETRY_COUNTDOWN, user_message="Job state is still settling. Retrying shortly.", ) diff --git a/apps/worker/app/services/document_ingestion/processing_run.py b/apps/worker/app/services/document_ingestion/processing_run.py index d3c31ca0a..086f59044 100644 --- a/apps/worker/app/services/document_ingestion/processing_run.py +++ b/apps/worker/app/services/document_ingestion/processing_run.py @@ -277,14 +277,14 @@ def _finalize_parse_job_success( source_file_name=str(source_file_name), ) try: - kb_dir_for_enrich = os.path.dirname(str(add_dir)) + document_root_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, + document_root_for_enrich, source_file=str(source_file_name), use_llm=summary_use_llm, ) diff --git a/apps/worker/app/services/document_ingestion/workspace.py b/apps/worker/app/services/document_ingestion/workspace.py index c6bd2ca85..579b975d3 100644 --- a/apps/worker/app/services/document_ingestion/workspace.py +++ b/apps/worker/app/services/document_ingestion/workspace.py @@ -58,7 +58,10 @@ def create_task_workspace(job_id: str) -> str: try: os.makedirs(temp_root, exist_ok=True) - return tempfile.mkdtemp(prefix=f"kb_task_{job_id}_", dir=temp_root) + return tempfile.mkdtemp( + prefix=f"document_ingestion_task_{job_id}_", + dir=temp_root, + ) except (OSError, PermissionError) as exc: raise FileSystemException( user_message="System error preparing temporary storage", diff --git a/apps/worker/tests/contract/README.md b/apps/worker/tests/contract/README.md index 725e982dc..d9f009a5c 100644 --- a/apps/worker/tests/contract/README.md +++ b/apps/worker/tests/contract/README.md @@ -18,9 +18,9 @@ These tests should avoid: - `app.core.tasks.stale_job_sweeper.expire_stale_jobs` verifies stale-job expiration, durable failure state, audit logging, and the Redis-backed duplicate-Beat lock -- `app.core.tasks.kb_tasks.upload_url_file_task` +- `app.core.tasks.document_ingestion_tasks.upload_url_file_task` verifies the cached storage target, Redis progress publication, and the stable `waiting-file` job state while mocking only outbound URL download and S3 boundaries -- `app.core.tasks.kb_tasks.parse_task` +- `app.core.tasks.document_ingestion_tasks.parse_task` verifies success publication, terminal skip handling, and failure cleanup/refund behavior while keeping billing, finalization, and retrieval publication real - `app.core.tasks.webhook_tasks.recover_orphaned_webhooks` verifies orphaned pending webhook recovery, durable QStash delivery-state persistence, and the Redis-backed duplicate-Beat lock diff --git a/apps/worker/tests/contract/test_parse_task_contract.py b/apps/worker/tests/contract/test_parse_task_contract.py index f0d45b529..1a1ce4d64 100644 --- a/apps/worker/tests/contract/test_parse_task_contract.py +++ b/apps/worker/tests/contract/test_parse_task_contract.py @@ -34,7 +34,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.core.tasks.document_ingestion_tasks as document_ingestion_tasks 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 @@ -45,7 +45,7 @@ def _load_parse_task_modules() -> tuple[Any, Any, Any, Engine, Any, Any, Any]: ) return ( - kb_tasks, + document_ingestion_tasks, parse_service, parse_job_service, get_sync_engine(), @@ -82,7 +82,7 @@ def _save_worker_task_cache( "s3_key": s3_key, "user_id": user_id, "webhook_enabled": False, - "job_type": "kb_management", + "job_type": "document_ingestion", "source_type": "file", }, ) @@ -107,21 +107,21 @@ def _find_task_workspaces(root: Path, job_id: str) -> list[Path]: return sorted( path for path in root.iterdir() - if path.is_dir() and path.name.startswith(f"kb_task_{job_id}_") + if path.is_dir() and path.name.startswith(f"document_ingestion_task_{job_id}_") ) def _bind_parse_task_to_current_module( monkeypatch: MonkeyPatch, *, - kb_tasks: Any, + document_ingestion_tasks: Any, ) -> None: monkeypatch.setitem( - kb_tasks.parse_task._orig_run.__globals__, + document_ingestion_tasks.parse_task._orig_run.__globals__, "_parse", - kb_tasks._parse, + document_ingestion_tasks._parse, ) - monkeypatch.setattr(kb_tasks.parse_task, "__trace__", None, raising=False) + monkeypatch.setattr(document_ingestion_tasks.parse_task, "__trace__", None, raising=False) @pytest.mark.parametrize( @@ -141,7 +141,7 @@ def test_should_parse_a_pending_file_job_and_persist_the_published_result_state( ) -> None: monkeypatch.setenv("BILLING_ENABLED", "true" if billing_enabled else "false") ( - kb_tasks, + document_ingestion_tasks, parse_service, parse_job_service, engine, @@ -185,7 +185,7 @@ def test_should_parse_a_pending_file_job_and_persist_the_published_result_state( sync_redis_service_factory=sync_redis_service_factory, ) - _bind_parse_task_to_current_module(monkeypatch, kb_tasks=kb_tasks) + _bind_parse_task_to_current_module(monkeypatch, document_ingestion_tasks=document_ingestion_tasks) monkeypatch.setattr(settings, "TMP_PATH", str(tmp_path)) monkeypatch.setattr(settings, "BILLING_ENABLED", billing_enabled) @@ -320,7 +320,7 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: monkeypatch.setattr(parse_service, "checkerboard_inject_parse", fake_checkerboard_inject_parse) monkeypatch.setattr(parse_job_service, "get_result_storage", lambda: FakeResultStorage()) - result = kb_tasks.parse_task.run(job_id, user_id, "kb_management") + result = document_ingestion_tasks.parse_task.run(job_id, user_id, "document_ingestion") expected_summary = "This document includes: 公司研究, 相关研报" expected_connect_to = [ @@ -610,7 +610,7 @@ def test_should_export_full_result_when_publication_deduplicates_existing_chunks ) -> None: monkeypatch.setenv("BILLING_ENABLED", "false") ( - kb_tasks, + document_ingestion_tasks, parse_service, parse_job_service, engine, @@ -800,7 +800,7 @@ def test_should_export_full_result_when_publication_deduplicates_existing_chunks sync_redis_service_factory=sync_redis_service_factory, ) - _bind_parse_task_to_current_module(monkeypatch, kb_tasks=kb_tasks) + _bind_parse_task_to_current_module(monkeypatch, document_ingestion_tasks=document_ingestion_tasks) monkeypatch.setattr(settings, "TMP_PATH", str(tmp_path)) monkeypatch.setattr(settings, "BILLING_ENABLED", False) @@ -907,7 +907,7 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: 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") + result = document_ingestion_tasks.parse_task.run(job_id, user_id, "document_ingestion") assert result["contents_count"] == 3 assert "images/duplicate.png" in captured_artifacts["zip_entries"] @@ -986,7 +986,7 @@ def test_should_initialize_billing_once_for_concurrent_parse_tasks( tmp_path: Path, ) -> None: ( - kb_tasks, + document_ingestion_tasks, parse_service, parse_job_service, engine, @@ -1034,7 +1034,7 @@ def test_should_initialize_billing_once_for_concurrent_parse_tasks( sync_redis_service_factory=sync_redis_service_factory, ) - _bind_parse_task_to_current_module(monkeypatch, kb_tasks=kb_tasks) + _bind_parse_task_to_current_module(monkeypatch, document_ingestion_tasks=document_ingestion_tasks) monkeypatch.setattr(settings, "TMP_PATH", str(tmp_path)) monkeypatch.setattr(settings, "BILLING_ENABLED", True) @@ -1107,7 +1107,7 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: 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")) + return dict(document_ingestion_tasks.parse_task.run(job_id, user_id, "document_ingestion")) with ThreadPoolExecutor(max_workers=len(job_ids)) as executor: results = list(executor.map(run_parse_task, job_ids)) @@ -1230,7 +1230,7 @@ def test_should_skip_parse_task_when_the_job_is_already_terminal( tmp_path: Path, ) -> None: ( - kb_tasks, + document_ingestion_tasks, parse_service, parse_job_service, engine, @@ -1270,7 +1270,7 @@ def test_should_skip_parse_task_when_the_job_is_already_terminal( sync_redis_service_factory=sync_redis_service_factory, ) - _bind_parse_task_to_current_module(monkeypatch, kb_tasks=kb_tasks) + _bind_parse_task_to_current_module(monkeypatch, document_ingestion_tasks=document_ingestion_tasks) 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,7 +1284,7 @@ def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: ), ) - result = kb_tasks.parse_task.run(job_id, user_id, "kb_management") + result = document_ingestion_tasks.parse_task.run(job_id, user_id, "document_ingestion") assert result == { "status": "skipped", @@ -1320,7 +1320,7 @@ def test_should_mark_the_job_failed_and_cleanup_the_workspace_when_parse_executi tmp_path: Path, ) -> None: ( - kb_tasks, + document_ingestion_tasks, parse_service, parse_job_service, engine, @@ -1360,7 +1360,7 @@ def test_should_mark_the_job_failed_and_cleanup_the_workspace_when_parse_executi sync_redis_service_factory=sync_redis_service_factory, ) - _bind_parse_task_to_current_module(monkeypatch, kb_tasks=kb_tasks) + _bind_parse_task_to_current_module(monkeypatch, document_ingestion_tasks=document_ingestion_tasks) monkeypatch.setattr(settings, "TMP_PATH", str(tmp_path)) def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: return { @@ -1392,8 +1392,8 @@ def fake_download_s3_file_to_temp( ), ) - result = kb_tasks.parse_task.apply( - args=[job_id, user_id, "kb_management"], + result = document_ingestion_tasks.parse_task.apply( + args=[job_id, user_id, "document_ingestion"], throw=False, ) diff --git a/apps/worker/tests/contract/test_url_upload_contract.py b/apps/worker/tests/contract/test_url_upload_contract.py index cdf4a3de4..4765472cd 100644 --- a/apps/worker/tests/contract/test_url_upload_contract.py +++ b/apps/worker/tests/contract/test_url_upload_contract.py @@ -14,7 +14,7 @@ def _load_upload_task_modules() -> tuple[Any, Any, Engine, Any, Any]: - import app.core.tasks.kb_tasks as kb_tasks + import app.core.tasks.document_ingestion_tasks as document_ingestion_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 ( @@ -23,7 +23,7 @@ def _load_upload_task_modules() -> tuple[Any, Any, Engine, Any, Any]: ) return ( - kb_tasks, + document_ingestion_tasks, url_upload_service, get_sync_engine(), SyncJobInfoRedisService, @@ -31,20 +31,37 @@ def _load_upload_task_modules() -> tuple[Any, Any, Engine, Any, Any]: ) +def _bind_upload_task_to_current_module( + monkeypatch: MonkeyPatch, + *, + document_ingestion_tasks: Any, +) -> None: + monkeypatch.setitem( + document_ingestion_tasks.upload_url_file_task._orig_run.__globals__, + "_upload_url_file", + document_ingestion_tasks._upload_url_file, + ) + monkeypatch.setattr( + document_ingestion_tasks.upload_url_file_task, + "__trace__", + None, + raising=False, + ) + + def test_should_upload_a_url_job_to_the_expected_storage_key_and_publish_progress( worker_contract_environment: None, monkeypatch: MonkeyPatch, tmp_path: Path, ) -> None: ( - kb_tasks, - _url_upload_service, + document_ingestion_tasks, + 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]}" job_id = f"job_url_upload_{uuid4().hex[:12]}" @@ -56,27 +73,50 @@ def test_should_upload_a_url_job_to_the_expected_storage_key_and_publish_progres 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))] monkeypatch.setattr( - JobFileStorage, - "download_file_from_url", - lambda self, _source_url, *, temp_dir=None: str(downloaded_path), + url_upload_service, + "download_source_url_to_temp", + lambda _source_url: str(downloaded_path), + ) + monkeypatch.setitem( + document_ingestion_tasks.upload_url_file.__globals__, + "download_source_url_to_temp", + lambda _source_url: str(downloaded_path), ) monkeypatch.setattr( - JobFileStorage, - "upload_source_file", - lambda self, local_path, storage_key: uploaded_calls.append( - (local_path, storage_key, self.uploads_bucket) + url_upload_service, + "upload_temp_file_to_source_storage", + lambda *, temp_file_path, s3_key: uploaded_calls.append( + (temp_file_path, s3_key, settings.S3_BUCKET_NAME) + ), + ) + monkeypatch.setitem( + document_ingestion_tasks.upload_url_file.__globals__, + "upload_temp_file_to_source_storage", + lambda *, temp_file_path, s3_key: uploaded_calls.append( + (temp_file_path, s3_key, settings.S3_BUCKET_NAME) ), ) monkeypatch.setattr( - JobFileStorage, - "verify_upload_exists", - lambda self, storage_key: {"exists": storage_key == s3_key, "size": 3}, + url_upload_service, + "verify_source_upload", + lambda storage_key: {"exists": storage_key == s3_key, "size": 3}, + ) + monkeypatch.setitem( + document_ingestion_tasks.upload_url_file.__globals__, + "verify_source_upload", + lambda storage_key: {"exists": storage_key == s3_key, "size": 3}, ) monkeypatch.setattr(socket, "getaddrinfo", resolve_public_address) + _bind_upload_task_to_current_module( + monkeypatch, + document_ingestion_tasks=document_ingestion_tasks, + ) downloaded_path.write_bytes(b"pdf") @@ -106,16 +146,16 @@ def resolve_public_address( "s3_key": s3_key, "user_id": user_id, "webhook_enabled": False, - "job_type": "kb_management", + "job_type": "document_ingestion", "source_type": "url", }, ) - result = kb_tasks.upload_url_file_task.run( + result = document_ingestion_tasks.upload_url_file_task.run( job_id, source_url, user_id, - "kb_management", + "document_ingestion", ) assert result == { diff --git a/apps/worker/tests/contract/test_worker_bootstrap_contract.py b/apps/worker/tests/contract/test_worker_bootstrap_contract.py index fa0ad1c05..c3701b0c6 100644 --- a/apps/worker/tests/contract/test_worker_bootstrap_contract.py +++ b/apps/worker/tests/contract/test_worker_bootstrap_contract.py @@ -2,6 +2,8 @@ import sys +from pytest import MonkeyPatch + def test_should_register_worker_task_modules_for_celery_consumers( worker_contract_environment: None, @@ -10,13 +12,15 @@ def test_should_register_worker_task_modules_for_celery_consumers( from shared.core.celery_app import celery_app expected_task_names: tuple[str, ...] = ( + "app.core.tasks.document_ingestion_tasks.upload_url_file_task", + "app.core.tasks.document_ingestion_tasks.parse_task", "app.core.tasks.kb_tasks.upload_url_file_task", "app.core.tasks.kb_tasks.parse_task", "app.core.tasks.stale_job_sweeper.expire_stale_jobs", "app.core.tasks.webhook_tasks.recover_orphaned_webhooks", ) task_module_names: tuple[str, ...] = ( - "app.core.tasks.kb_tasks", + "app.core.tasks.document_ingestion_tasks", "app.core.tasks.stale_job_sweeper", "app.core.tasks.webhook_tasks", ) @@ -31,3 +35,44 @@ def test_should_register_worker_task_modules_for_celery_consumers( for task_name in expected_task_names: assert task_name in celery_app.tasks + + +def test_should_consume_current_and_legacy_ingestion_queues( + worker_contract_environment: None, + monkeypatch: MonkeyPatch, +) -> None: + from app.core import worker_bootstrap + + worker_main_calls: list[list[str]] = [] + beat_commands: list[list[str]] = [] + + class FakeBeatProcess: + pass + + def record_beat_command(command: list[str]) -> FakeBeatProcess: + beat_commands.append(command) + return FakeBeatProcess() + + def record_worker_main(args: list[str]) -> None: + worker_main_calls.append(args) + + monkeypatch.setattr(worker_bootstrap.subprocess, "Popen", record_beat_command) + monkeypatch.setattr(worker_bootstrap.celery_app, "worker_main", record_worker_main) + + worker_bootstrap.run_worker() + + assert beat_commands != [] + assert len(worker_main_calls) == 1 + + worker_args = worker_main_calls[0] + queue_arg = worker_args[worker_args.index("-Q") + 1] + consumed_queues = set(queue_arg.split(",")) + + assert { + "document_ingestion_high", + "document_ingestion_medium", + "document_ingestion_low", + "kb_high", + "kb_medium", + "kb_low", + }.issubset(consumed_queues) diff --git a/apps/worker/tests/support/contract_database.py b/apps/worker/tests/support/contract_database.py index f586f1339..3cb31c307 100644 --- a/apps/worker/tests/support/contract_database.py +++ b/apps/worker/tests/support/contract_database.py @@ -39,7 +39,7 @@ def insert_contract_job( *, job_id: str, user_id: str, - job_type: str = "kb_management", + job_type: str = "document_ingestion", status: str = "pending", source_type: str = "file", file_path: str | None = None, diff --git a/packages/shared-python/shared/core/celery_app.py b/packages/shared-python/shared/core/celery_app.py index 8d4f2fbb1..83b6534bb 100644 --- a/packages/shared-python/shared/core/celery_app.py +++ b/packages/shared-python/shared/core/celery_app.py @@ -32,7 +32,11 @@ def get_unique_node_name() -> str: # Priority is handled via Celery's Redis transport priority support # (task_queue_max_priority + task_default_priority). celery_app.conf.task_queues = ( - # Knowledge-base queues (routed by priority in task invocation) + # Document Ingestion queues (routed by priority in task invocation) + Queue("document_ingestion_high", routing_key="document_ingestion.high"), + Queue("document_ingestion_medium", routing_key="document_ingestion.medium"), + Queue("document_ingestion_low", routing_key="document_ingestion.low"), + # Pre-rename ingestion queues kept until old broker messages have drained. Queue("kb_high", routing_key="kb.high"), Queue("kb_medium", routing_key="kb.medium"), Queue("kb_low", routing_key="kb.low"), @@ -87,7 +91,10 @@ def get_unique_node_name() -> str: beat_max_loop_interval=30, # Task routing task_routes={ - # Knowledge base tasks (default medium priority) + # Document Ingestion tasks (default medium priority) + "app.core.tasks.document_ingestion_tasks.*": { + "queue": "document_ingestion_medium" + }, "app.core.tasks.kb_tasks.*": {"queue": "kb_medium"}, # Webhook orphan recovery "app.core.tasks.webhook_tasks.recover_orphaned_webhooks": {"queue": "default"}, diff --git a/packages/shared-python/shared/core/celery_router.py b/packages/shared-python/shared/core/celery_router.py index 7c8fd86fb..eca089cbc 100644 --- a/packages/shared-python/shared/core/celery_router.py +++ b/packages/shared-python/shared/core/celery_router.py @@ -13,6 +13,20 @@ # Note: get_queue_for_job has been simplified and no longer needs direct DB access. +_DOCUMENT_INGESTION_TASK_TYPE = "document_ingestion" +_LEGACY_DOCUMENT_INGESTION_TASK_TYPES = frozenset( + { + "kb_encoding", + "kb_management", + } +) + + +def _normalize_task_type(task_type: str) -> str: + if task_type in _LEGACY_DOCUMENT_INGESTION_TASK_TYPES: + return _DOCUMENT_INGESTION_TASK_TYPE + return task_type + class TaskType(Enum): """Task-type enum.""" @@ -21,12 +35,11 @@ class TaskType(Enum): USER_AUTH = "user_auth" URGENT_DOCUMENT = "urgent_document" DOCUMENT_PROCESSING = "document_processing" - KB_ENCODING = "kb_encoding" + DOCUMENT_INGESTION = "document_ingestion" BATCH_PROCESSING = "batch_processing" ANALYTICS = "analytics" BACKUP = "backup" LOG_PROCESSING = "log_processing" - KB_MANAGEMENT = "kb_management" class UserLevel(Enum): @@ -75,12 +88,11 @@ def __init__(self): TaskType.USER_AUTH: 10, TaskType.URGENT_DOCUMENT: 10, TaskType.DOCUMENT_PROCESSING: 5, - TaskType.KB_ENCODING: 5, + TaskType.DOCUMENT_INGESTION: 6, TaskType.BATCH_PROCESSING: 5, TaskType.ANALYTICS: 1, TaskType.BACKUP: 1, TaskType.LOG_PROCESSING: 1, - TaskType.KB_MANAGEMENT: 6, } # User-level weights. @@ -171,7 +183,7 @@ def create_task_context( """ try: # Parse the task type. - task_type_enum = TaskType(task_type) + task_type_enum = TaskType(_normalize_task_type(task_type)) # Parse the user level. user_level = UserLevel(kwargs.get("user_level", "standard")) @@ -206,7 +218,7 @@ def get_queue_for_job(self, job_type: str, user_id: str) -> str: Resolve the queue name for a job based on job type and subscription. Args: - job_type: Task type such as kb_management or ai_query. + job_type: Task type such as document_ingestion or ai_query. user_id: User ID. Returns: @@ -214,8 +226,10 @@ def get_queue_for_job(self, job_type: str, user_id: str) -> str: """ try: # TODO: temporary simplified path to avoid async work here. - if job_type in ["kb_management", "kb_encoding"]: + if job_type in _LEGACY_DOCUMENT_INGESTION_TASK_TYPES: return "kb_low" + if job_type == _DOCUMENT_INGESTION_TASK_TYPE: + return "document_ingestion_low" elif job_type in ["ai_query", "user_auth", "urgent_document"]: return "ai_high_priority" elif job_type in ["document_processing"]: @@ -235,8 +249,10 @@ def get_queue_for_job(self, job_type: str, user_id: str) -> str: except Exception as e: logger.error(f"Failed to get queue for user {user_id}: {e}") # Default to a medium-priority queue on failure. - if job_type in ["kb_management", "kb_encoding"]: + if job_type in _LEGACY_DOCUMENT_INGESTION_TASK_TYPES: return "kb_medium" + if job_type == _DOCUMENT_INGESTION_TASK_TYPE: + return "document_ingestion_medium" elif job_type in ["ai_query", "user_auth", "urgent_document"]: return "ai_high_priority" else: @@ -305,7 +321,7 @@ def route_task(self, name, args, kwargs, options, task=None, **kwds): ) # Celery supports router list: try function router first, fallback to static dict -# 1. task_router.route_task: Dynamic routing (based on user subscription, for kb_tasks) +# 1. task_router.route_task: Dynamic routing for Document Ingestion tasks. # 2. static_routes: Static routing (webhook and other fixed routes, from celery_app.py) celery_app.conf.task_routes = [ task_router.route_task, # Dynamic routing priority diff --git a/packages/shared-python/shared/core/config/base.py b/packages/shared-python/shared/core/config/base.py index 0f098cd4e..d0f865eff 100644 --- a/packages/shared-python/shared/core/config/base.py +++ b/packages/shared-python/shared/core/config/base.py @@ -25,7 +25,7 @@ class BaseConfig(BaseSettings): description="Application version read from the APP_VERSION environment variable", ) APP_DESCRIPTION: str = Field( - default="AI-powered document parsing, retrieval, and knowledge access backend", + default="AI-powered document parsing and retrieval backend", description="Application description", ) diff --git a/packages/shared-python/shared/core/config/celery.py b/packages/shared-python/shared/core/config/celery.py index 99541fc62..e22053780 100644 --- a/packages/shared-python/shared/core/config/celery.py +++ b/packages/shared-python/shared/core/config/celery.py @@ -23,9 +23,13 @@ class CeleryConfig(BaseModel): ) # Task retry configuration - KB_TASK_MAX_RETRIES: int = Field(default=2, description="KB task max retries") - KB_TASK_RETRY_COUNTDOWN: int = Field( - default=120, description="KB task retry countdown (seconds)" + DOCUMENT_INGESTION_TASK_MAX_RETRIES: int = Field( + default=2, + description="Document Ingestion task max retries", + ) + DOCUMENT_INGESTION_TASK_RETRY_COUNTDOWN: int = Field( + default=120, + description="Document Ingestion task retry countdown (seconds)", ) PYMUPDF_MAX_CONCURRENT: int = Field( default=2, @@ -40,7 +44,7 @@ class CeleryConfig(BaseModel): "user_auth": 10, "urgent_document": 10, "document_processing": 5, - "kb_encoding": 5, + "document_ingestion": 6, "batch_processing": 5, "analytics": 1, "backup": 1, @@ -56,7 +60,7 @@ class CeleryConfig(BaseModel): "user_auth": "auth_queue", "urgent_document": "document_urgent", "document_processing": "document_processing", - "kb_encoding": "kb_encoding", + "document_ingestion": "document_ingestion_low", "batch_processing": "batch_processing", "analytics": "analytics_queue", "backup": "backup_queue", diff --git a/packages/shared-python/shared/core/tasks/__init__.py b/packages/shared-python/shared/core/tasks/__init__.py index 849c9f7b6..95e7349b9 100644 --- a/packages/shared-python/shared/core/tasks/__init__.py +++ b/packages/shared-python/shared/core/tasks/__init__.py @@ -1,5 +1,5 @@ """Shared Celery task definitions. -Knowledge-base tasks moved to the Worker service, and state-machine tasks moved -to the API service. Only generic Celery tasks remain here. +Document Ingestion tasks live in the Worker service, and state-machine tasks +live in the API service. Only generic Celery tasks remain here. """ diff --git a/packages/shared-python/shared/models/database/job.py b/packages/shared-python/shared/models/database/job.py index 76c9f1df1..39790cefd 100644 --- a/packages/shared-python/shared/models/database/job.py +++ b/packages/shared-python/shared/models/database/job.py @@ -51,7 +51,7 @@ class Job(Base): ) # Basic Job Info - job_type: Mapped[str] = mapped_column(String(50), nullable=False) # kb_management + job_type: Mapped[str] = mapped_column(String(50), nullable=False) # document_ingestion status: Mapped[str] = mapped_column( String(50), nullable=False, default="pending" ) # PRD Status: pending, waiting-file, running, converting, done, failed diff --git a/packages/shared-python/shared/services/redis/distributed_lock.py b/packages/shared-python/shared/services/redis/distributed_lock.py index 8933737f1..751309031 100644 --- a/packages/shared-python/shared/services/redis/distributed_lock.py +++ b/packages/shared-python/shared/services/redis/distributed_lock.py @@ -130,7 +130,7 @@ def __enter__(self) -> "RedisJobLock": internal_message=( f"Could not acquire processing lock for job {self._job_id}" ), - retry_after=settings.KB_TASK_RETRY_COUNTDOWN, + retry_after=settings.DOCUMENT_INGESTION_TASK_RETRY_COUNTDOWN, ) return self From 3b7b4f5dd40458790a720ee99b0db150c2e59fcc Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 14:40:33 +0000 Subject: [PATCH 73/78] refactor: replace parser kb dir with namespace --- AGENTS.md | 2 +- .../tests/contract/test_retrieval_contract.py | 26 ++++---- .../connect_builder/summary_builder.py | 16 ++--- .../document_ingestion/processing_run.py | 10 +-- .../document_parser/formats/docx/parser.py | 2 +- .../formats/fragment/parser.py | 34 ++++++---- .../document_parser/formats/image/parser.py | 14 ++--- .../orchestration/format_adapters.py | 2 +- .../orchestration/namespace_path_segment.py | 18 ++++++ .../orchestration/parse_input.py | 2 +- .../orchestration/parse_session.py | 62 ++++++++++--------- .../services/document_parser/parse_service.py | 8 +-- .../contract/test_excel_parser_contract.py | 35 ++++++++++- .../contract/test_parse_task_contract.py | 34 +++++----- .../shared/models/schemas/job.py | 3 - .../shared/services/ai/prompt_service.py | 2 +- .../shared/services/redis/key_builder.py | 20 ------ .../services/retrieval/agentic/prompts.py | 6 +- .../services/retrieval/execution/routes.py | 18 +++--- .../services/retrieval/search/lexical_text.py | 4 +- .../retrieval/workflow/orchestrator.py | 4 +- .../retrieval/workflow/plan_service.py | 8 +-- .../services/retrieval/workflow/planner.py | 10 +-- 23 files changed, 194 insertions(+), 146 deletions(-) create mode 100644 apps/worker/app/services/document_parser/orchestration/namespace_path_segment.py diff --git a/AGENTS.md b/AGENTS.md index 7017e6e82..efbc89049 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -608,7 +608,7 @@ Unlike legacy retrieval which relied on static `hydrate_mode` tags, hydration is ### Small Corpus Optimization When `total_chunks <= top_k`, skips the full pipeline and returns all chunks -directly (router: `small_kb_all`). +directly (router: `small_corpus_all`). ### Caching diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index 898a476ae..06cbc173b 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -153,16 +153,16 @@ async def fake_plan( self: object, *, query: str, - kb_total_docs: int = 0, - kb_total_chunks: int = 0, + corpus_total_docs: int = 0, + corpus_total_chunks: int = 0, ) -> QueryPlan: return QueryPlan( original_query=query, steps=[PlannedStep(id="request-policy", sub_query="policy marker")], final_strategy="concat_final_parts", reasoning_summary=( - f"request policy contract for {kb_total_docs} docs " - f"and {kb_total_chunks} chunks" + f"request policy contract for {corpus_total_docs} docs " + f"and {corpus_total_chunks} chunks" ), ) @@ -280,7 +280,7 @@ async def test_should_return_seeded_retrieval_results_for_the_authenticated_user assert response_json["namespace"] == "contract-retrieval" assert response_json["query"] == "alpha" - assert response_json["router_used"] == "small_kb_all" + assert response_json["router_used"] == "small_corpus_all" assert len(results) == 1 assert results[0]["chunk_type"] == "text" assert results[0]["content"] == "alpha contract retrieval content" @@ -763,8 +763,8 @@ async def fake_plan( self: object, *, query: str, - kb_total_docs: int = 0, - kb_total_chunks: int = 0, + corpus_total_docs: int = 0, + corpus_total_chunks: int = 0, ) -> QueryPlan: return QueryPlan( original_query=query, @@ -774,8 +774,8 @@ async def fake_plan( ], final_strategy="concat_final_parts", reasoning_summary=( - f"forced two-step contract plan for {kb_total_docs} docs " - f"and {kb_total_chunks} chunks" + f"forced two-step contract plan for {corpus_total_docs} docs " + f"and {corpus_total_chunks} chunks" ), ) @@ -874,8 +874,8 @@ async def fake_plan( self: object, *, query: str, - kb_total_docs: int = 0, - kb_total_chunks: int = 0, + corpus_total_docs: int = 0, + corpus_total_chunks: int = 0, ) -> QueryPlan: return QueryPlan( original_query=query, @@ -885,8 +885,8 @@ async def fake_plan( ], final_strategy="concat_final_parts", reasoning_summary=( - f"forced section identity contract plan for {kb_total_docs} docs " - f"and {kb_total_chunks} chunks" + f"forced section identity contract plan for {corpus_total_docs} docs " + f"and {corpus_total_chunks} chunks" ), ) diff --git a/apps/worker/app/services/connect_builder/summary_builder.py b/apps/worker/app/services/connect_builder/summary_builder.py index c52348112..7804bd6b6 100644 --- a/apps/worker/app/services/connect_builder/summary_builder.py +++ b/apps/worker/app/services/connect_builder/summary_builder.py @@ -7,7 +7,7 @@ Usage (standalone): from app.services.connect_builder.summary_builder import enrich_doc_nav_summaries - enrich_doc_nav_summaries(kb_dir, source_file="report.pdf") + enrich_doc_nav_summaries(document_workspace_dir, source_file="report.pdf") """ import json @@ -263,7 +263,7 @@ def _build_nav_top_summary( def enrich_doc_nav_summaries( - kb_dir: str, + document_workspace_dir: str, source_file: Optional[str] = None, force: bool = False, use_llm: bool = True, @@ -271,7 +271,7 @@ def enrich_doc_nav_summaries( """Enrich doc_nav.json with bottom-up recursive summaries. Args: - kb_dir: Absolute path to the KB directory. + document_workspace_dir: Absolute path to the temporary document workspace. source_file: If given, only process this file. Otherwise process all. force: If True, regenerate even if summaries already exist. use_llm: If True, use LLM for multi-child aggregation. @@ -287,13 +287,13 @@ def enrich_doc_nav_summaries( else: targets = [ entry - for entry in os.listdir(kb_dir) - if os.path.isdir(os.path.join(kb_dir, entry)) + for entry in os.listdir(document_workspace_dir) + if os.path.isdir(os.path.join(document_workspace_dir, entry)) and not entry.startswith(".") ] for file_name in targets: - file_dir = os.path.join(kb_dir, file_name) + file_dir = os.path.join(document_workspace_dir, file_name) doc_nav = _load_doc_nav(file_dir) if doc_nav is None: logger.debug(f"No {DOC_NAV_FILENAME} for {file_name}, skipping") @@ -334,7 +334,7 @@ def build_section_summary_lookup(file_dir: str) -> Dict[str, str]: """Build a flat {section_path: summary} dict from all nodes in doc_nav.json. Keys use the DocumentSection.section_path format produced by - ``section_path_from_chunk_path`` (strips kb_root + filename prefix, + ``section_path_from_chunk_path`` (strips namespace + filename prefix, joins remaining parts with ``" / "``). Traverses the full section tree at all depths. Used by the publication @@ -342,7 +342,7 @@ def build_section_summary_lookup(file_dir: str) -> Dict[str, str]: Args: file_dir: Absolute path to the file-level directory - (e.g. ~/.knowhere/{kb_id}/{source_file_name}/). + inside the task-scoped parse workspace. Returns: Dict mapping section_path → summary string (empty dict on any error). diff --git a/apps/worker/app/services/document_ingestion/processing_run.py b/apps/worker/app/services/document_ingestion/processing_run.py index 086f59044..ba9b49878 100644 --- a/apps/worker/app/services/document_ingestion/processing_run.py +++ b/apps/worker/app/services/document_ingestion/processing_run.py @@ -169,17 +169,17 @@ def _run_parse_job( filename=filename, doc_type=doc_type, ): + namespace = JobMetadataHelper.get_namespace( + job_context.job_metadata, + "default", + ) 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", - ), + namespace=namespace or "default", doc_type=doc_type, smart_title_parse=JobMetadataHelper.get_parsing_param( job_context.job_metadata, diff --git a/apps/worker/app/services/document_parser/formats/docx/parser.py b/apps/worker/app/services/document_parser/formats/docx/parser.py index 2f3a92551..93554ac89 100755 --- a/apps/worker/app/services/document_parser/formats/docx/parser.py +++ b/apps/worker/app/services/document_parser/formats/docx/parser.py @@ -634,7 +634,7 @@ def convert_doc2dics( ] pure_text = "\n".join(text_items).strip() know_id = gen_str_codes(pure_text) - # Use relative_root for path instead of absolute kb_dir + # Use relative_root for path instead of the absolute output directory. path_suffix = key if key.strip() else "" know_path = ( split_char.join([relative_root, path_suffix]) diff --git a/apps/worker/app/services/document_parser/formats/fragment/parser.py b/apps/worker/app/services/document_parser/formats/fragment/parser.py index 6ab680cc4..0d5dde56f 100644 --- a/apps/worker/app/services/document_parser/formats/fragment/parser.py +++ b/apps/worker/app/services/document_parser/formats/fragment/parser.py @@ -2,17 +2,19 @@ Fragment Parser - for user-pasted text content injection This parser handles .fragment files which represent user-pasted text content -that needs to be injected into the knowledge base without requiring a physical file. +that needs to be parsed without requiring a physical file. """ import os from typing import Any, Optional from app.services.document_parser.formats.markdown.parser import parse_md +from app.services.document_parser.orchestration.namespace_path_segment import ( + build_namespace_path_segment, +) from loguru import logger from openai.types.chat import ChatCompletionMessageParam -from shared.core.config import settings from app.services.common.file_utils import path_handle from shared.services.ai.openai_compatible_client_sync import get_openai_client @@ -54,7 +56,7 @@ def parse_fragment( fragment_content: str, filename: Optional[str] = None, output_dir: Optional[str] = None, - kb_dir: str = "Default_Root", + namespace: str = "default", base_llm_paras: Optional[dict[str, Any]] = None, **kwargs, ): @@ -65,13 +67,12 @@ def parse_fragment( fragment_content: The text content to parse filename: Optional filename, if not provided will be auto-generated output_dir: Base output directory - kb_dir: Knowledge base directory name + namespace: Retrieval namespace used for parser-relative paths base_llm_paras: LLM parameters for parsing Returns: tuple: (full_output_dir, relative_root, parsed_df) """ - split_char = settings.SPLIT_CHAR or "/" if output_dir is None: raise ValueError("output_dir is required for fragment parsing") @@ -91,13 +92,22 @@ def parse_fragment( logger.debug(f"Fragment filename: {filename}") # Build relative_root and full_output_dir - kb_dir_parts = kb_dir.split(split_char) - relative_root = "/".join(kb_dir_parts + [filename]) - full_output_dir = os.path.join(output_dir, relative_root.replace("/", os.sep)) - sanitized_output_dir = path_handle(full_output_dir, mode="sanitize") - if not isinstance(sanitized_output_dir, str): - raise ValueError("sanitized fragment output path must be a string") - full_output_dir = sanitized_output_dir + namespace_segment = build_namespace_path_segment(namespace) + filename_segment = path_handle(filename, mode="clean_single") + if not isinstance(filename_segment, str) or filename_segment in {"", ".", ".."}: + filename_segment = f"fragment_{os.urandom(4).hex()}.fragment" + relative_root = f"{namespace_segment}/{filename_segment}" + full_output_dir = os.path.realpath( + os.path.join(output_dir, namespace_segment, filename_segment) + ) + resolved_output_dir = os.path.realpath(output_dir) + if ( + os.path.commonpath([resolved_output_dir, full_output_dir]) + != resolved_output_dir + ): + raise ValueError( + f"Fragment output directory escaped workspace: {full_output_dir}" + ) os.makedirs(full_output_dir, exist_ok=True) logger.debug(f"Fragment relative_root: {relative_root}") diff --git a/apps/worker/app/services/document_parser/formats/image/parser.py b/apps/worker/app/services/document_parser/formats/image/parser.py index 4562a2135..357211437 100755 --- a/apps/worker/app/services/document_parser/formats/image/parser.py +++ b/apps/worker/app/services/document_parser/formats/image/parser.py @@ -97,12 +97,12 @@ def local_image_to_data_url(path, cut=True, min_size=None, max_size=None): return img_data_base64 -def process_img_path4read(paths_, kb_dir, cut): +def process_img_path4read(paths_, image_root_dir, cut): urls = [] for path_ in paths_: if not is_remote(path_): - kb_dir = Path(kb_dir).resolve() - url_ = local_image_to_data_url(kb_dir / path_, cut) + resolved_image_root = Path(image_root_dir).resolve() + url_ = local_image_to_data_url(resolved_image_root / path_, cut) if url_ is not None: urls.append(url_) else: @@ -112,7 +112,7 @@ def process_img_path4read(paths_, kb_dir, cut): def ask_image( client: OpenAICompatibleClientSync, - kb_dir, + image_root_dir, paths_, title_text="", task="summary-images", @@ -138,7 +138,7 @@ def ask_image( if not valid_paths: return None - urls_ = process_img_path4read(valid_paths, kb_dir, size_cut) + urls_ = process_img_path4read(valid_paths, image_root_dir, size_cut) if task in ("summary-images", "atlas-page-info"): image_model = settings.IMAGE_MODEL or "gpt-4-vision-preview" @@ -187,14 +187,14 @@ def ask_image( return None -def detect_summary_img_md(line, last_context, kb_dir, mode=False): +def detect_summary_img_md(line, last_context, image_root_dir, mode=False): client = _get_vision_client() imgs = [] img_paths = re.findall(MD_IMAGE_PATTERN, line, flags=re.IGNORECASE) for i, ip in enumerate(img_paths): if mode: try: - llm_resp = ask_image(client, kb_dir, paths_=[ip]) + llm_resp = ask_image(client, image_root_dir, paths_=[ip]) if llm_resp: from app.services.document_parser.formats.text.parser import ( split_title_summary, 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 55d0e4e74..48d05396d 100644 --- a/apps/worker/app/services/document_parser/orchestration/format_adapters.py +++ b/apps/worker/app/services/document_parser/orchestration/format_adapters.py @@ -28,7 +28,7 @@ def parse(self, session: ParseSession) -> ParseOutput: session.fragment_content, filename=session.filename, output_dir=session.output_dir, - kb_dir=session.kb_dir, + namespace=session.namespace, base_llm_paras=session.base_llm_paras, ) return ParseOutput(output_dir=full_output_dir, parsed_df=parsed_df) diff --git a/apps/worker/app/services/document_parser/orchestration/namespace_path_segment.py b/apps/worker/app/services/document_parser/orchestration/namespace_path_segment.py new file mode 100644 index 000000000..f07cfc32e --- /dev/null +++ b/apps/worker/app/services/document_parser/orchestration/namespace_path_segment.py @@ -0,0 +1,18 @@ +from __future__ import annotations + +from app.services.common.file_utils import path_handle + +_DEFAULT_NAMESPACE_SEGMENT = "default" + + +def build_namespace_path_segment(namespace: str | None) -> str: + """Map a retrieval namespace to one safe parser path segment.""" + raw_namespace = (namespace or _DEFAULT_NAMESPACE_SEGMENT).strip() + sanitized_namespace = path_handle(raw_namespace, mode="clean_single") + if not isinstance(sanitized_namespace, str): + return _DEFAULT_NAMESPACE_SEGMENT + + namespace_segment = sanitized_namespace.strip() + if namespace_segment in {"", ".", ".."}: + return _DEFAULT_NAMESPACE_SEGMENT + return namespace_segment diff --git a/apps/worker/app/services/document_parser/orchestration/parse_input.py b/apps/worker/app/services/document_parser/orchestration/parse_input.py index 89249b3d0..92c4749a1 100644 --- a/apps/worker/app/services/document_parser/orchestration/parse_input.py +++ b/apps/worker/app/services/document_parser/orchestration/parse_input.py @@ -22,7 +22,7 @@ class ParseInput: output_dir: str internal_output_filename: str job_id: str | None = None - kb_dir: str = "Default_Root" + namespace: str = "default" options: ParseOptions = field(default_factory=ParseOptions) base_url: str = "" fragment_content: str = "" 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 14670b8f8..1b35c3cf5 100644 --- a/apps/worker/app/services/document_parser/orchestration/parse_session.py +++ b/apps/worker/app/services/document_parser/orchestration/parse_session.py @@ -5,14 +5,16 @@ from typing import Any from app.services.document_parser.formats.atlas.classifier import classify_atlas_with_vlm -from app.services.document_parser.profiling.doc_profiler import profile_document +from app.services.document_parser.orchestration.namespace_path_segment import ( + build_namespace_path_segment, +) from app.services.document_parser.orchestration.parse_input import ParseInput +from app.services.document_parser.profiling.doc_profiler import profile_document from app.services.document_parser.support.stage_profiler import stage_timer from loguru import logger from shared.core.config import settings from shared.core.exceptions.domain_exceptions import ValidationException -from app.services.common.file_utils import path_handle PDF_PAGE_LIMIT = 600 @@ -27,7 +29,7 @@ class ParseSession: full_output_dir: str internal_output_filename: str job_id: str | None - kb_dir: str + namespace: str output_dir: str profile: Any relative_root: str @@ -52,7 +54,7 @@ def from_input( 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, + namespace=parse_input.namespace, output_dir=parse_input.output_dir, profile=profile, relative_root=relative_root, @@ -82,7 +84,7 @@ def build_parse_session(parse_input: ParseInput) -> ParseSession: relative_root, full_output_dir = _resolve_output_paths( filename=parse_input.filename, internal_output_filename=parse_input.internal_output_filename, - kb_dir=parse_input.kb_dir, + namespace=parse_input.namespace, output_dir=parse_input.output_dir, ) logger.debug(f"relative_root: {relative_root}") @@ -131,7 +133,7 @@ def build_parse_session(parse_input: ParseInput) -> ParseSession: _rename_atlas_output( filename=parse_input.filename, internal_output_filename=parse_input.internal_output_filename, - kb_dir=parse_input.kb_dir, + namespace=parse_input.namespace, output_dir=parse_input.output_dir, ) ) @@ -142,7 +144,7 @@ def build_parse_session(parse_input: ParseInput) -> ParseSession: output_dir=parse_input.output_dir, internal_output_filename=internal_output_filename, job_id=parse_input.job_id, - kb_dir=parse_input.kb_dir, + namespace=parse_input.namespace, options=parse_input.options, base_url=parse_input.base_url, fragment_content=parse_input.fragment_content, @@ -162,7 +164,7 @@ def _rename_atlas_output( *, filename: str, internal_output_filename: str, - kb_dir: str, + namespace: str, output_dir: str, ) -> tuple[str, str, str, str]: name_base, _ = os.path.splitext(filename) @@ -172,7 +174,7 @@ def _rename_atlas_output( relative_root, full_output_dir = _resolve_output_paths( filename=atlas_filename, internal_output_filename=atlas_internal_filename, - kb_dir=kb_dir, + namespace=namespace, output_dir=output_dir, ) return atlas_filename, atlas_internal_filename, relative_root, full_output_dir @@ -182,30 +184,34 @@ def _resolve_output_paths( *, filename: str, internal_output_filename: str, - kb_dir: str, + namespace: 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) + namespace_segment = build_namespace_path_segment(namespace) + relative_root = ( + f"{namespace_segment}/{filename}" if filename else namespace_segment + ) + internal_relative_root = ( + f"{namespace_segment}/{internal_output_filename}" + if internal_output_filename + else namespace_segment + ) full_output_dir = os.path.join( output_dir, - internal_relative_root.replace("/", os.sep), + namespace_segment, + internal_output_filename, ) - 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) + resolved_output_dir = os.path.realpath(output_dir) + resolved_full_output_dir = os.path.realpath(full_output_dir) + if ( + os.path.commonpath([resolved_output_dir, resolved_full_output_dir]) + != resolved_output_dir + ): + raise ValueError( + f"Parser output directory escaped task workspace: {full_output_dir}" + ) + os.makedirs(resolved_full_output_dir, exist_ok=True) logger.debug(f"internal_relative_root: {internal_relative_root}") - return relative_root, sanitized_output_dir + return relative_root, resolved_full_output_dir diff --git a/apps/worker/app/services/document_parser/parse_service.py b/apps/worker/app/services/document_parser/parse_service.py index 6266d5990..a9e9e7505 100644 --- a/apps/worker/app/services/document_parser/parse_service.py +++ b/apps/worker/app/services/document_parser/parse_service.py @@ -15,7 +15,7 @@ def checkerboard_parse_output( output_dir: str, internal_output_filename: str, job_id: str | None = None, - kb_dir: str = "Default_Root", + namespace: str = "default", llm_histories: int = 5, smart_title_parse: bool = True, summary_image: bool = True, @@ -34,7 +34,6 @@ def checkerboard_parse_output( filename=filename, internal_output_filename=internal_output_filename, job_id=job_id, - kb_dir=kb_dir, output_dir=output_dir, options=ParseOptions( add_frag_desc=add_frag_desc, @@ -48,6 +47,7 @@ def checkerboard_parse_output( ), base_url=base_url, fragment_content=fragment_content, + namespace=namespace, s3_key=s3_key, ) return run_parse_pipeline(parse_input) @@ -59,7 +59,7 @@ def checkerboard_inject_parse( output_dir: str, internal_output_filename: str, job_id: str | None = None, - kb_dir: str = "Default_Root", + namespace: str = "default", llm_histories: int = 5, smart_title_parse: bool = True, summary_image: bool = True, @@ -79,7 +79,7 @@ def checkerboard_inject_parse( output_dir=output_dir, internal_output_filename=internal_output_filename, job_id=job_id, - kb_dir=kb_dir, + namespace=namespace, llm_histories=llm_histories, smart_title_parse=smart_title_parse, summary_image=summary_image, diff --git a/apps/worker/tests/contract/test_excel_parser_contract.py b/apps/worker/tests/contract/test_excel_parser_contract.py index 5d80dd3cf..f7992a933 100644 --- a/apps/worker/tests/contract/test_excel_parser_contract.py +++ b/apps/worker/tests/contract/test_excel_parser_contract.py @@ -1,5 +1,6 @@ from __future__ import annotations +import os from pathlib import Path @@ -45,7 +46,7 @@ def test_xlsx_parser_contract_uses_stable_entrypoint_and_ignores_hidden_sheets( stopwords=[], ) - assert full_output_dir.endswith("Default_Root/budget.xlsx") + assert full_output_dir.endswith("default/budget.xlsx") assert parsed_df is not None assert parsed_df["type"].tolist() == ["table"] assert parsed_df["path"].tolist() == ["tables/table-Visible.html"] @@ -60,3 +61,35 @@ def test_xlsx_parser_contract_uses_stable_entrypoint_and_ignores_hidden_sheets( assert "10" in table_html_text assert "Secret" not in table_html_text assert "Hidden" not in table_html_text + + +def test_parser_maps_namespace_to_task_local_path_segment( + 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" + output_root = tmp_path / "parser-output" + _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(output_root), + internal_output_filename="budget.xlsx", + namespace="/tmp/../images", + summary_image=False, + summary_table=False, + summary_txt=False, + smart_title_parse=False, + stopwords=[], + ) + + assert ( + os.path.commonpath([str(output_root.resolve()), full_output_dir]) + == str(output_root.resolve()) + ) + assert full_output_dir.endswith("_tmp_.._images/budget.xlsx") + assert parsed_df is not None + assert parsed_df["path"].tolist() == ["tables/table-Visible.html"] diff --git a/apps/worker/tests/contract/test_parse_task_contract.py b/apps/worker/tests/contract/test_parse_task_contract.py index 1a1ce4d64..86b76f00b 100644 --- a/apps/worker/tests/contract/test_parse_task_contract.py +++ b/apps/worker/tests/contract/test_parse_task_contract.py @@ -28,7 +28,7 @@ def _build_pending_file_job_metadata(source_file_name: str) -> dict[str, Any]: "namespace": "worker-contract", "source_type": "file", "source_file_name": source_file_name, - "kb_dir": "Default_Root", + "parsing_params": {"kb_dir": "legacy-ignored"}, } return job_metadata @@ -210,7 +210,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: captured_artifacts["parse_kwargs"] = kwargs output_dir = ( Path(str(kwargs["output_dir"])) - / str(kwargs["kb_dir"]) + / str(kwargs["namespace"]) / str(kwargs["internal_output_filename"]) ) images_dir = output_dir / "images" @@ -225,7 +225,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: parsed_rows: list[dict[str, Any]] = [ { "content": text_content_with_refs, - "path": f"Default_Root/{file_root}/公司研究/自主可控加强，寒武纪或迎来营收快速放量周期", + "path": f"worker-contract/{file_root}/公司研究/自主可控加强，寒武纪或迎来营收快速放量周期", "type": "text", "length": len(text_content_with_refs), "keywords": "", @@ -246,7 +246,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: }, { "content": "chunk-2", - "path": f"Default_Root/{file_root}/相关研报/要点", + "path": f"worker-contract/{file_root}/相关研报/要点", "type": "text", "length": 7, "keywords": "", @@ -259,7 +259,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: }, { "content": "image caption", - "path": f"Default_Root/{file_root}/images/page-1.png", + "path": f"worker-contract/{file_root}/images/page-1.png", "type": "image", "length": 13, "keywords": "", @@ -272,7 +272,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: }, { "content": "table content", - "path": f"Default_Root/{file_root}/tables/table-1.html", + "path": f"worker-contract/{file_root}/tables/table-1.html", "type": "table", "length": 13, "keywords": "", @@ -361,7 +361,11 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: assert captured_artifacts["parse_kwargs"]["filename"] == source_file_name assert captured_artifacts["parse_kwargs"]["internal_output_filename"] == source_file_name assert Path(str(captured_artifacts["parse_kwargs"]["file_full_path"])).name == source_file_name - assert captured_artifacts["result_dir"].endswith("Default_Root/contract-parse.pdf") + assert captured_artifacts["parse_kwargs"]["namespace"] == "worker-contract" + assert "kb_dir" not in captured_artifacts["parse_kwargs"] + assert captured_artifacts["result_dir"].endswith( + "worker-contract/contract-parse.pdf" + ) assert captured_artifacts["doc_nav"]["file_name"] == source_file_name assert captured_artifacts["doc_nav"]["sections"][0]["title"] == "公司研究" assert captured_artifacts["manifest"]["HIERARCHY"] == { @@ -745,7 +749,7 @@ def test_should_export_full_result_when_publication_deduplicates_existing_chunks :result_id, 'text', 'already published text', - 'Default_Root/existing.pdf/Section/Duplicate text', + 'worker-contract/existing.pdf/Section/Duplicate text', NULL, CAST(:text_metadata AS JSON), 0, @@ -760,7 +764,7 @@ def test_should_export_full_result_when_publication_deduplicates_existing_chunks :result_id, 'image', 'already published image', - 'Default_Root/existing.pdf/images/duplicate.png', + 'worker-contract/existing.pdf/images/duplicate.png', 'images/duplicate.png', CAST(:image_metadata AS JSON), 1, @@ -825,7 +829,7 @@ def fake_download_s3_file_to_temp( def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: output_dir = ( Path(str(kwargs["output_dir"])) - / str(kwargs["kb_dir"]) + / str(kwargs["namespace"]) / str(kwargs["internal_output_filename"]) ) images_dir = output_dir / "images" @@ -837,7 +841,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: parsed_rows: list[dict[str, Any]] = [ { "content": "duplicate text", - "path": f"Default_Root/{file_root}/Section/Duplicate text", + "path": f"worker-contract/{file_root}/Section/Duplicate text", "type": "text", "length": 14, "keywords": "", @@ -850,7 +854,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: }, { "content": "duplicate image", - "path": f"Default_Root/{file_root}/images/duplicate.png", + "path": f"worker-contract/{file_root}/images/duplicate.png", "type": "image", "length": 15, "keywords": "", @@ -863,7 +867,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: }, { "content": "new text", - "path": f"Default_Root/{file_root}/Section/New text", + "path": f"worker-contract/{file_root}/Section/New text", "type": "text", "length": 8, "keywords": "", @@ -1064,7 +1068,7 @@ def fake_estimate_workload(file_path: str) -> Any: def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: output_dir = ( Path(str(kwargs["output_dir"])) - / str(kwargs["kb_dir"]) + / str(kwargs["namespace"]) / str(kwargs["internal_output_filename"]) ) output_dir.mkdir(parents=True, exist_ok=True) @@ -1074,7 +1078,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: parsed_rows: list[dict[str, Any]] = [ { "content": "chunk body", - "path": f"Default_Root/{file_root}/Section/Point", + "path": f"worker-contract/{file_root}/Section/Point", "type": "text", "length": 10, "keywords": "", diff --git a/packages/shared-python/shared/models/schemas/job.py b/packages/shared-python/shared/models/schemas/job.py index 20bad6763..692b810e3 100644 --- a/packages/shared-python/shared/models/schemas/job.py +++ b/packages/shared-python/shared/models/schemas/job.py @@ -17,9 +17,6 @@ class ParsingParams(BaseModel): model: Literal["base", "advanced"] = Field("base", description="Parsing model") ocr_enabled: bool = Field(False, description="Enable OCR") - kb_dir: Optional[str] = Field( - "Default_Root", description="Knowledge-base directory" - ) doc_type: Literal["auto", "pdf", "docx", "txt", "md"] = Field( "auto", description="Document type" ) diff --git a/packages/shared-python/shared/services/ai/prompt_service.py b/packages/shared-python/shared/services/ai/prompt_service.py index 9ea759675..339e00d7d 100755 --- a/packages/shared-python/shared/services/ai/prompt_service.py +++ b/packages/shared-python/shared/services/ai/prompt_service.py @@ -3,7 +3,7 @@ Only contains prompts required for document parsing workflow. Removed prompts for: -- RAG/Knowledge Base: talk-kb, merge-answers, judge-kb, rerank, connect-kb, detect-contradict +- Retrieval corpus: talk, merge-answers, judge, rerank, connect, detect-contradict - Document Generation: gen-titles-oneoff, gen-root-titles, gen-thoughts, reason-content-layout, rewrite-paras, rewrite-sentence, construct-table, reason-source - Table Filling: filling-tb-kv, filling-tb-ck diff --git a/packages/shared-python/shared/services/redis/key_builder.py b/packages/shared-python/shared/services/redis/key_builder.py index 62d3d5b28..02d0dce78 100644 --- a/packages/shared-python/shared/services/redis/key_builder.py +++ b/packages/shared-python/shared/services/redis/key_builder.py @@ -10,7 +10,6 @@ class RedisKeyType(Enum): USER = "user" TASK = "task" CONVERSATION = "conversation" - KNOWLEDGE_BASE = "kb" SESSION = "session" CACHE = "cache" QUEUE = "queue" @@ -106,24 +105,6 @@ def conversation_context(self, conversation_id: str) -> str: """Conversation context key.""" return self.build_key(RedisKeyType.CONVERSATION, conversation_id, "context") - # ==================== Knowledge Base Keys ==================== - - def kb_status(self, user_id: str) -> str: - """Knowledge-base status key.""" - return self.build_key(RedisKeyType.KNOWLEDGE_BASE, user_id, "status") - - def kb_vectors(self, user_id: str) -> str: - """Knowledge-base vectors key.""" - return self.build_key(RedisKeyType.KNOWLEDGE_BASE, user_id, "vectors") - - def kb_metadata(self, user_id: str) -> str: - """Knowledge-base metadata key.""" - return self.build_key(RedisKeyType.KNOWLEDGE_BASE, user_id, "metadata") - - def kb_index(self, user_id: str) -> str: - """Knowledge-base index key.""" - return self.build_key(RedisKeyType.KNOWLEDGE_BASE, user_id, "index") - # ==================== Session Keys ==================== def session_data(self, session_id: str) -> str: @@ -270,7 +251,6 @@ def get_key_ttl(self, key_type: RedisKeyType) -> int: RedisKeyType.USER: 86400, # 1 day (user_config cache). RedisKeyType.TASK: 86400, # 1 day. RedisKeyType.CONVERSATION: 3600 * 2, # 2 hours. - RedisKeyType.KNOWLEDGE_BASE: 86400 * 30, # 30 days. RedisKeyType.SESSION: 3600, # 1 hour. RedisKeyType.CACHE: 3600, # 1 hour. RedisKeyType.QUEUE: 86400, # 1 day. diff --git a/packages/shared-python/shared/services/retrieval/agentic/prompts.py b/packages/shared-python/shared/services/retrieval/agentic/prompts.py index 7c6295913..418ae46bc 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/prompts.py +++ b/packages/shared-python/shared/services/retrieval/agentic/prompts.py @@ -10,17 +10,17 @@ You are a document routing assistant. {budget_block} -Below is a knowledge base overview showing all available documents, +Below is a document corpus overview showing all available documents, their navigation summaries, chunk counts, and media counts. -=== Knowledge Base Overview === +=== Document Corpus 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 []. +If NO document in the corpus 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. """ diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index d7f13538a..b47c7351b 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -23,9 +23,9 @@ 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 + small_corpus_outcome = await _try_run_small_corpus_route(context) + if small_corpus_outcome is not None: + return small_corpus_outcome if _should_use_agentic_route(context.use_agentic): return await _run_agentic_route(context) @@ -33,7 +33,7 @@ async def run_retrieval_route( return await run_legacy_retrieval_route(context) -async def _try_run_small_kb_route( +async def _try_run_small_corpus_route( context: RetrievalRouteContext, ) -> RetrievalRouteOutcome | None: try: @@ -46,7 +46,7 @@ async def _try_run_small_kb_route( ) except Exception as exc: logger.warning( - f"Failed to count scoped chunks, skipping small KB optimization: {exc}" + f"Failed to count scoped chunks, skipping small corpus optimization: {exc}" ) total_chunk_count = context.top_k + 1 @@ -55,7 +55,7 @@ async def _try_run_small_kb_route( return None logger.info( - f" Small KB optimization: {total_chunk_count} chunks " + f" Small corpus optimization: {total_chunk_count} chunks " f"<= top_k={context.top_k}, returning all" ) all_rows = await load_all_scoped_chunks( @@ -69,7 +69,7 @@ async def _try_run_small_kb_route( filter_mode=context.filter_mode, ) logger.info( - f" small_kb load: loaded={len(all_rows)} rows after signal/exclude filters" + f" small_corpus load: loaded={len(all_rows)} rows after signal/exclude filters" ) assembled_rows = await assemble_retrieval_results( db=context.db, @@ -82,13 +82,13 @@ async def _try_run_small_kb_route( response = { "namespace": context.namespace, "query": context.query, - "router_used": "small_kb_all", + "router_used": "small_corpus_all", "results": results, } return RetrievalRouteOutcome( response=response, hit_stats_results=results, - completion_label="Small KB", + completion_label="Small corpus", completion_count=len(results), completion_detail="results", ) diff --git a/packages/shared-python/shared/services/retrieval/search/lexical_text.py b/packages/shared-python/shared/services/retrieval/search/lexical_text.py index 42d15366f..35ce5f8ae 100644 --- a/packages/shared-python/shared/services/retrieval/search/lexical_text.py +++ b/packages/shared-python/shared/services/retrieval/search/lexical_text.py @@ -60,13 +60,13 @@ def build_content_lexical_text(chunk: dict[str, Any]) -> Optional[str]: def section_path_from_chunk_path(source_path: Optional[str]) -> str: """Extract section hierarchy from chunk path. - Expected format: "/.ext/
//..." + Expected format: "/.ext/
//..." Returns " / "-joined section parts, or "Root" if no section hierarchy. """ if not source_path: return "Root" parts = split_section_path(source_path) - section_parts = parts[2:] # skip kb_root + filename + section_parts = parts[2:] # skip namespace + filename if not section_parts: return "Root" return normalize_section_path(" / ".join(section_parts)) diff --git a/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py b/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py index 3fec8dbda..234927d04 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py +++ b/packages/shared-python/shared/services/retrieval/workflow/orchestrator.py @@ -133,8 +133,8 @@ async def run_request( 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, + corpus_total_docs=total_docs, + corpus_total_chunks=total_chunks, ) wallet = BudgetWallet( diff --git a/packages/shared-python/shared/services/retrieval/workflow/plan_service.py b/packages/shared-python/shared/services/retrieval/workflow/plan_service.py index 11e7897f3..dec625979 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/plan_service.py +++ b/packages/shared-python/shared/services/retrieval/workflow/plan_service.py @@ -25,8 +25,8 @@ async def load_or_create( max_steps: int, wallet_total: int, per_retrieve: int, - kb_total_docs: int, - kb_total_chunks: int, + corpus_total_docs: int, + corpus_total_chunks: int, ) -> QueryPlan: try: cached = await get_cached_workflow_plan(user_id=user_id, namespace=namespace, query=query) @@ -44,8 +44,8 @@ async def load_or_create( ) plan = await planner.plan( query=query, - kb_total_docs=kb_total_docs, - kb_total_chunks=kb_total_chunks, + corpus_total_docs=corpus_total_docs, + corpus_total_chunks=corpus_total_chunks, ) try: await set_cached_workflow_plan( diff --git a/packages/shared-python/shared/services/retrieval/workflow/planner.py b/packages/shared-python/shared/services/retrieval/workflow/planner.py index 7f4937b80..d3a4294d8 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/planner.py +++ b/packages/shared-python/shared/services/retrieval/workflow/planner.py @@ -34,7 +34,7 @@ You are a retrieval workflow planner. Think step by step before answering. User query: {query} -Knowledge base inventory: {kb_total_docs} docs / {kb_total_chunks} chunks. +Document corpus inventory: {corpus_total_docs} docs / {corpus_total_chunks} chunks. Wallet status: total_budget={total_budget} tokens (planner_used={planner_used}). Decide whether the query needs decomposition into multiple sub-queries. @@ -80,16 +80,16 @@ async def plan( self, *, query: str, - kb_total_docs: int = 0, - kb_total_chunks: int = 0, + corpus_total_docs: int = 0, + corpus_total_chunks: int = 0, ) -> QueryPlan: if self._llm_fn is None: return QueryPlan.single_step(query, reason="planner_llm_unavailable") prompt = _PLANNER_PROMPT.format( query=query, - kb_total_docs=kb_total_docs, - kb_total_chunks=kb_total_chunks, + corpus_total_docs=corpus_total_docs, + corpus_total_chunks=corpus_total_chunks, total_budget=self._total_budget, planner_used=self._planner_used(), max_steps=self._max_steps, From 5835ac072559713902a6894977d144e3f4e04bd5 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 17:01:46 +0000 Subject: [PATCH 74/78] refactor: decouple parser paths from namespace --- AGENTS.md | 11 +- CONTEXT.md | 4 +- apps/api/app/api/v1/routes/demo.py | 16 +- apps/api/app/api/v1/routes/documents.py | 5 +- apps/api/app/api/v1/routes/retrieval.py | 12 +- .../demo_documents/tsla-q4-2025/chunks.json | 72 ++--- .../demo_documents/tsla-q4-2025/doc_nav.json | 294 +++++++++++++++++- apps/api/app/mcp/retrieval_server.py | 8 +- apps/api/app/services/demo/source_catalog.py | 6 +- .../app/services/demo/source_projection.py | 7 +- .../document_ingestion/scope_service.py | 7 +- .../test_chunk_document_path_contract.py | 247 +++++++++++++++ .../contract/test_job_creation_contract.py | 38 +++ .../connect_builder/summary_builder.py | 8 +- .../parse_result_package.py | 52 ++-- .../document_ingestion/processing_run.py | 18 +- .../formats/fragment/parser.py | 18 +- .../orchestration/format_adapters.py | 1 - .../orchestration/namespace_path_segment.py | 18 -- .../orchestration/parse_input.py | 1 - .../orchestration/parse_output.py | 13 +- .../orchestration/parse_session.py | 30 +- .../orchestration/path_segment.py | 19 ++ .../services/document_parser/parse_service.py | 45 --- .../contract/test_excel_parser_contract.py | 26 +- .../contract/test_parse_task_contract.py | 57 ++-- .../shared/models/schemas/job.py | 4 +- .../shared/models/schemas/job_metadata.py | 8 +- .../models/schemas/retrieval_namespace.py | 9 + .../shared/services/chunks/document_path.py | 156 ++++++++++ .../services/jobs/lifecycle/publication.py | 4 +- .../shared/services/retrieval/app_service.py | 3 +- .../services/retrieval/cache_service.py | 7 +- .../retrieval/execution/query_request.py | 3 +- .../services/retrieval/publication_service.py | 19 +- .../services/retrieval/search/lexical_text.py | 28 +- .../services/storage/zip_result_schema.py | 18 +- 37 files changed, 1011 insertions(+), 281 deletions(-) create mode 100644 apps/api/tests/contract/test_chunk_document_path_contract.py delete mode 100644 apps/worker/app/services/document_parser/orchestration/namespace_path_segment.py create mode 100644 apps/worker/app/services/document_parser/orchestration/path_segment.py create mode 100644 packages/shared-python/shared/models/schemas/retrieval_namespace.py create mode 100644 packages/shared-python/shared/services/chunks/document_path.py diff --git a/AGENTS.md b/AGENTS.md index efbc89049..6582747be 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -73,7 +73,7 @@ flowchart TB end subgraph PARSE["② Document Parsing (Worker)"] - Queue --> Router["parse_service.checkerboard_inject_parse"] + Queue --> Router["parse_service.checkerboard_parse_output"] Router --> Profiler["profiling.doc_profiler.profile_document"] Profiler --> PDF["formats.pdf.parser → MinerU"] Profiler --> DOCX["formats.docx.parser.parse_docx"] @@ -120,18 +120,15 @@ flowchart TB ### Entry Point `apps/worker/app/services/document_parser/parse_service.py` → -`checkerboard_inject_parse()` +`checkerboard_parse_output()` -This is the legacy tuple-compatible entry for all file types. Internally, -`checkerboard_parse_output()` and parser adapters use the typed `ParseOutput` -contract. The parser flow: +This is the typed `ParseOutput` entry for all file types. The parser flow: 1. **Profiles** the document via `profiling.doc_profiler.profile_document()` to detect file type, page count, and special categories (e.g. `atlas`). 2. **Routes** to the appropriate parser based on file extension. 3. **Post-processes**: cleans up unreferenced images, compresses PNG→JPG. -4. Returns typed parse output internally, with `(output_dir, parsed_df)` kept as - the compatibility shape for older callers. +4. Returns typed parse output with task-local artifact paths. ### Parser Routing Table diff --git a/CONTEXT.md b/CONTEXT.md index 24d47b7bf..53a329fb9 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -411,8 +411,8 @@ side effects. - `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 - legacy parser entrypoint; parser option shaping, format routing, rendered PDF +- Worker Document Parsing exposes `checkerboard_parse_output` as the stable + parser entrypoint; parser option shaping, format routing, rendered PDF transforms, typed Parse Output, and heading inference stay behind that entrypoint. - A Job and a Document are not the same thing. Jobs track intake and processing; diff --git a/apps/api/app/api/v1/routes/demo.py b/apps/api/app/api/v1/routes/demo.py index c8559efc7..4f8dbf6eb 100644 --- a/apps/api/app/api/v1/routes/demo.py +++ b/apps/api/app/api/v1/routes/demo.py @@ -9,11 +9,12 @@ from app.services.rate_limit.data_structures import CurrentUser from fastapi import APIRouter, Depends, Query from fastapi.responses import FileResponse -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, field_validator from sqlalchemy.ext.asyncio import AsyncSession from shared.core.database import get_db from shared.core.exceptions.domain_exceptions import NotFoundException +from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace router = APIRouter(tags=["Demo Documents"]) @@ -23,13 +24,22 @@ class DemoMaterializeRequest(BaseModel): """Request to copy selected canonical demo sources into a namespace.""" - namespace: str | None = Field(None, description="Target retrieval namespace") + namespace: str | None = Field( + None, + max_length=255, + description="Target retrieval namespace", + ) demo_source_ids: list[str] = Field( default_factory=list, min_length=1, description="Canonical demo source IDs to materialize", ) + @field_validator("namespace") + @classmethod + def normalize_namespace(cls, namespace: str | None) -> str: + return normalize_retrieval_namespace(namespace) + @router.get("/catalog") async def get_demo_catalog() -> dict[str, Any]: @@ -117,7 +127,7 @@ async def materialize_demo_sources( db: AsyncSession = Depends(get_db), ) -> dict[str, Any]: """Copy canonical demo sources into the authenticated user's namespace.""" - namespace = (payload.namespace or "default").strip() or "default" + namespace = normalize_retrieval_namespace(payload.namespace) try: materialized_sources = await _demo_service.materialize_sources( db, diff --git a/apps/api/app/api/v1/routes/documents.py b/apps/api/app/api/v1/routes/documents.py index 1cbcf7ad9..a77c5d122 100644 --- a/apps/api/app/api/v1/routes/documents.py +++ b/apps/api/app/api/v1/routes/documents.py @@ -12,6 +12,7 @@ from shared.core.database import get_db from shared.core.exceptions.domain_exceptions import NotFoundException +from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace router = APIRouter(tags=["Documents"]) @@ -41,11 +42,11 @@ async def _archive_document_response( @router.get("") async def list_documents( - namespace: str | None = Query(None), + namespace: str | None = Query(None, max_length=255), current_user: CurrentUser = Depends(with_current_user), db: AsyncSession = Depends(get_db), ): - effective_namespace = namespace or "default" + effective_namespace = normalize_retrieval_namespace(namespace) documents = await _document_service.list_documents( db, user_id=current_user.user_id, diff --git a/apps/api/app/api/v1/routes/retrieval.py b/apps/api/app/api/v1/routes/retrieval.py index 61037c4b3..e7272f6c4 100644 --- a/apps/api/app/api/v1/routes/retrieval.py +++ b/apps/api/app/api/v1/routes/retrieval.py @@ -11,6 +11,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from shared.core.database import get_db +from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace from shared.services.retrieval import run_retrieval_query router = APIRouter(tags=["Retrieval"]) @@ -23,7 +24,9 @@ class ExcludeSection(BaseModel): class RetrievalQueryRequest(BaseModel): namespace: str | None = Field( - None, description="Effective namespace; defaults to default" + None, + max_length=255, + description="Effective namespace; defaults to default", ) query: str top_k: int = 10 @@ -67,6 +70,11 @@ def validate_channels(cls, v: list[str]) -> list[str]: raise ValueError(f"Invalid channel: {ch}. Must be one of {valid}") return v + @field_validator("namespace") + @classmethod + def normalize_namespace(cls, namespace: str | None) -> str: + return normalize_retrieval_namespace(namespace) + class RetrievalQueryResponse(BaseModel): namespace: str @@ -86,7 +94,7 @@ async def query_retrieval( return await run_retrieval_query( db=db, user_id=current_user.user_id, - namespace=payload.namespace or "default", + namespace=normalize_retrieval_namespace(payload.namespace), query=payload.query, top_k=payload.top_k, exclude_document_ids=payload.exclude_document_ids, diff --git a/apps/api/app/data/demo_documents/tsla-q4-2025/chunks.json b/apps/api/app/data/demo_documents/tsla-q4-2025/chunks.json index 7e7147254..177ae25c9 100644 --- a/apps/api/app/data/demo_documents/tsla-q4-2025/chunks.json +++ b/apps/api/app/data/demo_documents/tsla-q4-2025/chunks.json @@ -25,7 +25,7 @@ "chunk_id": "68a6be7d-c587-5c73-abf2-56f4686e28e6", "type": "text", "content": "[tables/table-0 Tesla 2025 Results.html]", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->HIGHLIGHTS", + "path": "TSLA-Q4-2025-Update.pdf-->HIGHLIGHTS", "metadata": { "length": 83, "summary": "", @@ -138,7 +138,7 @@ "chunk_id": "60109008-6261-51e6-b202-093d904eb881", "type": "text", "content": "FINANCIAL SUMMARY\n(Unaudited)\n\n[tables/table-1 Q4 2025 Financials.html]\n\n(1) As a result of the adoption of the new crypto assets standard, the previously reported quarterly periods in 2024 have been recast.\n(2) Beginning in Q1'25, Adjusted EBITDA (non-GAAP) is presented net of digital assets gains and losses and all prior periods have been adjusted.\n(3) Beginning in Q1'25, Net income attributable to common stockholders (non-GAAP) is presented net of digital assets gains and losses and all prior periods have been adjusted.\n(4) Beginning in Q1'25, Capital expenditures is presented inclusive of purchases of energy generation and storage systems and all prior periods have been adjusted.\nFINANCIAL SUMMARY\n(Unaudited)\n\n[tables/table-2 Financial Data 2021-25.html]\n\n(1) Beginning in Q1'25, Adjusted EBITDA (non-GAAP) is presented net of digital assets gains and losses and all prior periods have been adjusted.\n(2) Beginning in Q1'25, Net income attributable to common stockholders (non-GAAP) is presented net of digital assets gains and losses and all prior periods have been adjusted.\n(3) Beginning in Q1'25, Capital expenditures is presented inclusive of purchases of energy generation and storage systems and all prior periods have been adjusted.\nOPERATIONAL SUMMARY\n(Unaudited)\n\n[tables/table-3 Tesla Q4-2025 Data.html]\n\nOPERATIONAL SUMMARY\n(Unaudited)\n\n[tables/table-4 Tesla 2021-2025 Data.html]", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY", + "path": "TSLA-Q4-2025-Update.pdf-->SUMMARY", "metadata": { "length": 1517, "summary": "The document presents unaudited financial and operational summaries for a company, covering quarterly data from Q4 2024 through Q4 2025 and annual data from 2021 to 2025. Key notes indicate significant accounting changes effective Q1 2025: Adjusted EBITDA and Net income attributable to common stockholders are now presented net of digital assets gains and losses, with all prior periods adjusted accordingly. Additionally, Capital expenditures now include purchases of energy generation and storage systems, requiring restatement of previous periods. The content references multiple tables detailing these metrics but does not display the specific numerical values.", @@ -240,7 +240,7 @@ "chunk_id": "60339310-5480-5ae2-8791-e6017bafb730", "type": "text", "content": "While automotive sales declined sequentially, gross margin (even when excluding the impact of regulatory credits) improved. The APAC region continued to show strength across multiple markets and set a record for deliveries in the quarter. We continued the rollout of Model Y variants across markets in Q4, including the standard and performance versions.\nPreparations continue in North America for the production ramps of Tesla Semi and Cybercab, both commencing 1H26, and production of the next-generation Roadster.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY-->Automotive", + "path": "TSLA-Q4-2025-Update.pdf-->SUMMARY-->Automotive", "metadata": { "length": 516, "summary": "", @@ -304,7 +304,7 @@ "chunk_id": "eb18e33e-b093-532e-b4df-bd05d63b0294", "type": "text", "content": "We achieved our highest quarterly energy storage deployments, driven by record Megapack deployments. Total gross profit rose, both sequentially and year-over-year, to a record \\$1.1 billion, marking the fifth consecutive record quarter. We plan to begin Megapack 3 and Megablock production at Megafactory Houston in 2026. In 2025, our global Powerwall network supported more than 89,000 Virtual Power Plant events across over 1 million installed units, allowing homeowners to save over \\$1 billion in electricity bills as Virtual Power Plant participation continues to scale rapidly.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY-->Energy generation and storage", + "path": "TSLA-Q4-2025-Update.pdf-->SUMMARY-->Energy generation and storage", "metadata": { "length": 583, "summary": "", @@ -395,7 +395,7 @@ "chunk_id": "60518074-bcbe-5b48-89a9-dfeb5d1b531b", "type": "text", "content": "We made further progress on the Optimus program in 2025. In Q1 of this year, we plan to unveil the Gen 3 version of Optimus, which will include major upgrades from version 2.5, including our latest hand design. The Gen 3 is our first design meant for mass production. Preparations are underway for the first production line, including supply chain readiness, with start of production planned before the end of 2026 and eventual planned capacity of 1 million robots per year.\nInstalled Annual Manufacturing Capacity\n\n[tables/table-5 Tesla Production.html]\n\nInstalled capacity ≠ current production rate and there may be limitations discovered as production rates approach capacity. Production rates depend on a variety of factors, including equipment uptime, component supply, downtime related to factory upgrades, regulatory considerations and other factors. Construction includes factory and infrastructure buildout as well as tool installation.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY-->Robotics", + "path": "TSLA-Q4-2025-Update.pdf-->SUMMARY-->Robotics", "metadata": { "length": 959, "summary": "", @@ -491,7 +491,7 @@ "chunk_id": "c49883a3-5834-5537-be53-35e58da70bf7", "type": "text", "content": "We are currently building Cortex 2 at Gigafactory Texas to further increase our AI training compute capacity. In the first half of 2026, we plan to more than double the size of onsite compute in Texas (in terms of H100 equivalents). We aim to maximize capital efficiency by scaling training compute judiciously, including when the training backlog gets too long or in anticipation of greater demand from our engineers to support our AI-related offerings.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY-->AI Training Compute", + "path": "TSLA-Q4-2025-Update.pdf-->SUMMARY-->AI Training Compute", "metadata": { "length": 454, "summary": "", @@ -545,7 +545,7 @@ "chunk_id": "fcbc3b02-dec6-5e91-849e-fccd3f22320e", "type": "text", "content": "Our lithium refinery commenced pilot production and is the first spodumene to lithium hydroxide refinery in North America, leveraging a simpler, cheaper and more environmentally friendly process. This refinery enables us to domestically produce critical minerals in support of energy storage, battery manufacturing and ultimately for EV growth.\nWe have begun to produce battery packs for certain Model Ys with our 4680 cells, unlocking an additional vector of supply to help navigate increasingly complex supply chain challenges caused by trade barriers and tariff risks. We now produce dry-electrode for 4680 cells with both anode and cathode made in Austin. We expect both domestic cathode material in Texas and LFP lines in Nevada to begin production in 2026.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY-->Battery", + "path": "TSLA-Q4-2025-Update.pdf-->SUMMARY-->Battery", "metadata": { "length": 762, "summary": "", @@ -667,7 +667,7 @@ "chunk_id": "34777c71-d716-5c39-b18a-bec0d0b64706", "type": "text", "content": "We continue to efficiently utilize our existing physical footprint in North America, with targeted augmentation to support the rollout of Robotaxi. While in the short-term, operational workstreams such as charging, cleaning and maintenance can be managed through our existing charging network, service centers and sales and delivery locations, we will have to add more capacity as the service expands. We added over 3,800 net new Supercharging stalls, growing the network 19% year-over-year.\nInstalled Annual Capacity\n\n[tables/table-6 Facility Status.html]\n\nInstalled capacity ≠ current production rate and there may be limitations discovered as production rates approach capacity. Production rates depend on a variety of factors, including equipment uptime, component supply, downtime related to factory upgrades, regulatory considerations and other factors. Construction includes factory and infrastructure buildout as well as tool installation.\n\n## Other Supporting Infrastructure Tesla AI Training Capacity Ramp (H100 equivalent GPUs)0\n[images/image-1-Capacity Growth Projection.jpg]\n\nTesla AI Training Capacity Ramp (H100 equivalent GPUs)", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY-->Other Supporting Infrastructure", + "path": "TSLA-Q4-2025-Update.pdf-->SUMMARY-->Other Supporting Infrastructure", "metadata": { "length": 1142, "summary": "", @@ -786,7 +786,7 @@ "chunk_id": "778a63c2-7955-514c-84d0-9c2cfd99a489", "type": "text", "content": "We continue to enhance FSD (Supervised) $^{1}$ via our end-to-end foundation model trained on both customer and Robotaxi real-world data with our latest version, v14. FSD (Supervised) $^{1}$ increasingly provides safety and convenience functionality that can relieve drivers of many tedious and potentially dangerous aspects of road travel, including giving access to personal transport for those who otherwise have difficulty driving. V14 offers unparalleled driver assistance to safely drive to the customer's destination, find a free parking spot and park at that location. Our global fleet can collect the equivalent of over 500 years of continuous driving data per day $^{2}$ , allowing us to safely deploy and scale capabilities that can handle the long and fat tail of corner cases across diverse geographies and driving environments.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY-->AI Software", + "path": "TSLA-Q4-2025-Update.pdf-->SUMMARY-->AI Software", "metadata": { "length": 841, "summary": "", @@ -876,7 +876,7 @@ "chunk_id": "164babde-7a64-5b81-9a5a-41443b221c40", "type": "text", "content": "Development of our in-house, custom designed AI5 and AI6 inference chips for autonomy progressed during the quarter, with production planned for 2027 and 2028, respectively. We are targeting a 50x improvement in performance for AI5 relative to AI4 thanks to 10x raw compute, 9x memory capacity and 5x hardened block quantization and softmax function (the latter enabling efficient low-precision computing without sacrificing model accuracy).", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY-->AI Inference Compute", + "path": "TSLA-Q4-2025-Update.pdf-->SUMMARY-->AI Inference Compute", "metadata": { "length": 441, "summary": "", @@ -970,7 +970,7 @@ "chunk_id": "32bfae6f-b2d4-51d3-94d9-82b9640828dd", "type": "text", "content": "The Robotaxi iOS app no longer has a waitlist in the areas we serve. Our vehicles keep getting better with our over-the-air updates, including: Grok (an AI companion) which now supports navigation commands (allowing users to find, add and edit navigation destinations hands-free); Tesla Photobooth which enables users to take photos in their car and download or share via the Tesla mobile app; Supercharger Site Maps which displays Supercharger layouts, nearby businesses and live availability details; Automatic HOV Lane Routing based on interior camera occupancy detection; Phone Left Behind Chime and SpaceX ISS Docking Simulator Game.\n\nWe continue to enhance FSD (Supervised) $^{1}$ via our end-to-end foundation model trained on both customer and Robotaxi real-world data with our latest version, v14. FSD (Supervised) $^{1}$ increasingly provides safety and convenience functionality that can relieve drivers of many tedious and potentially dangerous aspects of road travel, including giving access to personal transport for those who otherwise have difficulty driving. V14 offers unparalleled driver assistance to safely drive to the customer's destination, find a free parking spot and park at that location. Our global fleet can collect the equivalent of over 500 years of continuous driving data per day $^{2}$ , allowing us to safely deploy and scale capabilities that can handle the long and fat tail of corner cases across diverse geographies and driving environments. Cumulative Miles Driven with FSD (Supervised) $^{1}$ (billions)0\n[images/image-2-FSD Mileage Growth.jpg]\n\nCumulative Miles Driven with FSD (Supervised) $^{1}$ (billions)\n\nDevelopment of our in-house, custom designed AI5 and AI6 inference chips for autonomy progressed during the quarter, with production planned for 2027 and 2028, respectively. We are targeting a 50x improvement in performance for AI5 relative to AI4 thanks to 10x raw compute, 9x memory capacity and 5x hardened block quantization and softmax function (the latter enabling efficient low-precision computing without sacrificing model accuracy). Targeting Step-Function Improvement for our Next-Generation Inference Chip, AI50\n[images/image-3-Tesla Silicon Optimization.jpg]\n\nTargeting Step-Function Improvement for our Next-Generation Inference Chip, AI5\n(1) Active driver supervision required; does not make the vehicle autonomous\n(2) Calculated based on continuous hours of driving at an average of 30 miles per hour", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY-->Automotive and Other Software", + "path": "TSLA-Q4-2025-Update.pdf-->SUMMARY-->Automotive and Other Software", "metadata": { "length": 2444, "summary": "The Robotaxi iOS app has removed its waitlist in served areas and now includes features like Grok navigation, Tesla Photobooth, Supercharger maps, automatic HOV routing, phone left-behind chimes, and a SpaceX game. FSD (Supervised) v14 uses an end-to-end foundation model trained on vast real-world data to assist drivers with navigation, parking, and safety, though active supervision remains required. The company is developing custom AI5 and AI6 inference chips for 2027 and 2028, targeting significant performance improvements over previous generations to handle complex driving scenarios globally.", @@ -1205,7 +1205,7 @@ "chunk_id": "41533301-58f0-554f-916d-8254cc2700df", "type": "text", "content": "We began testing driverless Robotaxis in Austin in December and began removing the safety monitor from customer rides in January on a limited basis, which will unlock further expansion of our Robotaxi fleet and coverage area in the Austin-metro. Our Bay Area ride-hailing service began serving the San Jose Airport in October, with plans to expand to other major airports in the Bay Area upon receiving required permitting.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY-->Robotaxi", + "path": "TSLA-Q4-2025-Update.pdf-->SUMMARY-->Robotaxi", "metadata": { "length": 423, "summary": "", @@ -1263,7 +1263,7 @@ "chunk_id": "3968a25f-933e-5a7f-a118-07db2e581153", "type": "text", "content": "We launched FSD (Supervised) $^{1}$ in South Korea, where customers drove over 1 million kilometers using the software in just one month. While we continue to pursue regulatory approval in China and Europe, we began offering ride-along experiences to consumers in Italy, Germany, France and Switzerland.\nMonthly subscriptions to FSD (Supervised) $^{1}$ continued to grow sequentially and more than doubled in 2025. Starting this quarter, we are transitioning access to FSD (Supervised) $^{1}$ to monthly subscriptions only as we begin to sunset the up-front payment option.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY-->FSD (Supervised) $^{1}$", + "path": "TSLA-Q4-2025-Update.pdf-->SUMMARY-->FSD (Supervised) $^{1}$", "metadata": { "length": 573, "summary": "", @@ -1364,7 +1364,7 @@ "chunk_id": "adf89909-d366-51e1-b68c-459f9024191c", "type": "text", "content": "Services and Other gross profit of approximately \\$300 million was partly driven by Part Sales and Supercharging. We now offer Tesla Insurance in Florida, as we continue to expand our insurance product to new states. In certain states, customers receive a discount on their insurance premiums when using FSD (Supervised) $^{1}$ . The more you drive with FSD (Supervised) $^{1}$ enabled, the bigger the discount is on your insurance premium – helping, in certain cases, to completely offset the monthly subscription cost for FSD (Supervised) $^{1}$ .\n\n## FSD (Supervised) $^{1}$ Cumulative Paid Robotaxi Miles0\n[images/image-4-Growth Trend 2025.jpg]\n\nCumulative Paid Robotaxi Miles\n\n[tables/table-7 Autonomous Driving Status.html]\n\nPlanned Robotaxi Coverage", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY-->Automotive Services", + "path": "TSLA-Q4-2025-Update.pdf-->SUMMARY-->Automotive Services", "metadata": { "length": 742, "summary": "", @@ -1450,7 +1450,7 @@ "chunk_id": "cf91377e-176a-5768-855a-b859092f4695", "type": "text", "content": "On January 16, 2026, Tesla entered into an agreement to invest approximately \\$2 billion to acquire shares of Series E Preferred Stock of xAI as part of their recent publicly-disclosed financing round. Tesla’s investment was made on market terms consistent with those previously agreed to by other investors in the financing round. As set forth in Master Plan Part IV, Tesla is building products and services that bring AI into the physical world. Meanwhile, xAI is developing leading digital AI products and services, such as its large language model (Grok).\nIn that context, and as part of Tesla's broader strategy under Master Plan Part IV, Tesla and xAI also entered into a framework agreement in connection with the investment. Among other things, the framework agreement builds upon the existing relationship between Tesla and xAI by providing a framework for evaluating potential AI collaborations between the companies. Together, the investment and the related framework agreement are intended to enhance Tesla's ability to develop and deploy AI products and services into the physical world at scale. This investment is subject to customary regulatory conditions with the expectation to close in Q1'2026.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OTHER UPDATES", + "path": "TSLA-Q4-2025-Update.pdf-->OTHER UPDATES", "metadata": { "length": 1213, "summary": "", @@ -1552,7 +1552,7 @@ "chunk_id": "79b23d73-b353-5998-a646-93d6739ec465", "type": "text", "content": "", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK", "metadata": { "length": 0, "summary": "", @@ -1569,7 +1569,7 @@ "chunk_id": "8d4c49a6-c531-5fa6-a77d-ad028537fa52", "type": "text", "content": "We are focused on maximum capacity utilization at our factories. Deliveries and deployments will be impacted by aggregate demand for our products, supply chain readiness and allocation decisions between sale to customers or use for our owned and operated fleet.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Volume", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Volume", "metadata": { "length": 261, "summary": "", @@ -1609,7 +1609,7 @@ "chunk_id": "426e5119-3ef6-5176-9119-fc044eb90be7", "type": "text", "content": "We will manage the businesses such that we ensure a strong balance sheet, maintaining sufficient liquidity to fund our product roadmap, long-term capacity expansion plans – including further vertical integration – and other expenses.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Cash", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Cash", "metadata": { "length": 233, "summary": "", @@ -1649,7 +1649,7 @@ "chunk_id": "99f54247-8409-5df2-8299-39184863cd09", "type": "text", "content": "While we continue to execute on innovations to reduce the cost of manufacturing and operations, over time, we expect our hardware-related profits to be accompanied by an acceleration of AI, software and fleet-based profits.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Profit", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Profit", "metadata": { "length": 223, "summary": "", @@ -1686,7 +1686,7 @@ "chunk_id": "29e89187-3dbe-5298-87c3-c38d3cbe6883", "type": "text", "content": "We continue to evolve and augment our product lineup with a focus on cost, scale and future monetization opportunities via services powered by our AI software. We remain focused on growing our sales volumes through a differentiated and efficiently managed product portfolio, which includes leveraging and optimizing our existing production capacity before building new factories and production lines.\nCybercab, Tesla Semi and Megapack 3 are on schedule for volume production starting in 2026. First generation production lines for Optimus are being installed in anticipation of volume production.\nPHOTOS & CHARTS", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product", "metadata": { "length": 612, "summary": "", @@ -1772,7 +1772,7 @@ "chunk_id": "bbe9d89e-eebd-56b1-b7c0-7eda12dd3377", "type": "text", "content": "Cybercab, Tesla Semi and Megapack 3 are on schedule for volume production starting in 2026. First generation production lines for Optimus are being installed in anticipation of volume production. 0\n[images/image-5-Tesla Model Y Driving.jpg]", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product-->MODEL Y - 2025 BEST IN CLASS EURO NCAP $^{(1)}$ SMALL SUV", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product-->MODEL Y - 2025 BEST IN CLASS EURO NCAP $^{(1)}$ SMALL SUV", "metadata": { "length": 245, "summary": "", @@ -1835,7 +1835,7 @@ "chunk_id": "02f7e756-798f-5675-9dd9-6ce9123f92d4", "type": "text", "content": " 0\n[images/image-6-Red Tesla on Coastal Road.jpg]", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product-->MODEL 3 - 2025 BEST IN CLASS EURO NCAP $^{(1)}$ LARGE FAMILY CAR", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product-->MODEL 3 - 2025 BEST IN CLASS EURO NCAP $^{(1)}$ LARGE FAMILY CAR", "metadata": { "length": 66, "summary": "", @@ -1884,7 +1884,7 @@ "chunk_id": "494118dd-5788-526b-a139-407d1cfb0d20", "type": "text", "content": " 0\n[images/image-7-Tesla Interior Interface.jpg]", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product-->FSD (SUPERVISED) $^{1}$ – V14 OFFERS UNPARALLELED DRIVER ASSISTANCE", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product-->FSD (SUPERVISED) $^{1}$ – V14 OFFERS UNPARALLELED DRIVER ASSISTANCE", "metadata": { "length": 66, "summary": "", @@ -1933,7 +1933,7 @@ "chunk_id": "6ac7765a-4da5-573a-963b-f35ab3796f6f", "type": "text", "content": " 0\n[images/image-8-Tesla Interior.jpg]", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product-->DRIVERLESS ROBOTAXI - TESTING IN AUSTIN", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product-->DRIVERLESS ROBOTAXI - TESTING IN AUSTIN", "metadata": { "length": 66, "summary": "", @@ -2016,7 +2016,7 @@ "chunk_id": "4cfe7ce9-5f93-5816-9505-ba562683ee42", "type": "text", "content": " 0\n[images/image-9-Tesla Cybertruck in Snow.jpg]\n\n\n### DRIVERLESS ROBOTAXI - TESTING IN AUSTIN 0\n[images/image-10-Tesla Semi Trucks.jpg]\n\nTESLA SEMI - MEGACHARGER NETWORK PLANNED SITES FOR 2026\n\n### CYBERCAB - COLD WEATHER TESTING IN ALASKA 0\n[images/image-11-US Lightning Map.jpg]", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product-->CYBERCAB - COLD WEATHER TESTING IN ALASKA", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product-->CYBERCAB - COLD WEATHER TESTING IN ALASKA", "metadata": { "length": 318, "summary": "", @@ -2102,7 +2102,7 @@ "chunk_id": "07535539-a2f9-5b79-ac6b-d41cf785ec88", "type": "text", "content": " 0\n[images/image-12-Tesla Factory Milestone.jpg]", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product-->GIGAFACTORY SHANGHAI - 9 MILLIONTH VEHICLE PRODUCED (GLOBALLY)", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product-->GIGAFACTORY SHANGHAI - 9 MILLIONTH VEHICLE PRODUCED (GLOBALLY)", "metadata": { "length": 67, "summary": "", @@ -2202,7 +2202,7 @@ "chunk_id": "702bff82-1736-584d-956e-e7bf7eac4039", "type": "text", "content": " 0\n[images/image-13-Tesla Factory Milestone.jpg]\n\n\n### GIGAFACTORY SHANGHAI - 9 MILLIONTH VEHICLE PRODUCED (GLOBALLY) 0\n[images/image-14-Vehicle Delivery Trends.jpg]\n\n\n 0\n[images/image-15-Quarterly Cash Flow.jpg]\n\n\n 0\n[images/image-16-Financial Performance Chart.jpg]", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product-->GIGAFACTORY NEVADA - 6 MILLIONTH DRIVE UNIT PRODUCED", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->Product-->GIGAFACTORY NEVADA - 6 MILLIONTH DRIVE UNIT PRODUCED", "metadata": { "length": 327, "summary": "", @@ -2319,7 +2319,7 @@ "chunk_id": "5e17588f-ea71-56df-be53-135da93fb3a0", "type": "text", "content": " 0\n[images/image-17-Projected Vehicle Deliveries.jpg]\n\n\n 0\n[images/image-18-Cash Flow Trends.jpg]\n\n\n 0\n[images/image-19-Financial Performance Forecast.jpg]", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->KEY METRICS TRAILING 12 MONTHS (TTM) (Unaudited)", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->KEY METRICS TRAILING 12 MONTHS (TTM) (Unaudited)", "metadata": { "length": 207, "summary": "", @@ -2369,7 +2369,7 @@ "chunk_id": "c82a6fed-f085-5ec3-b10a-0381723fa3c9", "type": "text", "content": "Total quarterly revenue decreased 3% YoY to \\$24.9B. YoY, revenue was impacted by the following items $^{(1)}$ :\n- decrease in vehicle deliveries\n- lower regulatory credit revenue\n+ growth in Energy Generation and Storage\n+ growth in Services and Other\n+ positive FX impact of \\$0.3B $^{1}$\n\\+ growth in other automotive ancillary sales, partly driven by an increase in FSD subscriptions\n\\+ higher vehicle average selling price (ASP) (excl. FX impact $^{1}$ ), inclusive of mix impact", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->KEY METRICS TRAILING 12 MONTHS (TTM) (Unaudited)-->Revenue", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->KEY METRICS TRAILING 12 MONTHS (TTM) (Unaudited)-->Revenue", "metadata": { "length": 484, "summary": "", @@ -2428,7 +2428,7 @@ "chunk_id": "946fabf6-000f-557a-9d81-0b8c1a25aafb", "type": "text", "content": "Our quarterly operating income decreased 11% YoY to \\$1.4B, resulting in a 5.7% operating margin. YoY, operating income was primarily impacted by the following items $^{(1)}$ :\n- increase in SBC and Restructuring and Other charges\n- increase in operating expenses (excl. SBC and Restructuring and Other) driven by AI and other R&D projects and SG&A\n- higher average cost per vehicle due to lower fixed cost absorption for certain models and an increase in tariffs\n- decrease in vehicle deliveries\n- lower regulatory credit revenue\n+ higher vehicle average gross profit due to mix and pricing impacts\n+ growth in Energy Generation and Storage gross profit\n+ growth in Services and Other gross profit\n+ growth in other automotive ancillary sales, partly driven by an increase in FSD subscriptions", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->KEY METRICS TRAILING 12 MONTHS (TTM) (Unaudited)-->Profitability", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->KEY METRICS TRAILING 12 MONTHS (TTM) (Unaudited)-->Profitability", "metadata": { "length": 794, "summary": "", @@ -2502,7 +2502,7 @@ "chunk_id": "001d4b8a-b264-5cae-86a5-8c419eabbec0", "type": "text", "content": "Quarter-end cash, cash equivalents and investments was \\$44.1B. The sequential increase of \\$2.4B was primarily the result of positive free cash flow.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->KEY METRICS TRAILING 12 MONTHS (TTM) (Unaudited)-->Cash", + "path": "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->KEY METRICS TRAILING 12 MONTHS (TTM) (Unaudited)-->Cash", "metadata": { "length": 150, "summary": "", @@ -2683,7 +2683,7 @@ "chunk_id": "15c18264-8e1e-5db6-b3c9-70cd181d1f39", "type": "text", "content": "STATEMENT OF OPERATIONS\n(Unaudited)\n\n[tables/table-8 Q4 2024-Q4 2025 Rev.html]\n\nBALANCE SHEET\n(Unaudited)\n\n[tables/table-9 Balance Sheet 2024-25.html]\n\nSTATEMENT OF CASH FLOWS\n(Unaudited)\n\n[tables/table-10 Cash Flow Q4-25.html]\n\nRECONCILIATION OF GAAP TO NON-GAAP FINANCIAL INFORMATION (Unaudited)\n\n[tables/table-11 Q4 2024-Q4 2025.html]\n\nRECONCILIATION OF GAAP TO NON-GAAP FINANCIAL INFORMATION\n(Unaudited)\n\n[tables/table-12 Financial Metrics 2021-25.html]\n\nRECONCILIATION OF GAAP TO NON-GAAP FINANCIAL INFORMATION\n(Unaudited)\n\n[tables/table-13 Financial Data 2022-25.html]\n\n\n[tables/table-14 Financial Metrics 2023-25.html]\n\nTTM = Trailing twelve months\n(1) Beginning in Q1'25, Capital expenditures is presented inclusive of purchases of energy generation and storage systems and all prior periods have been adjusted.\n(2) As a result of the adoption of the new crypto assets standard, the previously reported quarterly periods in 2024 have been recast.\n(3) Beginning in Q1'25, Adjusted EBITDA (non-GAAP) is presented net of digital assets gains and losses and all prior periods have been adjusted.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->FINANCIAL STATEMENTS", + "path": "TSLA-Q4-2025-Update.pdf-->FINANCIAL STATEMENTS", "metadata": { "length": 1418, "summary": "", @@ -2816,7 +2816,7 @@ "chunk_id": "31babb3c-1abf-5f64-90c7-e8e8a5384bbd", "type": "text", "content": "Tesla will provide a live webcast of its fourth quarter 2025 financial results conference call beginning at 4:30 p.m. CT on January 28, 2026 at ir.tesla.com. This webcast will also be available for replay for approximately one year thereafter.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->WEBCAST INFORMATION", + "path": "TSLA-Q4-2025-Update.pdf-->WEBCAST INFORMATION", "metadata": { "length": 243, "summary": "", @@ -2857,7 +2857,7 @@ "chunk_id": "19e3f236-d15a-5b2a-8589-43d28c5316fe", "type": "text", "content": "When used in this update, certain terms have the following meanings. Our vehicle deliveries include only vehicles that have been transferred to end customers with all paperwork correctly completed. Our energy product deployment volume includes both customer units when installed and equipment sales at time of delivery. \"Net income attributable to common stockholders (non-GAAP)\" is equal to (i) net income attributable to common stockholders before (ii)(a) stock-based compensation expense, net of tax, (b) digital assets (gain) loss, net of tax and (c) release of valuation allowance on deferred tax assets. \"Adjusted EBITDA (non-GAAP)\" is equal to (i) net income attributable to common stockholders before (ii)(a) interest expense, (b) provision for (benefit from) income taxes, (c) depreciation, amortization and impairment, (d) stock-based compensation expense and (e) digital assets loss (gain), net. \"Free cash flow\" is operating cash flow less capital expenditures. Average cost per vehicle is cost of automotive sales divided by new vehicle deliveries (excluding operating leases). \"Days sales outstanding\" is equal to (i) average accounts receivable, net for the period divided by (ii) total revenues and multiplied by (iii) the number of days in the period. \"Days payable outstanding\" is equal to (i) average accounts payable for the period divided by (ii) total cost of revenues and multiplied by (iii) the number of days in the period. \"Days of supply\" is calculated by dividing new car ending inventory by the relevant period's deliveries and using trading days. Constant currency impacts are calculated by comparing actuals against current results converted into USD using average exchange rates from the prior period.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->CERTAIN TERMS", + "path": "TSLA-Q4-2025-Update.pdf-->CERTAIN TERMS", "metadata": { "length": 1733, "summary": "This passage defines key financial and operational terms used in a specific update. It clarifies that vehicle deliveries refer to units transferred to end customers with completed paperwork, while energy product deployment includes installed customer units and equipment sales. The text details non-GAAP measures: 'Net income attributable to common stockholders' adjusts for stock-based compensation, digital asset gains/losses, and valuation allowances; 'Adjusted EBITDA' excludes interest, taxes, depreciation, amortization, impairment, stock-based compensation, and digital asset impacts. 'Free cash flow' is defined as operating cash flow minus capital expenditures. Operational metrics include 'Average cost per vehicle' (automotive sales cost divided by new deliveries excluding leases), 'Days sales outstanding' (average receivables divided by revenue times days in period), 'Days payable outstanding' (average payables divided by cost of revenues times days in period), and 'Days of supply' (ending inventory divided by deliveries using trading days). Finally, constant currency impacts are calculated by comparing actuals against results converted to USD using prior period average exchange rates.", @@ -2982,7 +2982,7 @@ "chunk_id": "332edac0-294c-5e5b-8599-a0e52b25e53a", "type": "text", "content": "Consolidated financial information has been presented in accordance with GAAP as well as on a non-GAAP basis to supplement our consolidated financial results. Our non-GAAP financial measures include non-GAAP net income (loss) attributable to common stockholders, non-GAAP net income (loss) attributable to common stockholders on a diluted per share basis (calculated using weighted average shares for GAAP diluted net income (loss) attributable to common stockholders), Adjusted EBITDA margin, non-GAAP automotive gross margin and free cash flow. These non-GAAP financial measures also facilitate management's internal comparisons to Tesla's historical performance as well as comparisons to the operating results of other companies. Management believes that it is useful to supplement its GAAP financial statements with this non-GAAP information because management uses such information internally for its operating, budgeting and financial planning purposes. Management also believes that presentation of the non-GAAP financial measures provides useful information to our investors regarding our financial condition and results of operations, so that investors can see through the eyes of Tesla management regarding important financial metrics that Tesla uses to run the business and allowing investors to better understand Tesla's performance. Non-GAAP information is not prepared under a comprehensive set of accounting rules and therefore, should only be read in conjunction with financial information reported under U.S. GAAP when understanding Tesla's operating performance. A reconciliation between GAAP and non-GAAP financial information is provided above.", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->NON-GAAP FINANCIAL INFORMATION", + "path": "TSLA-Q4-2025-Update.pdf-->NON-GAAP FINANCIAL INFORMATION", "metadata": { "length": 1664, "summary": "Tesla presents consolidated financial information under both GAAP and non-GAAP standards to supplement its results. Key non-GAAP measures include net income attributable to common stockholders, diluted per share figures, Adjusted EBITDA margin, automotive gross margin, and free cash flow. These metrics aid internal management comparisons with historical performance and other companies, supporting operating, budgeting, and planning activities. Management believes these non-GAAP figures provide investors with a clearer view of Tesla's financial condition and operational results by reflecting the metrics used to run the business. However, since non-GAAP data is not prepared under comprehensive accounting rules, it should be read alongside U.S. GAAP information for a complete understanding of Tesla's performance. A reconciliation between GAAP and non-GAAP data is provided elsewhere.", @@ -3077,7 +3077,7 @@ "chunk_id": "34263854-525e-5a28-b9fd-a88d4813756e", "type": "text", "content": "Certain statements in this update, including, but not limited to, statements in the “Outlook” section; statements relating to the development, strategy, ramp, production and capacity, demand and market growth, cost, pricing and profitability, investment, deliveries, deployment, availability and other features and improvements and timing of existing and future Tesla products and services and supporting infrastructure; statements regarding operating margin, operating profits, spending and liquidity; and statements regarding expansion, improvements and/or ramp and related timing at our facilities are “forward-looking statements” within the meaning of the Private Securities Litigation Reform Act of 1995. Forward-looking statements are based on assumptions and management’s current expectations, involve certain risks and uncertainties, and are not guarantees. Future results may differ materially from those expressed in any forward-looking statement. The following important factors, without limitation, could cause actual results to differ materially from those in the forward-looking statements: the risk of delays in launching and/or manufacturing our products, services and features cost-effectively; our ability to build and/or grow our products and services, sales, delivery, installation, servicing and charging capabilities and effectively manage this growth; our ability to successfully and timely develop, introduce and scale, as well as our consumers’ demand for, products and services based on artificial intelligence, robotics and automation, electric vehicles, advanced driver assistance systems, and ride-hailing services generally and our vehicles and services specifically; the ability of suppliers to deliver components according to schedules, prices, quality and volumes acceptable to us, and our ability to manage such components effectively; any issues with lithium-ion cells or other components manufactured at our factories; our ability to ramp our factories in accordance with our plans; our ability to procure supply of battery cells, including through our own manufacturing; risks relating to international operations and expansion, including unfavorable and uncertain regulatory, political, economic, tax, tariff, export controls and labor conditions; any failures by Tesla products to perform as expected or if product recalls occur; the risk of product liability claims; competition in the automotive, transportation and energy product and services and robotics markets; our ability to maintain public credibility and confidence in our long-term business prospects; our ability to manage risks relating to our various product financing programs; the status of government and economic incentives for electric vehicles and energy products; our ability to attract, hire and retain key employees and qualified personnel; our ability to maintain the security of our information and production and product systems; our compliance with various regulations and laws applicable to our operations and products, which may evolve from time to time; risks relating to our indebtedness and financing strategies; and adverse foreign exchange movements. More information on potential factors that could affect our financial results is included from time to time in our Securities and Exchange Commission filings and reports, including the risks identified under the section captioned “Risk Factors” in our annual report on Form 10-K filed with the SEC on January 30, 2025 and subsequent quarterly reports on Form 10-Q. Tesla disclaims any obligation to update information contained in these forward-looking statements whether as a result of new information, future events or otherwise.\nTESLA", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->FORWARD-LOOKING STATEMENTS", + "path": "TSLA-Q4-2025-Update.pdf-->FORWARD-LOOKING STATEMENTS", "metadata": { "length": 3711, "summary": "This passage from Tesla outlines that various statements in the update, particularly regarding future outlooks, product development, production capacity, financial metrics, and facility expansions, constitute forward-looking statements under the Private Securities Litigation Reform Act of 1995. These statements reflect management's current expectations based on assumptions and are subject to risks and uncertainties, meaning actual results may differ materially. The text lists numerous specific risk factors that could impact outcomes, including manufacturing delays, supply chain challenges, regulatory hurdles, competition, product performance issues, employee retention, and foreign exchange fluctuations. Tesla advises investors to consult SEC filings, specifically the Form 10-K filed on January 30, 2025, for detailed risk disclosures. The company explicitly disclaims any obligation to update these forward-looking statements due to new information or future events.", diff --git a/apps/api/app/data/demo_documents/tsla-q4-2025/doc_nav.json b/apps/api/app/data/demo_documents/tsla-q4-2025/doc_nav.json index d27f1ba4c..d8a4cdc7e 100644 --- a/apps/api/app/data/demo_documents/tsla-q4-2025/doc_nav.json +++ b/apps/api/app/data/demo_documents/tsla-q4-2025/doc_nav.json @@ -6,15 +6,299 @@ "text_chunks": 36, "image_chunks": 19, "table_chunks": 15, - "max_depth": 1 + "max_depth": 3 }, "sections": [ { - "title": "Root", - "path": "Default_Root/TSLA-Q4-2025-Update.pdf-->HIGHLIGHTS", + "title": "HIGHLIGHTS", + "path": "TSLA-Q4-2025-Update.pdf/HIGHLIGHTS", "level": 1, "summary": "[tables/table-0 Tesla 2025 Results.html]", - "chunk_count": 36, + "chunk_count": 1, + "children": [] + }, + { + "title": "SUMMARY", + "path": "TSLA-Q4-2025-Update.pdf/SUMMARY", + "level": 1, + "summary": "The document presents unaudited financial and operational summaries for a company, covering quarterly data from Q4 2024 through Q4 2025 and annual data from 2021 to 2025. Key notes indicate significant accounting changes effective Q1 2025: Adjusted EBITDA and Net income attributable to common stockholders are now presented net of digital assets gains and losses, with all prior periods adjusted accordingly. Additionally, Capital expenditures now include purchases of energy generation and storage systems, requiring restatement of previous periods. The content references multiple tables detailing these metrics but does not display the specific numerical values.", + "chunk_count": 13, + "children": [ + { + "title": "Automotive", + "path": "TSLA-Q4-2025-Update.pdf/SUMMARY/Automotive", + "level": 2, + "summary": "While automotive sales declined sequentially, gross margin (even when excluding the impact of regulatory credits) improved. The APAC region continued to show strength across multiple markets and set a record for deliveries in the quarter. We continued the rollout of Model Y variants across markets in Q4, including the standard and performance versions. Preparations continue in North America for the production ramps of Tesla Semi and Cybercab, both commencing 1H26, and production of the next-generation Roadster.", + "chunk_count": 1, + "children": [] + }, + { + "title": "Energy generation and storage", + "path": "TSLA-Q4-2025-Update.pdf/SUMMARY/Energy generation and storage", + "level": 2, + "summary": "We achieved our highest quarterly energy storage deployments, driven by record Megapack deployments. Total gross profit rose, both sequentially and year-over-year, to a record \\$1.1 billion, marking the fifth consecutive record quarter. We plan to begin Megapack 3 and Megablock production at Megafactory Houston in 2026. In 2025, our global Powerwall network supported more than 89,000 Virtual Power Plant events across over 1 million installed units, allowing homeowners to save over \\$1 billion in electricity bills as Virtual Power Plant participation continues to scale rapidly.", + "chunk_count": 1, + "children": [] + }, + { + "title": "Robotics", + "path": "TSLA-Q4-2025-Update.pdf/SUMMARY/Robotics", + "level": 2, + "summary": "We made further progress on the Optimus program in 2025. In Q1 of this year, we plan to unveil the Gen 3 version of Optimus, which will include major upgrades from version 2.5, including our latest hand design. The Gen 3 is our first design meant for mass production. Preparations are underway for the first production line, including supply chain readiness, with start of production planned before the end of 2026 and eventual planned capacity of 1 million robots per year. Installed Annual Manufacturing Capacity [tables/table-5 Tesla Production.html] Installed capacity ≠ current production rate and there may be limitations discovered as production rates approach capacity. Production rates depend on a variety of factors, including equipment uptime, component supply, downtime related to factory upgrades, regulatory considerations and other factors. Construction includes factory and infrastructure buildout as well as tool installation.", + "chunk_count": 1, + "children": [] + }, + { + "title": "AI Training Compute", + "path": "TSLA-Q4-2025-Update.pdf/SUMMARY/AI Training Compute", + "level": 2, + "summary": "We are currently building Cortex 2 at Gigafactory Texas to further increase our AI training compute capacity. In the first half of 2026, we plan to more than double the size of onsite compute in Texas (in terms of H100 equivalents). We aim to maximize capital efficiency by scaling training compute judiciously, including when the training backlog gets too long or in anticipation of greater demand from our engineers to support our AI-related offerings.", + "chunk_count": 1, + "children": [] + }, + { + "title": "Battery", + "path": "TSLA-Q4-2025-Update.pdf/SUMMARY/Battery", + "level": 2, + "summary": "Our lithium refinery commenced pilot production and is the first spodumene to lithium hydroxide refinery in North America, leveraging a simpler, cheaper and more environmentally friendly process. This refinery enables us to domestically produce critical minerals in support of energy storage, battery manufacturing and ultimately for EV growth. We have begun to produce battery packs for certain Model Ys with our 4680 cells, unlocking an additional vector of supply to help navigate increasingly complex supply chain challenges caused by trade barriers and tariff risks. We now produce dry-electrode for 4680 cells with both anode and cathode made in Austin. We expect both domestic cathode material in Texas and LFP lines in Nevada to begin production in 2026.", + "chunk_count": 1, + "children": [] + }, + { + "title": "Other Supporting Infrastructure", + "path": "TSLA-Q4-2025-Update.pdf/SUMMARY/Other Supporting Infrastructure", + "level": 2, + "summary": "We continue to efficiently utilize our existing physical footprint in North America, with targeted augmentation to support the rollout of Robotaxi. While in the short-term, operational workstreams such as charging, cleaning and maintenance can be managed through our existing charging network, service centers and sales and delivery locations, we will have to add more capacity as the service expands. We added over 3,800 net new Supercharging stalls, growing the network 19% year-over-year. Installed Annual Capacity [tables/table-6 Facility Status.html] Installed capacity ≠ current production rate and there may be limitations discovered as production rates approach capacity. Production rates depend on a variety of factors, including equipment uptime, component supply, downtime related to factory upgrades, regulatory considerations and other factors. Construction includes factory and infrastructure buildout as well as tool installation. ## Other Supporting Infrastructure Tesla AI Training Capacity Ramp (H100 equivalent GPUs)0 [images/image-1-Capacity Growth Projection.jpg] Tesla AI Training Capacity Ramp (H100 equivalent GPUs)", + "chunk_count": 1, + "children": [] + }, + { + "title": "AI Software", + "path": "TSLA-Q4-2025-Update.pdf/SUMMARY/AI Software", + "level": 2, + "summary": "We continue to enhance FSD (Supervised) $^{1}$ via our end-to-end foundation model trained on both customer and Robotaxi real-world data with our latest version, v14. FSD (Supervised) $^{1}$ increasingly provides safety and convenience functionality that can relieve drivers of many tedious and potentially dangerous aspects of road travel, including giving access to personal transport for those who otherwise have difficulty driving. V14 offers unparalleled driver assistance to safely drive to the customer's destination, find a free parking spot and park at that location. Our global fleet can collect the equivalent of over 500 years of continuous driving data per day $^{2}$ , allowing us to safely deploy and scale capabilities that can handle the long and fat tail of corner cases across diverse geographies and driving environments.", + "chunk_count": 1, + "children": [] + }, + { + "title": "AI Inference Compute", + "path": "TSLA-Q4-2025-Update.pdf/SUMMARY/AI Inference Compute", + "level": 2, + "summary": "Development of our in-house, custom designed AI5 and AI6 inference chips for autonomy progressed during the quarter, with production planned for 2027 and 2028, respectively. We are targeting a 50x improvement in performance for AI5 relative to AI4 thanks to 10x raw compute, 9x memory capacity and 5x hardened block quantization and softmax function (the latter enabling efficient low-precision computing without sacrificing model accuracy).", + "chunk_count": 1, + "children": [] + }, + { + "title": "Automotive and Other Software", + "path": "TSLA-Q4-2025-Update.pdf/SUMMARY/Automotive and Other Software", + "level": 2, + "summary": "The Robotaxi iOS app has removed its waitlist in served areas and now includes features like Grok navigation, Tesla Photobooth, Supercharger maps, automatic HOV routing, phone left-behind chimes, and a SpaceX game. FSD (Supervised) v14 uses an end-to-end foundation model trained on vast real-world data to assist drivers with navigation, parking, and safety, though active supervision remains required. The company is developing custom AI5 and AI6 inference chips for 2027 and 2028, targeting significant performance improvements over previous generations to handle complex driving scenarios globally.", + "chunk_count": 1, + "children": [] + }, + { + "title": "Robotaxi", + "path": "TSLA-Q4-2025-Update.pdf/SUMMARY/Robotaxi", + "level": 2, + "summary": "We began testing driverless Robotaxis in Austin in December and began removing the safety monitor from customer rides in January on a limited basis, which will unlock further expansion of our Robotaxi fleet and coverage area in the Austin-metro. Our Bay Area ride-hailing service began serving the San Jose Airport in October, with plans to expand to other major airports in the Bay Area upon receiving required permitting.", + "chunk_count": 1, + "children": [] + }, + { + "title": "FSD (Supervised) $^{1}$", + "path": "TSLA-Q4-2025-Update.pdf/SUMMARY/FSD (Supervised) $^{1}$", + "level": 2, + "summary": "We launched FSD (Supervised) $^{1}$ in South Korea, where customers drove over 1 million kilometers using the software in just one month. While we continue to pursue regulatory approval in China and Europe, we began offering ride-along experiences to consumers in Italy, Germany, France and Switzerland. Monthly subscriptions to FSD (Supervised) $^{1}$ continued to grow sequentially and more than doubled in 2025. Starting this quarter, we are transitioning access to FSD (Supervised) $^{1}$ to monthly subscriptions only as we begin to sunset the up-front payment option.", + "chunk_count": 1, + "children": [] + }, + { + "title": "Automotive Services", + "path": "TSLA-Q4-2025-Update.pdf/SUMMARY/Automotive Services", + "level": 2, + "summary": "Services and Other gross profit of approximately \\$300 million was partly driven by Part Sales and Supercharging. We now offer Tesla Insurance in Florida, as we continue to expand our insurance product to new states. In certain states, customers receive a discount on their insurance premiums when using FSD (Supervised) $^{1}$ . The more you drive with FSD (Supervised) $^{1}$ enabled, the bigger the discount is on your insurance premium – helping, in certain cases, to completely offset the monthly subscription cost for FSD (Supervised) $^{1}$ . ## FSD (Supervised) $^{1}$ Cumulative Paid Robotaxi Miles0 [images/image-4-Growth Trend 2025.jpg] Cumulative Paid Robotaxi Miles [tables/table-7 Autonomous Driving Status.html] Planned Robotaxi Coverage", + "chunk_count": 1, + "children": [] + } + ] + }, + { + "title": "OTHER UPDATES", + "path": "TSLA-Q4-2025-Update.pdf/OTHER UPDATES", + "level": 1, + "summary": "On January 16, 2026, Tesla entered into an agreement to invest approximately \\$2 billion to acquire shares of Series E Preferred Stock of xAI as part of their recent publicly-disclosed financing round. Tesla’s investment was made on market terms consistent with those previously agreed to by other investors in the financing round. As set forth in Master Plan Part IV, Tesla is building products and services that bring AI into the physical world. Meanwhile, xAI is developing leading digital AI products and services, such as its large language model (Grok). In that context, and as part of Tesla's broader strategy under Master Plan Part IV, Tesla and xAI also entered into a framework agreement in connection with the investment. Among other things, the framework agreement builds upon the existing relationship between Tesla and xAI by providing a framework for evaluating potential AI collaborations between the companies. Together, the investment and the related framework agreement are intended to enhance Tesla's ability to develop and deploy AI products and services into the physical world at scale. This investment is subject to customary regulatory conditions with the expectation to close in Q1'2026.", + "chunk_count": 1, + "children": [] + }, + { + "title": "OUTLOOK", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK", + "level": 1, + "summary": "", + "chunk_count": 16, + "children": [ + { + "title": "Volume", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/Volume", + "level": 2, + "summary": "We are focused on maximum capacity utilization at our factories. Deliveries and deployments will be impacted by aggregate demand for our products, supply chain readiness and allocation decisions between sale to customers or use for our owned and operated fleet.", + "chunk_count": 1, + "children": [] + }, + { + "title": "Cash", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/Cash", + "level": 2, + "summary": "We will manage the businesses such that we ensure a strong balance sheet, maintaining sufficient liquidity to fund our product roadmap, long-term capacity expansion plans – including further vertical integration – and other expenses.", + "chunk_count": 1, + "children": [] + }, + { + "title": "Profit", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/Profit", + "level": 2, + "summary": "While we continue to execute on innovations to reduce the cost of manufacturing and operations, over time, we expect our hardware-related profits to be accompanied by an acceleration of AI, software and fleet-based profits.", + "chunk_count": 1, + "children": [] + }, + { + "title": "Product", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/Product", + "level": 2, + "summary": "We continue to evolve and augment our product lineup with a focus on cost, scale and future monetization opportunities via services powered by our AI software. We remain focused on growing our sales volumes through a differentiated and efficiently managed product portfolio, which includes leveraging and optimizing our existing production capacity before building new factories and production lines. Cybercab, Tesla Semi and Megapack 3 are on schedule for volume production starting in 2026. First generation production lines for Optimus are being installed in anticipation of volume production. PHOTOS & CHARTS", + "chunk_count": 8, + "children": [ + { + "title": "MODEL Y - 2025 BEST IN CLASS EURO NCAP $^{(1)}$ SMALL SUV", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/Product/MODEL Y - 2025 BEST IN CLASS EURO NCAP $^{(1)}$ SMALL SUV", + "level": 3, + "summary": "Cybercab, Tesla Semi and Megapack 3 are on schedule for volume production starting in 2026. First generation production lines for Optimus are being installed in anticipation of volume production. 0 [images/image-5-Tesla Model Y Driving.jpg]", + "chunk_count": 1, + "children": [] + }, + { + "title": "MODEL 3 - 2025 BEST IN CLASS EURO NCAP $^{(1)}$ LARGE FAMILY CAR", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/Product/MODEL 3 - 2025 BEST IN CLASS EURO NCAP $^{(1)}$ LARGE FAMILY CAR", + "level": 3, + "summary": " 0 [images/image-6-Red Tesla on Coastal Road.jpg]", + "chunk_count": 1, + "children": [] + }, + { + "title": "FSD (SUPERVISED) $^{1}$ – V14 OFFERS UNPARALLELED DRIVER ASSISTANCE", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/Product/FSD (SUPERVISED) $^{1}$ – V14 OFFERS UNPARALLELED DRIVER ASSISTANCE", + "level": 3, + "summary": " 0 [images/image-7-Tesla Interior Interface.jpg]", + "chunk_count": 1, + "children": [] + }, + { + "title": "DRIVERLESS ROBOTAXI - TESTING IN AUSTIN", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/Product/DRIVERLESS ROBOTAXI - TESTING IN AUSTIN", + "level": 3, + "summary": " 0 [images/image-8-Tesla Interior.jpg]", + "chunk_count": 1, + "children": [] + }, + { + "title": "CYBERCAB - COLD WEATHER TESTING IN ALASKA", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/Product/CYBERCAB - COLD WEATHER TESTING IN ALASKA", + "level": 3, + "summary": " 0 [images/image-9-Tesla Cybertruck in Snow.jpg] ### DRIVERLESS ROBOTAXI - TESTING IN AUSTIN 0 [images/image-10-Tesla Semi Trucks.jpg] TESLA SEMI - MEGACHARGER NETWORK PLANNED SITES FOR 2026 ### CYBERCAB - COLD WEATHER TESTING IN ALASKA 0 [images/image-11-US Lightning Map.jpg]", + "chunk_count": 1, + "children": [] + }, + { + "title": "GIGAFACTORY SHANGHAI - 9 MILLIONTH VEHICLE PRODUCED (GLOBALLY)", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/Product/GIGAFACTORY SHANGHAI - 9 MILLIONTH VEHICLE PRODUCED (GLOBALLY)", + "level": 3, + "summary": " 0 [images/image-12-Tesla Factory Milestone.jpg]", + "chunk_count": 1, + "children": [] + }, + { + "title": "GIGAFACTORY NEVADA - 6 MILLIONTH DRIVE UNIT PRODUCED", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/Product/GIGAFACTORY NEVADA - 6 MILLIONTH DRIVE UNIT PRODUCED", + "level": 3, + "summary": " 0 [images/image-13-Tesla Factory Milestone.jpg] ### GIGAFACTORY SHANGHAI - 9 MILLIONTH VEHICLE PRODUCED (GLOBALLY) 0 [images/image-14-Vehicle Delivery Trends.jpg] 0 [images/image-15-Quarterly Cash Flow.jpg] 0 [images/image-16-Financial Performance Chart.jpg]", + "chunk_count": 1, + "children": [] + } + ] + }, + { + "title": "KEY METRICS TRAILING 12 MONTHS (TTM) (Unaudited)", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/KEY METRICS TRAILING 12 MONTHS (TTM) (Unaudited)", + "level": 2, + "summary": " 0 [images/image-17-Projected Vehicle Deliveries.jpg] 0 [images/image-18-Cash Flow Trends.jpg] 0 [images/image-19-Financial Performance Forecast.jpg]", + "chunk_count": 4, + "children": [ + { + "title": "Revenue", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/KEY METRICS TRAILING 12 MONTHS (TTM) (Unaudited)/Revenue", + "level": 3, + "summary": "Total quarterly revenue decreased 3% YoY to \\$24.9B. YoY, revenue was impacted by the following items $^{(1)}$ : - decrease in vehicle deliveries - lower regulatory credit revenue + growth in Energy Generation and Storage + growth in Services and Other + positive FX impact of \\$0.3B $^{1}$ \\+ growth in other automotive ancillary sales, partly driven by an increase in FSD subscriptions \\+ higher vehicle average selling price (ASP) (excl. FX impact $^{1}$ ), inclusive of mix impact", + "chunk_count": 1, + "children": [] + }, + { + "title": "Profitability", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/KEY METRICS TRAILING 12 MONTHS (TTM) (Unaudited)/Profitability", + "level": 3, + "summary": "Our quarterly operating income decreased 11% YoY to \\$1.4B, resulting in a 5.7% operating margin. YoY, operating income was primarily impacted by the following items $^{(1)}$ : - increase in SBC and Restructuring and Other charges - increase in operating expenses (excl. SBC and Restructuring and Other) driven by AI and other R&D projects and SG&A - higher average cost per vehicle due to lower fixed cost absorption for certain models and an increase in tariffs - decrease in vehicle deliveries - lower regulatory credit revenue + higher vehicle average gross profit due to mix and pricing impacts + growth in Energy Generation and Storage gross profit + growth in Services and Other gross profit + growth in other automotive ancillary sales, partly driven by an increase in FSD subscriptions", + "chunk_count": 1, + "children": [] + }, + { + "title": "Cash", + "path": "TSLA-Q4-2025-Update.pdf/OUTLOOK/KEY METRICS TRAILING 12 MONTHS (TTM) (Unaudited)/Cash", + "level": 3, + "summary": "Quarter-end cash, cash equivalents and investments was \\$44.1B. The sequential increase of \\$2.4B was primarily the result of positive free cash flow.", + "chunk_count": 1, + "children": [] + } + ] + } + ] + }, + { + "title": "FINANCIAL STATEMENTS", + "path": "TSLA-Q4-2025-Update.pdf/FINANCIAL STATEMENTS", + "level": 1, + "summary": "STATEMENT OF OPERATIONS (Unaudited) [tables/table-8 Q4 2024-Q4 2025 Rev.html] BALANCE SHEET (Unaudited) [tables/table-9 Balance Sheet 2024-25.html] STATEMENT OF CASH FLOWS (Unaudited) [tables/table-10 Cash Flow Q4-25.html] RECONCILIATION OF GAAP TO NON-GAAP FINANCIAL INFORMATION (Unaudited) [tables/table-11 Q4 2024-Q4 2025.html] RECONCILIATION OF GAAP TO NON-GAAP FINANCIAL INFORMATION (Unaudited) [tables/table-12 Financial Metrics 2021-25.html] RECONCILIATION OF GAAP TO NON-GAAP FINANCIAL INFORMATION (Unaudited) [tables/table-13 Financial Data 2022-25.html] [tables/table-14 Financial Metrics 2023-25.html] TTM = Trailing twelve months (1) Beginning in Q1'25, Capital expenditures is presented inclusive of purchases of energy generation and storage systems and all prior periods have been adjusted. (2) As a result of the adoption of the new crypto assets standard, the previously reported quarterly periods in 2024 have been recast. (3) Beginning in Q1'25, Adjusted EBITDA (non-GAAP) is presented net of digital assets gains and losses and all prior periods have been adjusted.", + "chunk_count": 1, + "children": [] + }, + { + "title": "WEBCAST INFORMATION", + "path": "TSLA-Q4-2025-Update.pdf/WEBCAST INFORMATION", + "level": 1, + "summary": "Tesla will provide a live webcast of its fourth quarter 2025 financial results conference call beginning at 4:30 p.m. CT on January 28, 2026 at ir.tesla.com. This webcast will also be available for replay for approximately one year thereafter.", + "chunk_count": 1, + "children": [] + }, + { + "title": "CERTAIN TERMS", + "path": "TSLA-Q4-2025-Update.pdf/CERTAIN TERMS", + "level": 1, + "summary": "This passage defines key financial and operational terms used in a specific update. It clarifies that vehicle deliveries refer to units transferred to end customers with completed paperwork, while energy product deployment includes installed customer units and equipment sales. The text details non-GAAP measures: 'Net income attributable to common stockholders' adjusts for stock-based compensation, digital asset gains/losses, and valuation allowances; 'Adjusted EBITDA' excludes interest, taxes, depreciation, amortization, impairment, stock-based compensation, and digital asset impacts. 'Free cash flow' is defined as operating cash flow minus capital expenditures. Operational metrics include 'Average cost per vehicle' (automotive sales cost divided by new deliveries excluding leases), 'Days sales outstanding' (average receivables divided by revenue times days in period), 'Days payable outstanding' (average payables divided by cost of revenues times days in period), and 'Days of supply' (ending inventory divided by deliveries using trading days). Finally, constant currency impacts are calculated by comparing actuals against results converted to USD using prior period average exchange rates.", + "chunk_count": 1, + "children": [] + }, + { + "title": "NON-GAAP FINANCIAL INFORMATION", + "path": "TSLA-Q4-2025-Update.pdf/NON-GAAP FINANCIAL INFORMATION", + "level": 1, + "summary": "Tesla presents consolidated financial information under both GAAP and non-GAAP standards to supplement its results. Key non-GAAP measures include net income attributable to common stockholders, diluted per share figures, Adjusted EBITDA margin, automotive gross margin, and free cash flow. These metrics aid internal management comparisons with historical performance and other companies, supporting operating, budgeting, and planning activities. Management believes these non-GAAP figures provide investors with a clearer view of Tesla's financial condition and operational results by reflecting the metrics used to run the business. However, since non-GAAP data is not prepared under comprehensive accounting rules, it should be read alongside U.S. GAAP information for a complete understanding of Tesla's performance. A reconciliation between GAAP and non-GAAP data is provided elsewhere.", + "chunk_count": 1, + "children": [] + }, + { + "title": "FORWARD-LOOKING STATEMENTS", + "path": "TSLA-Q4-2025-Update.pdf/FORWARD-LOOKING STATEMENTS", + "level": 1, + "summary": "This passage from Tesla outlines that various statements in the update, particularly regarding future outlooks, product development, production capacity, financial metrics, and facility expansions, constitute forward-looking statements under the Private Securities Litigation Reform Act of 1995. These statements reflect management's current expectations based on assumptions and are subject to risks and uncertainties, meaning actual results may differ materially. The text lists numerous specific risk factors that could impact outcomes, including manufacturing delays, supply chain challenges, regulatory hurdles, competition, product performance issues, employee retention, and foreign exchange fluctuations. Tesla advises investors to consult SEC filings, specifically the Form 10-K filed on January 30, 2025, for detailed risk disclosures. The company explicitly disclaims any obligation to update these forward-looking statements due to new information or future events.", + "chunk_count": 1, "children": [] } ], @@ -160,4 +444,4 @@ } ] } -} \ No newline at end of file +} diff --git a/apps/api/app/mcp/retrieval_server.py b/apps/api/app/mcp/retrieval_server.py index b3f33a0c0..e9d9f0915 100644 --- a/apps/api/app/mcp/retrieval_server.py +++ b/apps/api/app/mcp/retrieval_server.py @@ -11,11 +11,11 @@ from sqlalchemy.ext.asyncio import AsyncSession from shared.core.database import get_db_context +from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace from shared.services.retrieval import run_retrieval_query DbFactory = Callable[[], AsyncContextManager[AsyncSession]] KNOWHERE_NAMESPACE_HEADER = "x-knowhere-namespace" -DEFAULT_NAMESPACE = "default" def create_public_mcp_transport_security() -> TransportSecuritySettings: @@ -44,10 +44,10 @@ def resolve_mcp_namespace(*, ctx: Context | None) -> str: try: request = get_mcp_request(ctx) except (RuntimeError, ValueError): - return DEFAULT_NAMESPACE + return normalize_retrieval_namespace(None) headers = getattr(request, "headers", {}) or {} - namespace = str(get_header(headers, KNOWHERE_NAMESPACE_HEADER) or "").strip() - return namespace or DEFAULT_NAMESPACE + namespace = get_header(headers, KNOWHERE_NAMESPACE_HEADER) + return normalize_retrieval_namespace(namespace) def to_mcp_query_response(response: dict[str, Any]) -> dict[str, Any]: diff --git a/apps/api/app/services/demo/source_catalog.py b/apps/api/app/services/demo/source_catalog.py index 1b0add285..2158f0889 100644 --- a/apps/api/app/services/demo/source_catalog.py +++ b/apps/api/app/services/demo/source_catalog.py @@ -65,7 +65,7 @@ class DemoSourceDefinition: citations=( DemoCitationDefinition( section_path=( - "Default_Root/TSLA-Q4-2025-Update.pdf-->OTHER UPDATES" + "TSLA-Q4-2025-Update.pdf-->OTHER UPDATES" ), description="xAI investment", content=( @@ -89,7 +89,7 @@ class DemoSourceDefinition: citations=( DemoCitationDefinition( section_path=( - "Default_Root/TSLA-Q4-2025-Update.pdf-->SUMMARY-->" + "TSLA-Q4-2025-Update.pdf-->SUMMARY-->" "Energy generation and storage" ), description="Storage deployment growth", @@ -113,7 +113,7 @@ class DemoSourceDefinition: citations=( DemoCitationDefinition( section_path=( - "Default_Root/TSLA-Q4-2025-Update.pdf-->OUTLOOK-->" + "TSLA-Q4-2025-Update.pdf-->OUTLOOK-->" "Product" ), description="2026 production plans", diff --git a/apps/api/app/services/demo/source_projection.py b/apps/api/app/services/demo/source_projection.py index c59511c17..46a5f8ffe 100644 --- a/apps/api/app/services/demo/source_projection.py +++ b/apps/api/app/services/demo/source_projection.py @@ -225,7 +225,7 @@ def _publication_path( source: _DemoSourceDefinition, raw_path: str | None, ) -> str: - prefix = f"Default_Root/{source.title}" + prefix = source.title raw = str(raw_path or "").strip() if not raw: return prefix @@ -238,8 +238,11 @@ def _publication_path( 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": + if parts and parts[0] == source.title: return raw + if len(parts) >= 2 and parts[0] == "Default_Root": + section_parts = parts[2:] if parts[1] == source.title else parts[1:] + return "/".join([prefix, *section_parts]) if section_parts else prefix return prefix diff --git a/apps/api/app/services/document_ingestion/scope_service.py b/apps/api/app/services/document_ingestion/scope_service.py index 34390799d..afc13ada0 100644 --- a/apps/api/app/services/document_ingestion/scope_service.py +++ b/apps/api/app/services/document_ingestion/scope_service.py @@ -17,6 +17,7 @@ ValidationException, ) from shared.models.database.job import Job +from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace _ACTIVE_JOB_STATUSES = ("waiting-file", "pending", "running", "converting") @@ -74,8 +75,9 @@ async def resolve_effective_document_scope( requested_namespace: Optional[str], repository: DocumentRepository | None = None, ) -> tuple[str, str]: + effective_requested_namespace = normalize_retrieval_namespace(requested_namespace) if not document_id: - return f"doc_{uuid.uuid4().hex[:12]}", requested_namespace or "default" + return f"doc_{uuid.uuid4().hex[:12]}", effective_requested_namespace document = await (repository or DocumentRepository()).get_document( db, @@ -88,7 +90,8 @@ async def resolve_effective_document_scope( resource_id=document_id, internal_message=f"Document not found for update flow: {document_id}", ) - if requested_namespace and requested_namespace != document.namespace: + has_requested_namespace = bool(str(requested_namespace or "").strip()) + if has_requested_namespace and effective_requested_namespace != document.namespace: raise ValidationException( user_message="namespace must match the existing document namespace", violations=[ diff --git a/apps/api/tests/contract/test_chunk_document_path_contract.py b/apps/api/tests/contract/test_chunk_document_path_contract.py new file mode 100644 index 000000000..d10a99c46 --- /dev/null +++ b/apps/api/tests/contract/test_chunk_document_path_contract.py @@ -0,0 +1,247 @@ +from typing import cast + +import pytest + +from shared.services.retrieval.search.lexical_text import section_path_from_chunk_path +from shared.services.storage.zip_result_schema import ZipResultSchemaBuilder + + +@pytest.mark.parametrize( + "chunk_path", + [ + "acme.com/report.pdf/Intro/Subsection", + "team/foo.fragment/Intro/Subsection", + "team/foo.atlas/Intro/Subsection", + "team/photo.png/Intro/Subsection", + "images/report.pdf/Intro/Subsection", + "tables/report.pdf/Intro/Subsection", + "client.pdf/report.pdf/Intro/Subsection", + ], +) +def test_should_read_legacy_namespace_chunk_paths_as_document_sections( + chunk_path: str, +) -> None: + source_file_name = chunk_path.split("/")[1] + + doc_nav = ZipResultSchemaBuilder().build_doc_nav( + [ + { + "chunk_id": "chunk_legacy_path", + "type": "text", + "content": "legacy path content", + "path": chunk_path, + "metadata": {"summary": "legacy path summary"}, + } + ], + source_file_name, + ) + + sections = cast(list[dict[str, object]], doc_nav["sections"]) + children = cast(list[dict[str, object]], sections[0]["children"]) + + assert ( + section_path_from_chunk_path( + chunk_path, + source_file_name=source_file_name, + ) + == "Intro / Subsection" + ) + assert sections[0]["title"] == "Intro" + assert sections[0]["path"] == "/".join(chunk_path.split("/")[:3]) + assert children[0]["title"] == "Subsection" + assert children[0]["path"] == chunk_path + + +@pytest.mark.parametrize( + "chunk_path", + [ + "images/photo.png", + "tables/table.html", + ], +) +def test_should_keep_media_resource_paths_at_root(chunk_path: str) -> None: + assert section_path_from_chunk_path(chunk_path) == "Root" + + +def test_should_not_treat_dotted_section_titles_as_legacy_document_files() -> None: + chunk_path = "report.pdf/1. Introduction/Details" + + doc_nav = ZipResultSchemaBuilder().build_doc_nav( + [ + { + "chunk_id": "chunk_numbered_section", + "type": "text", + "content": "numbered section content", + "path": chunk_path, + "metadata": {"summary": "numbered section summary"}, + } + ], + "report.pdf", + ) + + sections = cast(list[dict[str, object]], doc_nav["sections"]) + children = cast(list[dict[str, object]], sections[0]["children"]) + + assert section_path_from_chunk_path(chunk_path) == "1. Introduction / Details" + assert sections[0]["title"] == "1. Introduction" + assert sections[0]["path"] == "report.pdf/1. Introduction" + assert children[0]["title"] == "Details" + assert children[0]["path"] == chunk_path + + +def test_should_read_arrow_delimited_document_paths_as_section_paths() -> None: + chunk_path = "report.pdf-->Intro-->Subsection" + + doc_nav = ZipResultSchemaBuilder().build_doc_nav( + [ + { + "chunk_id": "chunk_arrow_delimited_path", + "type": "text", + "content": "arrow delimited content", + "path": chunk_path, + "metadata": {"summary": "arrow delimited summary"}, + } + ], + "report.pdf", + ) + + sections = cast(list[dict[str, object]], doc_nav["sections"]) + children = cast(list[dict[str, object]], sections[0]["children"]) + + assert ( + section_path_from_chunk_path( + chunk_path, + source_file_name="report.pdf", + ) + == "Intro / Subsection" + ) + assert sections[0]["title"] == "Intro" + assert sections[0]["path"] == "report.pdf/Intro" + assert children[0]["title"] == "Subsection" + assert children[0]["path"] == "report.pdf/Intro/Subsection" + + +def test_should_preserve_literal_arrow_text_in_slash_paths() -> None: + chunk_path = "report.pdf/Inputs --> Outputs/Details" + + doc_nav = ZipResultSchemaBuilder().build_doc_nav( + [ + { + "chunk_id": "chunk_literal_arrow_heading", + "type": "text", + "content": "literal arrow content", + "path": chunk_path, + "metadata": {"summary": "literal arrow summary"}, + } + ], + "report.pdf", + ) + + sections = cast(list[dict[str, object]], doc_nav["sections"]) + children = cast(list[dict[str, object]], sections[0]["children"]) + + assert ( + section_path_from_chunk_path( + chunk_path, + source_file_name="report.pdf", + ) + == "Inputs --> Outputs / Details" + ) + assert sections[0]["title"] == "Inputs --> Outputs" + assert sections[0]["path"] == "report.pdf/Inputs --> Outputs" + assert children[0]["title"] == "Details" + assert children[0]["path"] == chunk_path + + +def test_should_preserve_single_level_literal_arrow_headings() -> None: + chunk_path = "report.pdf/Inputs --> Outputs" + + doc_nav = ZipResultSchemaBuilder().build_doc_nav( + [ + { + "chunk_id": "chunk_single_arrow_heading", + "type": "text", + "content": "single arrow heading content", + "path": chunk_path, + "metadata": {"summary": "single arrow heading summary"}, + } + ], + "report.pdf", + ) + + sections = cast(list[dict[str, object]], doc_nav["sections"]) + + assert ( + section_path_from_chunk_path( + chunk_path, + source_file_name="report.pdf", + ) + == "Inputs --> Outputs" + ) + assert sections[0]["title"] == "Inputs --> Outputs" + assert sections[0]["path"] == chunk_path + assert sections[0]["children"] == [] + + +def test_should_preserve_filename_like_arrow_text_below_document_root() -> None: + chunk_path = "report.pdf/input.csv-->output/Details" + + doc_nav = ZipResultSchemaBuilder().build_doc_nav( + [ + { + "chunk_id": "chunk_filename_arrow_heading", + "type": "text", + "content": "filename arrow heading content", + "path": chunk_path, + "metadata": {"summary": "filename arrow heading summary"}, + } + ], + "report.pdf", + ) + + sections = cast(list[dict[str, object]], doc_nav["sections"]) + children = cast(list[dict[str, object]], sections[0]["children"]) + + assert ( + section_path_from_chunk_path( + chunk_path, + source_file_name="report.pdf", + ) + == "input.csv-->output / Details" + ) + assert sections[0]["title"] == "input.csv-->output" + assert sections[0]["path"] == "report.pdf/input.csv-->output" + assert children[0]["title"] == "Details" + assert children[0]["path"] == chunk_path + + +def test_should_preserve_filename_like_section_titles_for_new_paths() -> None: + chunk_path = "report.pdf/appendix.pdf/Details" + + doc_nav = ZipResultSchemaBuilder().build_doc_nav( + [ + { + "chunk_id": "chunk_filename_section", + "type": "text", + "content": "filename-like section content", + "path": chunk_path, + "metadata": {"summary": "filename-like section summary"}, + } + ], + "report.pdf", + ) + + sections = cast(list[dict[str, object]], doc_nav["sections"]) + children = cast(list[dict[str, object]], sections[0]["children"]) + + assert ( + section_path_from_chunk_path( + chunk_path, + source_file_name="report.pdf", + ) + == "appendix.pdf / Details" + ) + assert sections[0]["title"] == "appendix.pdf" + assert sections[0]["path"] == "report.pdf/appendix.pdf" + assert children[0]["title"] == "Details" + assert children[0]["path"] == chunk_path diff --git a/apps/api/tests/contract/test_job_creation_contract.py b/apps/api/tests/contract/test_job_creation_contract.py index 95bb57cb2..572488351 100644 --- a/apps/api/tests/contract/test_job_creation_contract.py +++ b/apps/api/tests/contract/test_job_creation_contract.py @@ -566,6 +566,44 @@ async def test_should_return_not_found_when_creating_a_job_for_an_archived_docum assert await _count_jobs() == 0 +@pytest.mark.asyncio +@pytest.mark.parametrize("namespace_value", [None, "", " "]) +async def test_should_inherit_existing_document_namespace_when_update_namespace_is_blank( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + namespace_value: str | None, +) -> None: + document_id = f"doc_contract_{uuid4().hex[:12]}" + existing_namespace = "contract.jobs.custom" + payload: dict[str, object] = { + "document_id": document_id, + "namespace": namespace_value, + "source_type": "file", + "file_name": "replacement-upload.pdf", + "data_id": "contract-document-update", + } + + async with developer_api_client_factory() as api_client: + await _insert_document( + document_id=document_id, + namespace=existing_namespace, + ) + response = await api_client.post("/api/v1/jobs", json=payload) + + assert response.status_code == 200 + + response_json: dict[str, object] = response.json() + job_id = cast(str, response_json["job_id"]) + job_metadata = cast(dict[str, object], (await _load_job_record(job_id))["job_metadata"]) + original_request = cast(dict[str, object], job_metadata["original_request"]) + + assert response_json["namespace"] == existing_namespace + assert job_metadata["document_id"] == document_id + assert job_metadata["namespace"] == existing_namespace + assert original_request["namespace"] == namespace_value + + @pytest.mark.asyncio async def test_should_create_a_waiting_file_job_for_a_url_source_and_enqueue_the_upload_worker( monkeypatch: MonkeyPatch, diff --git a/apps/worker/app/services/connect_builder/summary_builder.py b/apps/worker/app/services/connect_builder/summary_builder.py index 7804bd6b6..ae020e429 100644 --- a/apps/worker/app/services/connect_builder/summary_builder.py +++ b/apps/worker/app/services/connect_builder/summary_builder.py @@ -334,7 +334,7 @@ def build_section_summary_lookup(file_dir: str) -> Dict[str, str]: """Build a flat {section_path: summary} dict from all nodes in doc_nav.json. Keys use the DocumentSection.section_path format produced by - ``section_path_from_chunk_path`` (strips namespace + filename prefix, + ``section_path_from_chunk_path`` (strips the filename prefix, joins remaining parts with ``" / "``). Traverses the full section tree at all depths. Used by the publication @@ -354,12 +354,16 @@ def build_section_summary_lookup(file_dir: str) -> Dict[str, str]: return {} lookup: Dict[str, str] = {} + source_file_name = str(doc_nav.get("file_name") or "") def _walk(node: Dict[str, Any]) -> None: nav_path = node.get("path", "") summary = node.get("summary", "") if nav_path and summary: - section_path = section_path_from_chunk_path(nav_path) + section_path = section_path_from_chunk_path( + nav_path, + source_file_name=source_file_name, + ) if section_path and section_path != "Root": lookup[section_path] = summary for child in node.get("children", []): diff --git a/apps/worker/app/services/document_ingestion/parse_result_package.py b/apps/worker/app/services/document_ingestion/parse_result_package.py index f35c6414c..c26fdf186 100644 --- a/apps/worker/app/services/document_ingestion/parse_result_package.py +++ b/apps/worker/app/services/document_ingestion/parse_result_package.py @@ -4,6 +4,7 @@ from typing import Any import pandas as pd +from app.services.document_parser.orchestration.parse_output import ParseOutput from loguru import logger from shared.core.exceptions.domain_exceptions import WorkerHandlingException @@ -33,51 +34,48 @@ class GeneratedResultPackage: statistics: dict[str, Any] zip_size: int - @classmethod - def from_legacy_tuple( - cls, - zip_file_path: str, - checksum: dict[str, Any] | str | None, - statistics: dict[str, Any], - zip_size: int, - ) -> GeneratedResultPackage: - checksum_value = ( - str(checksum.get("value", "")) - if isinstance(checksum, dict) - else str(checksum or "") - ) - return cls( - zip_file_path=zip_file_path, - checksum_value=checksum_value, - statistics=statistics, - zip_size=zip_size, - ) - def build_parse_result_package( *, job_id: str, filename: str, - add_dir: str, - parsed_contents_df: pd.DataFrame | None, + parse_output: ParseOutput, ) -> ParseResultPackage: artifact = _build_parse_artifact( job_id=job_id, filename=filename, - add_dir=add_dir, - parsed_contents_df=parsed_contents_df, + parse_output=parse_output, ) chunks = dataframe_to_chunks(artifact.dataframe) return ParseResultPackage(artifact=artifact, chunks=chunks) +def build_generated_result_package( + zip_file_path: str, + checksum: dict[str, Any] | str | None, + statistics: dict[str, Any], + zip_size: int, +) -> GeneratedResultPackage: + checksum_value = ( + str(checksum.get("value", "")) + if isinstance(checksum, dict) + else str(checksum or "") + ) + return GeneratedResultPackage( + zip_file_path=zip_file_path, + checksum_value=checksum_value, + statistics=statistics, + zip_size=zip_size, + ) + + def _build_parse_artifact( *, job_id: str, filename: str, - add_dir: str, - parsed_contents_df: pd.DataFrame | None, + parse_output: ParseOutput, ) -> ParseArtifact: + parsed_contents_df = parse_output.parsed_df if parsed_contents_df is None: raise WorkerHandlingException( user_message="We could not extract content from your file", @@ -89,4 +87,4 @@ def _build_parse_artifact( f"No content returned from file parsing: job_id={job_id}, filename={filename}" ) - return ParseArtifact(add_dir=add_dir, dataframe=parsed_contents_df) + return ParseArtifact(add_dir=parse_output.output_dir, dataframe=parsed_contents_df) diff --git a/apps/worker/app/services/document_ingestion/processing_run.py b/apps/worker/app/services/document_ingestion/processing_run.py index ba9b49878..d5d405596 100644 --- a/apps/worker/app/services/document_ingestion/processing_run.py +++ b/apps/worker/app/services/document_ingestion/processing_run.py @@ -13,8 +13,8 @@ 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.parse_result_package import ( - GeneratedResultPackage, ParseArtifact, + build_generated_result_package, build_parse_result_package, ) from app.services.document_ingestion.processing_billing import ( @@ -169,17 +169,12 @@ def _run_parse_job( filename=filename, doc_type=doc_type, ): - namespace = JobMetadataHelper.get_namespace( - job_context.job_metadata, - "default", - ) - add_dir, parsed_contents_df = parse_service.checkerboard_inject_parse( + parse_output = parse_service.checkerboard_parse_output( file_full_path=local_temp_path, filename=filename, output_dir=output_dir, job_id=job_id, internal_output_filename=internal_parse_name, - namespace=namespace or "default", doc_type=doc_type, smart_title_parse=JobMetadataHelper.get_parsing_param( job_context.job_metadata, @@ -211,8 +206,8 @@ def _run_parse_job( 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}" + f"job_id={job_id}, output_dir={parse_output.output_dir}, " + f"chunks={parse_output.rows_count}" ) lifecycle_service.update_progress( @@ -223,8 +218,7 @@ def _run_parse_job( result_package = build_parse_result_package( job_id=job_id, filename=filename, - add_dir=add_dir, - parsed_contents_df=parsed_contents_df, + parse_output=parse_output, ) lifecycle_service.update_progress( @@ -319,7 +313,7 @@ def _finalize_parse_job_success( data_id = JobMetadataHelper.get_data_id(job_context.job_metadata) zip_service = ZipResultService() - generated_package = GeneratedResultPackage.from_legacy_tuple( + generated_package = build_generated_result_package( *zip_service.generate_zip_package( job_id=job_id, chunks=chunks, diff --git a/apps/worker/app/services/document_parser/formats/fragment/parser.py b/apps/worker/app/services/document_parser/formats/fragment/parser.py index 0d5dde56f..66a20069e 100644 --- a/apps/worker/app/services/document_parser/formats/fragment/parser.py +++ b/apps/worker/app/services/document_parser/formats/fragment/parser.py @@ -9,8 +9,8 @@ from typing import Any, Optional from app.services.document_parser.formats.markdown.parser import parse_md -from app.services.document_parser.orchestration.namespace_path_segment import ( - build_namespace_path_segment, +from app.services.document_parser.orchestration.path_segment import ( + build_parser_path_segment, ) from loguru import logger from openai.types.chat import ChatCompletionMessageParam @@ -56,7 +56,6 @@ def parse_fragment( fragment_content: str, filename: Optional[str] = None, output_dir: Optional[str] = None, - namespace: str = "default", base_llm_paras: Optional[dict[str, Any]] = None, **kwargs, ): @@ -67,7 +66,6 @@ def parse_fragment( fragment_content: The text content to parse filename: Optional filename, if not provided will be auto-generated output_dir: Base output directory - namespace: Retrieval namespace used for parser-relative paths base_llm_paras: LLM parameters for parsing Returns: @@ -92,13 +90,13 @@ def parse_fragment( logger.debug(f"Fragment filename: {filename}") # Build relative_root and full_output_dir - namespace_segment = build_namespace_path_segment(namespace) - filename_segment = path_handle(filename, mode="clean_single") - if not isinstance(filename_segment, str) or filename_segment in {"", ".", ".."}: - filename_segment = f"fragment_{os.urandom(4).hex()}.fragment" - relative_root = f"{namespace_segment}/{filename_segment}" + filename_segment = build_parser_path_segment( + filename, + default=f"fragment_{os.urandom(4).hex()}.fragment", + ) + relative_root = filename_segment full_output_dir = os.path.realpath( - os.path.join(output_dir, namespace_segment, filename_segment) + os.path.join(output_dir, filename_segment) ) resolved_output_dir = os.path.realpath(output_dir) if ( 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 48d05396d..c5b0972b7 100644 --- a/apps/worker/app/services/document_parser/orchestration/format_adapters.py +++ b/apps/worker/app/services/document_parser/orchestration/format_adapters.py @@ -28,7 +28,6 @@ def parse(self, session: ParseSession) -> ParseOutput: session.fragment_content, filename=session.filename, output_dir=session.output_dir, - namespace=session.namespace, base_llm_paras=session.base_llm_paras, ) return ParseOutput(output_dir=full_output_dir, parsed_df=parsed_df) diff --git a/apps/worker/app/services/document_parser/orchestration/namespace_path_segment.py b/apps/worker/app/services/document_parser/orchestration/namespace_path_segment.py deleted file mode 100644 index f07cfc32e..000000000 --- a/apps/worker/app/services/document_parser/orchestration/namespace_path_segment.py +++ /dev/null @@ -1,18 +0,0 @@ -from __future__ import annotations - -from app.services.common.file_utils import path_handle - -_DEFAULT_NAMESPACE_SEGMENT = "default" - - -def build_namespace_path_segment(namespace: str | None) -> str: - """Map a retrieval namespace to one safe parser path segment.""" - raw_namespace = (namespace or _DEFAULT_NAMESPACE_SEGMENT).strip() - sanitized_namespace = path_handle(raw_namespace, mode="clean_single") - if not isinstance(sanitized_namespace, str): - return _DEFAULT_NAMESPACE_SEGMENT - - namespace_segment = sanitized_namespace.strip() - if namespace_segment in {"", ".", ".."}: - return _DEFAULT_NAMESPACE_SEGMENT - return namespace_segment diff --git a/apps/worker/app/services/document_parser/orchestration/parse_input.py b/apps/worker/app/services/document_parser/orchestration/parse_input.py index 92c4749a1..98dd9fa2c 100644 --- a/apps/worker/app/services/document_parser/orchestration/parse_input.py +++ b/apps/worker/app/services/document_parser/orchestration/parse_input.py @@ -22,7 +22,6 @@ class ParseInput: output_dir: str internal_output_filename: str job_id: str | None = None - namespace: str = "default" options: ParseOptions = field(default_factory=ParseOptions) base_url: str = "" fragment_content: str = "" diff --git a/apps/worker/app/services/document_parser/orchestration/parse_output.py b/apps/worker/app/services/document_parser/orchestration/parse_output.py index 60084e6dd..8e70fd56d 100644 --- a/apps/worker/app/services/document_parser/orchestration/parse_output.py +++ b/apps/worker/app/services/document_parser/orchestration/parse_output.py @@ -7,19 +7,11 @@ @dataclass(frozen=True) class ParseOutput: - """Parser adapter output with a compatibility view for legacy callers.""" + """Parser adapter output.""" output_dir: str parsed_df: pd.DataFrame | None - @classmethod - def from_legacy_tuple( - cls, - output_dir: str, - parsed_df: pd.DataFrame | None, - ) -> ParseOutput: - return cls(output_dir=output_dir, parsed_df=parsed_df) - @property def rows_count(self) -> int: if self.parsed_df is None: @@ -28,6 +20,3 @@ def rows_count(self) -> int: def with_dataframe(self, parsed_df: pd.DataFrame | None) -> ParseOutput: return ParseOutput(output_dir=self.output_dir, parsed_df=parsed_df) - - def as_legacy_tuple(self) -> tuple[str, pd.DataFrame | None]: - return self.output_dir, self.parsed_df 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 1b35c3cf5..3fe956f1a 100644 --- a/apps/worker/app/services/document_parser/orchestration/parse_session.py +++ b/apps/worker/app/services/document_parser/orchestration/parse_session.py @@ -5,8 +5,8 @@ from typing import Any from app.services.document_parser.formats.atlas.classifier import classify_atlas_with_vlm -from app.services.document_parser.orchestration.namespace_path_segment import ( - build_namespace_path_segment, +from app.services.document_parser.orchestration.path_segment import ( + build_parser_path_segment, ) from app.services.document_parser.orchestration.parse_input import ParseInput from app.services.document_parser.profiling.doc_profiler import profile_document @@ -29,7 +29,6 @@ class ParseSession: full_output_dir: str internal_output_filename: str job_id: str | None - namespace: str output_dir: str profile: Any relative_root: str @@ -54,7 +53,6 @@ def from_input( full_output_dir=full_output_dir, internal_output_filename=parse_input.internal_output_filename, job_id=parse_input.job_id, - namespace=parse_input.namespace, output_dir=parse_input.output_dir, profile=profile, relative_root=relative_root, @@ -84,7 +82,6 @@ def build_parse_session(parse_input: ParseInput) -> ParseSession: relative_root, full_output_dir = _resolve_output_paths( filename=parse_input.filename, internal_output_filename=parse_input.internal_output_filename, - namespace=parse_input.namespace, output_dir=parse_input.output_dir, ) logger.debug(f"relative_root: {relative_root}") @@ -133,7 +130,6 @@ def build_parse_session(parse_input: ParseInput) -> ParseSession: _rename_atlas_output( filename=parse_input.filename, internal_output_filename=parse_input.internal_output_filename, - namespace=parse_input.namespace, output_dir=parse_input.output_dir, ) ) @@ -144,7 +140,6 @@ def build_parse_session(parse_input: ParseInput) -> ParseSession: output_dir=parse_input.output_dir, internal_output_filename=internal_output_filename, job_id=parse_input.job_id, - namespace=parse_input.namespace, options=parse_input.options, base_url=parse_input.base_url, fragment_content=parse_input.fragment_content, @@ -164,7 +159,6 @@ def _rename_atlas_output( *, filename: str, internal_output_filename: str, - namespace: str, output_dir: str, ) -> tuple[str, str, str, str]: name_base, _ = os.path.splitext(filename) @@ -174,7 +168,6 @@ def _rename_atlas_output( relative_root, full_output_dir = _resolve_output_paths( filename=atlas_filename, internal_output_filename=atlas_internal_filename, - namespace=namespace, output_dir=output_dir, ) return atlas_filename, atlas_internal_filename, relative_root, full_output_dir @@ -184,23 +177,18 @@ def _resolve_output_paths( *, filename: str, internal_output_filename: str, - namespace: str, output_dir: str, ) -> tuple[str, str]: - namespace_segment = build_namespace_path_segment(namespace) - relative_root = ( - f"{namespace_segment}/{filename}" if filename else namespace_segment - ) - internal_relative_root = ( - f"{namespace_segment}/{internal_output_filename}" - if internal_output_filename - else namespace_segment + filename_segment = build_parser_path_segment(filename) + internal_filename_segment = build_parser_path_segment( + internal_output_filename, + default=filename_segment, ) + relative_root = filename_segment full_output_dir = os.path.join( output_dir, - namespace_segment, - internal_output_filename, + internal_filename_segment, ) resolved_output_dir = os.path.realpath(output_dir) resolved_full_output_dir = os.path.realpath(full_output_dir) @@ -213,5 +201,5 @@ def _resolve_output_paths( ) os.makedirs(resolved_full_output_dir, exist_ok=True) - logger.debug(f"internal_relative_root: {internal_relative_root}") + logger.debug(f"internal_output_root: {internal_filename_segment}") return relative_root, resolved_full_output_dir diff --git a/apps/worker/app/services/document_parser/orchestration/path_segment.py b/apps/worker/app/services/document_parser/orchestration/path_segment.py new file mode 100644 index 000000000..ed03bf9eb --- /dev/null +++ b/apps/worker/app/services/document_parser/orchestration/path_segment.py @@ -0,0 +1,19 @@ +from __future__ import annotations + +import os + +from app.services.common.file_utils import path_handle + + +def build_parser_path_segment(value: str | None, default: str = "document") -> str: + """Map parser-owned names to one safe task-local path segment.""" + raw_value = str(value or "").strip() + raw_segment = os.path.basename(raw_value) if raw_value else default + sanitized_segment = path_handle(raw_segment, mode="clean_single") + if not isinstance(sanitized_segment, str): + return default + + segment = sanitized_segment.strip() + if segment in {"", ".", ".."}: + return default + return segment diff --git a/apps/worker/app/services/document_parser/parse_service.py b/apps/worker/app/services/document_parser/parse_service.py index a9e9e7505..ba3d9f6ec 100644 --- a/apps/worker/app/services/document_parser/parse_service.py +++ b/apps/worker/app/services/document_parser/parse_service.py @@ -1,7 +1,5 @@ """Stable parser seam backed by dedicated orchestration modules.""" -import pandas as pd - from app.services.document_parser.orchestration.parse_input import ParseInput, ParseOptions from app.services.document_parser.orchestration.parse_pipeline import ( ParsePipelineResult, @@ -15,7 +13,6 @@ def checkerboard_parse_output( output_dir: str, internal_output_filename: str, job_id: str | None = None, - namespace: str = "default", llm_histories: int = 5, smart_title_parse: bool = True, summary_image: bool = True, @@ -47,48 +44,6 @@ def checkerboard_parse_output( ), base_url=base_url, fragment_content=fragment_content, - namespace=namespace, s3_key=s3_key, ) return run_parse_pipeline(parse_input) - - -def checkerboard_inject_parse( - file_full_path: str, - filename: str, - output_dir: str, - internal_output_filename: str, - job_id: str | None = None, - namespace: str = "default", - 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 = "", - base_url: str = "", - fragment_content: str = "", - s3_key: str | None = None, -) -> tuple[str, pd.DataFrame | None]: - """Run the parser seam using the legacy tuple interface.""" - return checkerboard_parse_output( - file_full_path=file_full_path, - filename=filename, - output_dir=output_dir, - internal_output_filename=internal_output_filename, - job_id=job_id, - namespace=namespace, - 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, - add_frag_desc=add_frag_desc, - base_url=base_url, - fragment_content=fragment_content, - s3_key=s3_key, - ).as_legacy_tuple() diff --git a/apps/worker/tests/contract/test_excel_parser_contract.py b/apps/worker/tests/contract/test_excel_parser_contract.py index f7992a933..6f4935cf2 100644 --- a/apps/worker/tests/contract/test_excel_parser_contract.py +++ b/apps/worker/tests/contract/test_excel_parser_contract.py @@ -29,15 +29,16 @@ 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 + from app.services.document_parser.parse_service import checkerboard_parse_output workbook_path = tmp_path / "budget.xlsx" + output_root = tmp_path / "parser-output" _write_contract_workbook(workbook_path) - full_output_dir, parsed_df = checkerboard_inject_parse( + parse_output = checkerboard_parse_output( file_full_path=str(workbook_path), filename="budget.xlsx", - output_dir=str(tmp_path), + output_dir=str(output_root), internal_output_filename="budget.xlsx", summary_image=False, summary_table=False, @@ -46,7 +47,9 @@ def test_xlsx_parser_contract_uses_stable_entrypoint_and_ignores_hidden_sheets( stopwords=[], ) - assert full_output_dir.endswith("default/budget.xlsx") + full_output_dir = parse_output.output_dir + parsed_df = parse_output.parsed_df + assert full_output_dir.endswith("budget.xlsx") assert parsed_df is not None assert parsed_df["type"].tolist() == ["table"] assert parsed_df["path"].tolist() == ["tables/table-Visible.html"] @@ -63,22 +66,21 @@ def test_xlsx_parser_contract_uses_stable_entrypoint_and_ignores_hidden_sheets( assert "Hidden" not in table_html_text -def test_parser_maps_namespace_to_task_local_path_segment( +def test_parser_maps_document_name_to_task_local_path_segment( worker_contract_environment: None, tmp_path: Path, ) -> None: - from app.services.document_parser.parse_service import checkerboard_inject_parse + from app.services.document_parser.parse_service import checkerboard_parse_output workbook_path = tmp_path / "budget.xlsx" output_root = tmp_path / "parser-output" _write_contract_workbook(workbook_path) - full_output_dir, parsed_df = checkerboard_inject_parse( + parse_output = checkerboard_parse_output( file_full_path=str(workbook_path), - filename="budget.xlsx", + filename="/tmp/../images.xlsx", output_dir=str(output_root), - internal_output_filename="budget.xlsx", - namespace="/tmp/../images", + internal_output_filename="../../images.xlsx", summary_image=False, summary_table=False, summary_txt=False, @@ -86,10 +88,12 @@ def test_parser_maps_namespace_to_task_local_path_segment( stopwords=[], ) + full_output_dir = parse_output.output_dir + parsed_df = parse_output.parsed_df assert ( os.path.commonpath([str(output_root.resolve()), full_output_dir]) == str(output_root.resolve()) ) - assert full_output_dir.endswith("_tmp_.._images/budget.xlsx") + assert full_output_dir.endswith("images.xlsx") assert parsed_df is not None assert parsed_df["path"].tolist() == ["tables/table-Visible.html"] diff --git a/apps/worker/tests/contract/test_parse_task_contract.py b/apps/worker/tests/contract/test_parse_task_contract.py index 86b76f00b..1ca733f56 100644 --- a/apps/worker/tests/contract/test_parse_task_contract.py +++ b/apps/worker/tests/contract/test_parse_task_contract.py @@ -111,6 +111,12 @@ def _find_task_workspaces(root: Path, job_id: str) -> list[Path]: ) +def _build_fake_parse_output(*, output_dir: Path, rows: list[dict[str, Any]]) -> Any: + from app.services.document_parser.orchestration.parse_output import ParseOutput + + return ParseOutput(output_dir=str(output_dir), parsed_df=pd.DataFrame(rows)) + + def _bind_parse_task_to_current_module( monkeypatch: MonkeyPatch, *, @@ -206,11 +212,10 @@ def fake_download_s3_file_to_temp( shutil.copy2(_SAMPLE_PDF_PATH, downloaded_path) return str(downloaded_path) - def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: + def fake_checkerboard_parse_output(**kwargs: Any) -> Any: captured_artifacts["parse_kwargs"] = kwargs output_dir = ( Path(str(kwargs["output_dir"])) - / str(kwargs["namespace"]) / str(kwargs["internal_output_filename"]) ) images_dir = output_dir / "images" @@ -225,7 +230,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: parsed_rows: list[dict[str, Any]] = [ { "content": text_content_with_refs, - "path": f"worker-contract/{file_root}/公司研究/自主可控加强，寒武纪或迎来营收快速放量周期", + "path": f"{file_root}/公司研究/自主可控加强，寒武纪或迎来营收快速放量周期", "type": "text", "length": len(text_content_with_refs), "keywords": "", @@ -246,7 +251,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: }, { "content": "chunk-2", - "path": f"worker-contract/{file_root}/相关研报/要点", + "path": f"{file_root}/相关研报/要点", "type": "text", "length": 7, "keywords": "", @@ -259,7 +264,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: }, { "content": "image caption", - "path": f"worker-contract/{file_root}/images/page-1.png", + "path": f"{file_root}/images/page-1.png", "type": "image", "length": 13, "keywords": "", @@ -272,7 +277,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: }, { "content": "table content", - "path": f"worker-contract/{file_root}/tables/table-1.html", + "path": f"{file_root}/tables/table-1.html", "type": "table", "length": 13, "keywords": "", @@ -284,7 +289,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: "page_nums": "3", }, ] - return str(output_dir), pd.DataFrame(parsed_rows) + return _build_fake_parse_output(output_dir=output_dir, rows=parsed_rows) class FakeResultStorage: def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: @@ -317,7 +322,7 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: ) 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_service, "checkerboard_parse_output", fake_checkerboard_parse_output) monkeypatch.setattr(parse_job_service, "get_result_storage", lambda: FakeResultStorage()) result = document_ingestion_tasks.parse_task.run(job_id, user_id, "document_ingestion") @@ -361,11 +366,9 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: assert captured_artifacts["parse_kwargs"]["filename"] == source_file_name assert captured_artifacts["parse_kwargs"]["internal_output_filename"] == source_file_name assert Path(str(captured_artifacts["parse_kwargs"]["file_full_path"])).name == source_file_name - assert captured_artifacts["parse_kwargs"]["namespace"] == "worker-contract" + assert "namespace" not in captured_artifacts["parse_kwargs"] assert "kb_dir" not in captured_artifacts["parse_kwargs"] - assert captured_artifacts["result_dir"].endswith( - "worker-contract/contract-parse.pdf" - ) + assert captured_artifacts["result_dir"].endswith("contract-parse.pdf") assert captured_artifacts["doc_nav"]["file_name"] == source_file_name assert captured_artifacts["doc_nav"]["sections"][0]["title"] == "公司研究" assert captured_artifacts["manifest"]["HIERARCHY"] == { @@ -749,7 +752,7 @@ def test_should_export_full_result_when_publication_deduplicates_existing_chunks :result_id, 'text', 'already published text', - 'worker-contract/existing.pdf/Section/Duplicate text', + 'existing.pdf/Section/Duplicate text', NULL, CAST(:text_metadata AS JSON), 0, @@ -764,7 +767,7 @@ def test_should_export_full_result_when_publication_deduplicates_existing_chunks :result_id, 'image', 'already published image', - 'worker-contract/existing.pdf/images/duplicate.png', + 'existing.pdf/images/duplicate.png', 'images/duplicate.png', CAST(:image_metadata AS JSON), 1, @@ -826,10 +829,9 @@ def fake_download_s3_file_to_temp( shutil.copy2(_SAMPLE_PDF_PATH, downloaded_path) return str(downloaded_path) - def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: + def fake_checkerboard_parse_output(**kwargs: Any) -> Any: output_dir = ( Path(str(kwargs["output_dir"])) - / str(kwargs["namespace"]) / str(kwargs["internal_output_filename"]) ) images_dir = output_dir / "images" @@ -841,7 +843,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: parsed_rows: list[dict[str, Any]] = [ { "content": "duplicate text", - "path": f"worker-contract/{file_root}/Section/Duplicate text", + "path": f"{file_root}/Section/Duplicate text", "type": "text", "length": 14, "keywords": "", @@ -854,7 +856,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: }, { "content": "duplicate image", - "path": f"worker-contract/{file_root}/images/duplicate.png", + "path": f"{file_root}/images/duplicate.png", "type": "image", "length": 15, "keywords": "", @@ -867,7 +869,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: }, { "content": "new text", - "path": f"worker-contract/{file_root}/Section/New text", + "path": f"{file_root}/Section/New text", "type": "text", "length": 8, "keywords": "", @@ -879,7 +881,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: "page_nums": "3", }, ] - return str(output_dir), pd.DataFrame(parsed_rows) + return _build_fake_parse_output(output_dir=output_dir, rows=parsed_rows) class FakeResultStorage: def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: @@ -907,7 +909,7 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: _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_service, "checkerboard_parse_output", fake_checkerboard_parse_output) monkeypatch.setattr(parse_job_service, "get_result_storage", lambda: FakeResultStorage()) monkeypatch.setattr(parse_job_service, "cleanup_task_workspace", fake_cleanup_task_workspace) @@ -1065,10 +1067,9 @@ def fake_estimate_workload(file_path: str) -> Any: billing_start_barrier.wait(timeout=10) return WorkloadEstimate(page_count=1, method="contract_fake") - def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: + def fake_checkerboard_parse_output(**kwargs: Any) -> Any: output_dir = ( Path(str(kwargs["output_dir"])) - / str(kwargs["namespace"]) / str(kwargs["internal_output_filename"]) ) output_dir.mkdir(parents=True, exist_ok=True) @@ -1078,7 +1079,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: parsed_rows: list[dict[str, Any]] = [ { "content": "chunk body", - "path": f"worker-contract/{file_root}/Section/Point", + "path": f"{file_root}/Section/Point", "type": "text", "length": 10, "keywords": "", @@ -1090,7 +1091,7 @@ def fake_checkerboard_inject_parse(**kwargs: Any) -> tuple[str, pd.DataFrame]: "page_nums": "1", } ] - return str(output_dir), pd.DataFrame(parsed_rows) + return _build_fake_parse_output(output_dir=output_dir, rows=parsed_rows) class FakeResultStorage: def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: @@ -1107,7 +1108,7 @@ def upload(self, *, job_id: str, result_dir: str, zip_file_path: str) -> Any: "estimate_workload", fake_estimate_workload, ) - monkeypatch.setattr(parse_service, "checkerboard_inject_parse", fake_checkerboard_inject_parse) + monkeypatch.setattr(parse_service, "checkerboard_parse_output", fake_checkerboard_parse_output) monkeypatch.setattr(parse_job_service, "get_result_storage", lambda: FakeResultStorage()) def run_parse_task(job_id: str) -> dict[str, Any]: @@ -1282,7 +1283,7 @@ def fake_verify_s3_file_exists(storage_key: str) -> dict[str, Any]: _patch_verify_upload_exists(monkeypatch, fake_verify_s3_file_exists) monkeypatch.setattr( parse_service, - "checkerboard_inject_parse", + "checkerboard_parse_output", lambda **_kwargs: (_ for _ in ()).throw( AssertionError("terminal parse task should not invoke the parser") ), @@ -1385,7 +1386,7 @@ def 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", + "checkerboard_parse_output", lambda **_kwargs: (_ for _ in ()).throw(RuntimeError("parse failed")), ) monkeypatch.setattr( diff --git a/packages/shared-python/shared/models/schemas/job.py b/packages/shared-python/shared/models/schemas/job.py index 692b810e3..a2b686551 100644 --- a/packages/shared-python/shared/models/schemas/job.py +++ b/packages/shared-python/shared/models/schemas/job.py @@ -42,7 +42,9 @@ class JobCreate(BaseModel): """Request payload for creating a job.""" namespace: Optional[str] = Field( - None, description="Retrieval namespace; defaults to default" + None, + max_length=255, + description="Retrieval namespace; defaults to default", ) document_id: Optional[str] = Field( None, description="Existing document ID for update flows" diff --git a/packages/shared-python/shared/models/schemas/job_metadata.py b/packages/shared-python/shared/models/schemas/job_metadata.py index a4d2c7d96..71aa4d2dc 100644 --- a/packages/shared-python/shared/models/schemas/job_metadata.py +++ b/packages/shared-python/shared/models/schemas/job_metadata.py @@ -4,6 +4,8 @@ from pydantic import BaseModel, ConfigDict, Field +from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace + class JobMetadataBase(BaseModel): """Base schema for stored job metadata.""" @@ -39,9 +41,10 @@ class JobMetadataHelper: @staticmethod def create_from_request(request, **kwargs) -> Dict[str, Any]: """Build metadata from a JobCreate request without embedding user_config.""" + namespace = normalize_retrieval_namespace(request.namespace) metadata = { "original_request": request.model_dump(), - "namespace": request.namespace or "default", + "namespace": namespace, "document_id": request.document_id, "parsing_params": ( request.parsing_params.model_dump() if request.parsing_params else None @@ -115,7 +118,8 @@ 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) + namespace = JobMetadataHelper.get_string_field(metadata, "namespace", default) + return normalize_retrieval_namespace(namespace) if namespace is not None else None @staticmethod def get_document_id(metadata: Optional[Dict[str, Any]]) -> str | None: diff --git a/packages/shared-python/shared/models/schemas/retrieval_namespace.py b/packages/shared-python/shared/models/schemas/retrieval_namespace.py new file mode 100644 index 000000000..ef2be2349 --- /dev/null +++ b/packages/shared-python/shared/models/schemas/retrieval_namespace.py @@ -0,0 +1,9 @@ +from __future__ import annotations + +_DEFAULT_RETRIEVAL_NAMESPACE = "default" + + +def normalize_retrieval_namespace(namespace: str | None) -> str: + """Return the canonical namespace value used by jobs, retrieval, and caches.""" + normalized = str(namespace or "").strip() + return normalized or _DEFAULT_RETRIEVAL_NAMESPACE diff --git a/packages/shared-python/shared/services/chunks/document_path.py b/packages/shared-python/shared/services/chunks/document_path.py new file mode 100644 index 000000000..37baa9943 --- /dev/null +++ b/packages/shared-python/shared/services/chunks/document_path.py @@ -0,0 +1,156 @@ +"""Document path parsing for parser chunk paths.""" + +from __future__ import annotations + +import os + +_DOCUMENT_FILE_EXTENSIONS = { + ".csv", + ".atlas", + ".fragment", + ".gif", + ".doc", + ".docx", + ".htm", + ".html", + ".jpeg", + ".jpg", + ".json", + ".md", + ".markdown", + ".pdf", + ".png", + ".ppt", + ".pptx", + ".rtf", + ".webp", + ".txt", + ".xls", + ".xlsm", + ".xlsx", +} +_MEDIA_ROOT_SEGMENTS = {"images", "tables"} + + +def split_document_path( + path: str | None, + *, + source_file_name: str | None = None, +) -> tuple[list[str], list[str]]: + """Return ``(root_parts, section_parts)`` for new and legacy chunk paths.""" + parts = _split_path(path, source_file_name=source_file_name) + if not parts: + return [], [] + if parts[0] in _MEDIA_ROOT_SEGMENTS and not _is_legacy_namespace_path( + parts, + source_file_name=source_file_name, + ): + return parts[:1], [] + + document_index = _find_document_index(parts, source_file_name=source_file_name) + return parts[: document_index + 1], parts[document_index + 1 :] + + +def _split_path(path: str | None, *, source_file_name: str | None) -> list[str]: + raw = str(path or "").strip() + raw_segments = raw.split("/") + source_segment = _normalize_document_file_name(source_file_name) + parts: list[str] = [] + for index, segment in enumerate(raw_segments): + parts.extend( + _split_arrow_document_segment( + segment, + can_split=_can_split_arrow_document_segment( + index=index, + raw_segments=raw_segments, + segment=segment, + source_segment=source_segment, + ), + ) + ) + return parts + + +def _can_split_arrow_document_segment( + *, + index: int, + raw_segments: list[str], + segment: str, + source_segment: str, +) -> bool: + if index == 0: + return True + if index != 1: + return False + + first_segment = raw_segments[0].strip() if raw_segments else "" + if not _is_document_file_segment(first_segment): + return True + + arrow_document_segment = _normalize_document_file_name( + segment.split("-->", 1)[0] + ) + return bool(source_segment and arrow_document_segment == source_segment) + + +def _split_arrow_document_segment(segment: str, *, can_split: bool) -> list[str]: + normalized_segment = segment.strip() + if not normalized_segment: + return [] + if not can_split or "-->" not in normalized_segment: + return [normalized_segment] + + arrow_parts = [ + part.strip() + for part in normalized_segment.split("-->") + if part.strip() + ] + if arrow_parts and _is_document_file_segment(arrow_parts[0]): + return arrow_parts + return [normalized_segment] + + +def _find_document_index( + parts: list[str], + *, + source_file_name: str | None, +) -> int: + source_segment = _normalize_document_file_name(source_file_name) + if source_segment: + for index in range(min(2, len(parts))): + if _normalize_document_file_name(parts[index]) == source_segment: + return index + + if _is_legacy_namespace_path(parts, source_file_name=source_file_name): + return 1 + if _is_document_file_segment(parts[0]): + return 0 + return 0 + + +def _is_legacy_namespace_path( + parts: list[str], + *, + source_file_name: str | None, +) -> bool: + if len(parts) < 3 or not _is_document_file_segment(parts[1]): + return False + + source_segment = _normalize_document_file_name(source_file_name) + if source_segment: + return _normalize_document_file_name(parts[1]) == source_segment + return not _is_document_file_segment(parts[0]) + + +def _normalize_document_file_name(value: str | None) -> str: + if not value: + return "" + return os.path.basename(str(value).strip().replace("\\", "/")).lower() + + +def _is_document_file_segment(segment: str) -> bool: + normalized_segment = segment.lower().strip() + return any( + normalized_segment.endswith(extension) + for extension in _DOCUMENT_FILE_EXTENSIONS + ) diff --git a/packages/shared-python/shared/services/jobs/lifecycle/publication.py b/packages/shared-python/shared/services/jobs/lifecycle/publication.py index b6b2d69a1..a9dfe2f0a 100644 --- a/packages/shared-python/shared/services/jobs/lifecycle/publication.py +++ b/packages/shared-python/shared/services/jobs/lifecycle/publication.py @@ -10,6 +10,7 @@ from shared.models.database.document import DocumentSection from shared.models.database.job import Job from shared.models.schemas.job_metadata import JobMetadataHelper +from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace from shared.services.redis.redis_sync_service import SyncRedisServiceFactory from shared.services.retrieval.publication_service import RetrievalPublicationService @@ -88,7 +89,8 @@ def invalidate_cache_after_commit( redis_service = SyncRedisServiceFactory.get_service() user_id = cache_invalidation["user_id"] seen: set[str] = set() - for namespace in cache_invalidation["namespaces"]: + for raw_namespace in cache_invalidation["namespaces"]: + namespace = normalize_retrieval_namespace(str(raw_namespace)) if not namespace or namespace in seen: continue seen.add(namespace) diff --git a/packages/shared-python/shared/services/retrieval/app_service.py b/packages/shared-python/shared/services/retrieval/app_service.py index a5e3a55df..d8ec2ff22 100644 --- a/packages/shared-python/shared/services/retrieval/app_service.py +++ b/packages/shared-python/shared/services/retrieval/app_service.py @@ -4,6 +4,7 @@ from sqlalchemy.ext.asyncio import AsyncSession +from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace from shared.services.retrieval.execution.plan import ( run_retrieval_query as execute_retrieval_query, ) @@ -34,7 +35,7 @@ async def run_retrieval_query( return await execute_retrieval_query( db=db, user_id=user_id, - namespace=namespace, + namespace=normalize_retrieval_namespace(namespace), query=query, top_k=top_k, exclude_document_ids=exclude_document_ids, diff --git a/packages/shared-python/shared/services/retrieval/cache_service.py b/packages/shared-python/shared/services/retrieval/cache_service.py index bb165eb41..e081b2a8c 100644 --- a/packages/shared-python/shared/services/retrieval/cache_service.py +++ b/packages/shared-python/shared/services/retrieval/cache_service.py @@ -5,6 +5,7 @@ from loguru import logger +from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace from shared.services.redis import RedisServiceFactory _RETRIEVAL_CACHE_TTL_SECONDS = 300 @@ -13,6 +14,7 @@ def _namespace_version_key(*, user_id: str, namespace: str) -> str: + namespace = normalize_retrieval_namespace(namespace) return f"retrieval:version:{user_id}:{namespace}" @@ -75,6 +77,7 @@ def _query_cache_key( exclude_sections: list[dict[str, str]], **extra_params: Any, ) -> str: + namespace = normalize_retrieval_namespace(namespace) digest = _cache_shape_digest( query=query, top_k=top_k, @@ -110,7 +113,8 @@ async def invalidate_retrieval_cache_namespaces( *, user_id: str, namespaces: list[str] ) -> None: seen: set[str] = set() - for namespace in namespaces: + for raw_namespace in namespaces: + namespace = normalize_retrieval_namespace(raw_namespace) if not namespace or namespace in seen: continue seen.add(namespace) @@ -184,6 +188,7 @@ async def set_cached_retrieval_query_result( def _workflow_plan_cache_key(*, user_id: str, namespace: str, query: str) -> str: + namespace = normalize_retrieval_namespace(namespace) digest = hashlib.sha256(query.encode("utf-8")).hexdigest() return f"retrieval:workflow:plan:{user_id}:{namespace}:{digest}" diff --git a/packages/shared-python/shared/services/retrieval/execution/query_request.py b/packages/shared-python/shared/services/retrieval/execution/query_request.py index a233df233..493a31d53 100644 --- a/packages/shared-python/shared/services/retrieval/execution/query_request.py +++ b/packages/shared-python/shared/services/retrieval/execution/query_request.py @@ -5,6 +5,7 @@ from sqlalchemy.ext.asyncio import AsyncSession +from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace from shared.services.retrieval.execution.route_types import RetrievalRouteContext from shared.services.retrieval.settings import ( INTERNAL_RECALL_K_MULTIPLIER, @@ -55,7 +56,7 @@ def from_parameters( return cls( db=db, user_id=user_id, - namespace=namespace, + namespace=normalize_retrieval_namespace(namespace), query=str(query).strip(), top_k=top_k, exclude_document_ids=exclude_document_ids, diff --git a/packages/shared-python/shared/services/retrieval/publication_service.py b/packages/shared-python/shared/services/retrieval/publication_service.py index 40ae68045..1e7ffbe06 100644 --- a/packages/shared-python/shared/services/retrieval/publication_service.py +++ b/packages/shared-python/shared/services/retrieval/publication_service.py @@ -19,6 +19,7 @@ from shared.models.database.document import Document, DocumentChunk, DocumentSection from shared.models.database.job import Job from shared.models.database.job_result import JobResult +from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace from shared.services.retrieval.graph.service import DocumentGraphService, GraphScope from shared.services.retrieval.search.lexical_text import ( build_content_lexical_text, @@ -90,7 +91,7 @@ def _publish_document_state_for_job( ) -> Optional[Dict[str, Any]]: job_metadata = job.job_metadata or {} - namespace = job_metadata.get("namespace") or "default" + namespace = normalize_retrieval_namespace(job_metadata.get("namespace")) document_id = job_metadata.get("document_id") source_file_name = job_metadata.get("source_file_name") or job_metadata.get("file_name") @@ -132,7 +133,7 @@ def _publish_document_state_for_job( ) db.add(document) else: - namespace = namespace or document.namespace + namespace = normalize_retrieval_namespace(namespace or document.namespace) if self._is_stale_document_completion( db, document=document, @@ -154,7 +155,7 @@ def _publish_document_state_for_job( job_result = result.scalar_one_or_none() if job_result: job_result.document_id = document_id - namespace = namespace or document.namespace + namespace = normalize_retrieval_namespace(namespace or document.namespace) db.execute( delete(DocumentChunk) @@ -172,7 +173,10 @@ def _publish_document_state_for_job( for index, chunk in enumerate(deduped_chunks): chunk_metadata = chunk.get("metadata") or {} source_path = chunk_metadata.get("path") or chunk.get("path") - section_path = section_path_from_chunk_path(source_path) + section_path = section_path_from_chunk_path( + source_path, + source_file_name=source_file_name, + ) section = sections_by_path.get(section_path) if section is None: path_parts = [p for p in section_path.split(" / ") if p] @@ -222,7 +226,10 @@ def _publish_document_state_for_job( chunk_type=chunk.get("type") or chunk.get("chunk_type") or "text", content=chunk.get("content") or chunk.get("text"), content_lexical_text=build_content_lexical_text(chunk), - path_lexical_text=build_path_lexical_text(source_path), + path_lexical_text=build_path_lexical_text( + source_path, + source_file_name=source_file_name, + ), content_search_text=build_content_search_text( chunk, section_summary=section_summary ), @@ -269,7 +276,7 @@ def _publish_document_graph_for_job( ) -> None: metadata = job.job_metadata or {} - namespace = metadata.get("namespace") or "default" + namespace = normalize_retrieval_namespace(metadata.get("namespace")) document_id = metadata.get("document_id") if not document_id: document = db.execute( diff --git a/packages/shared-python/shared/services/retrieval/search/lexical_text.py b/packages/shared-python/shared/services/retrieval/search/lexical_text.py index 35ce5f8ae..b6cbc9c8c 100644 --- a/packages/shared-python/shared/services/retrieval/search/lexical_text.py +++ b/packages/shared-python/shared/services/retrieval/search/lexical_text.py @@ -6,6 +6,7 @@ from typing import Any, Optional +from shared.services.chunks.document_path import split_document_path from shared.utils.text_utils import tokenize_contents_for_retrieval @@ -57,23 +58,36 @@ def build_content_lexical_text(chunk: dict[str, Any]) -> Optional[str]: return "\n".join(lexical_parts) if lexical_parts else content -def section_path_from_chunk_path(source_path: Optional[str]) -> str: +def section_path_from_chunk_path( + source_path: Optional[str], + *, + source_file_name: str | None = None, +) -> str: """Extract section hierarchy from chunk path. - Expected format: "/.ext/
//..." + Expected format: ".ext/
//..." Returns " / "-joined section parts, or "Root" if no section hierarchy. """ if not source_path: return "Root" - parts = split_section_path(source_path) - section_parts = parts[2:] # skip namespace + filename + _, section_parts = split_document_path( + source_path, + source_file_name=source_file_name, + ) if not section_parts: return "Root" - return normalize_section_path(" / ".join(section_parts)) + return " / ".join(section_parts) -def build_path_lexical_text(source_path: Optional[str]) -> Optional[str]: - section_path = section_path_from_chunk_path(source_path) +def build_path_lexical_text( + source_path: Optional[str], + *, + source_file_name: str | None = None, +) -> Optional[str]: + section_path = section_path_from_chunk_path( + source_path, + source_file_name=source_file_name, + ) if not section_path: return None normalized_path = section_path.replace(" / ", " ") diff --git a/packages/shared-python/shared/services/storage/zip_result_schema.py b/packages/shared-python/shared/services/storage/zip_result_schema.py index 410bfa111..d4063b79d 100644 --- a/packages/shared-python/shared/services/storage/zip_result_schema.py +++ b/packages/shared-python/shared/services/storage/zip_result_schema.py @@ -10,6 +10,7 @@ normalize_connect_to_targets, parse_relationship_refs, ) +from shared.services.chunks.document_path import split_document_path from shared.utils.text_utils import truncate_content_preview from shared.utils.utc_now import utc_now_naive @@ -254,7 +255,10 @@ def build_doc_nav( } ) - sections = self._build_section_tree(text_chunks) + sections = self._build_section_tree( + text_chunks, + source_file_name=source_file_name, + ) stats["max_depth"] = _max_depth(sections) return { "version": "1.0", @@ -270,20 +274,24 @@ def build_doc_nav( def _build_section_tree( self, text_chunks: List[Dict[str, Any]], + *, + source_file_name: str, ) -> 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 [] + root_parts, section_parts = split_document_path( + path, + source_file_name=source_file_name, + ) 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, + "path": "/".join(root_parts) if root_parts else path, "summary": chunk.get("summary", ""), "chunk_count": 0, "_children_map": {}, @@ -294,7 +302,7 @@ def _build_section_tree( continue current_level = root_children - full_section_path_parts = parts[:2] + full_section_path_parts = list(root_parts) for index, part in enumerate(section_parts): full_section_path_parts.append(part) if part not in current_level: From bcbabdca2273577a3aad1b33dde04ea38c202f44 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 17:21:18 +0000 Subject: [PATCH 75/78] refactor: deepen document processing modules --- CONTEXT.md | 40 ++- apps/api/app/services/demo/source_catalog.py | 3 + apps/api/scripts/validate_demo_documents.py | 185 ++++++++++++ .../document_ingestion/parse_execution.py | 78 +++++ .../document_ingestion/processing_run.py | 274 ++---------------- .../document_ingestion/source_preparation.py | 66 +++++ .../success_finalization.py | 236 +++++++++++++++ .../services/document_ingestion/workspace.py | 30 ++ .../services/retrieval/publication_content.py | 178 ++++++++++++ .../services/retrieval/publication_models.py | 31 ++ .../services/retrieval/publication_service.py | 195 +++++-------- 11 files changed, 946 insertions(+), 370 deletions(-) create mode 100644 apps/api/scripts/validate_demo_documents.py create mode 100644 apps/worker/app/services/document_ingestion/parse_execution.py create mode 100644 apps/worker/app/services/document_ingestion/source_preparation.py create mode 100644 apps/worker/app/services/document_ingestion/success_finalization.py create mode 100644 packages/shared-python/shared/services/retrieval/publication_content.py create mode 100644 packages/shared-python/shared/services/retrieval/publication_models.py diff --git a/CONTEXT.md b/CONTEXT.md index 53a329fb9..9576fe215 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -75,10 +75,22 @@ 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. +### Worker Document Processing + +The worker-side Document Ingestion coordinator. It prepares task-local source +files, bills the workload estimate, invokes Worker Document Parsing, builds the +result package, uploads artifacts, and finalizes the Job. + +### Temporary Parse Workspace + +The task-local folder set used by Worker Document Processing: input files, +parser output files, and generated result packages. It is temporary worker +storage, not a retrieval Namespace or durable document scope. + ### Parse Output The stable parser adapter result with an output directory and optional parsed -DataFrame. It owns legacy tuple compatibility for older parser callers. +DataFrame. ### Parse Artifact @@ -140,6 +152,11 @@ jobs, and daily quota. The shared workflow that turns parsed chunks into Documents, Document Sections, Document Chunks, and document graph state. +### Publication Content + +The Publication module that replaces a single Document revision's Document +Sections and Document Chunks from parsed chunk rows. + ### Retrieval The query workflow that returns cited evidence from published documents. @@ -165,6 +182,12 @@ top-k, and data-type overrides while preserving the request policy. An API-owned canonical document shipped with the repository for demo and guest flows. +### Demo Source Validation + +The repeatable script workflow that validates Demo Source catalog metadata, +canonical chunks, citation projection, original file size, and regenerated +doc_nav.json output. + ### Demo Source Materialization The workflow that copies a Demo Source into a user's Namespace as normal Job, @@ -308,6 +331,9 @@ admission dependencies. - `app/api/v1/routes/retrieval.py` - `packages/shared-python/shared/services/retrieval/app_service.py` +- `packages/shared-python/shared/services/retrieval/publication_service.py` +- `packages/shared-python/shared/services/retrieval/publication_content.py` +- `packages/shared-python/shared/services/retrieval/publication_models.py` - `packages/shared-python/shared/services/retrieval/execution/*` - `packages/shared-python/shared/services/retrieval/search/*` - `packages/shared-python/shared/services/retrieval/hydration/*` @@ -323,6 +349,7 @@ admission dependencies. - `app/api/v1/routes/demo.py` - `app/services/demo/*` +- `apps/api/scripts/validate_demo_documents.py` ### Billing Workflow @@ -377,6 +404,17 @@ side effects. ## apps/worker Workflow Ownership +### Worker Document Processing + +- `app/services/document_ingestion/service.py` +- `app/services/document_ingestion/processing_run.py` +- `app/services/document_ingestion/source_preparation.py` +- `app/services/document_ingestion/parse_execution.py` +- `app/services/document_ingestion/success_finalization.py` +- `app/services/document_ingestion/workspace.py` +- `app/services/document_ingestion/parse_result_package.py` +- `app/services/document_ingestion/processing_billing.py` + ### Worker Document Parsing - `app/services/document_parser/parse_service.py` diff --git a/apps/api/app/services/demo/source_catalog.py b/apps/api/app/services/demo/source_catalog.py index 2158f0889..b3bdb71bb 100644 --- a/apps/api/app/services/demo/source_catalog.py +++ b/apps/api/app/services/demo/source_catalog.py @@ -133,6 +133,9 @@ class DemoSourceCatalog: def __init__(self, *, projection: DemoSourceProjection | None = None) -> None: self._projection = projection or DemoSourceProjection() + def list_sources(self) -> tuple[DemoSourceDefinition, ...]: + return _DEMO_SOURCE_DEFINITIONS + def get_catalog(self) -> dict[str, Any]: return { "sources": [ diff --git a/apps/api/scripts/validate_demo_documents.py b/apps/api/scripts/validate_demo_documents.py new file mode 100644 index 000000000..479e46556 --- /dev/null +++ b/apps/api/scripts/validate_demo_documents.py @@ -0,0 +1,185 @@ +from __future__ import annotations +# ruff: noqa: E402 + +import argparse +import json +import os +import sys +from dataclasses import dataclass +from pathlib import Path +from typing import Any + + +def _bootstrap_python_path() -> None: + current_dir = Path(__file__).resolve().parent + api_root = current_dir.parent + repo_root = api_root.parent.parent + shared_python_path = repo_root / "packages" / "shared-python" + + for path in (api_root, shared_python_path): + path_value = os.fspath(path) + if path_value not in sys.path: + sys.path.insert(0, path_value) + + +_bootstrap_python_path() + +from app.services.demo.source_catalog import DemoSourceCatalog, DemoSourceDefinition +from shared.services.storage.zip_result_schema import ZipResultSchemaBuilder + + +@dataclass(frozen=True) +class DemoValidationIssue: + source_id: str + message: str + + +def validate_demo_documents(*, write: bool = False) -> list[DemoValidationIssue]: + catalog = DemoSourceCatalog() + issues: list[DemoValidationIssue] = [] + for source in catalog.list_sources(): + issues.extend(_validate_source(catalog=catalog, source=source, write=write)) + return issues + + +def _validate_source( + *, + catalog: DemoSourceCatalog, + source: DemoSourceDefinition, + write: bool, +) -> list[DemoValidationIssue]: + issues: list[DemoValidationIssue] = [] + source_directory = catalog.source_directory(source) + chunks = _load_chunks(source_directory) + if len(chunks) != source.chunk_count: + issues.append( + DemoValidationIssue( + source_id=source.demo_source_id, + message=( + f"chunk_count mismatch: catalog={source.chunk_count}, " + f"chunks.json={len(chunks)}" + ), + ) + ) + + original_path = source_directory / "original.pdf" + if not original_path.is_file(): + issues.append( + DemoValidationIssue( + source_id=source.demo_source_id, + message="original.pdf is missing", + ) + ) + elif original_path.stat().st_size != source.size_bytes: + issues.append( + DemoValidationIssue( + source_id=source.demo_source_id, + message=( + f"size_bytes mismatch: catalog={source.size_bytes}, " + f"original.pdf={original_path.stat().st_size}" + ), + ) + ) + + _validate_catalog_projection(catalog=catalog, source=source, issues=issues) + _validate_doc_nav( + source=source, + source_directory=source_directory, + chunks=chunks, + write=write, + issues=issues, + ) + return issues + + +def _load_chunks(source_directory: Path) -> list[dict[str, Any]]: + chunks_path = source_directory / "chunks.json" + payload = json.loads(chunks_path.read_text(encoding="utf-8")) + raw_chunks = payload.get("chunks") if isinstance(payload, dict) else None + if not isinstance(raw_chunks, list): + return [] + return [chunk for chunk in raw_chunks if isinstance(chunk, dict)] + + +def _validate_catalog_projection( + *, + catalog: DemoSourceCatalog, + source: DemoSourceDefinition, + issues: list[DemoValidationIssue], +) -> None: + try: + catalog.get_catalog() + except Exception as exc: + issues.append( + DemoValidationIssue( + source_id=source.demo_source_id, + message=f"catalog projection failed: {exc}", + ) + ) + + +def _validate_doc_nav( + *, + source: DemoSourceDefinition, + source_directory: Path, + chunks: list[dict[str, Any]], + write: bool, + issues: list[DemoValidationIssue], +) -> None: + doc_nav_path = source_directory / "doc_nav.json" + expected_doc_nav = ZipResultSchemaBuilder().build_doc_nav(chunks, source.title) + expected_text = _serialize_json(expected_doc_nav) + if write: + doc_nav_path.write_text(expected_text, encoding="utf-8") + return + + if not doc_nav_path.is_file(): + issues.append( + DemoValidationIssue( + source_id=source.demo_source_id, + message="doc_nav.json is missing; run with --write", + ) + ) + return + + current_text = doc_nav_path.read_text(encoding="utf-8") + if current_text != expected_text: + issues.append( + DemoValidationIssue( + source_id=source.demo_source_id, + message="doc_nav.json is stale; run with --write", + ) + ) + + +def _serialize_json(payload: dict[str, Any]) -> str: + return json.dumps(payload, ensure_ascii=False, indent=2) + "\n" + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="Validate and optionally regenerate canonical demo documents.", + ) + parser.add_argument( + "--write", + action="store_true", + help="Regenerate doc_nav.json files from chunks.json before validating.", + ) + return parser + + +def main() -> int: + args = _build_parser().parse_args() + issues = validate_demo_documents(write=bool(args.write)) + if issues: + for issue in issues: + print(f"[FAIL] {issue.source_id}: {issue.message}") + return 1 + + action = "regenerated and validated" if args.write else "validated" + print(f"Demo documents {action}.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/apps/worker/app/services/document_ingestion/parse_execution.py b/apps/worker/app/services/document_ingestion/parse_execution.py new file mode 100644 index 000000000..594ff60fe --- /dev/null +++ b/apps/worker/app/services/document_ingestion/parse_execution.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +from app.services.document_ingestion.processing_context import ParseJobContext +from app.services.document_ingestion.source_preparation import PreparedSourceFile +from app.services.document_parser import parse_service +from app.services.document_parser.orchestration.parse_output import ParseOutput +from app.services.document_parser.support.stage_profiler import stage_timer +from loguru import logger + +from shared.models.schemas.job_metadata import JobMetadataHelper + + +def execute_document_parse( + *, + job_id: str, + job_context: ParseJobContext, + prepared_source: PreparedSourceFile, + output_dir: str, +) -> ParseOutput: + """Run the parser adapter for a prepared source file.""" + doc_type = JobMetadataHelper.get_parsing_param( + job_context.job_metadata, + "doc_type", + "auto", + ) + logger.info( + f"Start parse: job_id={job_id}, " + f"filename={prepared_source.source_file_name}, " + f"internal_filename={prepared_source.internal_parse_name}, type={doc_type}" + ) + + with stage_timer( + "worker.parse.document", + job_id=job_id, + filename=prepared_source.source_file_name, + doc_type=doc_type, + ): + parse_output = parse_service.checkerboard_parse_output( + file_full_path=prepared_source.local_file_path, + filename=prepared_source.source_file_name, + output_dir=output_dir, + job_id=job_id, + internal_output_filename=prepared_source.internal_parse_name, + 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}, output_dir={parse_output.output_dir}, " + f"chunks={parse_output.rows_count}" + ) + return parse_output diff --git a/apps/worker/app/services/document_ingestion/processing_run.py b/apps/worker/app/services/document_ingestion/processing_run.py index d5d405596..995990379 100644 --- a/apps/worker/app/services/document_ingestion/processing_run.py +++ b/apps/worker/app/services/document_ingestion/processing_run.py @@ -1,22 +1,14 @@ from __future__ import annotations -import os from datetime import datetime, timezone from typing import Any -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.parse_result_package import ( - ParseArtifact, - build_generated_result_package, build_parse_result_package, ) +from app.services.document_ingestion.parse_execution import execute_document_parse from app.services.document_ingestion.processing_billing import ( charge_parse_job_pages, record_processing_start, @@ -26,22 +18,21 @@ assert_source_file_within_size_limit, load_parse_job_context, ) +from app.services.document_ingestion.source_preparation import prepare_source_file +from app.services.document_ingestion.success_finalization import finalize_parse_success from app.services.document_ingestion.workspace import ( + TemporaryParseWorkspace, cleanup_task_workspace, - create_task_workspace, download_s3_file_to_temp, ) -from app.services.document_parser.support.stage_profiler import stage_timer from loguru import logger -from shared.models.schemas.job_metadata import JobMetadataHelper from shared.services.jobs.lifecycle.service import get_sync_job_lifecycle_service from shared.services.redis.distributed_lock import RedisJobLock from shared.services.redis.redis_sync_service import ( SyncRedisServiceFactory, ) from shared.services.storage.result_storage import get_result_storage -from shared.services.storage.zip_result_service import ZipResultService class DocumentProcessingRun: @@ -65,71 +56,37 @@ 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) + task_workspace = TemporaryParseWorkspace.create(job_id) try: result = _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, + task_workspace=task_workspace, ) finally: - cleanup_task_workspace(task_workspace_dir) + task_workspace.cleanup(cleanup_task_workspace) return result -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, + task_workspace: TemporaryParseWorkspace, ) -> dict[str, object]: lifecycle_service.update_progress(job_id, progress=10, message="Parsing document...") - 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}") - - from app.services.document_parser.support.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}" + prepared_source = prepare_source_file( + job_id=job_id, + job_context=job_context, + input_dir=task_workspace.input_dir, + download_source_file=download_s3_file_to_temp, ) - workload_estimate = PageEstimator.estimate_workload(local_temp_path) + workload_estimate = PageEstimator.estimate_workload(prepared_source.local_file_path) page_count = workload_estimate.page_count logger.info( "Workload estimation: " @@ -141,7 +98,7 @@ def _run_parse_job( processing_started_at = datetime.now(timezone.utc) billing_snapshot = charge_parse_job_pages( job_id=job_id, - filename=filename, + filename=prepared_source.source_file_name, job_user_id=job_context.job_user_id, workload_estimate=workload_estimate, ) @@ -153,61 +110,11 @@ def _run_parse_job( workload_estimate=workload_estimate, ) - 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", + parse_output = execute_document_parse( job_id=job_id, - filename=filename, - doc_type=doc_type, - ): - parse_output = parse_service.checkerboard_parse_output( - file_full_path=local_temp_path, - filename=filename, - output_dir=output_dir, - job_id=job_id, - internal_output_filename=internal_parse_name, - 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}, output_dir={parse_output.output_dir}, " - f"chunks={parse_output.rows_count}" + job_context=job_context, + prepared_source=prepared_source, + output_dir=task_workspace.output_dir, ) lifecycle_service.update_progress( @@ -217,7 +124,7 @@ def _run_parse_job( ) result_package = build_parse_result_package( job_id=job_id, - filename=filename, + filename=prepared_source.source_file_name, parse_output=parse_output, ) @@ -230,145 +137,12 @@ def _run_parse_job( f"Chunks prepared: job_id={job_id}, count={len(result_package.chunks)}" ) - return _finalize_parse_job_success( - parse_artifact=result_package.artifact, - chunks=result_package.chunks, + return finalize_parse_success( + result_package=result_package, job_context=job_context, job_id=job_id, lifecycle_service=lifecycle_service, processing_started_at=processing_started_at, - task_workspace_dir=task_workspace_dir, - ) - - -def _finalize_parse_job_success( - *, - parse_artifact: ParseArtifact, - chunks: list[dict[str, Any]], - job_context: ParseJobContext, - job_id: str, - lifecycle_service: Any, - processing_started_at: datetime, - task_workspace_dir: str, -) -> dict[str, object]: - source_file_name = JobMetadataHelper.get_source_file_name( - job_context.job_metadata, - ) 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) - - document_top_summary = "" - section_summaries: dict[str, str] = {} - add_dir = parse_artifact.add_dir - parsed_contents_df = parse_artifact.dataframe - 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: - document_root_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( - document_root_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_data_id(job_context.job_metadata) - zip_service = ZipResultService() - generated_package = build_generated_result_package( - *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, - ) - ) - - 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=generated_package.zip_file_path, - ) - result_s3_key = result_bundle.zip_key - stored_count = 0 - - finalization_response = lifecycle_service.finalize_job_success( - job_id=job_id, - chunks=chunks, - result_s3_key=result_s3_key, - checksum=generated_package.checksum_value, - zip_size=generated_package.zip_size, - stored_count=stored_count, - delivery_mode="url", - section_summaries=section_summaries, - ) - if finalization_response.get("status") != "success": - logger.error( - f"Worker processing finalization failed: job_id={job_id}, " - f"response={finalization_response}" - ) - return dict(finalization_response) - - lifecycle_service.update_progress(job_id, progress=100, message="Task complete!") - - logger.info( - f"Worker processing complete: job_id={job_id}, result_s3_key={result_s3_key}" + task_workspace_dir=task_workspace.root_dir, + result_storage_factory=get_result_storage, ) - - 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/source_preparation.py b/apps/worker/app/services/document_ingestion/source_preparation.py new file mode 100644 index 000000000..cc3c1bd79 --- /dev/null +++ b/apps/worker/app/services/document_ingestion/source_preparation.py @@ -0,0 +1,66 @@ +from __future__ import annotations + +import os +from collections.abc import Callable +from dataclasses import dataclass + +from app.services.document_ingestion.processing_context import ParseJobContext +from app.services.document_ingestion.workspace import download_s3_file_to_temp +from app.services.document_parser.support.internal_parse_name import ( + prepare_internal_parse_input, +) +from loguru import logger + +from shared.models.schemas.job_metadata import JobMetadataHelper + +DownloadSourceFile = Callable[[str, str, str], str] + + +@dataclass(frozen=True) +class PreparedSourceFile: + source_file_name: str + internal_parse_name: str + local_file_path: str + file_extension: str + + +def prepare_source_file( + *, + job_id: str, + job_context: ParseJobContext, + input_dir: str, + download_source_file: DownloadSourceFile = download_s3_file_to_temp, +) -> PreparedSourceFile: + """Download and normalize the source file before parser execution.""" + source_file_name = JobMetadataHelper.get_source_file_name( + job_context.job_metadata, + ) or os.path.basename(job_context.s3_key) + file_extension = ( + os.path.splitext(job_context.s3_key)[1].lower() if job_context.s3_key else "" + ) + + local_file_path = download_source_file( + job_context.s3_key, + file_extension, + input_dir, + ) + logger.info(f"File downloaded: job_id={job_id}, local_path={local_file_path}") + + prepared_parse_input = prepare_internal_parse_input( + local_file_path, + source_file_name, + fallback_ext=file_extension, + prefer_fallback_ext=True, + ) + logger.info( + f"File prepared for parsing: job_id={job_id}, " + f"internal_filename={prepared_parse_input.internal_filename}, " + f"local_path={prepared_parse_input.file_path}" + ) + + return PreparedSourceFile( + source_file_name=source_file_name, + internal_parse_name=prepared_parse_input.internal_filename, + local_file_path=prepared_parse_input.file_path, + file_extension=file_extension, + ) diff --git a/apps/worker/app/services/document_ingestion/success_finalization.py b/apps/worker/app/services/document_ingestion/success_finalization.py new file mode 100644 index 000000000..62397d56d --- /dev/null +++ b/apps/worker/app/services/document_ingestion/success_finalization.py @@ -0,0 +1,236 @@ +from __future__ import annotations + +import os +from collections.abc import Callable +from datetime import datetime, timezone +from typing import Any + +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.parse_result_package import ( + GeneratedResultPackage, + ParseArtifact, + ParseResultPackage, + build_generated_result_package, +) +from app.services.document_ingestion.processing_context import ParseJobContext +from loguru import logger + +from shared.models.schemas.job_metadata import JobMetadataHelper +from shared.services.storage.result_storage import ResultStorage, get_result_storage +from shared.services.storage.zip_result_service import ZipResultService + +ResultStorageFactory = Callable[[], ResultStorage] + + +def finalize_parse_success( + *, + result_package: ParseResultPackage, + job_context: ParseJobContext, + job_id: str, + lifecycle_service: Any, + processing_started_at: datetime, + task_workspace_dir: str, + result_storage_factory: ResultStorageFactory = get_result_storage, +) -> dict[str, object]: + """Package, upload, and publish a successful parser result.""" + source_file_name = _resolve_source_file_name(job_context) + document_top_summary, section_summaries = _enrich_document_navigation( + artifact=result_package.artifact, + chunks=result_package.chunks, + job_context=job_context, + source_file_name=source_file_name, + ) + _attach_document_top_summary(result_package.chunks, document_top_summary) + + lifecycle_service.update_progress( + job_id, + progress=80, + message="Generating ZIP package...", + ) + _record_processing_completion( + job_id=job_id, + job_context=job_context, + processing_started_at=processing_started_at, + ) + generated_package = _generate_result_package( + result_package=result_package, + job_context=job_context, + job_id=job_id, + source_file_name=source_file_name, + task_workspace_dir=task_workspace_dir, + ) + + lifecycle_service.update_progress( + job_id, + progress=90, + message="Uploading results to S3...", + ) + result_s3_key = _upload_result_package( + result_package=result_package, + generated_package=generated_package, + job_id=job_id, + result_storage_factory=result_storage_factory, + ) + stored_count = 0 + + finalization_response = lifecycle_service.finalize_job_success( + job_id=job_id, + chunks=result_package.chunks, + result_s3_key=result_s3_key, + checksum=generated_package.checksum_value, + zip_size=generated_package.zip_size, + stored_count=stored_count, + delivery_mode="url", + section_summaries=section_summaries, + ) + if finalization_response.get("status") != "success": + logger.error( + f"Worker processing finalization failed: job_id={job_id}, " + f"response={finalization_response}" + ) + return dict(finalization_response) + + lifecycle_service.update_progress(job_id, progress=100, message="Task complete!") + 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": result_package.artifact.contents_count, + "stored_count": stored_count, + "delivery_mode": "url", + "result_s3_key": result_s3_key, + } + + +def _resolve_source_file_name(job_context: ParseJobContext) -> str: + source_file_name = JobMetadataHelper.get_source_file_name( + job_context.job_metadata, + ) 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) + return str(source_file_name) + + +def _enrich_document_navigation( + *, + artifact: ParseArtifact, + chunks: list[dict[str, Any]], + job_context: ParseJobContext, + source_file_name: str, +) -> tuple[str, dict[str, str]]: + document_top_summary = "" + section_summaries: dict[str, str] = {} + add_dir = artifact.add_dir + parsed_contents_df = artifact.dataframe + 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=source_file_name, + ) + try: + document_root_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( + document_root_for_enrich, + source_file=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), source_file_name) + return document_top_summary, section_summaries + + +def _attach_document_top_summary( + chunks: list[dict[str, Any]], + document_top_summary: str, +) -> None: + if not document_top_summary: + return + + for chunk in chunks: + metadata = chunk.get("metadata") + if not isinstance(metadata, dict): + metadata = {} + chunk["metadata"] = metadata + metadata["document_top_summary"] = document_top_summary + + +def _record_processing_completion( + *, + job_id: str, + job_context: ParseJobContext, + processing_started_at: datetime, +) -> None: + 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) + + +def _generate_result_package( + *, + result_package: ParseResultPackage, + job_context: ParseJobContext, + job_id: str, + source_file_name: str, + task_workspace_dir: str, +) -> GeneratedResultPackage: + data_id = JobMetadataHelper.get_data_id(job_context.job_metadata) + zip_service = ZipResultService() + return build_generated_result_package( + *zip_service.generate_zip_package( + job_id=job_id, + chunks=result_package.chunks, + add_dir=str(result_package.artifact.add_dir) + if result_package.artifact.add_dir + else "", + source_file_name=source_file_name, + data_id=data_id, + job_metadata=job_context.job_metadata, + parsed_df=result_package.artifact.dataframe, + temp_dir=task_workspace_dir, + ) + ) + + +def _upload_result_package( + *, + result_package: ParseResultPackage, + generated_package: GeneratedResultPackage, + job_id: str, + result_storage_factory: ResultStorageFactory, +) -> str: + result_bundle = result_storage_factory().upload( + job_id=job_id, + result_dir=str(result_package.artifact.add_dir) + if result_package.artifact.add_dir + else "", + zip_file_path=generated_package.zip_file_path, + ) + return result_bundle.zip_key diff --git a/apps/worker/app/services/document_ingestion/workspace.py b/apps/worker/app/services/document_ingestion/workspace.py index 579b975d3..e8a0188dc 100644 --- a/apps/worker/app/services/document_ingestion/workspace.py +++ b/apps/worker/app/services/document_ingestion/workspace.py @@ -3,6 +3,8 @@ import os import shutil import tempfile +from collections.abc import Callable +from dataclasses import dataclass from loguru import logger @@ -14,6 +16,34 @@ ) from shared.services.storage.job_file_storage import JobFileStorage +CleanupTaskWorkspace = Callable[[str | None], bool] + + +@dataclass(frozen=True) +class TemporaryParseWorkspace: + """Task-local folders for parser input, parser output, and ZIP generation.""" + + root_dir: str + input_dir: str + output_dir: str + + @classmethod + def create(cls, job_id: str) -> "TemporaryParseWorkspace": + root_dir = create_task_workspace(job_id) + input_dir = os.path.join(root_dir, "input") + output_dir = os.path.join(root_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={root_dir}") + return cls(root_dir=root_dir, input_dir=input_dir, output_dir=output_dir) + + def cleanup( + self, + cleanup_workspace: CleanupTaskWorkspace | None = None, + ) -> bool: + resolved_cleanup = cleanup_workspace or cleanup_task_workspace + return resolved_cleanup(self.root_dir) + def cleanup_temp_file(file_path: str | None) -> None: """Best-effort cleanup for temp files created during parsing.""" diff --git a/packages/shared-python/shared/services/retrieval/publication_content.py b/packages/shared-python/shared/services/retrieval/publication_content.py new file mode 100644 index 000000000..1c5c3b4ef --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/publication_content.py @@ -0,0 +1,178 @@ +from __future__ import annotations + +from typing import Any +from uuid import uuid4 + +from sqlalchemy import delete +from sqlalchemy.orm import Session + +from shared.models.database.document import DocumentChunk, DocumentSection +from shared.services.retrieval.publication_models import DocumentPublicationScope +from shared.services.retrieval.search.lexical_text import ( + build_content_lexical_text, + build_content_search_text, + build_path_lexical_text, + build_path_search_text, + build_term_search_text, + section_path_from_chunk_path, +) + + +def replace_document_revision_content( + db: Session, + *, + scope: DocumentPublicationScope, + chunks: list[dict[str, Any]], +) -> None: + """Replace retrieval sections and chunks for one published document revision.""" + _delete_existing_revision_content(db, scope=scope) + section_publisher = DocumentSectionPublisher(db=db, scope=scope) + for index, chunk in enumerate(chunks): + chunk_metadata = _get_chunk_metadata(chunk) + source_path = _get_source_path(chunk=chunk, chunk_metadata=chunk_metadata) + section_path = section_path_from_chunk_path( + source_path, + source_file_name=scope.source_file_name, + ) + section = section_publisher.ensure_section(section_path) + db.add( + _build_document_chunk( + chunk=chunk, + chunk_metadata=chunk_metadata, + source_path=source_path, + section=section, + scope=scope, + fallback_sort_order=index, + ) + ) + + +class DocumentSectionPublisher: + def __init__(self, *, db: Session, scope: DocumentPublicationScope) -> None: + self._db = db + self._scope = scope + self._sections_by_path: dict[str, DocumentSection] = {} + + def ensure_section(self, section_path: str) -> DocumentSection: + existing_section = self._sections_by_path.get(section_path) + if existing_section is not None: + return existing_section + + path_parts = [part for part in section_path.split(" / ") if part] + for depth in range(1, len(path_parts) + 1): + ancestor_path = " / ".join(path_parts[:depth]) + if ancestor_path in self._sections_by_path: + continue + + ancestor_section = DocumentSection( + user_id=self._scope.user_id, + namespace=self._scope.namespace, + document_id=self._scope.document_id, + job_result_id=self._scope.job_result_id, + parent_section_id=self._get_parent_section_id(path_parts, depth), + section_path=ancestor_path, + section_title=path_parts[depth - 1], + section_level=depth, + section_metadata={}, + sort_order=len(self._sections_by_path), + ) + self._db.add(ancestor_section) + self._db.flush() + self._sections_by_path[ancestor_path] = ancestor_section + + return self._sections_by_path[section_path] + + def _get_parent_section_id( + self, + path_parts: list[str], + depth: int, + ) -> str | None: + if depth <= 1: + return None + parent_path = " / ".join(path_parts[: depth - 1]) + parent = self._sections_by_path.get(parent_path) + return parent.section_id if parent is not None else None + + +def _delete_existing_revision_content( + db: Session, + *, + scope: DocumentPublicationScope, +) -> None: + db.execute( + delete(DocumentChunk) + .where(DocumentChunk.document_id == scope.document_id) + .where(DocumentChunk.job_result_id == scope.job_result_id) + ) + db.execute( + delete(DocumentSection) + .where(DocumentSection.document_id == scope.document_id) + .where(DocumentSection.job_result_id == scope.job_result_id) + ) + + +def _build_document_chunk( + *, + chunk: dict[str, Any], + chunk_metadata: dict[str, Any], + source_path: str | None, + section: DocumentSection, + scope: DocumentPublicationScope, + fallback_sort_order: int, +) -> DocumentChunk: + section_summary = section.summary + section_path = section.section_path + section_title = section.section_title + path_text = f"{scope.source_file_name or ''} {section_path}".strip() + return DocumentChunk( + id=f"dchk_{uuid4().hex[:12]}", + chunk_id=str(chunk.get("chunk_id") or f"chunk_{uuid4().hex[:12]}"), + user_id=scope.user_id, + namespace=scope.namespace, + document_id=scope.document_id, + job_result_id=scope.job_result_id, + section_id=section.section_id, + chunk_type=chunk.get("type") or chunk.get("chunk_type") or "text", + content=chunk.get("content") or chunk.get("text"), + content_lexical_text=build_content_lexical_text(chunk), + path_lexical_text=build_path_lexical_text( + source_path, + source_file_name=scope.source_file_name, + ), + content_search_text=build_content_search_text( + chunk, + section_summary=section_summary, + ), + path_search_text=build_path_search_text( + source_file_name=scope.source_file_name, + section_path=section_path, + section_title=section_title, + section_summary=section_summary, + ), + term_search_text=build_term_search_text(chunk, path_text=path_text), + source_chunk_path=source_path, + file_path=chunk_metadata.get("file_path") or chunk.get("file_path"), + chunk_metadata=chunk_metadata, + sort_order=_get_sort_order(chunk, fallback_sort_order), + ) + + +def _get_chunk_metadata(chunk: dict[str, Any]) -> dict[str, Any]: + metadata = chunk.get("metadata") + return metadata if isinstance(metadata, dict) else {} + + +def _get_source_path( + *, + chunk: dict[str, Any], + chunk_metadata: dict[str, Any], +) -> str | None: + source_path = chunk_metadata.get("path") or chunk.get("path") + return str(source_path) if source_path is not None else None + + +def _get_sort_order(chunk: dict[str, Any], fallback_sort_order: int) -> int: + try: + return int(chunk.get("order", fallback_sort_order)) + except (TypeError, ValueError): + return fallback_sort_order diff --git a/packages/shared-python/shared/services/retrieval/publication_models.py b/packages/shared-python/shared/services/retrieval/publication_models.py new file mode 100644 index 000000000..87b628aa5 --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/publication_models.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + + +@dataclass(frozen=True) +class PublishedDocumentState: + user_id: str + namespace: str + document_id: str | None + skipped_all_duplicate: bool = False + + def to_dict(self) -> dict[str, Any]: + payload: dict[str, Any] = { + "user_id": self.user_id, + "namespace": self.namespace, + "document_id": self.document_id, + } + if self.skipped_all_duplicate: + payload["skipped_all_duplicate"] = True + return payload + + +@dataclass(frozen=True) +class DocumentPublicationScope: + user_id: str + namespace: str + document_id: str + job_result_id: str + source_file_name: str | None diff --git a/packages/shared-python/shared/services/retrieval/publication_service.py b/packages/shared-python/shared/services/retrieval/publication_service.py index 1e7ffbe06..c1f1a7ed5 100644 --- a/packages/shared-python/shared/services/retrieval/publication_service.py +++ b/packages/shared-python/shared/services/retrieval/publication_service.py @@ -13,21 +13,20 @@ from uuid import uuid4 from loguru import logger -from sqlalchemy import delete, select +from sqlalchemy import select from sqlalchemy.orm import Session -from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.models.database.document import Document from shared.models.database.job import Job from shared.models.database.job_result import JobResult from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace from shared.services.retrieval.graph.service import DocumentGraphService, GraphScope -from shared.services.retrieval.search.lexical_text import ( - build_content_lexical_text, - build_content_search_text, - build_path_lexical_text, - build_path_search_text, - build_term_search_text, - section_path_from_chunk_path, +from shared.services.retrieval.publication_content import ( + replace_document_revision_content, +) +from shared.services.retrieval.publication_models import ( + DocumentPublicationScope, + PublishedDocumentState, ) @@ -93,7 +92,9 @@ def _publish_document_state_for_job( job_metadata = job.job_metadata or {} namespace = normalize_retrieval_namespace(job_metadata.get("namespace")) document_id = job_metadata.get("document_id") - source_file_name = job_metadata.get("source_file_name") or job_metadata.get("file_name") + source_file_name = job_metadata.get("source_file_name") or job_metadata.get( + "file_name" + ) deduped_chunks = chunks @@ -103,14 +104,60 @@ def _publish_document_state_for_job( f"⏭️ All chunks are duplicates of existing documents. " f"Skipping document creation for job_id={job.job_id}." ) - return { - "user_id": str(job.user_id), - "namespace": namespace, - "document_id": None, - "skipped_all_duplicate": True, - } - - # ── Document upsert (original logic, but only for deduped chunks) ── + return PublishedDocumentState( + user_id=str(job.user_id), + namespace=namespace, + document_id=None, + skipped_all_duplicate=True, + ).to_dict() + + document = self._upsert_document_revision( + db, + job=job, + job_result_id=job_result_id, + document_id=str(document_id) if document_id else None, + namespace=namespace, + source_file_name=str(source_file_name) if source_file_name else None, + ) + if document is None: + return None + + self._bind_job_result_document( + db, + job_result_id=job_result_id, + document_id=document.document_id, + ) + namespace = normalize_retrieval_namespace(namespace or document.namespace) + scope = DocumentPublicationScope( + user_id=str(job.user_id), + namespace=namespace, + document_id=document.document_id, + job_result_id=job_result_id, + source_file_name=str(source_file_name) if source_file_name else None, + ) + replace_document_revision_content( + db, + scope=scope, + chunks=deduped_chunks, + ) + + db.flush() + return PublishedDocumentState( + user_id=str(job.user_id), + namespace=namespace, + document_id=document.document_id, + ).to_dict() + + def _upsert_document_revision( + self, + db: Session, + *, + job: Job, + job_result_id: str, + document_id: str | None, + namespace: str, + source_file_name: str | None, + ) -> Document | None: document = None if document_id: document = db.execute( @@ -133,14 +180,14 @@ def _publish_document_state_for_job( ) db.add(document) else: - namespace = normalize_retrieval_namespace(namespace or document.namespace) if self._is_stale_document_completion( db, document=document, job=job, ): logger.warning( - f"Skipping stale document publication: job_id={job.job_id}, document_id={document.document_id}" + "Skipping stale document publication: " + f"job_id={job.job_id}, document_id={document.document_id}" ) return None document.status = "active" @@ -150,109 +197,19 @@ def _publish_document_state_for_job( document.updated_at = utc_now_naive() db.flush() - document_id = document.document_id + return document + + def _bind_job_result_document( + self, + db: Session, + *, + job_result_id: str, + document_id: str, + ) -> None: result = db.execute(select(JobResult).where(JobResult.id == job_result_id)) job_result = result.scalar_one_or_none() if job_result: job_result.document_id = document_id - namespace = normalize_retrieval_namespace(namespace or document.namespace) - - db.execute( - delete(DocumentChunk) - .where(DocumentChunk.document_id == document_id) - .where(DocumentChunk.job_result_id == job_result_id) - ) - db.execute( - delete(DocumentSection) - .where(DocumentSection.document_id == document_id) - .where(DocumentSection.job_result_id == job_result_id) - ) - - # ── Insert only deduped (non-duplicate) chunks ────────────────── - sections_by_path: Dict[str, DocumentSection] = {} - for index, chunk in enumerate(deduped_chunks): - chunk_metadata = chunk.get("metadata") or {} - source_path = chunk_metadata.get("path") or chunk.get("path") - section_path = section_path_from_chunk_path( - source_path, - source_file_name=source_file_name, - ) - section = sections_by_path.get(section_path) - if section is None: - path_parts = [p for p in section_path.split(" / ") if p] - # Ensure all ancestor sections exist (top-down) - for depth in range(1, len(path_parts) + 1): - ancestor_path = " / ".join(path_parts[:depth]) - if ancestor_path in sections_by_path: - continue - ancestor_parent_id = None - if depth > 1: - parent_path = " / ".join(path_parts[:depth - 1]) - parent = sections_by_path.get(parent_path) - if parent is not None: - ancestor_parent_id = parent.section_id - ancestor_section = DocumentSection( - user_id=str(job.user_id), - namespace=namespace, - document_id=document_id, - job_result_id=job_result_id, - parent_section_id=ancestor_parent_id, - section_path=ancestor_path, - section_title=path_parts[depth - 1], - section_level=depth, - section_metadata={}, - sort_order=len(sections_by_path), - ) - db.add(ancestor_section) - db.flush() - sections_by_path[ancestor_path] = ancestor_section - section = sections_by_path[section_path] - - chunk_id = chunk.get("chunk_id") or f"chunk_{uuid4().hex[:12]}" - section_summary = section.summary if section else None - section_path_str = section.section_path if section else "Root" - section_title_str = section.section_title if section else None - path_text = f"{source_file_name or ''} {section_path_str}".strip() - - db.add( - DocumentChunk( - id=f"dchk_{uuid4().hex[:12]}", - chunk_id=chunk_id, - user_id=str(job.user_id), - namespace=namespace, - document_id=document_id, - job_result_id=job_result_id, - section_id=section.section_id, - chunk_type=chunk.get("type") or chunk.get("chunk_type") or "text", - content=chunk.get("content") or chunk.get("text"), - content_lexical_text=build_content_lexical_text(chunk), - path_lexical_text=build_path_lexical_text( - source_path, - source_file_name=source_file_name, - ), - content_search_text=build_content_search_text( - chunk, section_summary=section_summary - ), - path_search_text=build_path_search_text( - source_file_name=source_file_name, - section_path=section_path_str, - section_title=section_title_str, - section_summary=section_summary, - ), - term_search_text=build_term_search_text(chunk, path_text=path_text), - source_chunk_path=source_path, - file_path=chunk_metadata.get("file_path") or chunk.get("file_path"), - chunk_metadata=chunk_metadata, - sort_order=chunk.get("order", index), - ) - ) - - db.flush() - return { - "user_id": str(job.user_id), - "namespace": namespace, - "document_id": document_id, - } def publish_document_graph( self, From addb331015cdda4909acb0d0c087ba250d1302f4 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 17:47:29 +0000 Subject: [PATCH 76/78] refactor: deepen lifecycle and packaging modules --- CONTEXT.md | 7 + .../api/app/api/v1/routes/qstash_callbacks.py | 7 +- apps/api/app/api/v1/routes/retrieval.py | 2 +- apps/api/app/mcp/retrieval_server.py | 2 +- .../webhook/qstash_callback_service.py | 42 +- .../jobs/lifecycle/failure_finalizer.py | 2 +- .../jobs/lifecycle/post_commit_effects.py | 10 +- .../services/jobs/lifecycle/publication.py | 56 +- .../shared/services/jobs/lifecycle/service.py | 5 +- .../jobs/lifecycle/success_finalizer.py | 51 +- .../services/jobs/lifecycle/webhook_outbox.py | 14 +- .../shared/services/retrieval/__init__.py | 33 +- .../services/retrieval/publication_models.py | 17 +- .../services/retrieval/publication_service.py | 22 +- .../services/storage/zip_chunk_schema.py | 182 +++++++ .../services/storage/zip_doc_navigation.py | 183 +++++++ .../services/storage/zip_manifest_schema.py | 42 ++ .../services/storage/zip_package_writer.py | 112 ++++ .../services/storage/zip_result_resources.py | 321 +++++++++++ .../services/storage/zip_result_schema.py | 373 ++----------- .../services/storage/zip_result_service.py | 497 +++--------------- 21 files changed, 1128 insertions(+), 852 deletions(-) create mode 100644 packages/shared-python/shared/services/storage/zip_chunk_schema.py create mode 100644 packages/shared-python/shared/services/storage/zip_doc_navigation.py create mode 100644 packages/shared-python/shared/services/storage/zip_manifest_schema.py create mode 100644 packages/shared-python/shared/services/storage/zip_package_writer.py create mode 100644 packages/shared-python/shared/services/storage/zip_result_resources.py diff --git a/CONTEXT.md b/CONTEXT.md index 9576fe215..c8903a584 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -102,6 +102,13 @@ chunk conversion and result packaging. The generated ZIP bundle metadata used by terminal Job finalization, including ZIP path, checksum, statistics, and byte size. +### ZIP Result Packaging + +The storage-side module set that turns parsed chunks and parser artifacts into a +Generated Result Package. `ZipResultService` is the orchestration interface; +resource discovery, schema projection, and physical ZIP writing live in +separate modules. + ### Workload Estimate The worker-side estimate used for billing and processing metadata. It records diff --git a/apps/api/app/api/v1/routes/qstash_callbacks.py b/apps/api/app/api/v1/routes/qstash_callbacks.py index 041146adf..e2fba04e7 100644 --- a/apps/api/app/api/v1/routes/qstash_callbacks.py +++ b/apps/api/app/api/v1/routes/qstash_callbacks.py @@ -54,4 +54,9 @@ async def handle_qstash_failure(request: Request) -> Response: def _to_response(outcome: QStashCallbackOutcome) -> Response: - return Response(status_code=outcome.status_code, content=outcome.content) + content_by_kind = { + "processed": "OK", + "missing_event_id": "OK (no event_id)", + "event_not_found": "OK (event not found)", + } + return Response(status_code=200, content=content_by_kind[outcome.kind]) diff --git a/apps/api/app/api/v1/routes/retrieval.py b/apps/api/app/api/v1/routes/retrieval.py index e7272f6c4..5db937130 100644 --- a/apps/api/app/api/v1/routes/retrieval.py +++ b/apps/api/app/api/v1/routes/retrieval.py @@ -12,7 +12,7 @@ from shared.core.database import get_db from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace -from shared.services.retrieval import run_retrieval_query +from shared.services.retrieval.app_service import run_retrieval_query router = APIRouter(tags=["Retrieval"]) diff --git a/apps/api/app/mcp/retrieval_server.py b/apps/api/app/mcp/retrieval_server.py index e9d9f0915..e1bff80e8 100644 --- a/apps/api/app/mcp/retrieval_server.py +++ b/apps/api/app/mcp/retrieval_server.py @@ -12,7 +12,7 @@ from shared.core.database import get_db_context from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace -from shared.services.retrieval import run_retrieval_query +from shared.services.retrieval.app_service import run_retrieval_query DbFactory = Callable[[], AsyncContextManager[AsyncSession]] KNOWHERE_NAMESPACE_HEADER = "x-knowhere-namespace" diff --git a/apps/api/app/services/webhook/qstash_callback_service.py b/apps/api/app/services/webhook/qstash_callback_service.py index dfd16f49b..352e8fcd8 100644 --- a/apps/api/app/services/webhook/qstash_callback_service.py +++ b/apps/api/app/services/webhook/qstash_callback_service.py @@ -3,7 +3,7 @@ import json from dataclasses import dataclass from datetime import datetime, timezone -from typing import Any, Optional +from typing import Any, Literal from uuid import NAMESPACE_URL, uuid5 from loguru import logger @@ -15,10 +15,25 @@ from shared.models.database.webhook_log import WebhookLog +QStashCallbackOutcomeKind = Literal["processed", "missing_event_id", "event_not_found"] + + @dataclass(frozen=True) class QStashCallbackOutcome: - status_code: int - content: str + kind: QStashCallbackOutcomeKind + event_id: str | None = None + + @classmethod + def processed(cls, *, event_id: str) -> QStashCallbackOutcome: + return cls(kind="processed", event_id=event_id) + + @classmethod + def missing_event_id(cls) -> QStashCallbackOutcome: + return cls(kind="missing_event_id") + + @classmethod + def event_not_found(cls, *, event_id: str) -> QStashCallbackOutcome: + return cls(kind="event_not_found", event_id=event_id) def get_qstash_verification_url(callback_path: str, request_url: str) -> str: @@ -64,7 +79,7 @@ def handle_qstash_success_callback(raw_body: bytes) -> QStashCallbackOutcome: if not event_id: logger.warning("QStash callback: missing event_id, cannot correlate") - return QStashCallbackOutcome(status_code=200, content="OK (no event_id)") + return QStashCallbackOutcome.missing_event_id() retried = data.get("retried", 0) logger.info( @@ -86,7 +101,7 @@ def handle_qstash_failure_callback(raw_body: bytes) -> QStashCallbackOutcome: if not event_id: logger.warning("QStash failure callback: missing event_id, cannot correlate") - return QStashCallbackOutcome(status_code=200, content="OK (no event_id)") + return QStashCallbackOutcome.missing_event_id() retried = data.get("retried", 0) max_retries = data.get("maxRetries", 0) @@ -110,7 +125,7 @@ def extract_callback_data(body: bytes) -> dict[str, Any]: return {"raw": body.decode("utf-8", errors="replace")} -def find_event_id(data: dict[str, Any]) -> Optional[str]: +def find_event_id(data: dict[str, Any]) -> str | None: source_header = data.get("sourceHeader", {}) or {} event_id = normalize_header_value( source_header.get("X-Knowhere-Event-Id") @@ -124,7 +139,7 @@ def find_event_id(data: dict[str, Any]) -> Optional[str]: return event_id -def normalize_header_value(value: Any) -> Optional[str]: +def normalize_header_value(value: Any) -> str | None: if isinstance(value, list): if not value: return None @@ -141,7 +156,7 @@ def normalize_header_value(value: Any) -> Optional[str]: def build_callback_log_idempotency_key( - qstash_message_id: Optional[str], + qstash_message_id: str | None, event_id: str, ) -> str: if qstash_message_id: @@ -149,7 +164,7 @@ def build_callback_log_idempotency_key( return event_id -def get_response_status_code(value: Any) -> Optional[int]: +def get_response_status_code(value: Any) -> int | None: if value is None: return None @@ -159,7 +174,7 @@ def get_response_status_code(value: Any) -> Optional[int]: return None -def is_success_response_status(status_code: Optional[int]) -> bool: +def is_success_response_status(status_code: int | None) -> bool: return status_code is not None and 200 <= status_code < 300 @@ -208,10 +223,7 @@ def process_qstash_callback( if not event: logger.warning(f"QStash {log_label}: event {event_id} not found in DB") - return QStashCallbackOutcome( - status_code=200, - content="OK (event not found)", - ) + return QStashCallbackOutcome.event_not_found(event_id=event_id) now = datetime.now(timezone.utc).replace(tzinfo=None) event_status = resolve_event_status(event.status, callback_status) @@ -240,4 +252,4 @@ def process_qstash_callback( db.add(log) db.commit() - return QStashCallbackOutcome(status_code=200, content="OK") + return QStashCallbackOutcome.processed(event_id=event_id) diff --git a/packages/shared-python/shared/services/jobs/lifecycle/failure_finalizer.py b/packages/shared-python/shared/services/jobs/lifecycle/failure_finalizer.py index 41c2f4e64..b87b270fb 100644 --- a/packages/shared-python/shared/services/jobs/lifecycle/failure_finalizer.py +++ b/packages/shared-python/shared/services/jobs/lifecycle/failure_finalizer.py @@ -83,7 +83,7 @@ def finalize( return JobFailureFinalization( succeeded=True, post_commit_effects=PostCommitEffectPlan.from_failure( - webhook_event_id=webhook_event.id if webhook_event else None, + webhook_event_id=webhook_event.event_id if webhook_event else None, ), ) diff --git a/packages/shared-python/shared/services/jobs/lifecycle/post_commit_effects.py b/packages/shared-python/shared/services/jobs/lifecycle/post_commit_effects.py index 0575d5d42..e990ec5a4 100644 --- a/packages/shared-python/shared/services/jobs/lifecycle/post_commit_effects.py +++ b/packages/shared-python/shared/services/jobs/lifecycle/post_commit_effects.py @@ -1,15 +1,17 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Any -from shared.services.jobs.lifecycle.publication import SyncJobPublicationFinalizer +from shared.services.jobs.lifecycle.publication import ( + RetrievalCacheInvalidation, + SyncJobPublicationFinalizer, +) from shared.services.jobs.lifecycle.webhook_outbox import SyncJobWebhookOutbox @dataclass(frozen=True) class PostCommitEffectPlan: - retrieval_cache_invalidations: tuple[dict[str, Any], ...] = field( + retrieval_cache_invalidations: tuple[RetrievalCacheInvalidation, ...] = field( default_factory=tuple ) webhook_event_ids: tuple[str, ...] = field(default_factory=tuple) @@ -22,7 +24,7 @@ def none(cls) -> PostCommitEffectPlan: def from_success( cls, *, - cache_invalidation: dict[str, Any] | None, + cache_invalidation: RetrievalCacheInvalidation | None, webhook_event_id: str | None, ) -> PostCommitEffectPlan: return cls( diff --git a/packages/shared-python/shared/services/jobs/lifecycle/publication.py b/packages/shared-python/shared/services/jobs/lifecycle/publication.py index a9dfe2f0a..ee6c78e38 100644 --- a/packages/shared-python/shared/services/jobs/lifecycle/publication.py +++ b/packages/shared-python/shared/services/jobs/lifecycle/publication.py @@ -13,12 +13,23 @@ from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace from shared.services.redis.redis_sync_service import SyncRedisServiceFactory from shared.services.retrieval.publication_service import RetrievalPublicationService +from shared.services.retrieval.publication_models import ( + ExistingDocumentScope, + PublishedDocumentState, +) + + +@dataclass(frozen=True) +class RetrievalCacheInvalidation: + user_id: str + namespaces: tuple[str, ...] + job_id: str @dataclass(frozen=True) class JobPublicationOutcome: - published_document_state: dict[str, str] | None - cache_invalidation: dict[str, Any] | None + published_document_state: PublishedDocumentState | None + cache_invalidation: RetrievalCacheInvalidation | None class SyncJobPublicationFinalizer: @@ -57,7 +68,7 @@ def publish_result( if section_summaries: self._backfill_section_summaries( db, - document_id=published_document_state.get("document_id", ""), + document_id=published_document_state.document_id or "", job_result_id=job_result_id, section_summaries=section_summaries, ) @@ -80,16 +91,16 @@ def publish_result( def invalidate_cache_after_commit( self, - cache_invalidation: dict[str, Any] | None, + cache_invalidation: RetrievalCacheInvalidation | None, ) -> None: if not cache_invalidation: return try: redis_service = SyncRedisServiceFactory.get_service() - user_id = cache_invalidation["user_id"] + user_id = cache_invalidation.user_id seen: set[str] = set() - for raw_namespace in cache_invalidation["namespaces"]: + for raw_namespace in cache_invalidation.namespaces: namespace = normalize_retrieval_namespace(str(raw_namespace)) if not namespace or namespace in seen: continue @@ -98,7 +109,7 @@ def invalidate_cache_after_commit( except Exception as exc: logger.warning( "Failed to invalidate retrieval cache after publication " - f"(ignored): job_id={cache_invalidation.get('job_id')}, error={exc}" + f"(ignored): job_id={cache_invalidation.job_id}, error={exc}" ) def _backfill_section_summaries( @@ -136,28 +147,33 @@ def _build_cache_invalidation( db: Session, *, job_id: str, - published_document_state: dict[str, str] | None, - previous_document_scope: dict[str, str] | None, - ) -> dict[str, Any] | None: + published_document_state: PublishedDocumentState | None, + previous_document_scope: ExistingDocumentScope | None, + ) -> RetrievalCacheInvalidation | 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 = [ + namespaces: list[str] = [ 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} + if previous_document_scope: + namespaces.append(previous_document_scope.namespace) + if published_document_state: + namespaces.append(published_document_state.namespace) + + return RetrievalCacheInvalidation( + user_id=str(job.user_id), + namespaces=tuple(namespaces), + job_id=job_id, + ) def _should_publish_document_graph( - published_document_state: dict[str, str] | None, + published_document_state: PublishedDocumentState | None, ) -> bool: - return published_document_state is not None and not published_document_state.get( - "skipped_all_duplicate" + return ( + published_document_state is not None + and not published_document_state.skipped_all_duplicate ) diff --git a/packages/shared-python/shared/services/jobs/lifecycle/service.py b/packages/shared-python/shared/services/jobs/lifecycle/service.py index e0d82d832..3a611d1f8 100644 --- a/packages/shared-python/shared/services/jobs/lifecycle/service.py +++ b/packages/shared-python/shared/services/jobs/lifecycle/service.py @@ -79,9 +79,8 @@ def finalize_job_success( delivery_mode=delivery_mode, section_summaries=section_summaries, ), - should_commit=lambda finalization: finalization.response.get("status") - == "success", - build_response=lambda finalization: finalization.response, + should_commit=lambda finalization: finalization.response.should_commit(), + build_response=lambda finalization: finalization.response.to_dict(), build_effect_plan=lambda finalization: finalization.post_commit_effects, run_after_commit_effects=self._post_commit_effect_runner.run, ) diff --git a/packages/shared-python/shared/services/jobs/lifecycle/success_finalizer.py b/packages/shared-python/shared/services/jobs/lifecycle/success_finalizer.py index c325edefc..d05b41436 100644 --- a/packages/shared-python/shared/services/jobs/lifecycle/success_finalizer.py +++ b/packages/shared-python/shared/services/jobs/lifecycle/success_finalizer.py @@ -1,7 +1,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any +from typing import Any, Literal from loguru import logger from sqlalchemy.orm import Session @@ -13,9 +13,39 @@ from shared.services.jobs.lifecycle.webhook_outbox import SyncJobWebhookOutbox +@dataclass(frozen=True) +class JobSuccessResponse: + status: Literal["success", "failed"] + job_id: str + stored_count: int | None = None + reason: str | None = None + + @classmethod + def completed(cls, *, job_id: str, stored_count: int) -> JobSuccessResponse: + return cls(status="success", job_id=job_id, stored_count=stored_count) + + @classmethod + def state_transition_failed(cls, *, job_id: str) -> JobSuccessResponse: + return cls(status="failed", job_id=job_id, reason="state_transition_failed") + + def should_commit(self) -> bool: + return self.status == "success" + + def to_dict(self) -> dict[str, Any]: + response: dict[str, Any] = { + "status": self.status, + "job_id": self.job_id, + } + if self.stored_count is not None: + response["stored_count"] = self.stored_count + if self.reason: + response["reason"] = self.reason + return response + + @dataclass(frozen=True) class JobSuccessFinalization: - response: dict[str, Any] + response: JobSuccessResponse post_commit_effects: PostCommitEffectPlan @@ -82,11 +112,7 @@ def finalize( f"reason={transition_outcome.reason}" ) return JobSuccessFinalization( - response={ - "status": "failed", - "job_id": job_id, - "reason": "state_transition_failed", - }, + response=JobSuccessResponse.state_transition_failed(job_id=job_id), post_commit_effects=PostCommitEffectPlan.none(), ) @@ -96,13 +122,12 @@ def finalize( event_type="job.completed", ) return JobSuccessFinalization( - response={ - "status": "success", - "job_id": job_id, - "stored_count": stored_count, - }, + response=JobSuccessResponse.completed( + job_id=job_id, + stored_count=stored_count, + ), post_commit_effects=PostCommitEffectPlan.from_success( cache_invalidation=publication_outcome.cache_invalidation, - webhook_event_id=webhook_event.id if webhook_event else None, + webhook_event_id=webhook_event.event_id if webhook_event else None, ), ) diff --git a/packages/shared-python/shared/services/jobs/lifecycle/webhook_outbox.py b/packages/shared-python/shared/services/jobs/lifecycle/webhook_outbox.py index 05ea84333..5717d4fa4 100644 --- a/packages/shared-python/shared/services/jobs/lifecycle/webhook_outbox.py +++ b/packages/shared-python/shared/services/jobs/lifecycle/webhook_outbox.py @@ -1,5 +1,6 @@ from __future__ import annotations +from dataclasses import dataclass from datetime import datetime, timezone from typing import Any @@ -11,6 +12,11 @@ from shared.models.database.webhook import WebhookEvent, WebhookEventStatus +@dataclass(frozen=True) +class WebhookOutboxEvent: + event_id: str + + class SyncJobWebhookOutbox: """Create webhook events in-transaction and publish them after commit.""" @@ -21,7 +27,7 @@ def create_event( job_id: str, event_type: str, extra_payload: dict[str, Any] | None = None, - ) -> WebhookEvent | None: + ) -> WebhookOutboxEvent | None: result = db.execute(select(Job).where(Job.job_id == job_id)) job = result.scalar_one_or_none() @@ -54,12 +60,12 @@ def create_event( db.add(event) db.flush() logger.info(f"WebhookEvent created: event_id={event.id}, job_id={job_id}") - return event + return WebhookOutboxEvent(event_id=event.id) - def enqueue_after_commit(self, webhook_event: WebhookEvent | None) -> None: + def enqueue_after_commit(self, webhook_event: WebhookOutboxEvent | None) -> None: if not webhook_event: return - self.enqueue_event_id_after_commit(webhook_event.id) + self.enqueue_event_id_after_commit(webhook_event.event_id) def enqueue_event_id_after_commit(self, webhook_event_id: str | None) -> None: if not webhook_event_id: diff --git a/packages/shared-python/shared/services/retrieval/__init__.py b/packages/shared-python/shared/services/retrieval/__init__.py index 259fa6489..845269856 100644 --- a/packages/shared-python/shared/services/retrieval/__init__.py +++ b/packages/shared-python/shared/services/retrieval/__init__.py @@ -1,28 +1,7 @@ -from .app_service import merge_channels_rrf, run_retrieval_query -from .cache_service import ( - bump_retrieval_namespace_cache_version, - get_cached_retrieval_query_result, - get_retrieval_namespace_cache_version, - invalidate_retrieval_cache_namespaces, - set_cached_retrieval_query_result, -) -from .graph.query_service import GraphQueryService -from .graph.service import DocumentGraphService, GraphScope -from .stats.service import record_retrieval_hits -from .llm_adapter import create_retrieval_llm_fn, create_retrieval_planner_fn +"""Retrieval service package. -__all__ = [ - "create_retrieval_llm_fn", - "create_retrieval_planner_fn", - "run_retrieval_query", - "merge_channels_rrf", - "DocumentGraphService", - "GraphQueryService", - "GraphScope", - "record_retrieval_hits", - "bump_retrieval_namespace_cache_version", - "get_cached_retrieval_query_result", - "get_retrieval_namespace_cache_version", - "invalidate_retrieval_cache_namespaces", - "set_cached_retrieval_query_result", -] +Import concrete modules directly so importing one retrieval submodule does not +initialize unrelated retrieval dependencies. +""" + +from __future__ import annotations diff --git a/packages/shared-python/shared/services/retrieval/publication_models.py b/packages/shared-python/shared/services/retrieval/publication_models.py index 87b628aa5..8c653212c 100644 --- a/packages/shared-python/shared/services/retrieval/publication_models.py +++ b/packages/shared-python/shared/services/retrieval/publication_models.py @@ -1,7 +1,12 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any + + +@dataclass(frozen=True) +class ExistingDocumentScope: + document_id: str + namespace: str @dataclass(frozen=True) @@ -11,16 +16,6 @@ class PublishedDocumentState: document_id: str | None skipped_all_duplicate: bool = False - def to_dict(self) -> dict[str, Any]: - payload: dict[str, Any] = { - "user_id": self.user_id, - "namespace": self.namespace, - "document_id": self.document_id, - } - if self.skipped_all_duplicate: - payload["skipped_all_duplicate"] = True - return payload - @dataclass(frozen=True) class DocumentPublicationScope: diff --git a/packages/shared-python/shared/services/retrieval/publication_service.py b/packages/shared-python/shared/services/retrieval/publication_service.py index c1f1a7ed5..aa252ef8e 100644 --- a/packages/shared-python/shared/services/retrieval/publication_service.py +++ b/packages/shared-python/shared/services/retrieval/publication_service.py @@ -9,7 +9,7 @@ from __future__ import annotations from datetime import datetime, timezone -from typing import Any, Dict, List, Optional +from typing import Any from uuid import uuid4 from loguru import logger @@ -26,6 +26,7 @@ ) from shared.services.retrieval.publication_models import ( DocumentPublicationScope, + ExistingDocumentScope, PublishedDocumentState, ) @@ -42,7 +43,7 @@ def get_existing_document_scope( db: Session, *, job_id: str, - ) -> Optional[Dict[str, str]]: + ) -> ExistingDocumentScope | None: job = db.execute(select(Job).where(Job.job_id == job_id)).scalar_one_or_none() if not job: return None @@ -58,7 +59,10 @@ def get_existing_document_scope( if not document: return None - return {"document_id": document.document_id, "namespace": document.namespace} + return ExistingDocumentScope( + document_id=document.document_id, + namespace=document.namespace, + ) def publish_document_state( self, @@ -66,8 +70,8 @@ def publish_document_state( *, job_id: str, job_result_id: str, - chunks: List[Dict[str, Any]], - ) -> Optional[Dict[str, Any]]: + chunks: list[dict[str, Any]], + ) -> PublishedDocumentState | None: job = db.execute(select(Job).where(Job.job_id == job_id)).scalar_one_or_none() if not job: logger.warning(f"Job not found for document publication: {job_id}") @@ -86,8 +90,8 @@ def _publish_document_state_for_job( *, job: Job, job_result_id: str, - chunks: List[Dict[str, Any]], - ) -> Optional[Dict[str, Any]]: + chunks: list[dict[str, Any]], + ) -> PublishedDocumentState | None: job_metadata = job.job_metadata or {} namespace = normalize_retrieval_namespace(job_metadata.get("namespace")) @@ -109,7 +113,7 @@ def _publish_document_state_for_job( namespace=namespace, document_id=None, skipped_all_duplicate=True, - ).to_dict() + ) document = self._upsert_document_revision( db, @@ -146,7 +150,7 @@ def _publish_document_state_for_job( user_id=str(job.user_id), namespace=namespace, document_id=document.document_id, - ).to_dict() + ) def _upsert_document_revision( self, diff --git a/packages/shared-python/shared/services/storage/zip_chunk_schema.py b/packages/shared-python/shared/services/storage/zip_chunk_schema.py new file mode 100644 index 000000000..2dbf973d1 --- /dev/null +++ b/packages/shared-python/shared/services/storage/zip_chunk_schema.py @@ -0,0 +1,182 @@ +"""Chunk projection for Knowhere ZIP result packages.""" + +from __future__ import annotations + +from typing import Any + +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, +) + + +class ZipChunkSchemaBuilder: + 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", "") + normalized_type = _normalize_chunk_type(chunk_type) + 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: list[dict[str, Any]] = [] + for chunk in chunks: + chunk_id = str(chunk.get("chunk_id") or chunk.get("know_id")) + chunk_type_str = chunk.get("type", "") + normalized_type = _normalize_chunk_type(chunk_type_str) + image_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 = _base_chunk_metadata(existing_metadata, chunk, content) + + if chunk_type == "text": + metadata.update( + _format_text_metadata( + chunk=chunk, + chunk_type_str=chunk_type_str, + content=str(content), + existing_metadata=existing_metadata, + resource_target_map=resource_target_map, + ) + ) + elif chunk_type == "image": + if image_info: + metadata["file_path"] = image_info["file_path"] + metadata["keywords"] = existing_metadata.get("keywords") or chunk.get( + "keywords", [] + ) + metadata["tokens"] = [] + elif chunk_type == "table": + metadata["file_path"] = _resolve_table_file_path( + chunk=chunk, + chunk_id=chunk_id, + path=str(path), + existing_metadata=existing_metadata, + table_files_map=table_files_map, + ) + 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 _normalize_chunk_type(value: Any) -> str: + raw_type = str(value).strip() + return raw_type.split("\n", 1)[0].lower() + + +def _base_chunk_metadata( + existing_metadata: dict[str, Any], + chunk: dict[str, Any], + content: Any, +) -> dict[str, Any]: + 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 + return metadata + + +def _format_text_metadata( + *, + chunk: dict[str, Any], + chunk_type_str: Any, + content: str, + existing_metadata: dict[str, Any], + resource_target_map: dict[str, str], +) -> dict[str, Any]: + relationship_refs = parse_relationship_refs( + chunk.get("type_raw") or chunk_type_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, + ) + return { + "tokens": existing_metadata.get("tokens") or chunk.get("tokens", 0), + "keywords": existing_metadata.get("keywords") or chunk.get("keywords", []), + "connect_to": merge_connections(embed_connections, related_connections), + } + + +def _resolve_table_file_path( + *, + chunk: dict[str, Any], + chunk_id: str, + path: str, + existing_metadata: dict[str, Any], + table_files_map: dict[str, dict[str, Any]], +) -> Any: + file_path = existing_metadata.get("file_path") + if file_path: + return file_path + + table_info = table_files_map.get(chunk_id) + if table_info: + return table_info["file_path"] + + table_name = path.split("/")[-1] if "/" in path else f"table_{chunk_id}.html" + return f"tables/{table_name}" diff --git a/packages/shared-python/shared/services/storage/zip_doc_navigation.py b/packages/shared-python/shared/services/storage/zip_doc_navigation.py new file mode 100644 index 000000000..e4e2f8f13 --- /dev/null +++ b/packages/shared-python/shared/services/storage/zip_doc_navigation.py @@ -0,0 +1,183 @@ +"""Document navigation projection for Knowhere ZIP result packages.""" + +from __future__ import annotations + +from typing import Any + +from shared.services.chunks.document_path import split_document_path +from shared.utils.text_utils import truncate_content_preview + + +class ZipDocNavigationBuilder: + 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, + source_file_name=source_file_name, + ) + 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]], + *, + source_file_name: str, + ) -> list[dict[str, Any]]: + root_children: dict[str, dict[str, Any]] = {} + + for chunk in text_chunks: + path = chunk.get("path", "") + root_parts, section_parts = split_document_path( + path, + source_file_name=source_file_name, + ) + + if not section_parts: + key = "__root__" + if key not in root_children: + root_children[key] = { + "title": "Root", + "path": "/".join(root_parts) if root_parts 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 = list(root_parts) + 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[dict[str, Any]], 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[str, Any]], + level: int = 1, +) -> list[dict[str, Any]]: + result: list[dict[str, Any]] = [] + 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_manifest_schema.py b/packages/shared-python/shared/services/storage/zip_manifest_schema.py new file mode 100644 index 000000000..19afd84b9 --- /dev/null +++ b/packages/shared-python/shared/services/storage/zip_manifest_schema.py @@ -0,0 +1,42 @@ +"""Manifest projection for Knowhere ZIP result packages.""" + +from __future__ import annotations + +from typing import Any + +from shared.utils.utc_now import utc_now_naive + + +class ZipManifestBuilder: + def generate_manifest( + self, + *, + job_id: str, + data_id: str | None, + source_file_name: str, + statistics: dict[str, Any], + job_metadata: dict[str, Any], + hierarchy: dict[str, Any] | None = 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 {}, + } diff --git a/packages/shared-python/shared/services/storage/zip_package_writer.py b/packages/shared-python/shared/services/storage/zip_package_writer.py new file mode 100644 index 000000000..42c78532e --- /dev/null +++ b/packages/shared-python/shared/services/storage/zip_package_writer.py @@ -0,0 +1,112 @@ +"""Physical ZIP writing for Knowhere result packages.""" + +from __future__ import annotations + +import hashlib +import json +import os +import tempfile +import zipfile +from dataclasses import dataclass +from typing import Any + +from loguru import logger + +from shared.services.storage.zip_result_resources import ZipResourceFileInfo + + +@dataclass(frozen=True) +class ZipPackageWriteRequest: + job_id: str + add_dir: str + formatted_chunks: list[dict[str, Any]] + image_files: tuple[ZipResourceFileInfo, ...] + table_files: tuple[ZipResourceFileInfo, ...] + doc_nav: dict[str, Any] | None + manifest: dict[str, Any] + temp_dir: str | None + + +@dataclass(frozen=True) +class ZipPackageArtifact: + zip_file_path: str + checksum: dict[str, str] + zip_size: int + + +class ZipPackageWriter: + """Write a prepared result package to a ZIP file.""" + + def write(self, request: ZipPackageWriteRequest) -> ZipPackageArtifact: + effective_temp_dir = request.temp_dir or tempfile.gettempdir() + os.makedirs(effective_temp_dir, exist_ok=True) + zip_file_path = os.path.join(effective_temp_dir, f"result_{request.job_id}.zip") + + with zipfile.ZipFile(zip_file_path, "w", zipfile.ZIP_DEFLATED) as zip_file: + chunks_json = json.dumps( + {"chunks": request.formatted_chunks}, + ensure_ascii=False, + indent=2, + ) + zip_file.writestr("chunks.json", chunks_json.encode("utf-8")) + + self._write_optional_file(zip_file, request.add_dir, "full.md") + if self._write_optional_file( + zip_file, + request.add_dir, + "toc_hierarchies.json", + ): + logger.info("Added toc_hierarchies.json to ZIP") + + self._write_resource_files(zip_file, request.image_files, label="Image") + self._write_resource_files(zip_file, request.table_files, label="Table") + + if request.doc_nav is not None: + doc_nav_json = json.dumps(request.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") + + manifest_json = json.dumps(request.manifest, ensure_ascii=False, indent=2) + zip_file.writestr("manifest.json", manifest_json.encode("utf-8")) + + checksum_value = _calculate_zip_checksum(zip_file_path) + zip_size = os.path.getsize(zip_file_path) + return ZipPackageArtifact( + zip_file_path=zip_file_path, + checksum={"algorithm": "sha256", "value": checksum_value}, + zip_size=zip_size, + ) + + def _write_optional_file( + self, + zip_file: zipfile.ZipFile, + add_dir: str, + filename: str, + ) -> bool: + file_path = os.path.join(add_dir, filename) + if not os.path.exists(file_path): + return False + zip_file.write(file_path, filename) + return True + + def _write_resource_files( + self, + zip_file: zipfile.ZipFile, + resources: tuple[ZipResourceFileInfo, ...], + *, + label: str, + ) -> None: + for resource in resources: + source_path = resource["source_path"] + if os.path.exists(source_path): + zip_file.write(source_path, resource["zip_path"]) + else: + logger.warning(f"{label} file not found: {source_path}") + + +def _calculate_zip_checksum(zip_file_path: str) -> str: + sha256_hash = hashlib.sha256() + with open(zip_file_path, "rb") as file_obj: + for byte_block in iter(lambda: file_obj.read(4096), b""): + sha256_hash.update(byte_block) + return sha256_hash.hexdigest().lower() diff --git a/packages/shared-python/shared/services/storage/zip_result_resources.py b/packages/shared-python/shared/services/storage/zip_result_resources.py new file mode 100644 index 000000000..8c1851dd6 --- /dev/null +++ b/packages/shared-python/shared/services/storage/zip_result_resources.py @@ -0,0 +1,321 @@ +"""Resource discovery for Knowhere ZIP result packages.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import Any + +from PIL import Image + +from shared.core.exceptions.domain_exceptions import StorageServiceException + +ZipResourceFileInfo = dict[str, Any] + + +@dataclass(frozen=True) +class ZipPackageResources: + image_files: tuple[ZipResourceFileInfo, ...] + table_files: tuple[ZipResourceFileInfo, ...] + + @property + def image_files_map(self) -> dict[str, ZipResourceFileInfo]: + return {str(image["id"]): image for image in self.image_files} + + @property + def table_files_map(self) -> dict[str, ZipResourceFileInfo]: + return {str(table["id"]): table for table in self.table_files} + + +class ZipResourceCollector: + """Resolve chunk resource references to files that belong in the result ZIP.""" + + def collect( + self, + *, + chunks: list[dict[str, Any]], + add_dir: str, + ) -> ZipPackageResources: + images_dir = os.path.join(add_dir, "images") + tables_dir = os.path.join(add_dir, "tables") + return ZipPackageResources( + image_files=tuple(self._collect_image_files(chunks, images_dir)), + table_files=tuple(self._collect_table_files(chunks, tables_dir)), + ) + + def _collect_image_files( + self, + chunks: list[dict[str, Any]], + images_dir: str, + ) -> list[ZipResourceFileInfo]: + image_files: list[ZipResourceFileInfo] = [] + if not os.path.exists(images_dir): + has_image_chunks = any(chunk.get("type", "") == "image" for chunk in chunks) + if has_image_chunks: + raise StorageServiceException( + internal_message=( + "Image directory not found for ZIP packaging: " + f"images_dir={images_dir}" + ), + operation="collect_image_files", + ) + return image_files + + image_files_map = _collect_files_by_name(images_dir) + + for chunk in chunks: + chunk_id = chunk.get("chunk_id") or chunk.get("know_id") + chunk_type = chunk.get("type", "") + + if chunk_type != "image": + continue + + metadata = chunk.get("metadata", {}) + candidate_names: list[str] = [] + metadata_file_path = "" + if metadata and isinstance(metadata, dict): + metadata_file_path = str(metadata.get("file_path") or "").strip() + _add_candidate(candidate_names, metadata.get("file_path")) + _add_candidate(candidate_names, metadata.get("original_name")) + + if metadata_file_path.startswith("images/"): + original_name = os.path.basename(metadata_file_path) + else: + original_name = None + + original_path = chunk.get("path", "") + if original_path: + _add_candidate(candidate_names, original_path) + normalized_path = original_path.replace("-->", "/") + original_name = os.path.basename(normalized_path) + + source_path, matched_name, ext = _resolve_image_source_path( + image_files_map, + candidate_names, + ) + if matched_name: + original_name = matched_name + + if not source_path: + raise StorageServiceException( + internal_message=( + "Cannot resolve image file for ZIP packaging: " + f"chunk_id={chunk_id}, candidates={candidate_names}" + ), + operation="collect_image_files", + ) + + width = None + height = None + try: + with Image.open(source_path) as img: + width, height = img.size + except Exception: + pass + + file_size = os.path.getsize(source_path) + zip_path = _resolve_image_zip_path( + metadata=metadata, + original_name=original_name, + chunk_id=str(chunk_id), + ext=ext, + ) + + image_files.append( + { + "id": str(chunk_id), + "file_path": zip_path, + "original_name": original_name or f"image_{chunk_id}.{ext}", + "size_bytes": file_size, + "format": ext.lower(), + "width": width, + "height": height, + "source_path": source_path, + "zip_path": zip_path, + } + ) + + return image_files + + def _collect_table_files( + self, + chunks: list[dict[str, Any]], + tables_dir: str, + ) -> list[ZipResourceFileInfo]: + table_files: list[ZipResourceFileInfo] = [] + if not os.path.exists(tables_dir): + return table_files + + table_files_map = { + filename: file_path + for filename, file_path in _collect_files_by_name(tables_dir).items() + if filename.endswith(".html") + } + + for chunk in chunks: + chunk_id = chunk.get("chunk_id") or chunk.get("know_id") + chunk_type = chunk.get("type", "") + + if chunk_type != "table": + continue + + metadata = chunk.get("metadata", {}) + candidate_names: list[str] = [] + if metadata and isinstance(metadata, dict): + _add_candidate(candidate_names, metadata.get("file_path")) + _add_candidate(candidate_names, metadata.get("original_name")) + + original_path = chunk.get("path", "") + if original_path: + normalized_path = original_path.replace("-->", "/") + original_name = os.path.basename(normalized_path) + _add_candidate(candidate_names, original_path) + else: + original_name = None + + source_path, matched_name = _resolve_table_source_path( + table_files_map, + candidate_names, + ) + if matched_name: + original_name = matched_name + + if not source_path: + raise StorageServiceException( + internal_message=( + "Cannot resolve table file for ZIP packaging: " + f"chunk_id={chunk_id}, candidates={candidate_names}" + ), + operation="collect_table_files", + ) + + file_size = os.path.getsize(source_path) + zip_path = _resolve_table_zip_path( + metadata=metadata, + original_name=original_name, + chunk_id=str(chunk_id), + ) + + table_files.append( + { + "id": str(chunk_id), + "file_path": zip_path, + "original_name": original_name or f"table_{chunk_id}.html", + "size_bytes": file_size, + "format": "html", + "source_path": source_path, + "zip_path": zip_path, + } + ) + + return table_files + + +def _collect_files_by_name(directory_path: str) -> dict[str, str]: + files: dict[str, str] = {} + for filename in os.listdir(directory_path): + file_path = os.path.join(directory_path, filename) + if os.path.isfile(file_path): + files[filename] = file_path + return files + + +def _add_candidate(candidates: list[str], value: str | None) -> None: + if not value: + return + candidate = os.path.basename(str(value).strip().replace("-->", "/")) + if not candidate: + return + if candidate.startswith("[") and candidate.endswith("]"): + candidate = candidate[1:-1].strip() + if candidate and candidate not in candidates: + candidates.append(candidate) + + +def _resolve_image_source_path( + image_files_map: dict[str, str], + candidates: list[str], +) -> tuple[str | None, str | None, str]: + for candidate in candidates: + if candidate in image_files_map: + _, ext = os.path.splitext(candidate) + return image_files_map[candidate], candidate, ext.lstrip(".") or "jpg" + + stem, _ = os.path.splitext(candidate) + if stem: + stem_matches = [ + filename + for filename in image_files_map + if os.path.splitext(filename)[0] == stem + ] + if len(stem_matches) == 1: + matched = stem_matches[0] + _, matched_ext = os.path.splitext(matched) + return ( + image_files_map[matched], + matched, + matched_ext.lstrip(".") or "jpg", + ) + + return None, None, "jpg" + + +def _resolve_table_source_path( + table_files_map: dict[str, str], + candidates: list[str], +) -> tuple[str | None, str | None]: + for candidate in candidates: + if candidate in table_files_map: + return table_files_map[candidate], candidate + + stem, _ = os.path.splitext(candidate) + if stem: + stem_matches = [ + filename + for filename in table_files_map + if os.path.splitext(filename)[0] == stem + ] + if len(stem_matches) == 1: + matched = stem_matches[0] + return table_files_map[matched], matched + + return None, None + + +def _resolve_image_zip_path( + *, + metadata: Any, + original_name: str | None, + chunk_id: str, + ext: str, +) -> str: + if metadata and isinstance(metadata, dict): + zip_file_path = metadata.get("file_path") + if zip_file_path and zip_file_path.startswith("images/"): + return zip_file_path + if original_name: + return f"images/{original_name}" + return f"images/{chunk_id}.{ext}" + + if original_name: + return f"images/{original_name}" + return f"images/{chunk_id}.{ext}" + + +def _resolve_table_zip_path( + *, + metadata: Any, + original_name: str | None, + chunk_id: str, +) -> str: + if metadata and isinstance(metadata, dict): + zip_file_path = metadata.get("file_path") + if zip_file_path and zip_file_path.startswith("tables/"): + return zip_file_path + if original_name: + return f"tables/{original_name}" + return f"tables/{chunk_id}.html" + + if original_name: + return f"tables/{original_name}" + return f"tables/{chunk_id}.html" diff --git a/packages/shared-python/shared/services/storage/zip_result_schema.py b/packages/shared-python/shared/services/storage/zip_result_schema.py index d4063b79d..e444475a7 100644 --- a/packages/shared-python/shared/services/storage/zip_result_schema.py +++ b/packages/shared-python/shared/services/storage/zip_result_schema.py @@ -1,353 +1,72 @@ -"""Schema projection for Knowhere ZIP result packages.""" +"""Schema projection facade for Knowhere ZIP result packages.""" + from __future__ import annotations -from typing import Any, Dict, List, Optional +from typing import Any -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.services.chunks.document_path import split_document_path -from shared.utils.text_utils import truncate_content_preview -from shared.utils.utc_now import utc_now_naive +from shared.services.storage.zip_chunk_schema import ZipChunkSchemaBuilder +from shared.services.storage.zip_doc_navigation import ZipDocNavigationBuilder +from shared.services.storage.zip_manifest_schema import ZipManifestBuilder 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 + def __init__( + self, + *, + chunk_schema: ZipChunkSchemaBuilder | None = None, + doc_navigation: ZipDocNavigationBuilder | None = None, + manifest_builder: ZipManifestBuilder | None = None, + ) -> None: + self._chunk_schema = chunk_schema or ZipChunkSchemaBuilder() + self._doc_navigation = doc_navigation or ZipDocNavigationBuilder() + self._manifest_builder = manifest_builder or ZipManifestBuilder() - return { - "total_chunks": total_chunks, - "text_chunks": text_chunks, - "image_chunks": image_chunks, - "table_chunks": table_chunks, - "total_pages": None, - } + def calculate_statistics(self, chunks: list[dict[str, Any]]) -> dict[str, Any]: + return self._chunk_schema.calculate_statistics(chunks) 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: list[dict[str, Any]], + image_files_map: dict[str, dict[str, Any]], + table_files_map: dict[str, dict[str, Any]], + ) -> list[dict[str, Any]]: + return self._chunk_schema.format_chunks( chunks, - image_files_map=image_files_map, - table_files_map=table_files_map, + image_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], + data_id: str | None, 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 {}, - } + statistics: dict[str, Any], + job_metadata: dict[str, Any], + hierarchy: dict[str, Any] | None = None, + ) -> dict[str, Any]: + return self._manifest_builder.generate_manifest( + job_id=job_id, + data_id=data_id, + source_file_name=source_file_name, + statistics=statistics, + job_metadata=job_metadata, + hierarchy=hierarchy, + ) 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 + sections: list[dict[str, Any]], + ) -> dict[str, Any]: + return self._doc_navigation.build_hierarchy_dict(sections) def build_doc_nav( self, - formatted_chunks: List[Dict[str, Any]], + 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, - source_file_name=source_file_name, - ) - 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]], - *, - source_file_name: str, - ) -> List[Dict[str, Any]]: - root_children: Dict[str, dict] = {} - - for chunk in text_chunks: - path = chunk.get("path", "") - root_parts, section_parts = split_document_path( - path, - source_file_name=source_file_name, - ) - - if not section_parts: - key = "__root__" - if key not in root_children: - root_children[key] = { - "title": "Root", - "path": "/".join(root_parts) if root_parts 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 = list(root_parts) - 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, - } + ) -> dict[str, Any]: + return self._doc_navigation.build_doc_nav( + formatted_chunks, + source_file_name, ) - 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 033c2db25..b2fee6e17 100644 --- a/packages/shared-python/shared/services/storage/zip_result_service.py +++ b/packages/shared-python/shared/services/storage/zip_result_service.py @@ -1,17 +1,14 @@ """ -ZIP Result Package Generation Service -Generates ZIP packages according to Knowhere-API-ZIP-Spec.md specification +ZIP Result Package Generation Service. + +Generates ZIP packages according to the Knowhere API ZIP result format. """ -import hashlib -import json -import os -import tempfile -import zipfile -from typing import Any, Dict, List, Optional, Tuple +from __future__ import annotations + +from typing import Any from loguru import logger -from PIL import Image import pandas as pd @@ -19,446 +16,116 @@ KnowhereException, StorageServiceException, ) +from shared.services.storage.zip_package_writer import ( + ZipPackageWriter, + ZipPackageWriteRequest, +) +from shared.services.storage.zip_result_resources import ZipResourceCollector from shared.services.storage.zip_result_schema import ZipResultSchemaBuilder class ZipResultService: - """ZIP Result Package Generation Service""" + """Generate a result ZIP while preserving the worker-facing package contract.""" - def __init__(self): - self._schema = ZipResultSchemaBuilder() + def __init__( + self, + *, + schema_builder: ZipResultSchemaBuilder | None = None, + resource_collector: ZipResourceCollector | None = None, + package_writer: ZipPackageWriter | None = None, + ) -> None: + self._schema = schema_builder or ZipResultSchemaBuilder() + self._resources = resource_collector or ZipResourceCollector() + self._writer = package_writer or ZipPackageWriter() def generate_zip_package( self, job_id: str, - chunks: List[Dict[str, Any]], + chunks: list[dict[str, Any]], add_dir: str, source_file_name: str, - data_id: Optional[str], - job_metadata: Dict[str, Any], - parsed_df: Optional["pd.DataFrame"] = None, - temp_dir: Optional[str] = None, - ) -> Tuple[str, Dict[str, str], Dict[str, Any], int]: + data_id: str | None, + job_metadata: dict[str, Any], + parsed_df: pd.DataFrame | None = None, + temp_dir: str | None = None, + ) -> tuple[str, dict[str, str], dict[str, Any], int]: """ - Generate ZIP result package - - Args: - job_id: Job ID - chunks: List of chunks data - add_dir: Parsed directory path (contains images/ and tables/ directories) - source_file_name: Source file name - data_id: User-defined ID - job_metadata: Job metadata - parsed_df: Optional, parsed DataFrame (legacy, unused after doc_nav.json migration) - temp_dir: Optional directory for the generated ZIP file + Generate ZIP result package. Returns: Tuple[zip_file_path, checksum, statistics, zip_size]: - zip_file_path: ZIP file path - checksum: {"algorithm": "sha256", "value": "..."} - - statistics: {"total_chunks": int, "text_chunks": int, "image_chunks": int, "table_chunks": int, "total_pages": Optional[int]} + - statistics: chunk and page statistics persisted with the Job Result - zip_size: ZIP file size in bytes """ try: - # Create temporary ZIP file - effective_temp_dir = temp_dir or tempfile.gettempdir() - os.makedirs(effective_temp_dir, exist_ok=True) - zip_file_path = os.path.join(effective_temp_dir, f"result_{job_id}.zip") - - # Collect image and table file info (must be done before formatting chunks as file info is needed) - images_dir = os.path.join(add_dir, "images") - tables_dir = os.path.join(add_dir, "tables") - image_files_info = self._collect_image_files(chunks, images_dir) - table_files_info = self._collect_table_files(chunks, tables_dir) - - # Create image and table file mappings (chunk_id -> file_info) - image_files_map = {img["id"]: img for img in image_files_info} - table_files_map = {tb["id"]: tb for tb in table_files_info} - - # Convert chunks data format (using file info) + resources = self._resources.collect(chunks=chunks, add_dir=add_dir) formatted_chunks = self._schema.format_chunks( - chunks, image_files_map, table_files_map + chunks, + resources.image_files_map, + resources.table_files_map, ) statistics = self._schema.calculate_statistics(formatted_chunks) - doc_nav: Dict[str, Any] = {} - hierarchy: Dict[str, Any] = {} - - # Create ZIP package - with zipfile.ZipFile(zip_file_path, "w", zipfile.ZIP_DEFLATED) as zip_file: - # 1. Generate chunks.json (full version) - chunks_json = json.dumps( - {"chunks": formatted_chunks}, ensure_ascii=False, indent=2 - ) - zip_file.writestr("chunks.json", chunks_json.encode("utf-8")) - - # 2. Try to add full.md (if exists) - full_md_path = os.path.join(add_dir, "full.md") - if os.path.exists(full_md_path): - zip_file.write(full_md_path, "full.md") - - # 2b. Try to add toc_hierarchies.json (if exists) - toc_path = os.path.join(add_dir, "toc_hierarchies.json") - if os.path.exists(toc_path): - zip_file.write(toc_path, "toc_hierarchies.json") - logger.info("Added toc_hierarchies.json to ZIP") - - # 3. Add image files - for img_info in image_files_info: - source_path = img_info["source_path"] - if os.path.exists(source_path): - zip_file.write( - source_path, - img_info["zip_path"], - ) - else: - logger.warning(f"Image file not found: {source_path}") - - # 4. Add table files - for table_info in table_files_info: - source_path = table_info["source_path"] - if os.path.exists(source_path): - zip_file.write( - source_path, - table_info["zip_path"], - ) - else: - logger.warning(f"Table file not found: {source_path}") - - # 5. Generate doc_nav.json — unified navigation file - try: - 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") - except Exception as e: - logger.warning(f"generate doc_nav.json fail {e}") - - # 6. Generate manifest.json (checksum not included, stored in database) - manifest = self._schema.generate_manifest( + doc_nav, hierarchy = self._build_navigation_outputs( + formatted_chunks=formatted_chunks, + source_file_name=source_file_name, + ) + manifest = self._schema.generate_manifest( + job_id=job_id, + data_id=data_id, + source_file_name=source_file_name, + statistics=statistics, + job_metadata=job_metadata, + hierarchy=hierarchy, + ) + artifact = self._writer.write( + ZipPackageWriteRequest( job_id=job_id, - data_id=data_id, - source_file_name=source_file_name, - statistics=statistics, - job_metadata=job_metadata, - hierarchy=hierarchy, + add_dir=add_dir, + formatted_chunks=formatted_chunks, + image_files=resources.image_files, + table_files=resources.table_files, + doc_nav=doc_nav, + manifest=manifest, + temp_dir=temp_dir, ) - manifest_json = json.dumps(manifest, ensure_ascii=False, indent=2) - zip_file.writestr("manifest.json", manifest_json.encode("utf-8")) - - # Calculate ZIP package SHA-256 - checksum_value = self._calculate_zip_checksum(zip_file_path) - checksum = {"algorithm": "sha256", "value": checksum_value} - - # Get ZIP file size - zip_size = os.path.getsize(zip_file_path) + ) logger.info( - f"ZIP package generated successfully: job_id={job_id}, size={zip_size}, checksum={checksum_value[:16]}..." + "ZIP package generated successfully: " + f"job_id={job_id}, size={artifact.zip_size}, " + f"checksum={artifact.checksum['value'][:16]}..." ) - return zip_file_path, checksum, statistics, zip_size + return ( + artifact.zip_file_path, + artifact.checksum, + statistics, + artifact.zip_size, + ) except KnowhereException: raise - except Exception as e: - logger.error(f"Failed to generate ZIP package: {e}") + except Exception as exc: + logger.error(f"Failed to generate ZIP package: {exc}") raise StorageServiceException( - internal_message=f"Failed to generate ZIP package: {str(e)}", + internal_message=f"Failed to generate ZIP package: {str(exc)}", operation="generate_zip_package", - original_exception=e, + original_exception=exc, ) - def _collect_image_files( - self, chunks: List[Dict[str, Any]], images_dir: str - ) -> List[Dict[str, Any]]: - """Collect image file information""" - image_files = [] - if not os.path.exists(images_dir): - has_image_chunks = any(chunk.get("type", "") == "image" for chunk in chunks) - if has_image_chunks: - raise StorageServiceException( - internal_message=( - "Image directory not found for ZIP packaging: " - f"images_dir={images_dir}" - ), - operation="collect_image_files", - ) - return image_files - - # Get all image files - image_files_map = {} - for filename in os.listdir(images_dir): - file_path = os.path.join(images_dir, filename) - if os.path.isfile(file_path): - image_files_map[filename] = file_path - - def add_candidate(candidates: List[str], value: Optional[str]) -> None: - if not value: - return - candidate = os.path.basename(str(value).strip().replace("-->", "/")) - if not candidate: - return - if candidate.startswith("[") and candidate.endswith("]"): - candidate = candidate[1:-1].strip() - if candidate and candidate not in candidates: - candidates.append(candidate) - - def resolve_source_path( - candidates: List[str], chunk_id: str - ) -> Tuple[Optional[str], Optional[str], str]: - for candidate in candidates: - if candidate in image_files_map: - _, ext = os.path.splitext(candidate) - return ( - image_files_map[candidate], - candidate, - ext.lstrip(".") or "jpg", - ) - - stem, ext = os.path.splitext(candidate) - if stem: - stem_matches = [ - filename - for filename in image_files_map - if os.path.splitext(filename)[0] == stem - ] - if len(stem_matches) == 1: - matched = stem_matches[0] - _, matched_ext = os.path.splitext(matched) - return ( - image_files_map[matched], - matched, - matched_ext.lstrip(".") or "jpg", - ) - - return None, None, "jpg" - - # Match images from chunks - for chunk in chunks: - chunk_id = chunk.get("chunk_id") or chunk.get("know_id") - chunk_type = chunk.get("type", "") - - if chunk_type != "image": - continue - - metadata = chunk.get("metadata", {}) - candidate_names: List[str] = [] - metadata_file_path = "" - if metadata and isinstance(metadata, dict): - metadata_file_path = str(metadata.get("file_path") or "").strip() - add_candidate(candidate_names, metadata.get("file_path")) - add_candidate(candidate_names, metadata.get("original_name")) - - # Try to get original filename from chunk's path field - if metadata_file_path.startswith("images/"): - original_name = os.path.basename(metadata_file_path) - else: - original_name = None - - original_path = chunk.get("path", "") - if original_path: - add_candidate(candidate_names, original_path) - # Normalize path separators: replace --> with /, then extract filename - normalized_path = original_path.replace("-->", "/") - original_name = os.path.basename(normalized_path) - - source_path, matched_name, ext = resolve_source_path( - candidate_names, str(chunk_id) - ) - if matched_name: - original_name = matched_name - - if not source_path: - raise StorageServiceException( - internal_message=( - "Cannot resolve image file for ZIP packaging: " - f"chunk_id={chunk_id}, candidates={candidate_names}" - ), - operation="collect_image_files", - ) - - # Get image dimensions - width = None - height = None - try: - with Image.open(source_path) as img: - width, height = img.size - except Exception: - pass - - file_size = os.path.getsize(source_path) - - # Priority: use file_path from metadata, then original_name, finally chunk_id - if metadata and isinstance(metadata, dict): - # metadata.file_path format: "images/xxx.jpg" - zip_file_path = metadata.get("file_path") - if zip_file_path and zip_file_path.startswith("images/"): - # Use complete path from metadata - zip_path = zip_file_path - # Extract filename as original_name - if not original_name: - original_name = metadata.get( - "original_name" - ) or os.path.basename(zip_file_path) - else: - # If metadata has no file_path, use original_name or chunk_id - if original_name: - zip_path = f"images/{original_name}" - else: - zip_path = f"images/{chunk_id}.{ext}" - else: - # If no metadata, use original_name or chunk_id - if original_name: - zip_path = f"images/{original_name}" - else: - zip_path = f"images/{chunk_id}.{ext}" - - image_files.append( - { - "id": str(chunk_id), - "file_path": zip_path, - "original_name": original_name or f"image_{chunk_id}.{ext}", - "size_bytes": file_size, - "format": ext.lower(), - "width": width, - "height": height, - "source_path": source_path, - "zip_path": zip_path, - } - ) - - return image_files - - def _collect_table_files( - self, chunks: List[Dict[str, Any]], tables_dir: str - ) -> List[Dict[str, Any]]: - """Collect table file information""" - table_files = [] - if not os.path.exists(tables_dir): - return table_files - - # Get all table files - table_files_map = {} - for filename in os.listdir(tables_dir): - file_path = os.path.join(tables_dir, filename) - if os.path.isfile(file_path) and filename.endswith(".html"): - table_files_map[filename] = file_path - - def add_candidate(candidates: List[str], value: Optional[str]) -> None: - if not value: - return - candidate = os.path.basename(str(value).strip().replace("-->", "/")) - if not candidate: - return - if candidate.startswith("[") and candidate.endswith("]"): - candidate = candidate[1:-1].strip() - if candidate and candidate not in candidates: - candidates.append(candidate) - - def resolve_source_path( - candidates: List[str], - ) -> Tuple[Optional[str], Optional[str]]: - for candidate in candidates: - if candidate in table_files_map: - return table_files_map[candidate], candidate - - stem, _ = os.path.splitext(candidate) - if stem: - stem_matches = [ - filename - for filename in table_files_map - if os.path.splitext(filename)[0] == stem - ] - if len(stem_matches) == 1: - matched = stem_matches[0] - return table_files_map[matched], matched - - return None, None - - # Match tables from chunks - for chunk in chunks: - chunk_id = chunk.get("chunk_id") or chunk.get("know_id") - chunk_type = chunk.get("type", "") - - if chunk_type != "table": - continue - - metadata = chunk.get("metadata", {}) - candidate_names: List[str] = [] - if metadata and isinstance(metadata, dict): - add_candidate(candidate_names, metadata.get("file_path")) - add_candidate(candidate_names, metadata.get("original_name")) - - # Try to get original filename from chunk's path field - original_path = chunk.get("path", "") - if original_path: - # Normalize path separators: replace --> with /, then extract filename - normalized_path = original_path.replace("-->", "/") - original_name = os.path.basename(normalized_path) - add_candidate(candidate_names, original_path) - else: - original_name = None - - source_path, matched_name = resolve_source_path(candidate_names) - if matched_name: - original_name = matched_name - - if not source_path: - raise StorageServiceException( - internal_message=( - "Cannot resolve table file for ZIP packaging: " - f"chunk_id={chunk_id}, candidates={candidate_names}" - ), - operation="collect_table_files", - ) - - file_size = os.path.getsize(source_path) - - # Priority: use file_path from metadata, then original_name, finally chunk_id - if metadata and isinstance(metadata, dict): - # metadata.file_path format: "tables/xxx.html" - zip_file_path = metadata.get("file_path") - if zip_file_path and zip_file_path.startswith("tables/"): - # Use complete path from metadata - zip_path = zip_file_path - # Extract filename as original_name - if not original_name: - original_name = metadata.get( - "original_name" - ) or os.path.basename(zip_file_path) - else: - # If metadata has no file_path, use original_name or chunk_id - if original_name: - zip_path = f"tables/{original_name}" - else: - zip_path = f"tables/{chunk_id}.html" - else: - # If no metadata, use original_name or chunk_id - if original_name: - zip_path = f"tables/{original_name}" - else: - zip_path = f"tables/{chunk_id}.html" - - table_files.append( - { - "id": str(chunk_id), - "file_path": zip_path, - "original_name": original_name or f"table_{chunk_id}.html", - "size_bytes": file_size, - "format": "html", - "source_path": source_path, - "zip_path": zip_path, - } - ) - - return table_files - - def _calculate_zip_checksum(self, zip_file_path: str) -> str: - """Calculate SHA-256 checksum of ZIP file""" - sha256_hash = hashlib.sha256() - with open(zip_file_path, "rb") as f: - for byte_block in iter(lambda: f.read(4096), b""): - sha256_hash.update(byte_block) - return sha256_hash.hexdigest().lower() + def _build_navigation_outputs( + self, + *, + formatted_chunks: list[dict[str, Any]], + source_file_name: str, + ) -> tuple[dict[str, Any] | None, dict[str, Any]]: + try: + doc_nav = self._schema.build_doc_nav(formatted_chunks, source_file_name) + hierarchy = self._schema.build_hierarchy_dict(doc_nav.get("sections", [])) + return doc_nav, hierarchy + except Exception as exc: + logger.warning(f"generate doc_nav.json fail {exc}") + return None, {} From 6d4f57bd893ed8058f7a065672288142e4082341 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 18:06:41 +0000 Subject: [PATCH 77/78] fix: address advanced security review comments --- apps/worker/app/core/worker_bootstrap.py | 16 ++++- .../formats/excel/structure_parser.py | 2 - .../formats/markdown/deferred_task.py | 7 +++ .../formats/markdown/image_asset.py | 4 +- .../orchestration/format_adapters.py | 2 +- .../providers/mineru/task_polling.py | 15 +++-- .../document_parser/tables/table_parser.py | 18 ------ .../services/retrieval/agentic/prompts.py | 61 ++++++++++--------- .../services/retrieval/stats/recorder.py | 2 +- .../services/storage/zip_result_resources.py | 7 ++- 10 files changed, 71 insertions(+), 63 deletions(-) diff --git a/apps/worker/app/core/worker_bootstrap.py b/apps/worker/app/core/worker_bootstrap.py index 44a40c60c..5a0827a22 100644 --- a/apps/worker/app/core/worker_bootstrap.py +++ b/apps/worker/app/core/worker_bootstrap.py @@ -106,6 +106,18 @@ def run_worker() -> None: node_name = f"celery@{hostname}-{pid}" log_level = os.getenv("LOG_LEVEL", "INFO").lower() concurrency = settings.WORKER_CONCURRENCY + worker_queues = ",".join( + [ + "document_ingestion_high", + "document_ingestion_medium", + "document_ingestion_low", + "kb_high", + "kb_medium", + "kb_low", + "ai_high_priority", + "default", + ] + ) celery_args = [ "worker", @@ -114,9 +126,7 @@ def run_worker() -> None: f"--loglevel={log_level}", f"--hostname={node_name}", "-Q", - "document_ingestion_high,document_ingestion_medium," - "document_ingestion_low,kb_high,kb_medium,kb_low," - "ai_high_priority,default", + worker_queues, "--without-gossip", "--without-mingle", ] diff --git a/apps/worker/app/services/document_parser/formats/excel/structure_parser.py b/apps/worker/app/services/document_parser/formats/excel/structure_parser.py index 3ff1ea6b8..bfe3553de 100644 --- a/apps/worker/app/services/document_parser/formats/excel/structure_parser.py +++ b/apps/worker/app/services/document_parser/formats/excel/structure_parser.py @@ -551,8 +551,6 @@ def _split_sheet_recursive( 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), ( diff --git a/apps/worker/app/services/document_parser/formats/markdown/deferred_task.py b/apps/worker/app/services/document_parser/formats/markdown/deferred_task.py index da935abaa..f2cb82409 100644 --- a/apps/worker/app/services/document_parser/formats/markdown/deferred_task.py +++ b/apps/worker/app/services/document_parser/formats/markdown/deferred_task.py @@ -31,3 +31,10 @@ class TextDeferredSummaryTask: MarkdownDeferredSummaryTask: TypeAlias = ( ImageDeferredSummaryTask | TableDeferredSummaryTask | TextDeferredSummaryTask ) + +__all__ = [ + "ImageDeferredSummaryTask", + "MarkdownDeferredSummaryTask", + "TableDeferredSummaryTask", + "TextDeferredSummaryTask", +] diff --git a/apps/worker/app/services/document_parser/formats/markdown/image_asset.py b/apps/worker/app/services/document_parser/formats/markdown/image_asset.py index 88a684e56..b71ace46a 100644 --- a/apps/worker/app/services/document_parser/formats/markdown/image_asset.py +++ b/apps/worker/app/services/document_parser/formats/markdown/image_asset.py @@ -178,8 +178,8 @@ def _build_duplicate_image_asset( ) try: source_path.unlink() - except OSError: - pass + except OSError as exc: + logger.debug(f"Failed to remove duplicate image source {source_path}: {exc}") logger.debug("Skipped duplicate image") return MarkdownImageAsset( content_item=cache_entry["img_content"], 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 c5b0972b7..6dc3409c2 100644 --- a/apps/worker/app/services/document_parser/orchestration/format_adapters.py +++ b/apps/worker/app/services/document_parser/orchestration/format_adapters.py @@ -14,7 +14,7 @@ def document_format(self) -> object: def parse(self, session: ParseSession) -> ParseOutput: """Parse a document session into a stable parser output object.""" - ... + raise NotImplementedError @dataclass(frozen=True) diff --git a/apps/worker/app/services/document_parser/providers/mineru/task_polling.py b/apps/worker/app/services/document_parser/providers/mineru/task_polling.py index 689ac69ab..fa6ee7775 100644 --- a/apps/worker/app/services/document_parser/providers/mineru/task_polling.py +++ b/apps/worker/app/services/document_parser/providers/mineru/task_polling.py @@ -158,10 +158,17 @@ def poll_mineru_task( if state == "running": if "extract_progress" in status: try: - _progress = ( - status["extract_progress"]["extracted_pages"] - / status["extract_progress"]["total_pages"] - ) + extracted_pages = status["extract_progress"][ + "extracted_pages" + ] + total_pages = status["extract_progress"]["total_pages"] + progress = extracted_pages / total_pages + polling_logger.bind( + token_id=lease.token_id, + extracted_pages=extracted_pages, + total_pages=total_pages, + progress_pct=round(progress * 100, 1), + ).debug("MinerU parsing in progress") except (KeyError, ZeroDivisionError): polling_logger.bind(token_id=lease.token_id).info( "MinerU parsing in progress" diff --git a/apps/worker/app/services/document_parser/tables/table_parser.py b/apps/worker/app/services/document_parser/tables/table_parser.py index 7cdafe6f3..ef69e1bfc 100755 --- a/apps/worker/app/services/document_parser/tables/table_parser.py +++ b/apps/worker/app/services/document_parser/tables/table_parser.py @@ -1,24 +1,6 @@ from __future__ import annotations import pandas as pd -from app.services.document_parser.tables.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.tables.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( diff --git a/packages/shared-python/shared/services/retrieval/agentic/prompts.py b/packages/shared-python/shared/services/retrieval/agentic/prompts.py index 418ae46bc..e48e82901 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/prompts.py +++ b/packages/shared-python/shared/services/retrieval/agentic/prompts.py @@ -135,30 +135,21 @@ def extract(data: dict) -> dict: return {"action": action, "tools": tools, "selections": selections} - try: - data = json.loads(text) - if isinstance(data, dict): - return extract(data) - except (ValueError, json.JSONDecodeError): - pass + data = _parse_json_object(text) + if data is not None: + return extract(data) 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 + data = _parse_json_object(fence_match.group(1).strip()) + if data is not None: + return extract(data) 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 + data = _parse_json_object(brace_match.group()) + if data is not None: + return extract(data) return default @@ -192,23 +183,33 @@ def parse_json_array(text: str) -> list[str]: 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 + result = _parse_json_array(text) + if result is not None: + return result 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 + result = _parse_json_array(match.group()) + if result is not None: + return result return [] +def _parse_json_object(raw_value: str) -> dict[str, Any] | None: + try: + result = json.loads(raw_value) + except (ValueError, json.JSONDecodeError): + return None + return result if isinstance(result, dict) else None + + +def _parse_json_array(raw_value: str) -> list[Any] | None: + try: + result = json.loads(raw_value) + except (ValueError, json.JSONDecodeError): + return None + return result if isinstance(result, list) else None + + def normalize_confidence(value: Any) -> float | None: if value is None: return None diff --git a/packages/shared-python/shared/services/retrieval/stats/recorder.py b/packages/shared-python/shared/services/retrieval/stats/recorder.py index 05b018678..967086a4b 100644 --- a/packages/shared-python/shared/services/retrieval/stats/recorder.py +++ b/packages/shared-python/shared/services/retrieval/stats/recorder.py @@ -17,7 +17,7 @@ def _finalize_retrieval_hit_stats_task(task: asyncio.Task[None]) -> None: try: task.result() except asyncio.CancelledError: - pass + logger.debug('Retrieval hit stats task was cancelled') except Exception as exc: logger.warning(f'Failed to record retrieval hit stats (ignored): {exc}') diff --git a/packages/shared-python/shared/services/storage/zip_result_resources.py b/packages/shared-python/shared/services/storage/zip_result_resources.py index 8c1851dd6..e84e0cb5a 100644 --- a/packages/shared-python/shared/services/storage/zip_result_resources.py +++ b/packages/shared-python/shared/services/storage/zip_result_resources.py @@ -6,6 +6,7 @@ from dataclasses import dataclass from typing import Any +from loguru import logger from PIL import Image from shared.core.exceptions.domain_exceptions import StorageServiceException @@ -110,8 +111,10 @@ def _collect_image_files( try: with Image.open(source_path) as img: width, height = img.size - except Exception: - pass + except Exception as exc: + logger.debug( + f"Failed to read image dimensions for ZIP resource {source_path}: {exc}" + ) file_size = os.path.getsize(source_path) zip_path = _resolve_image_zip_path( From 96655cf2fc5cf90cae5ff6043e0ba480f4ca1f1b Mon Sep 17 00:00:00 2001 From: suguanYang Date: Mon, 18 May 2026 18:31:59 +0000 Subject: [PATCH 78/78] fix: clear remaining codeql alerts --- .../document_parser/formats/atlas/parser.py | 4 +-- .../document_parser/formats/docx/parser.py | 8 ++--- .../document_parser/formats/image/parser.py | 8 ++--- .../formats/pdf/pymupdf_subprocess.py | 8 ++--- .../structure/metadata_extractor.py | 10 +++--- .../shared/models/database/job.py | 18 +++------- .../shared/models/database/job_result.py | 7 ++-- .../models/database/job_state_audit_log.py | 7 ++-- .../models/database/job_state_history.py | 7 ++-- .../shared/models/database/webhook.py | 10 ++---- .../shared/models/database/webhook_log.py | 9 ++--- .../services/ai/iloveapi_quota_manager.py | 4 +-- .../ai/openai_compatible_client_sync.py | 8 +++-- .../services/retrieval/agentic/core/trace.py | 14 +++++--- .../services/retrieval/agentic/policy.py | 22 +++++++----- .../services/retrieval/workflow/planner.py | 6 ++-- .../shared/services/storage/result_storage.py | 17 +++++---- .../shared/testing/contract_runtime.py | 36 ++++++++++++------- 18 files changed, 101 insertions(+), 102 deletions(-) diff --git a/apps/worker/app/services/document_parser/formats/atlas/parser.py b/apps/worker/app/services/document_parser/formats/atlas/parser.py index 6c5815c29..a3210d67e 100644 --- a/apps/worker/app/services/document_parser/formats/atlas/parser.py +++ b/apps/worker/app/services/document_parser/formats/atlas/parser.py @@ -143,8 +143,8 @@ def _vlm_extract_page_info(output_dir: str, img_name: str) -> str: if vlm_path != img_path: try: os.remove(vlm_path) - except OSError: - pass + except OSError as exc: + logger.debug(f"Failed to remove compressed atlas image {vlm_path}: {exc}") if result: return result.strip() diff --git a/apps/worker/app/services/document_parser/formats/docx/parser.py b/apps/worker/app/services/document_parser/formats/docx/parser.py index 93554ac89..59c7b297c 100755 --- a/apps/worker/app/services/document_parser/formats/docx/parser.py +++ b/apps/worker/app/services/document_parser/formats/docx/parser.py @@ -531,12 +531,8 @@ def parse_docx( outline_level = outline_dic.get(ele_num, -1) if outline_level > 0: # logger.debug('Found a title: ', text, ' current level: ', outline_level) - try: - last_heading = headings_stack[-1]["heading"] - if last_heading == text: - continue - except Exception: - pass + if headings_stack and headings_stack[-1].get("heading") == text: + continue while headings_stack and headings_stack[-1]["level"] >= outline_level: headings_stack.pop() diff --git a/apps/worker/app/services/document_parser/formats/image/parser.py b/apps/worker/app/services/document_parser/formats/image/parser.py index 357211437..9dd3808ac 100755 --- a/apps/worker/app/services/document_parser/formats/image/parser.py +++ b/apps/worker/app/services/document_parser/formats/image/parser.py @@ -319,20 +319,18 @@ def parse_image( if inferred_suffix.lower() in {".png", ".jpg", ".jpeg", ".gif", ".webp"}: img_name = img_stem or img_name if auto_rename: - update_img_path = os.path.join(img_dir, f"{img_name}{img_suffix}") + target_img_path = os.path.join(img_dir, f"{img_name}{img_suffix}") if os.path.exists(img_path): - if img_path != update_img_path: - os.rename(img_path, update_img_path) + if img_path != target_img_path: + os.rename(img_path, target_img_path) # Store the relative filename for path construction final_img_name = f"{img_name}{img_suffix}" else: logger.warning( f"Image file missing before rename, keeping original name: {filename}" ) - update_img_path = img_path final_img_name = filename else: - update_img_path = img_path final_img_name = filename except KnowhereException: raise diff --git a/apps/worker/app/services/document_parser/formats/pdf/pymupdf_subprocess.py b/apps/worker/app/services/document_parser/formats/pdf/pymupdf_subprocess.py index 88bec26af..fdcf64d35 100644 --- a/apps/worker/app/services/document_parser/formats/pdf/pymupdf_subprocess.py +++ b/apps/worker/app/services/document_parser/formats/pdf/pymupdf_subprocess.py @@ -121,13 +121,13 @@ def _close_result_queue(result_queue: MultiprocessingQueue) -> None: """Release parent-side queue resources once the child result is no longer needed.""" try: result_queue.close() - except Exception: - pass + except Exception as exc: + logger.debug(f"Failed to close PyMuPDF result queue: {exc}") try: result_queue.join_thread() - except Exception: - pass + except Exception as exc: + logger.debug(f"Failed to join PyMuPDF result queue thread: {exc}") def _run_worker_in_spawned_process( diff --git a/apps/worker/app/services/document_parser/structure/metadata_extractor.py b/apps/worker/app/services/document_parser/structure/metadata_extractor.py index 3322f15ac..52d595424 100644 --- a/apps/worker/app/services/document_parser/structure/metadata_extractor.py +++ b/apps/worker/app/services/document_parser/structure/metadata_extractor.py @@ -124,8 +124,8 @@ def extract_md_headings(md_lines: List[str]) -> List[dict]: if text: text_key = normalize_content(text) occurrence_counter[text_key] += 1 - except Exception: - pass + except Exception as exc: + logger.debug(f"Failed to collect table metadata from Markdown: {exc}") continue if is_heading: @@ -212,8 +212,10 @@ def process_block(block, page_idx): "type": "table", } ) - except Exception: - pass + except Exception as exc: + logger.debug( + f"Failed to collect table metadata from layout span: {exc}" + ) continue # Regular text span diff --git a/packages/shared-python/shared/models/database/job.py b/packages/shared-python/shared/models/database/job.py index 39790cefd..8a4fb6bf3 100644 --- a/packages/shared-python/shared/models/database/job.py +++ b/packages/shared-python/shared/models/database/job.py @@ -6,8 +6,7 @@ from datetime import datetime -# Forward references avoid circular imports. -from typing import TYPE_CHECKING, Any, Dict, Optional +from typing import Any, Dict, Optional from uuid import uuid4 from sqlalchemy import ( @@ -28,13 +27,6 @@ from shared.utils.utc_now import utc_now_naive -if TYPE_CHECKING: - from shared.models.database.job_result import JobResult - from shared.models.database.job_state_audit_log import JobStateAuditLog - from shared.models.database.job_state_history import JobStateHistory - from shared.models.database.webhook_log import WebhookLog - - class Job(Base): """Job Model - User API Business Task""" @@ -113,22 +105,22 @@ class Job(Base): # Relationships — default to noload to prevent implicit SELECTs. # Use explicit selectinload() in queries that need related data. - state_history: Mapped[list["JobStateHistory"]] = relationship( + state_history: Mapped[list[Any]] = relationship( "JobStateHistory", back_populates="job", cascade="all, delete-orphan", lazy="noload", ) - state_audit_logs: Mapped[list["JobStateAuditLog"]] = relationship( + state_audit_logs: Mapped[list[Any]] = relationship( "JobStateAuditLog", back_populates="job", cascade="all, delete-orphan", lazy="noload", ) - webhook_logs: Mapped[list["WebhookLog"]] = relationship( + webhook_logs: Mapped[list[Any]] = relationship( "WebhookLog", back_populates="job", cascade="all, delete-orphan", lazy="noload" ) - job_result: Mapped[Optional["JobResult"]] = relationship( + job_result: Mapped[Optional[Any]] = relationship( "JobResult", back_populates="job", uselist=False, lazy="noload" ) diff --git a/packages/shared-python/shared/models/database/job_result.py b/packages/shared-python/shared/models/database/job_result.py index c6d4d9e73..f88f3b98b 100644 --- a/packages/shared-python/shared/models/database/job_result.py +++ b/packages/shared-python/shared/models/database/job_result.py @@ -3,7 +3,7 @@ from __future__ import annotations from datetime import datetime -from typing import TYPE_CHECKING, Any, Dict, List, Optional +from typing import Any, Dict, List, Optional from uuid import uuid4 from sqlalchemy import JSON, DateTime, ForeignKey, Index, Integer, String, Text @@ -12,9 +12,6 @@ from shared.core.database import Base from shared.utils.utc_now import utc_now_naive -if TYPE_CHECKING: - from shared.models.database.job import Job - class JobResult(Base): """Primary job-result table.""" @@ -57,7 +54,7 @@ class JobResult(Base): nullable=False, ) - job: Mapped["Job"] = relationship( + job: Mapped[Any] = relationship( "Job", back_populates="job_result", lazy="joined" ) chunks: Mapped[List["JobChunk"]] = relationship( diff --git a/packages/shared-python/shared/models/database/job_state_audit_log.py b/packages/shared-python/shared/models/database/job_state_audit_log.py index 970b2af57..c2edc4ade 100644 --- a/packages/shared-python/shared/models/database/job_state_audit_log.py +++ b/packages/shared-python/shared/models/database/job_state_audit_log.py @@ -3,7 +3,7 @@ from __future__ import annotations from datetime import datetime -from typing import TYPE_CHECKING, Any, Dict, Optional +from typing import Any, Dict, Optional from sqlalchemy import JSON, DateTime, ForeignKey, Index, Integer, String from sqlalchemy.orm import Mapped, mapped_column, relationship @@ -11,9 +11,6 @@ from shared.core.database import Base from shared.utils.utc_now import utc_now_naive -if TYPE_CHECKING: - from shared.models.database.job import Job - class JobStateAuditLog(Base): """Job state-transition audit log.""" @@ -58,7 +55,7 @@ class JobStateAuditLog(Base): ) # Relationships. - job: Mapped["Job"] = relationship( + job: Mapped[Any] = relationship( "Job", back_populates="state_audit_logs", lazy="select" ) diff --git a/packages/shared-python/shared/models/database/job_state_history.py b/packages/shared-python/shared/models/database/job_state_history.py index b480ae3f7..a564f160b 100644 --- a/packages/shared-python/shared/models/database/job_state_history.py +++ b/packages/shared-python/shared/models/database/job_state_history.py @@ -3,7 +3,7 @@ from __future__ import annotations from datetime import datetime -from typing import TYPE_CHECKING, Any, Dict, Optional +from typing import Any, Dict, Optional from uuid import uuid4 from sqlalchemy import JSON, DateTime, ForeignKey, Index, String @@ -12,9 +12,6 @@ from shared.core.database import Base from shared.utils.utc_now import utc_now_naive -if TYPE_CHECKING: - from shared.models.database.job import Job - class JobStateHistory(Base): """Job state-history model for individual state-machine transitions.""" @@ -44,7 +41,7 @@ class JobStateHistory(Base): ) # Relationships. - job: Mapped["Job"] = relationship("Job", back_populates="state_history") + job: Mapped[Any] = relationship("Job", back_populates="state_history") # Indexes. __table_args__ = ( diff --git a/packages/shared-python/shared/models/database/webhook.py b/packages/shared-python/shared/models/database/webhook.py index 09f7f4dd1..73dc31dc8 100644 --- a/packages/shared-python/shared/models/database/webhook.py +++ b/packages/shared-python/shared/models/database/webhook.py @@ -8,7 +8,7 @@ from __future__ import annotations from datetime import datetime -from typing import TYPE_CHECKING, Any, Dict, Optional +from typing import Any, Dict, Optional from uuid import uuid4 from sqlalchemy import DateTime, ForeignKey, Index, Integer, String @@ -18,10 +18,6 @@ from shared.core.database import Base from shared.utils.utc_now import utc_now_naive -if TYPE_CHECKING: - from shared.models.database.job import Job - from shared.models.database.webhook_log import WebhookLog - class WebhookEventStatus: """Webhook event status constants.""" @@ -83,8 +79,8 @@ class WebhookEvent(Base): ) # Relationships - job: Mapped["Job"] = relationship("Job", lazy="select") - deliveries: Mapped[list["WebhookLog"]] = relationship( + job: Mapped[Any] = relationship("Job", lazy="select") + deliveries: Mapped[list[Any]] = relationship( "WebhookLog", back_populates="event", cascade="all, delete-orphan", # how about the manual trigger? diff --git a/packages/shared-python/shared/models/database/webhook_log.py b/packages/shared-python/shared/models/database/webhook_log.py index 2beb6dca2..cd664680a 100644 --- a/packages/shared-python/shared/models/database/webhook_log.py +++ b/packages/shared-python/shared/models/database/webhook_log.py @@ -7,7 +7,7 @@ from __future__ import annotations from datetime import datetime -from typing import TYPE_CHECKING, Any, Dict, Optional +from typing import Any, Dict, Optional from uuid import uuid4 from sqlalchemy import JSON, DateTime, ForeignKey, Index, Integer, String, Text @@ -16,9 +16,6 @@ from shared.core.database import Base from shared.utils.utc_now import utc_now_naive -if TYPE_CHECKING: - from shared.models.database.job import Job - from shared.models.database.webhook import WebhookEvent class WebhookLog(Base): """Webhook Log Model - Records webhook delivery history.""" @@ -75,8 +72,8 @@ class WebhookLog(Base): ) # Relationships - job: Mapped["Job"] = relationship("Job", back_populates="webhook_logs") - event: Mapped[Optional["WebhookEvent"]] = relationship( + job: Mapped[Any] = relationship("Job", back_populates="webhook_logs") + event: Mapped[Optional[Any]] = relationship( "WebhookEvent", back_populates="deliveries", foreign_keys=[event_id] ) diff --git a/packages/shared-python/shared/services/ai/iloveapi_quota_manager.py b/packages/shared-python/shared/services/ai/iloveapi_quota_manager.py index 8cea78bb5..510719008 100644 --- a/packages/shared-python/shared/services/ai/iloveapi_quota_manager.py +++ b/packages/shared-python/shared/services/ai/iloveapi_quota_manager.py @@ -124,8 +124,8 @@ def parse_tokens_from_settings() -> List[TokenConfig]: ), ) ) - except json.JSONDecodeError: - pass + except json.JSONDecodeError as exc: + logger.debug(f"Invalid I Love API token JSON ignored: {exc}") if not specs and legacy_pub and legacy_sec: specs.append( diff --git a/packages/shared-python/shared/services/ai/openai_compatible_client_sync.py b/packages/shared-python/shared/services/ai/openai_compatible_client_sync.py index 60ea29ae3..79a00e97a 100644 --- a/packages/shared-python/shared/services/ai/openai_compatible_client_sync.py +++ b/packages/shared-python/shared/services/ai/openai_compatible_client_sync.py @@ -390,11 +390,13 @@ def _parse_retry_after(exc: openai.RateLimitError) -> int: """Extract Retry-After seconds from a RateLimitError, with sane bounds.""" try: if hasattr(exc, "response") and exc.response is not None: - header_value = exc.response.headers.get("retry-after") or exc.response.headers.get("Retry-After") + header_value = exc.response.headers.get( + "retry-after" + ) or exc.response.headers.get("Retry-After") if header_value: return max(1, min(int(header_value), 120)) - except (ValueError, TypeError, AttributeError): - pass + except (ValueError, TypeError, AttributeError) as parse_error: + logger.debug(f"Could not parse retry-after header: {parse_error}") return settings.ALI_TOKEN_COOLDOWN_SECONDS diff --git a/packages/shared-python/shared/services/retrieval/agentic/core/trace.py b/packages/shared-python/shared/services/retrieval/agentic/core/trace.py index 08b98e642..bc6ccaa02 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/core/trace.py +++ b/packages/shared-python/shared/services/retrieval/agentic/core/trace.py @@ -105,8 +105,11 @@ async def create_run(self) -> None: logger.debug(f'agentic trace: failed to create run {self._run_id}: {e}') try: await self._db.rollback() - except Exception: - pass + except Exception as rollback_error: + logger.debug( + f'agentic trace: failed to roll back create failure ' + f'{self._run_id}: {rollback_error}' + ) def record_step( self, @@ -213,5 +216,8 @@ async def complete( logger.debug(f'agentic trace: failed to complete run {self._run_id}: {e}') try: await self._db.rollback() - except Exception: - pass + except Exception as rollback_error: + logger.debug( + f'agentic trace: failed to roll back completion failure ' + f'{self._run_id}: {rollback_error}' + ) diff --git a/packages/shared-python/shared/services/retrieval/agentic/policy.py b/packages/shared-python/shared/services/retrieval/agentic/policy.py index 6517fc861..6c201b074 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/policy.py +++ b/packages/shared-python/shared/services/retrieval/agentic/policy.py @@ -27,19 +27,23 @@ def _parse_answer_response(text: str) -> dict[str, Any] | None: """Extract a JSON answer object from LLM response text.""" text = text.strip() - try: - return json.loads(text) - except (json.JSONDecodeError, ValueError): - pass + parsed = _load_json_object(text) + if parsed is not None: + return parsed match = re.search(r'\{.*\}', text, re.DOTALL) if match: - try: - return json.loads(match.group()) - except (json.JSONDecodeError, ValueError): - pass + return _load_json_object(match.group()) return None +def _load_json_object(raw_value: str) -> dict[str, Any] | None: + try: + parsed = json.loads(raw_value) + except (json.JSONDecodeError, ValueError): + return None + return parsed if isinstance(parsed, dict) else None + + def _looks_like_json_wrapper(text: str) -> bool: """Detect malformed JSON-ish answer wrappers without exposing them.""" stripped = text.strip() @@ -78,7 +82,7 @@ def _budget_line_parts(budget_snapshot: dict | None, pool_name: str) -> dict[str min(100, round(remaining_int * 100 / capacity_int)), ) except (TypeError, ValueError): - pass + remaining_pct = 'unknown' return { 'status': pool.get('status', 'HEALTHY'), 'remaining': remaining, diff --git a/packages/shared-python/shared/services/retrieval/workflow/planner.py b/packages/shared-python/shared/services/retrieval/workflow/planner.py index d3a4294d8..82b574677 100644 --- a/packages/shared-python/shared/services/retrieval/workflow/planner.py +++ b/packages/shared-python/shared/services/retrieval/workflow/planner.py @@ -205,10 +205,10 @@ def _extract_json_object(text: str) -> dict[str, Any]: text = text.strip() try: parsed = json.loads(text) - if isinstance(parsed, dict): - return parsed except (ValueError, json.JSONDecodeError): - pass + parsed = None + if isinstance(parsed, dict): + return parsed match = re.search(r"\{.*\}", text, re.DOTALL) if not match: raise ValueError("no JSON object found") diff --git a/packages/shared-python/shared/services/storage/result_storage.py b/packages/shared-python/shared/services/storage/result_storage.py index 650019385..e16b8dab1 100644 --- a/packages/shared-python/shared/services/storage/result_storage.py +++ b/packages/shared-python/shared/services/storage/result_storage.py @@ -1,11 +1,13 @@ from __future__ import annotations import os -from dataclasses import dataclass from collections.abc import Iterator +from dataclasses import dataclass from pathlib import Path from typing import Protocol +from loguru import logger + from shared.services.storage.job_file_storage import JobFileStorage from shared.services.storage.storage_adapter import StorageAdapter @@ -24,13 +26,16 @@ class UploadedResultBundle: class ResultStorage(Protocol): def upload( self, *, job_id: str, result_dir: str, zip_file_path: str - ) -> UploadedResultBundle: ... + ) -> UploadedResultBundle: + raise NotImplementedError def generate_artifact_url( self, *, job_id: str, artifact_ref: str, expires_in: int = 3600 - ) -> str | None: ... + ) -> str | None: + raise NotImplementedError - def normalize_artifact_ref(self, artifact_ref: str | None) -> str | None: ... + def normalize_artifact_ref(self, artifact_ref: str | None) -> str | None: + raise NotImplementedError class JobResultStorage: @@ -156,8 +161,8 @@ def _is_excluded_dir(self, dir_name: str) -> bool: def _cleanup_file(self, file_path: Path) -> None: try: file_path.unlink(missing_ok=True) - except Exception: - pass + except Exception as exc: + logger.debug(f"Failed to clean up result file {file_path}: {exc}") def get_result_storage() -> ResultStorage: diff --git a/packages/shared-python/shared/testing/contract_runtime.py b/packages/shared-python/shared/testing/contract_runtime.py index 2865226a3..ce029a8c6 100644 --- a/packages/shared-python/shared/testing/contract_runtime.py +++ b/packages/shared-python/shared/testing/contract_runtime.py @@ -7,6 +7,7 @@ import subprocess import sys from collections.abc import Awaitable, Callable +from dataclasses import dataclass from pathlib import Path from types import ModuleType from typing import Any, Protocol, cast @@ -77,23 +78,36 @@ CONTRACT_DEVELOPER_USER_EMAIL: str = "local-dev-user@knowhere.local" CONTRACT_DEVELOPER_USER_TIER: str = "tier_5" CONTRACT_DEVELOPER_API_KEY_NAME: str = "contract-developer-api-key" -_contract_storage_prepared: bool = False _contract_storage_database_url: str | None = None _contract_fake_redis_server: fakeredis.FakeServer = fakeredis.FakeServer() +@dataclass +class ContractStorageRuntimeState: + is_prepared: bool = False + + +_contract_storage_runtime_state: ContractStorageRuntimeState = ( + ContractStorageRuntimeState() +) + + class PostgreSQLProcess(Protocol): @property - def host(self) -> str: ... + def host(self) -> str: + raise NotImplementedError @property - def port(self) -> int: ... + def port(self) -> int: + raise NotImplementedError @property - def user(self) -> str: ... + def user(self) -> str: + raise NotImplementedError @property - def password(self) -> str | None: ... + def password(self) -> str | None: + raise NotImplementedError def _ensure_import_paths() -> None: @@ -113,13 +127,12 @@ def _ensure_test_directories() -> None: def _reset_contract_storage_state(database_url: str) -> None: global _contract_storage_database_url - global _contract_storage_prepared if _contract_storage_database_url == database_url: return _contract_storage_database_url = database_url - _contract_storage_prepared = False + _contract_storage_runtime_state.is_prepared = False def _ensure_contract_postgresql_port(database_url: URL) -> None: @@ -596,7 +609,6 @@ def drop_contract_database( postgresql_process: PostgreSQLProcess | None = None, ) -> None: global _contract_storage_database_url - global _contract_storage_prepared contract_database_url = get_contract_database_url(postgresql_process) contract_database_name = make_url(contract_database_url).database @@ -611,7 +623,7 @@ def drop_contract_database( if _contract_storage_database_url == contract_database_url: _contract_storage_database_url = None - _contract_storage_prepared = False + _contract_storage_runtime_state.is_prepared = False def _run_contract_migrations() -> None: @@ -679,16 +691,14 @@ async def _create_contract_engine() -> AsyncEngine: async def prepare_contract_storage() -> None: - global _contract_storage_prepared - _ensure_import_paths() - if not _contract_storage_prepared: + if not _contract_storage_runtime_state.is_prepared: _recreate_contract_database() _initialize_contract_database() _run_contract_migrations() _assert_contract_schema_ready() - _contract_storage_prepared = True + _contract_storage_runtime_state.is_prepared = True await reset_contract_database() await reset_contract_redis()