|
|
| |
|
|
| import os |
| from langchain_community.vectorstores import Chroma |
| from langchain_community.embeddings import HuggingFaceEmbeddings |
| from langchain.text_splitter import CharacterTextSplitter |
| from modules import common_utils,file_utils |
| from modules import app_logger |
| |
|
|
| from modules import app_constants |
|
|
| app_logger = app_logger.app_logger |
|
|
| TEMP_DIR = app_constants.WORKSPACE_DIRECTORY + "tmp" |
| DB_DIR = app_constants.WORKSPACE_DIRECTORY + "db" |
|
|
| processed_files_record = os.path.join(app_constants.WORKSPACE_DIRECTORY, app_constants.PROCESSED_DOCS) |
|
|
| def load_documents_from_jsonl(file_path, loader_class): |
| try: |
| loader = loader_class(file_path, json_lines=True, text_content=False, jq_schema='.') |
| return loader.load() |
| except Exception as e: |
| app_logger.error(f"Error loading documents from JSONL file {file_path}: {e}") |
| return None |
|
|
| def update_processed_files_record(file_md5,module, file_path): |
| try: |
| with open(processed_files_record, 'a') as file: |
| file.write(f"{file_md5},{module},{file_path}\n") |
| except Exception as e: |
| app_logger.error(f"Error updating processed files record: {e}") |
|
|
| def is_file_processed(file_md5): |
| if os.path.exists(processed_files_record): |
| with open(processed_files_record, 'r') as file: |
| for line in file: |
| md5, _ = line.strip().split(',', 1) |
| if md5 == file_md5: |
| return True |
| return False |
|
|
| def get_chroma_index(file_path, current_page="nav_playbooks", is_persistent=True): |
| app_logger.info(f"Starting get_chroma_index for {file_path}") |
| file_md5 = file_utils.compute_md5(file_path) |
| if is_file_processed(file_md5): |
| app_logger.info(f"File {file_path} has already been processed. Skipping.") |
| db = None |
| return False |
| |
| _, file_extension = os.path.splitext(file_path) |
| loader_class = app_constants.DOCUMENT_MAP.get(file_extension.lower(), None) |
|
|
| if not loader_class: |
| app_logger.error(f"No suitable loader found for file type {file_extension}") |
| return None, False |
|
|
| embedding_model = app_constants.EMBEDDING_MODEL_NAME |
| chunk_size = app_constants.CHUNK_SIZE |
| chunk_overlap = app_constants.CHUNK_OVERLAP |
|
|
| storage_dir = DB_DIR if is_persistent else TEMP_DIR |
|
|
| base_filename = f"{current_page}_chroma_db" if is_persistent else f"{os.path.splitext(os.path.basename(file_path))[0]}_chroma_db" |
| sanitized_base_filename = file_utils.sanitize_filename(base_filename) |
| chroma_persist_directory = os.path.join(storage_dir, sanitized_base_filename) |
|
|
|
|
| embeddings = HuggingFaceEmbeddings(model_name=embedding_model) |
| try: |
| if file_extension.lower() == '.jsonl': |
| documents = load_documents_from_jsonl(file_path, loader_class) |
| else: |
| loader = loader_class(file_path) |
| documents = loader.load() |
|
|
| if not documents: |
| app_logger.error(f"No documents loaded from {file_path}.") |
| db = None |
| return False |
|
|
| text_splitter = CharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap) |
| docs = text_splitter.split_documents(documents) |
|
|
| if not docs: |
| app_logger.error(f"No documents to process after splitting from {file_path}.") |
| db = None |
| return False |
|
|
| db = Chroma.from_documents(docs, embeddings, persist_directory=chroma_persist_directory, client_settings=app_constants.CHROMA_SETTINGS) |
| update_processed_files_record(file_md5,current_page, file_path) |
| app_logger.info("Created index and saved to disk") |
| db.persist() |
| except Exception as e: |
| app_logger.error(f"Error in get_chroma_index for {file_path}: {e}") |
| db = None |
| return False |
| app_logger.info("Completed get_chroma_index operation") |
| db = None |
| return True |
