79 lines
2.9 KiB
Python
79 lines
2.9 KiB
Python
import logging
|
|
import sys
|
|
|
|
from core.config import settings
|
|
from ingestion.providers.sharepoint_provider import SharePointProvider
|
|
from extraction.ocr_service import OCRService
|
|
from chunking.markdown_chunker import MarkdownChunker
|
|
from indexing.vector_store import VectorStore
|
|
|
|
logging.basicConfig(level=logging.INFO, format="%(levelname)s:%(name)s:%(message)s")
|
|
logger = logging.getLogger("RAGPipeline")
|
|
|
|
def run_pipeline():
|
|
logger.info("=== BẮT ĐẦU TEST TOÀN BỘ ĐƯỜNG ỐNG RAG ===")
|
|
|
|
# Ép buộc dùng localhost cho OpenSearch khi chạy trực tiếp trên WSL
|
|
if settings.opensearch_host == "opensearch":
|
|
settings.opensearch_host = "localhost"
|
|
|
|
# 1. Tầng Ingestion
|
|
logger.info("\n--- BƯỚC 1: Lấy file từ SharePoint ---")
|
|
provider = SharePointProvider()
|
|
items, _ = provider.fetch_changes({})
|
|
|
|
target_item = None
|
|
for item in items:
|
|
if item.get("name", "").lower().endswith(".pdf"):
|
|
target_item = item
|
|
break
|
|
|
|
if not target_item:
|
|
logger.error("Không tìm thấy file PDF nào trên SharePoint để test!")
|
|
sys.exit(1)
|
|
|
|
logger.info(f"Đã chọn file: {target_item['name']}. Đang tải...")
|
|
pdf_bytes = provider.download_file(target_item)
|
|
logger.info(f"Tải thành công {len(pdf_bytes)} bytes.")
|
|
|
|
# 2. Tầng Extraction (VLM)
|
|
logger.info("\n--- BƯỚC 2: OCR / VLM Trích xuất Markdown ---")
|
|
ocr = OCRService()
|
|
pages = ocr.process_pdf_bytes(pdf_bytes)
|
|
|
|
if not pages:
|
|
logger.error("VLM không trích xuất được nội dung nào!")
|
|
sys.exit(1)
|
|
|
|
logger.info(f"VLM đã trích xuất thành công {len(pages)} trang.")
|
|
|
|
# 3. Tầng Chunking
|
|
logger.info("\n--- BƯỚC 3: Băm nhỏ văn bản (Semantic Chunking) ---")
|
|
chunker = MarkdownChunker(max_chunk_size=1000, overlap=100)
|
|
|
|
# Tạo metadata giả lập để lưu vào Chunk
|
|
metadata = {
|
|
"item_id": target_item["id"],
|
|
"name": target_item["name"],
|
|
"web_url": "https://285pdg.sharepoint.com/...",
|
|
"site_id": settings.sharepoint_site_id
|
|
}
|
|
|
|
chunks = chunker.chunk_document(pages, metadata)
|
|
logger.info(f"Đã băm thành {len(chunks)} chunks độc lập.")
|
|
if chunks:
|
|
logger.info(f"Ví dụ Chunk đầu tiên:\n[ID: {chunks[0].chunk_id}] {chunks[0].text[:150]}...")
|
|
|
|
# 4. Tầng Vector Database (OpenSearch)
|
|
logger.info("\n--- BƯỚC 4: Mã hóa Vector & Indexing ---")
|
|
try:
|
|
vector_db = VectorStore(index_name="poc_sharepoint_docs")
|
|
vector_db.embed_and_index(chunks)
|
|
logger.info("🎉 CHÚC MỪNG! DỮ LIỆU ĐÃ NẰM TRONG OPENSEARCH SẴN SÀNG ĐỂ CHAT!")
|
|
except Exception as e:
|
|
logger.error(f"LỖI trong quá trình Embedding / Indexing: {e}")
|
|
logger.warning("Gợi ý: Hãy chắc chắn Docker OpenSearch đang chạy trên cổng 9200!")
|
|
|
|
if __name__ == "__main__":
|
|
run_pipeline()
|