Spaces:
Running
Running
| import os | |
| import logging | |
| import shutil | |
| import json | |
| import hashlib | |
| from typing import Optional | |
| from rag_components import KnowledgeRAG | |
| from utils import download_and_unzip_gdrive_folder, prepare_rag_source_directory, validate_rag_source_csv | |
| from config import ( | |
| GDRIVE_SOURCES_ENABLED, GDRIVE_FOLDER_ID_OR_URL, RAG_SOURCES_DIR, FAQ_SOURCES_DIR, | |
| RAG_STORAGE_PARENT_DIR, RAG_FAISS_INDEX_SUBDIR_NAME, RAG_LOAD_INDEX_ON_STARTUP, | |
| RAG_EMBEDDING_MODEL_NAME, RAG_EMBEDDING_USE_GPU, | |
| RAG_CHUNK_SIZE, RAG_CHUNK_OVERLAP, | |
| RAG_RERANKER_MODEL_NAME, RAG_RERANKER_ENABLED, | |
| RAG_CHUNKED_SOURCES_FILENAME | |
| ) | |
| logger = logging.getLogger(__name__) | |
| def _approved_source_signature(source_dir: str, expected_kind: str) -> dict: | |
| """Tie each saved FAISS index to the approved CSV contents and filter version.""" | |
| expected_name = "faqs.csv" if expected_kind == "faq" else "quiz.csv" | |
| file_path = os.path.join(source_dir, expected_name) | |
| allowed, _ = validate_rag_source_csv(file_path, expected_name, "faq" if expected_kind == "faq" else "quiz") | |
| signature = {"source_filter_version": 1, "kind": expected_kind, "files": {}} | |
| if allowed: | |
| digest = hashlib.sha256() | |
| with open(file_path, "rb") as handle: | |
| for block in iter(lambda: handle.read(1024 * 1024), b""): | |
| digest.update(block) | |
| signature["files"][expected_name] = digest.hexdigest() | |
| if expected_kind == "url": | |
| url_text_path = os.path.join(source_dir, "url_data.txt") | |
| if os.path.isfile(url_text_path): | |
| digest = hashlib.sha256() | |
| with open(url_text_path, "rb") as handle: | |
| for block in iter(lambda: handle.read(1024 * 1024), b""): | |
| digest.update(block) | |
| signature["files"]["url_data.txt"] = digest.hexdigest() | |
| return signature | |
| def initialize_and_get_rag_system(force_rebuild: bool = False, source_dir_override: Optional[str] = None, storage_dir_override: Optional[str] = None) -> Optional[KnowledgeRAG]: | |
| logger.info("[RAG_SYSTEM_INIT] Initializing...") | |
| source_dir_to_use = source_dir_override if source_dir_override else RAG_SOURCES_DIR | |
| source_kind = ("faq" if source_dir_override is not None and os.path.abspath(source_dir_override) == os.path.abspath(FAQ_SOURCES_DIR) | |
| else "url" if source_dir_override is not None else "quiz") | |
| storage_dir_to_use = storage_dir_override if storage_dir_override else RAG_STORAGE_PARENT_DIR | |
| # GDrive Logic | |
| if GDRIVE_SOURCES_ENABLED and not source_dir_override and GDRIVE_FOLDER_ID_OR_URL: | |
| logger.info("[RAG_SYSTEM_INIT] Downloading sources from GDrive...") | |
| import tempfile | |
| staged_sources = tempfile.mkdtemp(prefix="rag_gdrive_download_") | |
| try: | |
| if download_and_unzip_gdrive_folder(GDRIVE_FOLDER_ID_OR_URL, staged_sources): | |
| if os.path.exists(RAG_SOURCES_DIR): | |
| shutil.rmtree(RAG_SOURCES_DIR) | |
| shutil.move(staged_sources, RAG_SOURCES_DIR) | |
| else: | |
| logger.error("[RAG_SYSTEM_INIT] Google Drive folder download failed; local sources retained.") | |
| finally: | |
| if os.path.isdir(staged_sources): | |
| shutil.rmtree(staged_sources) | |
| if source_kind == "quiz": | |
| prepare_rag_source_directory(RAG_SOURCES_DIR, FAQ_SOURCES_DIR) | |
| faiss_index_path = os.path.join(storage_dir_to_use, RAG_FAISS_INDEX_SUBDIR_NAME) | |
| manifest_path = os.path.join(storage_dir_to_use, "approved_sources_manifest.json") | |
| current_signature = _approved_source_signature(source_dir_to_use, source_kind) | |
| manifest_matches = False | |
| try: | |
| with open(manifest_path, "r", encoding="utf-8") as handle: | |
| manifest_matches = json.load(handle) == current_signature | |
| except (OSError, ValueError): | |
| pass | |
| if not manifest_matches: | |
| logger.info("[SOURCE_FILTER] %s source changed or prior index unverified; rebuilding to exclude disallowed files.", source_kind.upper()) | |
| force_rebuild = True | |
| pre_chunked_file = os.path.join(storage_dir_to_use, RAG_CHUNKED_SOURCES_FILENAME) | |
| if force_rebuild: | |
| if os.path.exists(faiss_index_path): | |
| logger.info("[RAG_SYSTEM_INIT] Force rebuild: deleting old index.") | |
| shutil.rmtree(faiss_index_path) | |
| if os.path.exists(pre_chunked_file): | |
| logger.info("[RAG_SYSTEM_INIT] Force rebuild: deleting old pre-chunked file.") | |
| os.remove(pre_chunked_file) | |
| if os.path.exists(manifest_path): | |
| os.remove(manifest_path) | |
| try: | |
| rag = KnowledgeRAG( | |
| index_storage_dir=storage_dir_to_use, | |
| embedding_model_name=RAG_EMBEDDING_MODEL_NAME, | |
| use_gpu_for_embeddings=RAG_EMBEDDING_USE_GPU, | |
| chunk_size=RAG_CHUNK_SIZE, | |
| chunk_overlap=RAG_CHUNK_OVERLAP, | |
| reranker_model_name=RAG_RERANKER_MODEL_NAME, | |
| enable_reranker=RAG_RERANKER_ENABLED, | |
| ) | |
| rag.allowed_source_kind = source_kind | |
| loaded = False | |
| if RAG_LOAD_INDEX_ON_STARTUP and not force_rebuild: | |
| if rag.chunk_config_has_changed(): | |
| logger.warning("[RAG_SYSTEM_INIT] Chunk config changed — forcing index rebuild.") | |
| if os.path.exists(faiss_index_path): | |
| shutil.rmtree(faiss_index_path) | |
| if os.path.exists(pre_chunked_file): | |
| os.remove(pre_chunked_file) | |
| else: | |
| try: | |
| rag.load_index_from_disk() | |
| loaded = True | |
| except Exception as e: | |
| logger.warning(f"[RAG_SYSTEM_INIT] Load failed ({e}). Building new.") | |
| if not loaded: | |
| if not current_signature["files"]: | |
| logger.warning( | |
| f"[RAG_SYSTEM_INIT] No supported source files found in {source_dir_to_use}. " | |
| "System empty for this index." | |
| ) | |
| else: | |
| rag.build_index_from_source_files(source_dir_to_use) | |
| os.makedirs(storage_dir_to_use, exist_ok=True) | |
| with open(manifest_path, "w", encoding="utf-8") as handle: | |
| json.dump(current_signature, handle, sort_keys=True) | |
| logger.info("[RAG_SYSTEM_INIT] Complete.") | |
| return rag | |
| except Exception as e: | |
| logger.critical(f"[RAG_SYSTEM_INIT] FATAL: {e}", exc_info=True) | |
| return None |