| 1 | import glob |
| 2 | import os |
| 3 | import hashlib |
| 4 | from typing import Any, Dict, Literal, TypedDict |
| 5 | from langchain_community.document_loaders import ( |
| 6 | CSVLoader, |
| 7 | PyPDFLoader, |
| 8 | TextLoader, |
| 9 | UnstructuredHTMLLoader, |
| 10 | ) |
| 11 | from helpers.log import LogItem |
| 12 | from helpers.print_style import PrintStyle |
| 13 | |
| 14 | text_loader_kwargs = {"autodetect_encoding": True} |
| 15 | |
| 16 | |
| 17 | class KnowledgeImport(TypedDict): |
| 18 | file: str |
| 19 | checksum: str |
| 20 | ids: list[str] |
| 21 | state: Literal["changed", "original", "removed"] |
| 22 | documents: list[Any] |
| 23 | |
| 24 | |
| 25 | def calculate_checksum(file_path: str) -> str: |
| 26 | hasher = hashlib.md5() |
| 27 | with open(file_path, "rb") as f: |
| 28 | buf = f.read() |
| 29 | hasher.update(buf) |
| 30 | return hasher.hexdigest() |
| 31 | |
| 32 | |
| 33 | def load_knowledge( |
| 34 | log_item: LogItem | None, |
| 35 | knowledge_dir: str, |
| 36 | index: Dict[str, KnowledgeImport], |
| 37 | metadata: dict[str, Any] = {}, |
| 38 | filename_pattern: str = "**/*", |
| 39 | recursive: bool = True, |
| 40 | ) -> Dict[str, KnowledgeImport]: |
| 41 | """ |
| 42 | Load knowledge files from a directory with change detection and metadata enhancement. |
| 43 | |
| 44 | This function now includes enhanced error handling and compatibility with the |
| 45 | intelligent memory consolidation system. |
| 46 | """ |
| 47 | |
| 48 | # Mapping file extensions to corresponding loader classes |
| 49 | # Note: Using TextLoader for JSON and MD to avoid parsing issues with consolidation |
| 50 | file_types_loaders = { |
| 51 | "txt": TextLoader, |
| 52 | "pdf": PyPDFLoader, |
| 53 | "csv": CSVLoader, |
| 54 | "html": UnstructuredHTMLLoader, |
| 55 | "json": TextLoader, # Use TextLoader for better consolidation compatibility |
| 56 | "md": TextLoader, # Use TextLoader for better consolidation compatibility |
| 57 | } |
| 58 | |
| 59 | cnt_files = 0 |
| 60 | cnt_docs = 0 |
| 61 | |
| 62 | # Validate and create knowledge directory if needed |
| 63 | if not knowledge_dir: |
| 64 | if log_item: |
| 65 | log_item.stream(progress="\nNo knowledge directory specified") |
| 66 | PrintStyle(font_color="yellow").print("No knowledge directory specified") |
| 67 | return index |
| 68 | |
| 69 | if not os.path.exists(knowledge_dir): |
| 70 | try: |
| 71 | os.makedirs(knowledge_dir, exist_ok=True) |
| 72 | # Verify the directory was actually created and is accessible |
| 73 | if not os.path.exists(knowledge_dir) or not os.access(knowledge_dir, os.R_OK): |
| 74 | error_msg = f"Knowledge directory {knowledge_dir} was created but is not accessible" |
| 75 | if log_item: |
| 76 | log_item.stream(progress=f"\n{error_msg}") |
| 77 | PrintStyle(font_color="red").print(error_msg) |
| 78 | return index |
| 79 | |
| 80 | if log_item: |
| 81 | log_item.stream(progress=f"\nCreated knowledge directory: {knowledge_dir}") |
| 82 | PrintStyle(font_color="green").print(f"Created knowledge directory: {knowledge_dir}") |
| 83 | except (OSError, PermissionError) as e: |
| 84 | error_msg = f"Failed to create knowledge directory {knowledge_dir}: {e}" |
| 85 | if log_item: |
| 86 | log_item.stream(progress=f"\n{error_msg}") |
| 87 | PrintStyle(font_color="red").print(error_msg) |
| 88 | return index |
| 89 | |
| 90 | # Final accessibility check for existing directories |
| 91 | if not os.access(knowledge_dir, os.R_OK): |
| 92 | error_msg = f"Knowledge directory {knowledge_dir} exists but is not readable" |
| 93 | if log_item: |
| 94 | log_item.stream(progress=f"\n{error_msg}") |
| 95 | PrintStyle(font_color="red").print(error_msg) |
| 96 | return index |
| 97 | |
| 98 | # Fetch all files in the directory with specified extensions |
| 99 | try: |
| 100 | kn_files = glob.glob(os.path.join(knowledge_dir, filename_pattern), recursive=recursive) |
| 101 | kn_files = [f for f in kn_files if os.path.isfile(f) and not os.path.basename(f).startswith('.')] |
| 102 | except Exception as e: |
| 103 | PrintStyle(font_color="red").print(f"Error scanning knowledge directory {knowledge_dir}: {e}") |
| 104 | if log_item: |
| 105 | log_item.stream(progress=f"\nError scanning directory: {e}") |
| 106 | return index |
| 107 | |
| 108 | if kn_files: |
| 109 | PrintStyle.standard( |
| 110 | f"Found {len(kn_files)} knowledge files in {knowledge_dir}, processing..." |
| 111 | ) |
| 112 | if log_item: |
| 113 | log_item.stream( |
| 114 | progress=f"\nFound {len(kn_files)} knowledge files in {knowledge_dir}, processing...", |
| 115 | ) |
| 116 | |
| 117 | for file_path in kn_files: |
| 118 | try: |
| 119 | # Get file extension safely |
| 120 | file_parts = os.path.basename(file_path).split('.') |
| 121 | if len(file_parts) < 2: |
| 122 | continue # Skip files without extensions |
| 123 | |
| 124 | ext = file_parts[-1].lower() |
| 125 | if ext not in file_types_loaders: |
| 126 | continue # Skip unsupported file types |
| 127 | |
| 128 | checksum = calculate_checksum(file_path) |
| 129 | if not checksum: |
| 130 | continue # Skip files with checksum errors |
| 131 | |
| 132 | file_key = file_path |
| 133 | |
| 134 | # Load existing data from the index or create a new entry |
| 135 | file_data: KnowledgeImport = index.get(file_key, { |
| 136 | "file": file_key, |
| 137 | "checksum": "", |
| 138 | "ids": [], |
| 139 | "state": "changed", |
| 140 | "documents": [] |
| 141 | }) |
| 142 | |
| 143 | # Check if file has changed |
| 144 | if file_data.get("checksum") == checksum: |
| 145 | file_data["state"] = "original" |
| 146 | else: |
| 147 | file_data["state"] = "changed" |
| 148 | |
| 149 | # Process changed files |
| 150 | if file_data["state"] == "changed": |
| 151 | file_data["checksum"] = checksum |
| 152 | loader_cls = file_types_loaders[ext] |
| 153 | |
| 154 | try: |
| 155 | loader = loader_cls( |
| 156 | file_path, |
| 157 | **( |
| 158 | text_loader_kwargs |
| 159 | if ext in ["txt", "csv", "html", "md"] |
| 160 | else {} |
| 161 | ), |
| 162 | ) |
| 163 | documents = loader.load_and_split() |
| 164 | |
| 165 | # Enhanced metadata for better consolidation compatibility |
| 166 | enhanced_metadata = { |
| 167 | **metadata, |
| 168 | "source_file": os.path.basename(file_path), |
| 169 | "source_path": file_path, |
| 170 | "file_type": ext, |
| 171 | "knowledge_source": True, # Flag to distinguish from conversation memories |
| 172 | "import_timestamp": None, # Will be set when inserted into memory |
| 173 | } |
| 174 | |
| 175 | # Apply metadata to all documents |
| 176 | for doc in documents: |
| 177 | doc.metadata = {**doc.metadata, **enhanced_metadata} |
| 178 | |
| 179 | file_data["documents"] = documents |
| 180 | cnt_files += 1 |
| 181 | cnt_docs += len(documents) |
| 182 | |
| 183 | except Exception as e: |
| 184 | PrintStyle(font_color="red").print(f"Error loading {file_path}: {e}") |
| 185 | if log_item: |
| 186 | log_item.stream(progress=f"\nError loading {os.path.basename(file_path)}: {e}") |
| 187 | continue |
| 188 | |
| 189 | # Update the index |
| 190 | index[file_key] = file_data |
| 191 | |
| 192 | except Exception as e: |
| 193 | PrintStyle(font_color="red").print(f"Error processing {file_path}: {e}") |
| 194 | continue |
| 195 | |
| 196 | # Mark removed files |
| 197 | current_files = set(kn_files) |
| 198 | for file_key, file_data in list(index.items()): |
| 199 | if file_key not in current_files and not file_data.get("state"): |
| 200 | index[file_key]["state"] = "removed" |
| 201 | |
| 202 | # Log results |
| 203 | if cnt_files > 0 or cnt_docs > 0: |
| 204 | PrintStyle.standard(f"Processed {cnt_docs} documents from {cnt_files} files.") |
| 205 | if log_item: |
| 206 | log_item.stream( |
| 207 | progress=f"\nProcessed {cnt_docs} documents from {cnt_files} files." |
| 208 | ) |
| 209 | |
| 210 | return index |