from langchain_huggingface import HuggingFaceEmbeddings import config import prompt_templete from langchain_core.prompts import ChatPromptTemplate from langchain_core.runnables import RunnableLambda, RunnablePassthrough from langchain_core.documents import Document import logging from langchain_core.output_parsers import StrOutputParser from typing import List,Any,Dict from langchain_weaviate.vectorstores import WeaviateVectorStore from langchain_google_genai import ChatGoogleGenerativeAI from utils.process_data import filter_and_serialize_complex_metadata import weaviate import weaviate.classes.config as wvc_config from weaviate.exceptions import WeaviateQueryException from google.api_core.exceptions import ResourceExhausted, PermissionDenied import time import json import re from operator import itemgetter logger = logging.getLogger(__name__) WEAVIATE_SCHEMA_CONFIG: List[Dict[str, Any]] = [ # Tên trường, Kiểu dữ liệu trong Weaviate, Có nên vector hóa trường này không? {"name": "source", "dataType": wvc_config.DataType.TEXT,"index_searchable": False, "vectorize": False}, {"name": "title", "dataType": wvc_config.DataType.TEXT, "index_searchable": True, "tokenization": wvc_config.Tokenization.WORD, "vectorize": True}, {"name": "field", "dataType": wvc_config.DataType.TEXT,"index_searchable": True, "vectorize": True}, {"name": "so_hieu", "dataType": wvc_config.DataType.TEXT, "index_searchable": False,"vectorize": False}, {"name": "loai_van_ban", "dataType": wvc_config.DataType.TEXT, "index_searchable": True,"vectorize": True}, {"name": "ten_van_ban", "dataType": wvc_config.DataType.TEXT,"index_searchable": True, "tokenization": wvc_config.Tokenization.WORD, "vectorize": True}, {"name": "co_quan_ban_hanh", "dataType": wvc_config.DataType.TEXT, "index_searchable": False,"vectorize": False}, {"name": "ngay_ban_hanh_str", "dataType": wvc_config.DataType.TEXT,"index_searchable": False, "vectorize": False}, {"name": "nam_ban_hanh", "dataType": wvc_config.DataType.INT,"index_searchable": True, "vectorize": False}, {"name": "phan_code", "dataType": wvc_config.DataType.TEXT,"index_searchable": False, "vectorize": False}, {"name": "chuong_code", "dataType": wvc_config.DataType.TEXT, "index_searchable": False,"vectorize": False}, {"name": "muc_code", "dataType": wvc_config.DataType.TEXT,"index_searchable": False, "vectorize": False}, {"name": "dieu_code", "dataType": wvc_config.DataType.TEXT,"index_searchable": False, "vectorize": False}, {"name": "entity_type", "dataType": wvc_config.DataType.TEXT,"index_searchable": True, "vectorize": False}, {"name": "penalties", "dataType": wvc_config.DataType.TEXT,"index_searchable": False, "vectorize": False}, {"name": "cross_references", "dataType": wvc_config.DataType.TEXT, "index_searchable": False, "vectorize": False}, ] # Hàm get_huggingface_embeddings giữ nguyên def get_huggingface_embeddings(model_name: str, device: str = 'cpu'): logger.info(f"🔸Đang khởi tạo model embedding: {model_name} trên thiết bị {device}...") model_kwargs = { 'device': device, 'trust_remote_code': True # thêm để đảm bảo load được những model custom } encode_kwargs = { 'batch_size': 32, # kích thước batch cho embedding 'normalize_embeddings': True # normalize để cosine similarity chuẩn } try: embeddings = HuggingFaceEmbeddings( model_name=model_name, model_kwargs=model_kwargs, encode_kwargs=encode_kwargs ) logger.info("🔸Khởi tạo model embedding thành công.") return embeddings except Exception as e: logger.error(f"🔸Lỗi khi khởi tạo model embedding: {e}") raise Exception(f"Khởi tạo model embedding thất bại: {str(e)}") # Begin New def create_weaviate_schema_if_not_exists(client: weaviate.WeaviateClient, collection_name: str): """ CẢI TIẾN: Tạo schema với cấu hình chi tiết cho filtering và hybrid search. """ if client.collections.exists(collection_name): logger.info(f"✅ Schema for collection '{collection_name}' already exists.") return logger.info(f"🔸 Schema for collection '{collection_name}' not found. Creating...") try: properties = [] for prop_config in WEAVIATE_SCHEMA_CONFIG: properties.append( wvc_config.Property( name=prop_config["name"], data_type=prop_config["dataType"], # Bỏ qua vector hóa nếu vectorize=False hoặc không được định nghĩa skip_vectorization=not prop_config.get("vectorize", False), # Kích hoạt tokenization cho các trường cần tìm kiếm từ khóa tokenization=prop_config.get("tokenization") ) ) # Thêm trường 'text' chính, tối ưu cho cả vector và keyword search properties.append( wvc_config.Property( name="text", data_type=wvc_config.DataType.TEXT, skip_vectorization=False, # Luôn vector hóa nội dung chính tokenization=wvc_config.Tokenization.WORD # Cho phép tìm kiếm BM25 trên nội dung ) ) client.collections.create( name=collection_name, properties=properties, # Kích hoạt inverted index (bắt buộc cho filtering và BM25) inverted_index_config=wvc_config.Configure.inverted_index( index_null_state=True, index_property_length=True, index_timestamps=True, bm25_b=0.75, # Tham số BM25, có thể điều chỉnh bm25_k1=1.2 # Tham số K1 cho BM25 ), vectorizer_config=wvc_config.Configure.Vectorizer.none(), vector_index_config=wvc_config.Configure.VectorIndex.hnsw( distance_metric=wvc_config.VectorDistances.COSINE ) ) logger.info(f"✅ Successfully created schema for collection '{collection_name}'.") except WeaviateQueryException as e: logger.error(f"❌ Error creating schema: {e}", exc_info=True) raise def ingest_chunks_with_native_batching(client: weaviate.WeaviateClient, collection_name: str, chunks: List[Document], embeddings_model): """Sử dụng API batch gốc của Weaviate, an toàn và hiệu suất cao.""" logger.info(f"🚀 Bắt đầu quá trình ingestion cho {len(chunks)} chunks...") texts_to_embed = [chunk.page_content for chunk in chunks] logger.info(f"🧠 Đang tạo embeddings cho {len(texts_to_embed)} chunks...") start_embed_time = time.time() chunk_vectors = embeddings_model.embed_documents(texts_to_embed) logger.info(f"⏱️ Thời gian tạo embedding: {time.time() - start_embed_time:.2f} giây.") # 3. CẢI TIẾN: Đảm bảo chỉ ingest các thuộc tính hợp lệ valid_property_names = {prop["name"] for prop in WEAVIATE_SCHEMA_CONFIG} valid_property_names.add("text") # Thêm trường 'text' with client.batch.dynamic() as batch: for i, chunk in enumerate(chunks): if not isinstance(chunk,Document) or not hasattr(chunk, 'id') or not chunk.id: logger.warning(f"Bỏ qua chunk ở vị trí {i} do không hợp lệ (sai type hoặc thiếu ID).") continue properties = {"text": chunk.page_content} # Lọc metadata để chỉ giữ lại các key hợp lệ đã định nghĩa trong schema filtered_metadata = { k: v for k, v in chunk.metadata.items() if k in valid_property_names } properties.update(filtered_metadata) batch.add_object( collection=collection_name, properties=properties, uuid=chunk.id, vector=chunk_vectors[i] ) logger.info(f"✅ Batching hoàn tất. Đã gửi {len(chunks)} objects.") if batch.number_errors > 0: logger.error(f"❌ Có {batch.number_errors} lỗi xảy ra trong quá trình batching.") # Log ra 5 lỗi đầu tiên để dễ gỡ lỗi for i, error_msg in enumerate(batch.errors): if i >= 5: break logger.error(f" - Lỗi {i+1}: {error_msg}") # End new def create_or_load_vectorstore(embeddings, weaviate_url, collection_name, weaviate_client, chunks=None): vectorstore = None if not embeddings: logger.error("🔸Không có model embedding để tạo/tải vector store.") return None logger.info(f"🔸Truy cập Weaviate tại: {weaviate_url} với collection: {collection_name}") try: # Kết nối tới Weaviate client = weaviate_client if not client: logger.error("🔸Không thể kết nối tới Weaviate.") return None # Tên collection cần kiểm tra collection_name = config.WEAVIATE_COLLECTION_NAME # Kiểm tra xem collection có tồn tại không collection_exists = client.collections.exists(collection_name) logger.info(f"Collection {collection_name} exists: {collection_exists}") if chunks is not None and not collection_exists: logger.info(f"🔸Tạo Weaviate collection mới từ {len(chunks)} chunks...") # Kiểm tra mẫu dữ liệu đầu tiên logger.info(f"🔸Chunk đầu tiên:\n{chunks[0].metadata}") logger.info(f"🔸Nội dung:\n{chunks[0].page_content[:500]}...") # Lọc metadata để đảm bảo tương thích với Weaviate chunks = filter_and_serialize_complex_metadata(chunks) logger.info(f"🔸Metadata chunk đầu tiên sau khi lọc/serialize:\n{chunks[0].metadata}") # KIỂM TRA TYPE if chunks: logger.info(f"Type của chunk đầu tiên: {type(chunks[0])}") # Kiểm tra xem có phải là langchain Document không from langchain_core.documents import Document as LangchainDocument is_langchain_doc = isinstance(chunks[0], LangchainDocument) logger.info(f"Chunk đầu tiên có phải là langchain_core.documents.Document không? {is_langchain_doc}") if not is_langchain_doc: logger.error("!!! LỖI NGHIÊM TRỌNG: Chunks không phải là instance của langchain_core.documents.Document") # In ra các attribute của object để xem nó là gì try: logger.error(f"Attributes của chunk[0]: {dir(chunks[0])}") if hasattr(chunks[0], "metadata"): logger.error(f"Metadata của chunk[0] (nếu có): {chunks[0].metadata}") if hasattr(chunks[0], "page_content"): logger.error(f"Page_content của chunk[0] (nếu có): {chunks[0].page_content[:100]}") except: pass # Bỏ qua nếu không thể dir() return None # Dừng ở đây nếu type sai # Tạo vectorstore max_batch_size = 1000 # Kích thước batch an toàn total_chunks = len(chunks) logger.info("🔸Đang nhúng dữ liệu...") # Tạo collection mới vectorstore = WeaviateVectorStore.from_documents( documents=chunks[:1], # Khởi tạo với 1 tài liệu để tạo schema embedding=embeddings, client=client, index_name=collection_name, text_key="text", # Tên trường văn bản trong tài liệu # by_texts=False # Nếu dùng ids thì không cần by_texts, nhưng để rõ ràng ) # Thêm tài liệu theo batch for i in range(1, total_chunks, max_batch_size): end_idx = min(i + max_batch_size, total_chunks) current_batch = chunks[i:end_idx] logger.info(f"🔸Đang xử lý batch {i//max_batch_size + 1}/{(total_chunks-1)//max_batch_size + 1}: từ {i} đến {end_idx-1}") try: vectorstore.add_documents(current_batch) logger.info(f"🔸Đã thêm batch {i//max_batch_size + 1} thành công") except Exception as batch_error: logger.error(f"🔸Lỗi khi xử lý batch từ {i} đến {end_idx-1}: {str(batch_error)}") # Thử với batch nhỏ hơn smaller_batch_size = max_batch_size // 2 if smaller_batch_size >= 10: logger.info(f"🔸Thử lại với batch size nhỏ hơn: {smaller_batch_size}") for j in range(i, end_idx, smaller_batch_size): end_j = min(j + smaller_batch_size, end_idx) smaller_batch = chunks[j:end_j] try: vectorstore.add_documents(smaller_batch) logger.info(f"🔸Đã thêm batch nhỏ từ {j} đến {end_j-1} thành công") except Exception as small_batch_error: logger.error(f"🔸Vẫn lỗi với batch nhỏ hơn từ {j} đến {end_j-1}: {str(small_batch_error)}") else: logger.error(f"🔸Batch size đã quá nhỏ, không thể giảm thêm. Bỏ qua batch này.") logger.info(f"🔸Tạo Weaviate collection thành công: {collection_name}") elif collection_exists: logger.info(f"🔸Tải Weaviate collection đã tồn tại: {collection_name}") vectorstore = WeaviateVectorStore( client=client, index_name=collection_name, embedding=embeddings, text_key="text", attributes=[ # Liệt kê TẤT CẢ các metadata bạn cần để retriever hoạt động "nam_ban_hanh", "title", "source", "field", "loai_van_ban", "so_hieu", "ten_van_ban", "ngay_ban_hanh_str", "co_quan_ban_hanh", "entity_type", # Các trường serialize thành JSON cũng cần được liệt kê nếu muốn lấy về "cross_references", "penalties" ] ) logger.info("🔸Tải Weaviate collection thành công.") else: logger.error(f"🔸Collection '{collection_name}' không tồn tại và không có dữ liệu chunks để tạo mới.") return None logger.info("🔸Vectorstore sẵn sàng.") return vectorstore except Exception as e: if client: client.close() logger.info("🔸Đã đóng kết nối tới Weaviate.") logger.error(f"🔸Lỗi khi tạo/tải Weaviate vector store: {e}") return None def get_google_llm(google_api_key): logger.info("🔸Đang khởi tạo LLM từ Google Generative AI...") if not google_api_key: logger.error("🔸Google API Key không được cung cấp.") return None try: def create_chat_google(): return ChatGoogleGenerativeAI( model="gemini-2.5-flash", google_api_key=google_api_key, temperature=0.0, # Điều chỉnh nhiệt độ nếu cần, 0.1-0.3 thường tốt cho RAG safety_settings={}, ) llm = create_chat_google() logger.info("🔸Khởi tạo Google Generative AI LLM thành công.") return llm except Exception as e: logger.error(f"🔸Lỗi khi khởi tạo Google Generative AI LLM: {e}") return None def create_llm_from_google_key_list(google_api_keys: List[str]): """ Khởi tạo một LLM duy nhất có khả năng tự động fallback qua một danh sách các API key của Google. Khi một key hết hạn mức (lỗi ResourceExhausted), nó sẽ tự động thử key tiếp theo. :param google_api_keys: Một list chứa các chuỗi API key của Google. :return: Một đối tượng LLM của LangChain, hoặc None nếu có lỗi. """ if not google_api_keys or not isinstance(google_api_keys, list): logger.error("❌ Danh sách API key không hợp lệ hoặc bị rỗng.") return None logger.info(f"🔸 Đang khởi tạo chuỗi LLM từ {len(google_api_keys)} API key của Google...") try: # --- 1. Tạo một danh sách các instance LLM, mỗi cái với một key khác nhau --- llm_instances = [ ChatGoogleGenerativeAI( model="gemini-2.5-flash-preview-05-20", google_api_key=key, temperature=0.0, # Điều chỉnh nhiệt độ nếu cần, 0.1-0.3 thường tốt cho RAG safety_settings={}, ) for key in google_api_keys ] # --- 2. Nếu chỉ có một key, không cần fallback --- if len(llm_instances) == 1: logger.info("✅ Chỉ có một API key được cung cấp. Không cấu hình fallback.") return llm_instances[0] # --- 3. Dùng LLM đầu tiên làm LLM chính, phần còn lại làm fallback --- primary_llm = llm_instances[0] fallback_llms = llm_instances[1:] logger.info(f"▶️ LLM chính sẽ dùng key: '...{google_api_keys[0][-4:]}'") for i, llm in enumerate(fallback_llms): logger.info(f"↪️ Fallback {i+1} sẽ dùng key: '...{google_api_keys[i+1][-4:]}'") # --- 4. Kết hợp chúng lại --- llm_with_fallbacks = primary_llm.with_fallbacks( fallbacks=fallback_llms, exceptions_to_handle=(ResourceExhausted, PermissionDenied) ) logger.info("✅ Đã tạo thành công chuỗi LLM với cơ chế fallback giữa các key Google!") return llm_with_fallbacks except Exception as e: logger.error(f"❌ Lỗi nghiêm trọng khi tạo chuỗi LLM từ danh sách key: {e}", exc_info=True) return None #new update def _extract_and_parse_json(text_with_json: str) -> dict: """ Hàm trợ giúp để tìm và trích xuất khối JSON đầu tiên từ một chuỗi văn bản. """ json_match = re.search(r'\{.*\}', text_with_json, re.DOTALL) if json_match: json_str = json_match.group(0) try: return json.loads(json_str) except json.JSONDecodeError: logger.error(f"Không thể phân tích chuỗi JSON được trích xuất: {json_str}") raise else: logger.error(f"Không tìm thấy khối JSON nào trong output: {text_with_json}") raise ValueError("Không tìm thấy đối tượng JSON trong output của LLM") def _extract_final_answer(rag_output_with_thinking: str) -> str: """ Hàm trợ giúp để trích xuất câu trả lời cuối cùng từ output của QA_PROMPT_TEMPLATE. """ start_tag = "[BEGIN FINAL ANSWER]" end_tag = "[END FINAL ANSWER]" start_index = rag_output_with_thinking.find(start_tag) end_index = rag_output_with_thinking.find(end_tag) if start_index != -1 and end_index != -1: return rag_output_with_thinking[start_index + len(start_tag):end_index].strip() logger.warning("Không tìm thấy thẻ đánh dấu trả lời trong output của RAG. Trả về toàn bộ output.") return rag_output_with_thinking def create_qa_chain( llm: any, retriever: any, process_input_llm: any = None ): """ PHIÊN BẢN CUỐI CÙNG: Tạo ra một RAG chain hoàn chỉnh, có khả năng xử lý đa ngôn ngữ bằng cách bảo toàn dữ liệu đầu vào gốc. """ if not all([llm, retriever]): logger.error("🔸 Thiếu LLM hoặc Retriever chính để tạo QA Chain.") return None try: logger.info("🔸 Bắt đầu tạo QA Chain phiên bản Tối Ưu Nhất...") preprocessing_llm = process_input_llm or llm # ----- 1. KHAI BÁO PROMPTS ----- unified_preprocessing_prompt = ChatPromptTemplate.from_template( prompt_templete.UNIFIED_PREPROCESSING_PROMPT ) qa_rag_prompt = ChatPromptTemplate.from_template( prompt_templete.QA_PROMPT_TEMPLATE ) general_response_prompt = ChatPromptTemplate.from_template( prompt_templete.GENERAL_RESPONSE_PROMPT ) # ----- 2. ĐỊNH NGHĨA CÁC NHÁNH XỬ LÝ ----- # --- Nhánh A: LEGAL RAG (Nhận dict đầy đủ) --- legal_rag_chain = ( RunnablePassthrough.assign( context=itemgetter("rewritten_question") | retriever ).assign( # Lấy câu trả lời đã được dọn dẹp answer=( qa_rag_prompt | llm | StrOutputParser() | RunnableLambda(_extract_final_answer) ) ) # Chỉ trả về 2 key quan trọng nhất | (lambda x: {"answer": x["answer"], "context": x["context"]}) ).with_config({"run_name": "LegalRAGChain"}) # --- Nhánh B: GENERAL RESPONSE (Nhận dict đầy đủ) --- general_response_chain = ( general_response_prompt | llm | StrOutputParser() | (lambda answer_str: {"answer": answer_str, "context": []}) ).with_config({"run_name": "GeneralResponseChain"}) # ----- 3. BỘ ĐỊNH TUYẾN (ROUTER) ----- def route(info: dict): classification = info.get("classification") logger.info(f"➡️ Định tuyến truy vấn với phân loại: '{classification}'") if classification == "legal_rag": return legal_rag_chain else: return general_response_chain # ----- 4. KẾT HỢP THÀNH FULL CHAIN ----- # <--- THAY ĐỔI QUAN TRỌNG BẮT ĐẦU TỪ ĐÂY ---> # 4.1. Chuỗi con để thực hiện tiền xử lý và trả về JSON preprocessing_logic = ( unified_preprocessing_prompt | preprocessing_llm | StrOutputParser() | RunnableLambda(_extract_and_parse_json) ) # 4.2. Xây dựng chuỗi chính để BẢO TOÀN và HỢP NHẤT dữ liệu full_chain = ( # Bắt đầu với một Passthrough để giữ lại dữ liệu gốc (input, chat_history) RunnablePassthrough.assign( # Chạy chuỗi tiền xử lý và gán kết quả của nó vào một key mới là `processed` processed=preprocessing_logic ) | RunnableLambda( # Hàm này sẽ "làm phẳng" dict trên thành một dict duy nhất # để các nhánh sau có thể truy cập tất cả các key lambda x: { "input": x["input"], "chat_history": x.get("chat_history", []), "classification": x["processed"]["classification"], "rewritten_question": x["processed"]["rewritten_question"] } ) # 4.3. Chạy bộ định tuyến với dict đã được làm phẳng | RunnableLambda( # `info_dict` bây giờ chứa tất cả các key cần thiết lambda info_dict: route(info_dict).invoke(info_dict) ) ) # <--- THAY ĐỔI QUAN TRỌNG KẾT THÚC TẠI ĐÂY ---> logger.info("✅ Tạo thành công QA Chain phiên bản TỐI ƯU NHẤT.") return full_chain except Exception as e: logger.error(f"❌ Lỗi khi tạo QA Chain: {e}", exc_info=True) return None