| 1 | from typing import Any, List, Sequence |
| 2 | from langchain.storage import InMemoryByteStore, LocalFileStore |
| 3 | from langchain.embeddings import CacheBackedEmbeddings |
| 4 | from helpers import guids |
| 5 | |
| 6 | # from langchain_chroma import Chroma |
| 7 | from langchain_community.vectorstores import FAISS |
| 8 | |
| 9 | # faiss needs to be patched for python 3.12 on arm #TODO remove once not needed |
| 10 | from helpers import faiss_monkey_patch |
| 11 | import faiss |
| 12 | |
| 13 | |
| 14 | from langchain_community.docstore.in_memory import InMemoryDocstore |
| 15 | from langchain_community.vectorstores.utils import ( |
| 16 | DistanceStrategy, |
| 17 | ) |
| 18 | from langchain_core.embeddings import Embeddings |
| 19 | |
| 20 | import os, json, hashlib, re |
| 21 | |
| 22 | import numpy as np |
| 23 | |
| 24 | from helpers.print_style import PrintStyle |
| 25 | from helpers import files, plugins, projects |
| 26 | from helpers.localization import Localization |
| 27 | from langchain_core.documents import Document |
| 28 | from . import knowledge_import |
| 29 | from helpers.log import Log, LogItem |
| 30 | from enum import Enum |
| 31 | from agent import Agent, AgentContext |
| 32 | import models |
| 33 | import logging |
| 34 | from simpleeval import simple_eval |
| 35 | |
| 36 | |
| 37 | # Raise the log level so WARNING messages aren't shown |
| 38 | logging.getLogger("langchain_core.vectorstores.base").setLevel(logging.ERROR) |
| 39 | |
| 40 | |
| 41 | class MyFaiss(FAISS): |
| 42 | # override aget_by_ids |
| 43 | def get_by_ids(self, ids: Sequence[str], /) -> List[Document]: |
| 44 | # return all self.docstore._dict[id] in ids |
| 45 | return [self.docstore._dict[id] for id in (ids if isinstance(ids, list) else [ids]) if id in self.docstore._dict] # type: ignore |
| 46 | |
| 47 | async def aget_by_ids(self, ids: Sequence[str], /) -> List[Document]: |
| 48 | return self.get_by_ids(ids) |
| 49 | |
| 50 | def get_all_docs(self): |
| 51 | return self.docstore._dict # type: ignore |
| 52 | |
| 53 | |
| 54 | class Memory: |
| 55 | |
| 56 | class Area(Enum): |
| 57 | MAIN = "main" |
| 58 | FRAGMENTS = "fragments" |
| 59 | SOLUTIONS = "solutions" |
| 60 | |
| 61 | index: dict[str, "MyFaiss"] = {} |
| 62 | |
| 63 | @staticmethod |
| 64 | def _get_embedding_config(agent=None): |
| 65 | from plugins._model_config.helpers.model_config import get_embedding_model_config_object |
| 66 | return get_embedding_model_config_object(agent) |
| 67 | |
| 68 | @staticmethod |
| 69 | async def get(agent: Agent): |
| 70 | memory_subdir = get_agent_memory_subdir(agent) |
| 71 | if Memory.index.get(memory_subdir) is None: |
| 72 | log_item = agent.context.log.log( |
| 73 | type="util", |
| 74 | heading=f"Initializing VectorDB in '/{memory_subdir}'", |
| 75 | ) |
| 76 | db, created = Memory.initialize( |
| 77 | log_item, |
| 78 | Memory._get_embedding_config(agent), |
| 79 | memory_subdir, |
| 80 | False, |
| 81 | ) |
| 82 | Memory.index[memory_subdir] = db |
| 83 | wrap = Memory(db, memory_subdir=memory_subdir) |
| 84 | knowledge_subdirs = get_knowledge_subdirs_by_memory_subdir( |
| 85 | memory_subdir, agent.config.knowledge_subdirs or [] |
| 86 | ) |
| 87 | if knowledge_subdirs: |
| 88 | await wrap.preload_knowledge(log_item, knowledge_subdirs, memory_subdir) |
| 89 | return wrap |
| 90 | else: |
| 91 | return Memory( |
| 92 | db=Memory.index[memory_subdir], |
| 93 | memory_subdir=memory_subdir, |
| 94 | ) |
| 95 | |
| 96 | @staticmethod |
| 97 | async def get_by_subdir( |
| 98 | memory_subdir: str, |
| 99 | log_item: LogItem | None = None, |
| 100 | preload_knowledge: bool = True, |
| 101 | ): |
| 102 | if not Memory.index.get(memory_subdir): |
| 103 | import initialize |
| 104 | |
| 105 | agent_config = initialize.initialize_agent() |
| 106 | model_config = Memory._get_embedding_config() |
| 107 | db, _created = Memory.initialize( |
| 108 | log_item=log_item, |
| 109 | model_config=model_config, |
| 110 | memory_subdir=memory_subdir, |
| 111 | in_memory=False, |
| 112 | ) |
| 113 | wrap = Memory(db, memory_subdir=memory_subdir) |
| 114 | if preload_knowledge: |
| 115 | knowledge_subdirs = get_knowledge_subdirs_by_memory_subdir( |
| 116 | memory_subdir, agent_config.knowledge_subdirs or [] |
| 117 | ) |
| 118 | if knowledge_subdirs: |
| 119 | await wrap.preload_knowledge( |
| 120 | log_item, knowledge_subdirs, memory_subdir |
| 121 | ) |
| 122 | Memory.index[memory_subdir] = db |
| 123 | return Memory(db=Memory.index[memory_subdir], memory_subdir=memory_subdir) |
| 124 | |
| 125 | @staticmethod |
| 126 | async def reload(agent: Agent): |
| 127 | memory_subdir = get_agent_memory_subdir(agent) |
| 128 | if Memory.index.get(memory_subdir): |
| 129 | del Memory.index[memory_subdir] |
| 130 | return await Memory.get(agent) |
| 131 | |
| 132 | @staticmethod |
| 133 | def initialize( |
| 134 | log_item: LogItem | None, |
| 135 | model_config: models.ModelConfig, |
| 136 | memory_subdir: str, |
| 137 | in_memory=False, |
| 138 | ) -> tuple[MyFaiss, bool]: |
| 139 | |
| 140 | PrintStyle.standard("Initializing VectorDB...") |
| 141 | |
| 142 | if log_item: |
| 143 | log_item.stream(progress="\nInitializing VectorDB") |
| 144 | |
| 145 | em_dir = files.get_abs_path( |
| 146 | "tmp/memory/embeddings" |
| 147 | ) # just caching, no need to parameterize |
| 148 | db_dir = abs_db_dir(memory_subdir) |
| 149 | |
| 150 | # make sure embeddings and database directories exist |
| 151 | os.makedirs(db_dir, exist_ok=True) |
| 152 | |
| 153 | if in_memory: |
| 154 | store = InMemoryByteStore() |
| 155 | else: |
| 156 | os.makedirs(em_dir, exist_ok=True) |
| 157 | store = LocalFileStore(em_dir) |
| 158 | |
| 159 | embeddings_model = models.get_embedding_model( |
| 160 | model_config.provider, |
| 161 | model_config.name, |
| 162 | **model_config.build_kwargs(), |
| 163 | ) |
| 164 | embeddings_model_id = files.safe_file_name( |
| 165 | model_config.provider + "_" + model_config.name |
| 166 | ) |
| 167 | |
| 168 | # here we setup the embeddings model with the chosen cache storage |
| 169 | embedder = CacheBackedEmbeddings.from_bytes_store( |
| 170 | embeddings_model, store, namespace=embeddings_model_id |
| 171 | ) |
| 172 | |
| 173 | # initial DB and docs variables |
| 174 | db: MyFaiss | None = None |
| 175 | docs: dict[str, Document] | None = None |
| 176 | |
| 177 | created = False |
| 178 | |
| 179 | # if db folder exists and is not empty: |
| 180 | if os.path.exists(db_dir) and files.exists(db_dir, "index.faiss"): |
| 181 | if not Memory._verify_index_hash(db_dir): |
| 182 | PrintStyle(font_color="yellow").print( |
| 183 | f"FAISS index hash mismatch in '{db_dir}' — index will be rebuilt." |
| 184 | ) |
| 185 | else: |
| 186 | db = MyFaiss.load_local( |
| 187 | folder_path=db_dir, |
| 188 | embeddings=embedder, |
| 189 | allow_dangerous_deserialization=True, |
| 190 | distance_strategy=DistanceStrategy.COSINE, |
| 191 | # normalize_L2=True, |
| 192 | relevance_score_fn=Memory._cosine_normalizer, |
| 193 | ) # type: ignore |
| 194 | |
| 195 | # if there is a mismatch in embeddings used, re-index the whole DB |
| 196 | emb_ok = False |
| 197 | emb_set_file = files.get_abs_path(db_dir, "embedding.json") |
| 198 | if files.exists(emb_set_file): |
| 199 | embedding_set = json.loads(files.read_file(emb_set_file)) |
| 200 | if ( |
| 201 | embedding_set["model_provider"] == model_config.provider |
| 202 | and embedding_set["model_name"] == model_config.name |
| 203 | ): |
| 204 | # model matches |
| 205 | emb_ok = True |
| 206 | |
| 207 | # re-index - create new DB and insert existing docs |
| 208 | if db and not emb_ok: |
| 209 | docs = db.get_all_docs() |
| 210 | db = None |
| 211 | |
| 212 | # DB not loaded, create one |
| 213 | if not db: |
| 214 | index = faiss.IndexFlatIP(len(embedder.embed_query("example"))) |
| 215 | |
| 216 | db = MyFaiss( |
| 217 | embedding_function=embedder, |
| 218 | index=index, |
| 219 | docstore=InMemoryDocstore(), |
| 220 | index_to_docstore_id={}, |
| 221 | distance_strategy=DistanceStrategy.COSINE, |
| 222 | # normalize_L2=True, |
| 223 | relevance_score_fn=Memory._cosine_normalizer, |
| 224 | ) |
| 225 | |
| 226 | # insert docs if reindexing |
| 227 | if docs: |
| 228 | PrintStyle.standard("Indexing memories...") |
| 229 | if log_item: |
| 230 | log_item.stream(progress="\nIndexing memories") |
| 231 | db.add_documents(documents=list(docs.values()), ids=list(docs.keys())) |
| 232 | |
| 233 | # save DB |
| 234 | Memory._save_db_file(db, memory_subdir) |
| 235 | # save meta file |
| 236 | meta_file_path = files.get_abs_path(db_dir, "embedding.json") |
| 237 | files.write_file( |
| 238 | meta_file_path, |
| 239 | json.dumps( |
| 240 | { |
| 241 | "model_provider": model_config.provider, |
| 242 | "model_name": model_config.name, |
| 243 | } |
| 244 | ), |
| 245 | ) |
| 246 | |
| 247 | created = True |
| 248 | |
| 249 | return db, created |
| 250 | |
| 251 | def __init__( |
| 252 | self, |
| 253 | db: MyFaiss, |
| 254 | memory_subdir: str, |
| 255 | ): |
| 256 | self.db = db |
| 257 | self.memory_subdir = memory_subdir |
| 258 | |
| 259 | async def preload_knowledge( |
| 260 | self, log_item: LogItem | None, kn_dirs: list[str], memory_subdir: str |
| 261 | ): |
| 262 | if log_item: |
| 263 | log_item.update(heading="Preloading knowledge...") |
| 264 | |
| 265 | # db abs path |
| 266 | db_dir = abs_db_dir(memory_subdir) |
| 267 | |
| 268 | # Load the index file if it exists |
| 269 | index_path = files.get_abs_path(db_dir, "knowledge_import.json") |
| 270 | |
| 271 | # make sure directory exists |
| 272 | if not os.path.exists(db_dir): |
| 273 | os.makedirs(db_dir) |
| 274 | |
| 275 | index: dict[str, knowledge_import.KnowledgeImport] = {} |
| 276 | if os.path.exists(index_path): |
| 277 | with open(index_path, "r") as f: |
| 278 | index = json.load(f) |
| 279 | |
| 280 | # preload knowledge folders |
| 281 | index = self._preload_knowledge_folders(log_item, kn_dirs, index) |
| 282 | |
| 283 | for file in index: |
| 284 | if index[file]["state"] in ["changed", "removed"] and index[file].get( |
| 285 | "ids", [] |
| 286 | ): # for knowledge files that have been changed or removed and have IDs |
| 287 | await self.delete_documents_by_ids( |
| 288 | index[file]["ids"] |
| 289 | ) # remove original version |
| 290 | if index[file]["state"] == "changed": |
| 291 | index[file]["ids"] = await self.insert_documents( |
| 292 | index[file]["documents"] |
| 293 | ) # insert new version |
| 294 | |
| 295 | # remove index where state="removed" |
| 296 | index = {k: v for k, v in index.items() if v["state"] != "removed"} |
| 297 | |
| 298 | # strip state and documents from index and save it |
| 299 | for file in index: |
| 300 | if "documents" in index[file]: |
| 301 | del index[file]["documents"] # type: ignore |
| 302 | if "state" in index[file]: |
| 303 | del index[file]["state"] # type: ignore |
| 304 | with open(index_path, "w") as f: |
| 305 | json.dump(index, f) |
| 306 | |
| 307 | def _preload_knowledge_folders( |
| 308 | self, |
| 309 | log_item: LogItem | None, |
| 310 | kn_dirs: list[str], |
| 311 | index: dict[str, knowledge_import.KnowledgeImport], |
| 312 | ): |
| 313 | # load knowledge folders, subfolders by area |
| 314 | for kn_dir in kn_dirs: |
| 315 | # everything in the root of the knowledge goes to main |
| 316 | index = knowledge_import.load_knowledge( |
| 317 | log_item, |
| 318 | abs_knowledge_dir(kn_dir), |
| 319 | index, |
| 320 | {"area": Memory.Area.MAIN.value}, |
| 321 | filename_pattern="*", |
| 322 | recursive=False, |
| 323 | ) |
| 324 | # subdirectories go to their folders |
| 325 | for area in Memory.Area: |
| 326 | index = knowledge_import.load_knowledge( |
| 327 | log_item, |
| 328 | # files.get_abs_path("knowledge", kn_dir, area.value), |
| 329 | abs_knowledge_dir(kn_dir, area.value), |
| 330 | index, |
| 331 | {"area": area.value}, |
| 332 | recursive=True, |
| 333 | ) |
| 334 | |
| 335 | return index |
| 336 | |
| 337 | def get_document_by_id(self, id: str) -> Document | None: |
| 338 | return self.db.get_by_ids(id)[0] |
| 339 | |
| 340 | async def embed_query(self, query: str) -> list[float]: |
| 341 | return await self.db.embedding_function.aembed_query(query) |
| 342 | |
| 343 | async def search_similarity_threshold( |
| 344 | self, |
| 345 | query: str, |
| 346 | limit: int, |
| 347 | threshold: float, |
| 348 | filter: str = "", |
| 349 | embedding: list[float] | None = None, |
| 350 | ): |
| 351 | comparator = Memory._get_comparator(filter) if filter else None |
| 352 | |
| 353 | if embedding is not None: |
| 354 | docs_and_scores = await self.db.asimilarity_search_with_score_by_vector( |
| 355 | embedding, |
| 356 | k=limit, |
| 357 | filter=comparator, |
| 358 | ) |
| 359 | return [ |
| 360 | doc |
| 361 | for doc, score in docs_and_scores |
| 362 | if Memory._cosine_normalizer(score) >= threshold |
| 363 | ] |
| 364 | |
| 365 | return await self.db.asearch( |
| 366 | query, |
| 367 | search_type="similarity_score_threshold", |
| 368 | k=limit, |
| 369 | score_threshold=threshold, |
| 370 | filter=comparator, |
| 371 | ) |
| 372 | |
| 373 | async def search_similarity_threshold_with_scores( |
| 374 | self, query: str, limit: int, threshold: float, filter: str = "" |
| 375 | ) -> list[tuple[Document, float]]: |
| 376 | comparator = Memory._get_comparator(filter) if filter else None |
| 377 | |
| 378 | return await self.db.asimilarity_search_with_relevance_scores( |
| 379 | query, |
| 380 | k=limit, |
| 381 | score_threshold=threshold, |
| 382 | filter=comparator, |
| 383 | ) |
| 384 | |
| 385 | async def delete_documents_by_query( |
| 386 | self, |
| 387 | query: str, |
| 388 | threshold: float, |
| 389 | filter: str = "", |
| 390 | *, |
| 391 | include_exact: bool = False, |
| 392 | cascade: bool = False, |
| 393 | ): |
| 394 | k = 100 |
| 395 | tot = 0 |
| 396 | removed = [] |
| 397 | removed_ids: set[str] = set() |
| 398 | |
| 399 | while True: |
| 400 | # Perform similarity search with score |
| 401 | docs = await self.search_similarity_threshold( |
| 402 | query, limit=k, threshold=threshold, filter=filter |
| 403 | ) |
| 404 | removed += docs |
| 405 | |
| 406 | # Extract document IDs and filter based on score |
| 407 | # document_ids = [result[0].metadata["id"] for result in docs if result[1] < score_limit] |
| 408 | document_ids = [result.metadata["id"] for result in docs] |
| 409 | removed_ids.update(str(doc_id) for doc_id in document_ids) |
| 410 | |
| 411 | # Delete documents with IDs over the threshold score |
| 412 | if document_ids: |
| 413 | # fnd = self.db.get(where={"id": {"$in": document_ids}}) |
| 414 | # if fnd["ids"]: self.db.delete(ids=fnd["ids"]) |
| 415 | # tot += len(fnd["ids"]) |
| 416 | await self.db.adelete(ids=document_ids) |
| 417 | tot += len(document_ids) |
| 418 | |
| 419 | # If fewer than K document IDs, break the loop |
| 420 | if len(document_ids) < k: |
| 421 | break |
| 422 | |
| 423 | if include_exact: |
| 424 | exact_docs = self._find_exact_query_docs(query, filter, removed_ids) |
| 425 | if exact_docs: |
| 426 | exact_ids = [doc.metadata["id"] for doc in exact_docs] |
| 427 | await self.db.adelete(ids=exact_ids) |
| 428 | removed += exact_docs |
| 429 | removed_ids.update(str(doc_id) for doc_id in exact_ids) |
| 430 | tot += len(exact_ids) |
| 431 | |
| 432 | if cascade and removed_ids: |
| 433 | related_docs = self._find_related_docs_by_ids(removed_ids) |
| 434 | if related_docs: |
| 435 | related_ids = [doc.metadata["id"] for doc in related_docs] |
| 436 | await self.db.adelete(ids=related_ids) |
| 437 | removed += related_docs |
| 438 | removed_ids.update(str(doc_id) for doc_id in related_ids) |
| 439 | tot += len(related_ids) |
| 440 | |
| 441 | if tot: |
| 442 | self._save_db() # persist |
| 443 | return removed |
| 444 | |
| 445 | async def delete_documents_by_ids( |
| 446 | self, ids: list[str], *, cascade: bool = False, filter: str = "" |
| 447 | ): |
| 448 | # aget_by_ids is not yet implemented in faiss, need to do a workaround |
| 449 | rem_docs = await self.db.aget_by_ids( |
| 450 | ids |
| 451 | ) # existing docs to remove (prevents error) |
| 452 | rem_ids = [doc.metadata["id"] for doc in rem_docs] |
| 453 | |
| 454 | if cascade: |
| 455 | related_docs = self._find_related_docs_by_ids(set(ids) | set(rem_ids)) |
| 456 | if related_docs: |
| 457 | existing = {doc.metadata["id"] for doc in rem_docs} |
| 458 | rem_docs.extend( |
| 459 | doc for doc in related_docs if doc.metadata["id"] not in existing |
| 460 | ) |
| 461 | |
| 462 | if rem_docs: |
| 463 | rem_ids = [doc.metadata["id"] for doc in rem_docs] # ids to remove |
| 464 | await self.db.adelete(ids=rem_ids) |
| 465 | |
| 466 | if rem_docs: |
| 467 | self._save_db() # persist |
| 468 | return rem_docs |
| 469 | |
| 470 | async def insert_text(self, text, metadata: dict = {}): |
| 471 | doc = Document(text, metadata=metadata) |
| 472 | ids = await self.insert_documents([doc]) |
| 473 | return ids[0] |
| 474 | |
| 475 | async def insert_documents(self, docs: list[Document]): |
| 476 | ids = [self._generate_doc_id() for _ in range(len(docs))] |
| 477 | timestamp = self.get_timestamp() |
| 478 | |
| 479 | if ids: |
| 480 | for doc, id in zip(docs, ids): |
| 481 | doc.metadata["id"] = id # add ids to documents metadata |
| 482 | doc.metadata["timestamp"] = timestamp # add timestamp |
| 483 | if not doc.metadata.get("area", ""): |
| 484 | doc.metadata["area"] = Memory.Area.MAIN.value |
| 485 | |
| 486 | await self.db.aadd_documents(documents=docs, ids=ids) |
| 487 | self._save_db() # persist |
| 488 | return ids |
| 489 | |
| 490 | async def update_documents(self, docs: list[Document]): |
| 491 | ids = [doc.metadata["id"] for doc in docs] |
| 492 | await self.db.adelete(ids=ids) # delete originals |
| 493 | ins = await self.db.aadd_documents(documents=docs, ids=ids) # add updated |
| 494 | self._save_db() # persist |
| 495 | return ins |
| 496 | |
| 497 | def _save_db(self): |
| 498 | Memory._save_db_file(self.db, self.memory_subdir) |
| 499 | |
| 500 | def _generate_doc_id(self): |
| 501 | while True: |
| 502 | doc_id = guids.generate_id(10) # random ID |
| 503 | if not self.db.get_by_ids(doc_id): # check if exists |
| 504 | return doc_id |
| 505 | |
| 506 | def _find_exact_query_docs( |
| 507 | self, query: str, filter: str, skip_ids: set[str] |
| 508 | ) -> list[Document]: |
| 509 | needle = _normalize_memory_match_text(query) |
| 510 | if len(needle) < 3: |
| 511 | return [] |
| 512 | |
| 513 | docs: list[Document] = [] |
| 514 | comparator = Memory._get_comparator(filter) if filter else None |
| 515 | for doc in self.db.get_all_docs().values(): |
| 516 | doc_id = str(doc.metadata.get("id", "")) |
| 517 | if not doc_id or doc_id in skip_ids: |
| 518 | continue |
| 519 | if comparator and not comparator(doc.metadata): |
| 520 | continue |
| 521 | haystack = _normalize_memory_match_text( |
| 522 | f"{doc.page_content}\n{json.dumps(doc.metadata, sort_keys=True, default=str)}" |
| 523 | ) |
| 524 | if needle in haystack: |
| 525 | docs.append(doc) |
| 526 | return docs |
| 527 | |
| 528 | def _find_related_docs_by_ids( |
| 529 | self, ids: set[str], filter: str = "" |
| 530 | ) -> list[Document]: |
| 531 | ids = {str(doc_id) for doc_id in ids if str(doc_id)} |
| 532 | if not ids: |
| 533 | return [] |
| 534 | |
| 535 | docs: list[Document] = [] |
| 536 | comparator = Memory._get_comparator(filter) if filter else None |
| 537 | for doc in self.db.get_all_docs().values(): |
| 538 | doc_id = str(doc.metadata.get("id", "")) |
| 539 | if not doc_id or doc_id in ids: |
| 540 | continue |
| 541 | if comparator and not comparator(doc.metadata): |
| 542 | continue |
| 543 | if _metadata_references_any(doc.metadata, ids): |
| 544 | docs.append(doc) |
| 545 | return docs |
| 546 | |
| 547 | @staticmethod |
| 548 | def _save_db_file(db: MyFaiss, memory_subdir: str): |
| 549 | abs_dir = abs_db_dir(memory_subdir) |
| 550 | db.save_local(folder_path=abs_dir) |
| 551 | Memory._write_index_hash(abs_dir) |
| 552 | |
| 553 | @staticmethod |
| 554 | def _write_index_hash(abs_dir: str) -> None: |
| 555 | faiss_path = os.path.join(abs_dir, "index.faiss") |
| 556 | hash_path = os.path.join(abs_dir, "index.faiss.sha256") |
| 557 | try: |
| 558 | h = hashlib.sha256() |
| 559 | with open(faiss_path, "rb") as f: |
| 560 | for chunk in iter(lambda: f.read(65536), b""): |
| 561 | h.update(chunk) |
| 562 | with open(hash_path, "w") as f: |
| 563 | f.write(h.hexdigest()) |
| 564 | except Exception as e: |
| 565 | PrintStyle(font_color="yellow").print(f"Warning: could not write FAISS hash: {e}") |
| 566 | |
| 567 | @staticmethod |
| 568 | def _verify_index_hash(abs_dir: str) -> bool: |
| 569 | faiss_path = os.path.join(abs_dir, "index.faiss") |
| 570 | hash_path = os.path.join(abs_dir, "index.faiss.sha256") |
| 571 | if not os.path.exists(hash_path): |
| 572 | return True |
| 573 | try: |
| 574 | with open(hash_path, "r") as f: |
| 575 | stored = f.read().strip() |
| 576 | h = hashlib.sha256() |
| 577 | with open(faiss_path, "rb") as f: |
| 578 | for chunk in iter(lambda: f.read(65536), b""): |
| 579 | h.update(chunk) |
| 580 | return h.hexdigest() == stored |
| 581 | except Exception as e: |
| 582 | PrintStyle(font_color="yellow").print(f"Warning: FAISS hash check failed: {e}") |
| 583 | return True |
| 584 | |
| 585 | @staticmethod |
| 586 | def _get_comparator(condition: str): |
| 587 | _FILTER_SAFE = re.compile( |
| 588 | r"^[a-zA-Z0-9_\-\.\ \t'\"=<>!()\[\],:\+]+$" |
| 589 | ) |
| 590 | if len(condition) > 512 or not _FILTER_SAFE.match(condition): |
| 591 | PrintStyle.error( |
| 592 | f"Memory filter rejected (unsafe characters or too long): {condition!r}" |
| 593 | ) |
| 594 | return lambda _data: False |
| 595 | |
| 596 | def comparator(data: dict[str, Any]): |
| 597 | try: |
| 598 | result = simple_eval(condition, names=data, functions={}) |
| 599 | return result |
| 600 | except Exception as e: |
| 601 | PrintStyle.error(f"Error evaluating condition: {e}") |
| 602 | return False |
| 603 | |
| 604 | return comparator |
| 605 | |
| 606 | @staticmethod |
| 607 | def _score_normalizer(val: float) -> float: |
| 608 | res = 1 - 1 / (1 + np.exp(val)) |
| 609 | return res |
| 610 | |
| 611 | @staticmethod |
| 612 | def _cosine_normalizer(val: float) -> float: |
| 613 | res = (1 + val) / 2 |
| 614 | res = max( |
| 615 | 0, min(1, res) |
| 616 | ) # float precision can cause values like 1.0000000596046448 |
| 617 | return float(res) # native float, not numpy scalar (JSON serializable) |
| 618 | |
| 619 | @staticmethod |
| 620 | def format_docs_plain(docs: list[Document]) -> list[str]: |
| 621 | result = [] |
| 622 | for doc in docs: |
| 623 | text = "" |
| 624 | for k, v in doc.metadata.items(): |
| 625 | text += f"{k}: {v}\n" |
| 626 | text += f"Content: {doc.page_content}" |
| 627 | result.append(text) |
| 628 | return result |
| 629 | |
| 630 | @staticmethod |
| 631 | def get_timestamp(): |
| 632 | return Localization.get().now_iso(timespec="seconds") |
| 633 | |
| 634 | |
| 635 | def get_custom_knowledge_subdir_abs(agent: Agent) -> str: |
| 636 | for dir in agent.config.knowledge_subdirs: |
| 637 | if dir != "default": |
| 638 | if dir == "custom": |
| 639 | return files.get_abs_path("usr/knowledge") |
| 640 | return files.get_abs_path("usr/knowledge", dir) |
| 641 | raise Exception("No custom knowledge subdir set") |
| 642 | |
| 643 | |
| 644 | def reload(): |
| 645 | # clear the memory index, this will force all DBs to reload |
| 646 | Memory.index = {} |
| 647 | |
| 648 | |
| 649 | def _normalize_memory_match_text(value: str) -> str: |
| 650 | return " ".join(str(value or "").casefold().split()) |
| 651 | |
| 652 | |
| 653 | def _metadata_references_any(value: Any, ids: set[str]) -> bool: |
| 654 | if isinstance(value, dict): |
| 655 | return any(_metadata_references_any(item, ids) for item in value.values()) |
| 656 | if isinstance(value, (list, tuple, set)): |
| 657 | return any(_metadata_references_any(item, ids) for item in value) |
| 658 | text = str(value or "").strip() |
| 659 | if not text: |
| 660 | return False |
| 661 | if text in ids: |
| 662 | return True |
| 663 | return any(doc_id in text.split(",") for doc_id in ids) |
| 664 | |
| 665 | |
| 666 | def abs_db_dir(memory_subdir: str) -> str: |
| 667 | # patch for projects, this way we don't need to re-work the structure of memory subdirs |
| 668 | if memory_subdir.startswith("projects/"): |
| 669 | from helpers.projects import get_project_meta |
| 670 | |
| 671 | return files.get_abs_path(get_project_meta(memory_subdir[9:]), "memory") |
| 672 | # standard subdirs |
| 673 | return files.get_abs_path("usr/memory", memory_subdir) |
| 674 | |
| 675 | |
| 676 | def abs_knowledge_dir(knowledge_subdir: str, *sub_dirs: str) -> str: |
| 677 | # patch for projects, this way we don't need to re-work the structure of knowledge subdirs |
| 678 | if knowledge_subdir.startswith("projects/"): |
| 679 | from helpers.projects import get_project_meta |
| 680 | |
| 681 | return files.get_abs_path( |
| 682 | get_project_meta(knowledge_subdir[9:]), "knowledge", *sub_dirs |
| 683 | ) |
| 684 | # standard subdirs |
| 685 | if knowledge_subdir == "default": |
| 686 | return files.get_abs_path("knowledge", *sub_dirs) |
| 687 | if knowledge_subdir == "custom": |
| 688 | return files.get_abs_path("usr/knowledge", *sub_dirs) |
| 689 | return files.get_abs_path("usr/knowledge", knowledge_subdir, *sub_dirs) |
| 690 | |
| 691 | |
| 692 | def get_memory_subdir_abs(agent: Agent) -> str: |
| 693 | subdir = get_agent_memory_subdir(agent) |
| 694 | return abs_db_dir(subdir) |
| 695 | |
| 696 | |
| 697 | def get_agent_memory_subdir(agent: Agent) -> str: |
| 698 | config = plugins.get_plugin_config("_memory", agent) |
| 699 | |
| 700 | if not config: |
| 701 | return "default" |
| 702 | |
| 703 | # Check if project isolation is enabled and we are in a project |
| 704 | if config.get("project_memory_isolation", True): |
| 705 | project_name = projects.get_context_project_name(agent.context) |
| 706 | if project_name: |
| 707 | return "projects/" + project_name |
| 708 | |
| 709 | # Fallback to configured subdir or default |
| 710 | return config.get("agent_memory_subdir", "") or "default" |
| 711 | |
| 712 | |
| 713 | def get_context_memory_subdir(context: AgentContext) -> str: |
| 714 | agent = context.get_agent() |
| 715 | return get_agent_memory_subdir(agent) |
| 716 | |
| 717 | |
| 718 | def get_existing_memory_subdirs() -> list[str]: |
| 719 | try: |
| 720 | from helpers.projects import ( |
| 721 | get_project_meta, |
| 722 | get_projects_parent_folder, |
| 723 | ) |
| 724 | |
| 725 | # Get subdirectories from memory folder |
| 726 | subdirs = files.get_subdirectories("usr/memory") |
| 727 | |
| 728 | project_subdirs = files.get_subdirectories(get_projects_parent_folder()) |
| 729 | for project_subdir in project_subdirs: |
| 730 | if files.exists( |
| 731 | get_project_meta(project_subdir), "memory", "index.faiss" |
| 732 | ): |
| 733 | subdirs.append(f"projects/{project_subdir}") |
| 734 | |
| 735 | # Ensure 'default' is always available |
| 736 | if "default" not in subdirs: |
| 737 | subdirs.insert(0, "default") |
| 738 | |
| 739 | return subdirs |
| 740 | except Exception as e: |
| 741 | PrintStyle.error(f"Failed to get memory subdirectories: {str(e)}") |
| 742 | return ["default"] |
| 743 | |
| 744 | |
| 745 | def get_knowledge_subdirs_by_memory_subdir( |
| 746 | memory_subdir: str, default: list[str] |
| 747 | ) -> list[str]: |
| 748 | if memory_subdir.startswith("projects/"): |
| 749 | from helpers.projects import get_project_meta |
| 750 | |
| 751 | default.append(get_project_meta(memory_subdir[9:], "knowledge")) |
| 752 | return default |