update: retriever

services/document_service.py

@@ -2,25 +2,25 @@ from pathlib import Path
 from llama_parse import LlamaParse
 import docx
 import pypandoc
-import shutil
+from datetime import datetime, timezone
+from rag_components import get_huggingface_embeddings
+from io import BytesIO
 import os
 import logging
 from langchain_core.documents import Document
 import config
 from db.weaviateDB import connect_to_weaviate
-import utils.utils as utils
+from db.mongoDB import mongo_db
+
+from fastapi.concurrency import run_in_threadpool
 logger = logging.getLogger(__name__)
 
 from rag_components import create_weaviate_schema_if_not_exists, ingest_chunks_with_native_batching
 from utils.process_data import hierarchical_split_law_document,extract_document_metadata,clean_document_text,infer_field, infer_entity_type, filter_and_serialize_complex_metadata
 
 class ApiKeyManager:
-    """
-    Quản lý một danh sách các API key.
-    Cung cấp key tiếp theo trong danh sách và cho phép xoay vòng.
-    """
+    """Quản lý một danh sách các API key."""
     def __init__(self, api_key_string: str):
-        # Tách chuỗi thành danh sách các key, loại bỏ khoảng trắng thừa
         self.keys = [key.strip() for key in api_key_string.split(',') if key.strip()]
         if not self.keys:
             raise ValueError("Chuỗi API key không hợp lệ hoặc rỗng.")
@@ -39,105 +39,171 @@ class ApiKeyManager:
         logger.warning(f"Chuyển sang sử dụng API key tiếp theo (index: {self.current_key_index}).")
         return self.get_key()
 
-# Khởi tạo một instance của Key Manager với chuỗi key từ config
-# Đặt nó ở cấp module để nó duy trì trạng thái qua các lần gọi hàm
-llama_key_manager = ApiKeyManager(config.LLAMA_CLOUD_API_KEYS)
+    def reset(self):
+        """Reset lại index để bắt đầu từ key đầu tiên cho lần xử lý mới."""
+        self.current_key_index = 0
+        logger.info("Key Manager đã được reset.")
 
+llama_key_manager = ApiKeyManager(config.LLAMA_CLOUD_API_KEYS)
 
-def convert_to_text_content(source_path: str) -> str:
-    source_file = Path(source_path)
-    file_extension = source_file.suffix.lower()
-    logger.info(f"Đang trích xuất nội dung từ: {source_file.name}")
+# --- SỬA LẠI HÀM NÀY ĐỂ NHẬN STREAM ---
+def convert_to_text_content(source_stream: BytesIO, original_filename: str) -> str:
+    """Trích xuất nội dung text từ một stream trong bộ nhớ."""
+    file_extension = Path(original_filename).suffix.lower()
+    logger.info(f"Extracting content from: {original_filename}")
     content = ""
 
-    if file_extension == ".pdf":
-        # Bắt đầu vòng lặp để thử các API key
-        while (current_key := llama_key_manager.get_key()) is not None:
-            try:
-                logger.info(f"Đang thử chuyển đổi PDF bằng key index: {llama_key_manager.current_key_index}...")
-                parser = LlamaParse(
-                    api_key=current_key,
-                    result_type="text",
-                    verbose=True,
-                    language="vi"
-                )
-                # parser.load_data có thể ném ra lỗi nếu key không hợp lệ/hết hạn
-                documents = parser.load_data([str(source_file)])
-
-                if documents and documents[0].text.strip():
-                    content = documents[0].text
-                    logger.info(f"✅ Chuyển đổi PDF thành công bằng key index: {llama_key_manager.current_key_index}.")
-                    # Nếu thành công, thoát khỏi vòng lặp
-                    break
-                else:
-                    # Trường hợp hiếm: không có lỗi nhưng nội dung rỗng
-                    raise ValueError("LlamaParse trả về nội dung rỗng.")
-
-            except Exception as e:
-                logger.error(f"❌ Lỗi với key index {llama_key_manager.current_key_index}: {e}")
-                # Nếu có lỗi, thử key tiếp theo trong lần lặp tới
-                if llama_key_manager.get_next_key() is None:
-                    # Nếu đã hết key để thử
-                    logger.critical("Đã thử hết tất cả các API key nhưng đều thất bại.")
-                    raise Exception("Không thể chuyển đổi file PDF sau khi đã thử tất cả các API key.") from e
-                # Nếu còn key, vòng lặp while sẽ tiếp tục với key mới
-
-        # Sau vòng lặp, nếu không có content, nghĩa là đã có lỗi nghiêm trọng
-        if not content:
-             raise ValueError("Không thể trích xuất nội dung PDF sau khi thử các key.")
+    source_stream.seek(0)
 
