"""Adapt legacy processors to committed document-indexing phases.""" from concurrent.futures import ThreadPoolExecutor from typing import cast from flask import Flask, current_app from sqlalchemy.orm import Session, sessionmaker from werkzeug.exceptions import HTTPException from werkzeug.local import LocalProxy from core.errors.error import ( AppInvokeQuotaExceededError, InvokeRateLimitError, LLMError, ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError, ) from core.model_manager import ModelManager from core.plugin.impl.exc import PluginDaemonError from core.rag.datasource.keyword.jieba.jieba import Jieba from core.rag.datasource.vdb.vector_factory import Vector from core.rag.embedding.token_counter import calculate_segment_token_counts from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType from core.rag.index_processor.index_processor import IndexProcessorFactory from core.rag.models.document import Document from extensions.application_services.data_sources import build_data_source_credentials from extensions.ext_redis import redis_client from extensions.otel import propagate_context from graphon.model_runtime.entities.model_entities import ModelType from libs import helper from models.enums import DataSourceType from repositories.knowledge.document_repository import SQLAlchemyDocumentRepository from repositories.knowledge.segment_repository import SQLAlchemySegmentRepository from repositories.knowledge.upload_file_repository import SQLAlchemyKnowledgeUploadRepository from services.knowledge.indexing.adapters.sources import ( CompositeStoredSourceResolver, FileSourceAdapter, NotionSourceResolver, WebsiteSourceAdapter, ) from services.knowledge.indexing.errors import DocumentIsPausedError from services.knowledge.indexing.execution import DocumentIndexingService, IndexingDocument from services.knowledge.resource_scope import DocumentRef from services.vector_space_admission_service import VectorSpaceAdmissionService class IndexingExecutionAdapter: def __init__( self, *, session_factory: sessionmaker[Session], documents: SQLAlchemyDocumentRepository, segments: SQLAlchemySegmentRepository, uploads: SQLAlchemyKnowledgeUploadRepository, sources: CompositeStoredSourceResolver, ) -> None: self._session_factory = session_factory self._documents = documents self._segments = segments self._uploads = uploads self._sources = sources def describe_error(self, error: Exception) -> str: """Preserve provider and extractor descriptions without their transport prefixes.""" if isinstance( error, ( LLMError, ProviderTokenNotInitError, QuotaExceededError, AppInvokeQuotaExceededError, ModelCurrentlyNotSupportError, InvokeRateLimitError, PluginDaemonError, HTTPException, ), ): return str(error.description) return str(error) def check_paused(self, ref: DocumentRef) -> None: if redis_client.get(f"document_{ref.document_id}_is_paused"): raise DocumentIsPausedError(f"Document paused, document id: {ref.document_id}") def extract(self, document: IndexingDocument) -> list[Document]: if not document.processing_rule: raise ValueError("no process rule found") setting = self._sources.resolve_for_indexing(document.source) processor = IndexProcessorFactory(document.doc_form).init_index_processor() # Legacy extractors persist downloaded images and Notion source metadata. # Their session is independent of document-state write transactions. with self._session_factory() as session: texts = processor.extract(setting, process_rule_mode=document.processing_rule["mode"], session=session) session.commit() for text in texts: text.metadata.update(document_id=document.ref.document_id, dataset_id=document.ref.dataset.dataset_id) return texts def transform(self, document: IndexingDocument, texts: list[Document]) -> list[Document]: dataset, _ = self._documents.get_indexing_models(document.ref) user = self._documents.get_indexing_user(document.ref) model = None if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: manager = ModelManager.for_tenant(tenant_id=dataset.tenant_id) if dataset.embedding_model_provider: model = manager.get_model_instance( tenant_id=dataset.tenant_id, provider=dataset.embedding_model_provider, model_type=ModelType.TEXT_EMBEDDING, model=dataset.embedding_model, ) else: model = manager.get_default_model_instance( tenant_id=dataset.tenant_id, model_type=ModelType.TEXT_EMBEDDING ) processor = IndexProcessorFactory(document.doc_form).init_index_processor() with self._session_factory() as session: chunks = processor.transform( texts, user, embedding_model_instance=model, process_rule=document.processing_rule, tenant_id=dataset.tenant_id, doc_language=document.doc_language, session=session, ) session.commit() return chunks def count_tokens(self, document: IndexingDocument, chunks: list[Document]) -> list[int]: dataset, _ = self._documents.get_indexing_models(document.ref) return calculate_segment_token_counts(dataset=dataset, documents=chunks) def ensure_admission(self, document: IndexingDocument, chunks: list[Document]) -> None: dataset, _ = self._documents.get_indexing_models(document.ref) with self._session_factory() as session: VectorSpaceAdmissionService().ensure_document_can_be_indexed( dataset=dataset, document_id=document.ref.document_id, doc_form=document.doc_form, documents=chunks, include_summaries=document.need_summary, session=session, ) def load(self, document: IndexingDocument, chunks: list[Document]) -> None: dataset, _ = self._documents.get_indexing_models(document.ref) app = cast(LocalProxy[Flask], current_app)._get_current_object() if dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: groups: list[list[Document]] = [[] for _ in range(10)] for chunk in chunks: groups[int(helper.generate_text_hash(chunk.page_content), 16) % 10].append(chunk) with ThreadPoolExecutor(max_workers=10) as executor: futures = [ executor.submit( propagate_context(self._process_chunk), app, document.ref, group, False, ) for group in groups if group ] for future in futures: future.result() elif ( dataset.indexing_technique == IndexTechniqueType.ECONOMY and document.doc_form != IndexStructureType.PARENT_CHILD_INDEX ): # A future propagates keyword failures; Thread.join() used to hide # them and incorrectly let the document become completed. with ThreadPoolExecutor(max_workers=1) as executor: executor.submit(propagate_context(self._process_chunk), app, document.ref, chunks, True).result() def _process_chunk(self, app: Flask, ref: DocumentRef, chunks: list[Document], keywords: bool) -> None: with app.app_context(): self.check_paused(ref) # Workers load their own inputs and own their sessions. No ORM # instance or session crosses the executor boundary. dataset, document = self._documents.get_indexing_models(ref) if keywords: # Keep the read/modify/write lock while each repository call # closes its transaction before keyword extraction or storage I/O. Jieba(dataset).update_texts( chunks, read=lambda: self._segments.get_keyword_table(ref.dataset), write=lambda storage_type, data, selected: self._segments.save_keyword_table( ref.dataset, storage_type=storage_type, data=data, keywords=selected ), lock=redis_client.lock(f"keyword_indexing_lock_{dataset.id}", timeout=600), ) elif dataset.indexing_technique == IndexTechniqueType.HIGH_QUALITY: with self._session_factory() as session: vector_type = Vector.resolve_vector_type(dataset, session=session) attachments = ( [attachment for chunk in chunks for attachment in chunk.attachments or []] if dataset.is_multimodal else [] ) files = self._uploads.get_files( workspace_id=ref.dataset.tenant_id, file_ids=[attachment.metadata["doc_id"] for attachment in attachments], ) # All database inputs are materialized before initializing the # vector backend, embedding text/images, or reading object storage. vector = Vector(dataset, session=None, vector_type=vector_type) if document.doc_form == IndexStructureType.PARENT_CHILD_INDEX: for chunk in chunks: if chunk.children: vector.create([Document.model_validate(child.model_dump()) for child in chunk.children]) else: vector.create(chunks) if attachments: vector.create_multimodal(attachments, upload_files=files) self.check_paused(ref) self._segments.complete_indexing_segments(ref, [chunk.metadata["doc_id"] for chunk in chunks]) def build_document_indexing_service( *, session_factory: sessionmaker[Session], enforce_vector_space_admission: bool = False, ) -> DocumentIndexingService: """Assemble a fresh indexing use case from explicit database dependencies.""" documents = SQLAlchemyDocumentRepository(session_factory=session_factory) segments = SQLAlchemySegmentRepository(session_factory=session_factory) uploads = SQLAlchemyKnowledgeUploadRepository(session_factory=session_factory) credentials = build_data_source_credentials(database_client=session_factory) sources = CompositeStoredSourceResolver( adapters={ DataSourceType.UPLOAD_FILE: FileSourceAdapter(uploads=uploads), DataSourceType.NOTION_IMPORT: NotionSourceResolver( actor_credentials=credentials.actor, stored_credentials=credentials.stored ), DataSourceType.WEBSITE_CRAWL: WebsiteSourceAdapter(), } ) return DocumentIndexingService( documents=documents, segments=segments, backend=IndexingExecutionAdapter( session_factory=session_factory, documents=documents, segments=segments, uploads=uploads, sources=sources ), enforce_vector_space_admission=enforce_vector_space_admission, )