import logging import os import hashlib import tempfile import traceback import json from app.db.session import SessionLocal from io import BytesIO from typing import Optional, List, Dict, Any from fastapi import UploadFile from langchain_community.document_loaders import ( PyPDFLoader, Docx2txtLoader, UnstructuredMarkdownLoader, TextLoader ) from langchain.text_splitter import RecursiveCharacterTextSplitter from langchain_core.documents import Document as LangchainDocument from pydantic import BaseModel from sqlalchemy.orm import Session from app.core.config import settings from app.core.minio import get_minio_client from app.models.knowledge import ProcessingTask, Document, DocumentChunk from app.services.chunk_record import ChunkRecord from minio.error import MinioException from minio.commonconfig import CopySource from app.services.vector_store import VectorStoreFactory from app.services.embedding.embedding_factory import EmbeddingsFactory class UploadResult(BaseModel): file_path: str file_name: str file_size: int content_type: str file_hash: str class TextChunk(BaseModel): content: str metadata: Optional[Dict] = None class PreviewResult(BaseModel): chunks: List[TextChunk] total_chunks: int def _estimate_token_count(text: str) -> int: # Lightweight estimation without adding tokenizer dependencies. return len(text) def _build_enriched_chunk_metadata( *, source_metadata: Optional[Dict[str, Any]], chunk_id: str, file_name: str, file_path: str, kb_id: int, document_id: int, chunk_index: int, chunk_text: str, ) -> Dict[str, Any]: source_metadata = source_metadata or {} token_count = _estimate_token_count(chunk_text) return { **source_metadata, "source": file_name, "chunk_id": chunk_id, "file_name": file_name, "file_path": file_path, "kb_id": kb_id, "document_id": document_id, "chunk_index": chunk_index, "chunk_text": chunk_text, "token_count": token_count, "language": source_metadata.get("language", "zh"), "source_type": "document", "mission_phase": source_metadata.get("mission_phase"), "section_title": source_metadata.get("section_title"), "publish_time": source_metadata.get("publish_time"), # Keep graph-linked fields for future graph/vector federation. "extracted_entities": source_metadata.get("extracted_entities", []), "extracted_entity_types": source_metadata.get("extracted_entity_types", []), "extracted_relations": source_metadata.get("extracted_relations", []), "graph_node_ids": source_metadata.get("graph_node_ids", []), "graph_edge_ids": source_metadata.get("graph_edge_ids", []), "community_ids": source_metadata.get("community_ids", []), } def _sanitize_metadata_for_vector_store(metadata: Optional[Dict[str, Any]]) -> Dict[str, Any]: """Normalize metadata to satisfy Chroma's strict metadata constraints.""" if not metadata: return {} sanitized: Dict[str, Any] = {} scalar_types = (str, int, float, bool) for key, value in metadata.items(): if value is None: continue if isinstance(value, scalar_types): sanitized[key] = value continue if isinstance(value, list): primitive_items = [item for item in value if isinstance(item, scalar_types)] if primitive_items: sanitized[key] = primitive_items elif value: sanitized[key] = json.dumps(value, ensure_ascii=False) continue if isinstance(value, dict): sanitized[key] = json.dumps(value, ensure_ascii=False) continue sanitized[key] = str(value) return sanitized async def process_document(file_path: str, file_name: str, kb_id: int, document_id: int, chunk_size: int = 1000, chunk_overlap: int = 200) -> None: """Process document and store in vector database with incremental updates""" logger = logging.getLogger(__name__) try: preview_result = await preview_document(file_path, chunk_size, chunk_overlap) # Initialize embeddings logger.info("Initializing OpenAI embeddings...") embeddings = EmbeddingsFactory.create() logger.info(f"Initializing vector store with collection: kb_{kb_id}") vector_store = VectorStoreFactory.create( store_type=settings.VECTOR_STORE_TYPE, collection_name=f"kb_{kb_id}", embedding_function=embeddings, ) # Initialize chunk record manager chunk_manager = ChunkRecord(kb_id) # Get existing chunk hashes for this file existing_hashes = chunk_manager.list_chunks(file_name) # Prepare new chunks new_chunks = [] current_hashes = set() documents_to_update = [] for i, chunk in enumerate(preview_result.chunks): # Calculate chunk hash chunk_hash = hashlib.sha256( (chunk.content + str(chunk.metadata)).encode() ).hexdigest() current_hashes.add(chunk_hash) # Skip if chunk hasn't changed if chunk_hash in existing_hashes: continue # Create unique ID for the chunk chunk_id = hashlib.sha256( f"{kb_id}:{file_name}:{chunk_hash}".encode() ).hexdigest() metadata = _build_enriched_chunk_metadata( source_metadata=chunk.metadata, chunk_id=chunk_id, file_name=file_name, file_path=file_path, kb_id=kb_id, document_id=document_id, chunk_index=i, chunk_text=chunk.content, ) vector_metadata = _sanitize_metadata_for_vector_store(metadata) new_chunks.append({ "id": chunk_id, "kb_id": kb_id, "document_id": document_id, "file_name": file_name, "metadata": metadata, "hash": chunk_hash }) # Prepare document for vector store doc = LangchainDocument( page_content=chunk.content, metadata=vector_metadata ) documents_to_update.append(doc) # Add new chunks to database and vector store if new_chunks: logger.info(f"Adding {len(new_chunks)} new/updated chunks") chunk_manager.add_chunks(new_chunks) vector_store.add_documents(documents_to_update) if settings.GRAPHRAG_ENABLED: try: from app.services.graph.graphrag_adapter import GraphRAGAdapter graph_adapter = GraphRAGAdapter() source_texts = [doc.page_content for doc in documents_to_update if doc.page_content.strip()] await graph_adapter.ingest_texts(kb_id, source_texts) logger.info("GraphRAG ingestion completed in incremental processing") except Exception as graph_exc: logger.error(f"GraphRAG ingestion failed in incremental processing: {graph_exc}") # Delete removed chunks chunks_to_delete = chunk_manager.get_deleted_chunks(current_hashes, file_name) if chunks_to_delete: logger.info(f"Removing {len(chunks_to_delete)} deleted chunks") chunk_manager.delete_chunks(chunks_to_delete) vector_store.delete(chunks_to_delete) logger.info("Document processing completed successfully") except Exception as e: logger.error(f"Error processing document: {str(e)}") raise async def upload_document(file: UploadFile, kb_id: int) -> UploadResult: """Step 1: Upload document to MinIO""" content = await file.read() file_size = len(content) file_hash = hashlib.sha256(content).hexdigest() # Clean and normalize filename file_name = "".join(c for c in file.filename if c.isalnum() or c in ('-', '_', '.')).strip() object_path = f"kb_{kb_id}/{file_name}" content_types = { ".pdf": "application/pdf", ".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document", ".md": "text/markdown", ".txt": "text/plain" } _, ext = os.path.splitext(file_name) content_type = content_types.get(ext.lower(), "application/octet-stream") # Upload to MinIO minio_client = get_minio_client() try: minio_client.put_object( bucket_name=settings.MINIO_BUCKET_NAME, object_name=object_path, data=BytesIO(content), length=file_size, content_type=content_type ) except Exception as e: logging.error(f"Failed to upload file to MinIO: {str(e)}") raise return UploadResult( file_path=object_path, file_name=file_name, file_size=file_size, content_type=content_type, file_hash=file_hash ) async def preview_document(file_path: str, chunk_size: int = 1000, chunk_overlap: int = 200) -> PreviewResult: """Step 2: Generate preview chunks""" # Get file from MinIO minio_client = get_minio_client() _, ext = os.path.splitext(file_path) ext = ext.lower() # Download to temp file with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as temp_file: minio_client.fget_object( bucket_name=settings.MINIO_BUCKET_NAME, object_name=file_path, file_path=temp_file.name ) temp_path = temp_file.name try: # Select appropriate loader if ext == ".pdf": loader = PyPDFLoader(temp_path) elif ext == ".docx": loader = Docx2txtLoader(temp_path) elif ext == ".md": loader = UnstructuredMarkdownLoader(temp_path) else: # Default to text loader loader = TextLoader(temp_path) # Load and split the document documents = loader.load() text_splitter = RecursiveCharacterTextSplitter( chunk_size=chunk_size, chunk_overlap=chunk_overlap ) chunks = text_splitter.split_documents(documents) # Convert to preview format preview_chunks = [ TextChunk( content=chunk.page_content, metadata=chunk.metadata ) for chunk in chunks ] return PreviewResult( chunks=preview_chunks, total_chunks=len(chunks) ) finally: os.unlink(temp_path) async def process_document_background( temp_path: str, file_name: str, kb_id: int, task_id: int, db: Session = None, chunk_size: int = 1000, chunk_overlap: int = 200 ) -> None: """Process document in background""" logger = logging.getLogger(__name__) logger.info(f"Starting background processing for task {task_id}, file: {file_name}") # if we don't pass in db, create a new database session if db is None: db = SessionLocal() should_close_db = True else: should_close_db = False task = db.query(ProcessingTask).get(task_id) if not task: logger.error(f"Task {task_id} not found") return minio_client = None local_temp_path = None try: logger.info(f"Task {task_id}: Setting status to processing") task.status = "processing" db.commit() # 1. 从临时目录下载文件 minio_client = get_minio_client() try: local_temp_path = f"/tmp/temp_{task_id}_{file_name}" # 使用系统临时目录 logger.info(f"Task {task_id}: Downloading file from MinIO: {temp_path} to {local_temp_path}") minio_client.fget_object( bucket_name=settings.MINIO_BUCKET_NAME, object_name=temp_path, file_path=local_temp_path ) logger.info(f"Task {task_id}: File downloaded successfully") except MinioException as e: # Idempotent fallback: temp object may already be consumed by another task. # If the final document is already created, treat current task as completed. if "NoSuchKey" in str(e) and task.document_upload: existing_document = db.query(Document).filter( Document.knowledge_base_id == kb_id, Document.file_name == file_name, Document.file_hash == task.document_upload.file_hash, ).first() if existing_document: logger.warning( f"Task {task_id}: Temp object missing but document already exists, " f"marking task as completed (document_id={existing_document.id})" ) task.status = "completed" task.document_id = existing_document.id task.error_message = None task.document_upload.status = "completed" task.document_upload.error_message = None db.commit() return error_msg = f"Failed to download temp file: {str(e)}" logger.error(f"Task {task_id}: {error_msg}") raise Exception(error_msg) try: # 2. 加载和分块文档 _, ext = os.path.splitext(file_name) ext = ext.lower() logger.info(f"Task {task_id}: Loading document with extension {ext}") # 选择合适的加载器 if ext == ".pdf": loader = PyPDFLoader(local_temp_path) elif ext == ".docx": loader = Docx2txtLoader(local_temp_path) elif ext == ".md": loader = UnstructuredMarkdownLoader(local_temp_path) else: # 默认使用文本加载器 loader = TextLoader(local_temp_path) logger.info(f"Task {task_id}: Loading document content") documents = loader.load() logger.info(f"Task {task_id}: Document loaded successfully") logger.info(f"Task {task_id}: Splitting document into chunks") text_splitter = RecursiveCharacterTextSplitter( chunk_size=chunk_size, chunk_overlap=chunk_overlap ) chunks = text_splitter.split_documents(documents) logger.info(f"Task {task_id}: Document split into {len(chunks)} chunks") # 3. 创建向量存储 logger.info(f"Task {task_id}: Initializing vector store") embeddings = EmbeddingsFactory.create() vector_store = VectorStoreFactory.create( store_type=settings.VECTOR_STORE_TYPE, collection_name=f"kb_{kb_id}", embedding_function=embeddings, ) # 4. 将临时文件移动到永久目录 permanent_path = f"kb_{kb_id}/{file_name}" try: logger.info(f"Task {task_id}: Moving file to permanent storage") # 复制到永久目录 source = CopySource(settings.MINIO_BUCKET_NAME, temp_path) minio_client.copy_object( bucket_name=settings.MINIO_BUCKET_NAME, object_name=permanent_path, source=source ) logger.info(f"Task {task_id}: File moved to permanent storage") # 删除临时文件 logger.info(f"Task {task_id}: Removing temporary file from MinIO") minio_client.remove_object( bucket_name=settings.MINIO_BUCKET_NAME, object_name=temp_path ) logger.info(f"Task {task_id}: Temporary file removed") except MinioException as e: error_msg = f"Failed to move file to permanent storage: {str(e)}" logger.error(f"Task {task_id}: {error_msg}") raise Exception(error_msg) # 5. 创建文档记录 logger.info(f"Task {task_id}: Creating document record") document = Document( file_name=file_name, file_path=permanent_path, file_hash=task.document_upload.file_hash, file_size=task.document_upload.file_size, content_type=task.document_upload.content_type, knowledge_base_id=kb_id ) db.add(document) db.flush() db.refresh(document) logger.info(f"Task {task_id}: Document record created with ID {document.id}") # 6. 存储文档块 logger.info(f"Task {task_id}: Storing document chunks") for i, chunk in enumerate(chunks): # 为每个 chunk 生成唯一的 ID chunk_id = hashlib.sha256( f"{kb_id}:{file_name}:{chunk.page_content}".encode() ).hexdigest() metadata = _build_enriched_chunk_metadata( source_metadata=chunk.metadata, chunk_id=chunk_id, file_name=file_name, file_path=permanent_path, kb_id=kb_id, document_id=document.id, chunk_index=i, chunk_text=chunk.page_content, ) chunk.metadata = metadata doc_chunk = DocumentChunk( id=chunk_id, # 添加 ID 字段 document_id=document.id, kb_id=kb_id, file_name=file_name, chunk_metadata={ "page_content": chunk.page_content, **metadata }, hash=hashlib.sha256( (chunk.page_content + str(metadata)).encode() ).hexdigest() ) db.add(doc_chunk) if i > 0 and i % 100 == 0: logger.info(f"Task {task_id}: Stored {i} chunks") db.flush() # 7. 添加到向量存储 logger.info(f"Task {task_id}: Adding chunks to vector store") vector_chunks = [ LangchainDocument( page_content=chunk.page_content, metadata=_sanitize_metadata_for_vector_store(chunk.metadata), ) for chunk in chunks ] vector_store.add_documents(vector_chunks) # 移除 persist() 调用，因为新版本不需要 logger.info(f"Task {task_id}: Chunks added to vector store") if settings.GRAPHRAG_ENABLED: try: from app.services.graph.graphrag_adapter import GraphRAGAdapter logger.info(f"Task {task_id}: Starting GraphRAG ingestion") graph_adapter = GraphRAGAdapter() source_texts = [doc.page_content for doc in documents if doc.page_content.strip()] await graph_adapter.ingest_texts(kb_id, source_texts) logger.info(f"Task {task_id}: GraphRAG ingestion completed") except Exception as graph_exc: logger.error(f"Task {task_id}: GraphRAG ingestion failed: {graph_exc}") # 8. 更新任务状态 logger.info(f"Task {task_id}: Updating task status to completed") task.status = "completed" task.document_id = document.id # 更新为新创建的文档ID # 9. 更新上传记录状态 upload = task.document_upload # 直接通过关系获取 if upload: logger.info(f"Task {task_id}: Updating upload record status to completed") upload.status = "completed" db.commit() logger.info(f"Task {task_id}: Processing completed successfully") finally: # 清理本地临时文件 try: if os.path.exists(local_temp_path): logger.info(f"Task {task_id}: Cleaning up local temp file") os.remove(local_temp_path) logger.info(f"Task {task_id}: Local temp file cleaned up") except Exception as e: logger.warning(f"Task {task_id}: Failed to clean up local temp file: {str(e)}") except Exception as e: logger.error(f"Task {task_id}: Error processing document: {str(e)}") logger.error(f"Task {task_id}: Stack trace: {traceback.format_exc()}") db.rollback() failed_task = db.query(ProcessingTask).get(task_id) if failed_task: failed_task.status = "failed" failed_task.error_message = str(e) if failed_task.document_upload: failed_task.document_upload.status = "failed" failed_task.document_upload.error_message = str(e) db.commit() # 清理临时文件 try: logger.info(f"Task {task_id}: Cleaning up temporary file after error") if minio_client is not None: minio_client.remove_object( bucket_name=settings.MINIO_BUCKET_NAME, object_name=temp_path ) logger.info(f"Task {task_id}: Temporary file cleaned up after error") except: logger.warning(f"Task {task_id}: Failed to clean up temporary file after error") finally: # if we create the db session, we need to close it if should_close_db and db: db.close()