+    if file_extension == ".pdf":
+        # Do LlamaParse cần đường dẫn file, chúng ta sẽ ghi stream ra file tạm MỘT LẦN
+        # và tái sử dụng đường dẫn này trong vòng lặp thử key.
+
+        # Tạo tên file tạm duy nhất để tránh xung đột khi xử lý song song
+
+        temp_pdf_path = f"/tmp/{original_filename}"
+
+        try:
+            with open(temp_pdf_path, "wb") as f:
+                f.write(source_stream.getvalue())
+
+            # Reset key manager trước khi bắt đầu để đảm bảo nó luôn thử từ key đầu tiên
+            llama_key_manager.reset()
+
+            # Bắt đầu vòng lặp để thử các API key
+            while (current_key := llama_key_manager.get_key()) is not None:
+                try:
+                    logger.info(f"Đang thử chuyển đổi PDF '{original_filename}' bằng key index: {llama_key_manager.current_key_index}...")
+                    parser = LlamaParse(
+                        api_key=current_key,
+                        result_type="text",
+                        verbose=True, # Giữ để debug
+                        language="vi"
+                    )
+                    # Sử dụng đường dẫn file tạm đã tạo
+                    documents = parser.load_data([temp_pdf_path])
+
+                    if documents and documents[0].text.strip():
+                        content = documents[0].text
+                        logger.info(f"✅ Chuyển đổi PDF thành công bằng key index: {llama_key_manager.current_key_index}.")
+                        break # Thành công, thoát khỏi vòng lặp
+                    else:
+                        raise ValueError("LlamaParse trả về nội dung rỗng.")
+
+                except Exception as e:
+                    logger.error(f"❌ Lỗi với key index {llama_key_manager.current_key_index} cho file '{original_filename}': {e}")
+                    if llama_key_manager.get_next_key() is None:
+                        logger.critical("Đã thử hết tất cả các API key nhưng đều thất bại cho file PDF.")
+                        raise Exception(f"Không thể chuyển đổi file '{original_filename}' sau khi đã thử tất cả các API key.") from e
+
+            if not content:
+                raise ValueError(f"Không thể trích xuất nội dung từ PDF '{original_filename}' sau khi thử các key.")
+
+        finally:
+            # Luôn dọn dẹp file tạm, dù thành công hay thất bại
+            if os.path.exists(temp_pdf_path):
+                os.remove(temp_pdf_path)
+                logger.debug(f"Đã dọn dẹp file tạm: {temp_pdf_path}")
 
     elif file_extension == ".docx":
-        doc = docx.Document(source_path)
+        # docx có thể đọc trực tiếp từ stream
+        doc = docx.Document(source_stream)
         content = '\n'.join([para.text for para in doc.paragraphs])
+
     elif file_extension == ".doc":
-        content = pypandoc.convert_file(source_path, 'plain', format='doc')
+        # pypandoc cần file trên đĩa
+
+        temp_doc_path = f"/tmp/{original_filename}"
+        try:
+            with open(temp_doc_path, "wb") as f:
+                f.write(source_stream.getvalue())
+            content = pypandoc.convert_file(temp_doc_path, 'plain', format='doc')
+        finally:
+             if os.path.exists(temp_doc_path):
+                os.remove(temp_doc_path)
+                logger.debug(f"Đã dọn dẹp file tạm: {temp_doc_path}")
+
     else:
         raise ValueError(f"Định dạng file không được hỗ trợ: {file_extension}")
 
     if not content.strip():
-        # Lỗi này chỉ nên xảy ra với docx/doc hoặc nếu LlamaParse thất bại một cách thầm lặng
-        raise ValueError("Nội dung trích xuất bị rỗng.")
+        raise ValueError(f"Nội dung trích xuất từ '{original_filename}' bị rỗng.")
 
-    logger.info(f"✅ Trích xuất nội dung thành công từ {source_file.name}.")
+    logger.info(f"✅ Trích xuất nội dung thành công từ stream của file: {original_filename}.")
     return content
 
-def full_process_and_ingest_pipeline(filepath: str, file_hash: str, embedding_model):
-    filename = os.path.basename(filepath)
-    logger.info(f"BACKGROUND TASK: Starting full pipeline for: {filename} (Hash: {file_hash[:10]}...)")
+
+
+
+async def full_process_and_ingest_pipeline(raw_content: str, filename: str, file_hash: str):
+    """
+    Pipeline xử lý nền đã được tối ưu hoàn toàn.
+    """
+    logger.info(f"BACKGROUND TASK: Starting NLP and Ingestion for: {filename} (Hash: {file_hash[:10]}...)")
     weaviate_client = None
     try:
-        raw_content = convert_to_text_content(filepath)
-
-        doc_metadata = extract_document_metadata(raw_content, filename)
+        embeddings_model = get_huggingface_embeddings(
+            config.EMBEDDING_MODEL_NAME
+        )
+        # Giai đoạn 1: Xử lý NLP (CPU-bound)
+        doc_metadata = await run_in_threadpool(extract_document_metadata, raw_content, filename)
         doc_metadata["source"] = filename
-        cleaned_content = clean_document_text(raw_content)
-        doc_metadata["field"] = infer_field(cleaned_content, doc_metadata.get("ten_van_ban"))
-        doc_metadata["entity_type"] = infer_entity_type(cleaned_content, doc_metadata.get("field", ""))
+
+        cleaned_content = await run_in_threadpool(clean_document_text, raw_content)
+        doc_metadata["field"] = await run_in_threadpool(infer_field, cleaned_content, doc_metadata.get("ten_van_ban"))
+        doc_metadata["entity_type"] = await run_in_threadpool(infer_entity_type, cleaned_content, doc_metadata.get("field", ""))
 
         doc_to_split = Document(page_content=cleaned_content, metadata=doc_metadata)
-        chunks_from_file = hierarchical_split_law_document(doc_to_split)
+        chunks_from_file = await run_in_threadpool(hierarchical_split_law_document, doc_to_split)
 
         if not chunks_from_file:
             raise ValueError("File did not yield any chunks after processing.")
 
         processed_chunks = filter_and_serialize_complex_metadata(chunks_from_file)
 
-        weaviate_client = connect_to_weaviate()
-        embeddings_model = embedding_model
-        collection_name = config.WEAVIATE_COLLECTION_NAME
-        create_weaviate_schema_if_not_exists(weaviate_client, collection_name)
+        # Giai đoạn 2: Ingest vào Weaviate (I/O-bound và CPU-bound)
+        weaviate_client = connect_to_weaviate(run_diagnostics=False)
+
+
+        await run_in_threadpool(create_weaviate_schema_if_not_exists, weaviate_client, config.WEAVIATE_COLLECTION_NAME)
+        await run_in_threadpool(
+            ingest_chunks_with_native_batching,
+            weaviate_client, config.WEAVIATE_COLLECTION_NAME, processed_chunks, embeddings_model
+        )
 
-        ingest_chunks_with_native_batching(weaviate_client, collection_name, processed_chunks, embeddings_model)
+        # Giai đoạn 3: Ghi log thành công
+        await log_processed_hash(file_hash, filename)
+        logger.info(f"✅✅✅ All tasks completed for '{filename}'.")
 
-        utils.log_processed_hash(file_hash)
-        logger.info(f"✅ Successfully ingested '{filename}'.")
-        shutil.move(filepath, os.path.join(config.PROCESSED_FILES_FOLDER, filename))
-        logger.info(f"Moved '{filename}' to processed folder.")
     except Exception as e:
         logger.error(f"❌ FAILED pipeline for '{filename}': {e}", exc_info=True)
-        shutil.move(filepath, os.path.join(config.FAILED_FILES_FOLDER, filename))
-        logger.info(f"Moved '{filename}' to failed folder.")
+        await log_failed_process(file_hash, filename, str(e))
     finally:
         if weaviate_client and weaviate_client.is_connected():
-            weaviate_client.close()
+            weaviate_client.close()
+
+
+async def log_processed_hash(file_hash: str, filename: str, status: str = "SUCCESS", error_message: str = None):
+    """Ghi lại trạng thái xử lý (thành công hoặc thất bại) vào MongoDB."""
+    try:
+        record = {
+            "file_hash": file_hash,
+            "original_filename": filename,
+            "processed_at": datetime.now(timezone.utc),
+            "status": status,
+        }
+        if error_message:
+            record["error_message"] = error_message
+
+        await mongo_db.processed_documents.insert_one(record)
+    except Exception as e:
+        logger.error(f"Could not write process record to MongoDB for hash {file_hash}: {e}")
+
+# Wrapper cho việc log lỗi
+async def log_failed_process(file_hash: str, filename: str, error_message: str):
+    await log_processed_hash(file_hash, filename, status="FAILED", error_message=error_message)
+
+# Hàm kiểm tra trùng lặp
+async def check_if_hash_exists(file_hash: str) -> bool:
+    count = await mongo_db.processed_documents.count_documents({"file_hash": file_hash, "status": "SUCCESS"})
+    return count > 0