| 1 | from typing import Any, List, Sequence |
| 2 | from langchain_community.vectorstores import FAISS |
| 3 | |
| 4 | # faiss needs to be patched for python 3.12 on arm #TODO remove once not needed |
| 5 | from helpers import faiss_monkey_patch |
| 6 | import faiss |
| 7 | |
| 8 | |
| 9 | from langchain_core.documents import Document |
| 10 | from langchain.storage import InMemoryByteStore |
| 11 | from langchain_community.docstore.in_memory import InMemoryDocstore |
| 12 | from langchain_community.vectorstores.utils import ( |
| 13 | DistanceStrategy, |
| 14 | ) |
| 15 | from langchain.embeddings import CacheBackedEmbeddings |
| 16 | from simpleeval import simple_eval |
| 17 | |
| 18 | from agent import Agent |
| 19 | from helpers import guids |
| 20 | |
| 21 | |
| 22 | class MyFaiss(FAISS): |
| 23 | # override aget_by_ids |
| 24 | def get_by_ids(self, ids: Sequence[str], /) -> List[Document]: |
| 25 | # return all self.docstore._dict[id] in ids |
| 26 | return [self.docstore._dict[id] for id in (ids if isinstance(ids, list) else [ids]) if id in self.docstore._dict] # type: ignore |
| 27 | |
| 28 | async def aget_by_ids(self, ids: Sequence[str], /) -> List[Document]: |
| 29 | return self.get_by_ids(ids) |
| 30 | |
| 31 | def get_all_docs(self) -> dict[str, Document]: |
| 32 | return self.docstore._dict # type: ignore |
| 33 | |
| 34 | |
| 35 | class VectorDB: |
| 36 | |
| 37 | _cached_embeddings: dict[str, CacheBackedEmbeddings] = {} |
| 38 | |
| 39 | @staticmethod |
| 40 | def _get_embeddings(agent: Agent, cache: bool = True): |
| 41 | model = agent.get_embedding_model() |
| 42 | if not cache: |
| 43 | return model # return raw embeddings if cache is False |
| 44 | namespace = getattr( |
| 45 | model, |
| 46 | "model_name", |
| 47 | "default", |
| 48 | ) |
| 49 | if namespace not in VectorDB._cached_embeddings: |
| 50 | store = InMemoryByteStore() |
| 51 | VectorDB._cached_embeddings[namespace] = ( |
| 52 | CacheBackedEmbeddings.from_bytes_store( |
| 53 | model, |
| 54 | store, |
| 55 | namespace=namespace, |
| 56 | ) |
| 57 | ) |
| 58 | return VectorDB._cached_embeddings[namespace] |
| 59 | |
| 60 | def __init__(self, agent: Agent, cache: bool = True): |
| 61 | self.agent = agent |
| 62 | self.cache = cache # store cache preference |
| 63 | self.embeddings = self._get_embeddings(agent, cache=cache) |
| 64 | self.index = faiss.IndexFlatIP(len(self.embeddings.embed_query("example"))) |
| 65 | |
| 66 | self.db = MyFaiss( |
| 67 | embedding_function=self.embeddings, |
| 68 | index=self.index, |
| 69 | docstore=InMemoryDocstore(), |
| 70 | index_to_docstore_id={}, |
| 71 | distance_strategy=DistanceStrategy.COSINE, |
| 72 | # normalize_L2=True, |
| 73 | relevance_score_fn=cosine_normalizer, |
| 74 | ) |
| 75 | |
| 76 | async def search_by_similarity_threshold( |
| 77 | self, query: str, limit: int, threshold: float, filter: str = "" |
| 78 | ): |
| 79 | comparator = get_comparator(filter) if filter else None |
| 80 | |
| 81 | return await self.db.asearch( |
| 82 | query, |
| 83 | search_type="similarity_score_threshold", |
| 84 | k=limit, |
| 85 | score_threshold=threshold, |
| 86 | filter=comparator, |
| 87 | ) |
| 88 | |
| 89 | async def search_by_metadata(self, filter: str, limit: int = 0) -> list[Document]: |
| 90 | comparator = get_comparator(filter) |
| 91 | all_docs = self.db.get_all_docs() |
| 92 | result = [] |
| 93 | for doc in all_docs.values(): |
| 94 | if comparator(doc.metadata): |
| 95 | result.append(doc) |
| 96 | # stop if limit reached and limit > 0 |
| 97 | if limit > 0 and len(result) >= limit: |
| 98 | break |
| 99 | return result |
| 100 | |
| 101 | async def insert_documents(self, docs: list[Document]): |
| 102 | ids = [guids.generate_id() for _ in range(len(docs))] |
| 103 | |
| 104 | if ids: |
| 105 | for doc, id in zip(docs, ids): |
| 106 | doc.metadata["id"] = id # add ids to documents metadata |
| 107 | |
| 108 | self.db.add_documents(documents=docs, ids=ids) |
| 109 | return ids |
| 110 | |
| 111 | async def delete_documents_by_ids(self, ids: list[str]): |
| 112 | # aget_by_ids is not yet implemented in faiss, need to do a workaround |
| 113 | rem_docs = await self.db.aget_by_ids( |
| 114 | ids |
| 115 | ) # existing docs to remove (prevents error) |
| 116 | if rem_docs: |
| 117 | rem_ids = [doc.metadata["id"] for doc in rem_docs] # ids to remove |
| 118 | await self.db.adelete(ids=rem_ids) |
| 119 | return rem_docs |
| 120 | |
| 121 | |
| 122 | def format_docs_plain(docs: list[Document]) -> list[str]: |
| 123 | result = [] |
| 124 | for doc in docs: |
| 125 | text = "" |
| 126 | for k, v in doc.metadata.items(): |
| 127 | text += f"{k}: {v}\n" |
| 128 | text += f"Content: {doc.page_content}" |
| 129 | result.append(text) |
| 130 | return result |
| 131 | |
| 132 | |
| 133 | def cosine_normalizer(val: float) -> float: |
| 134 | res = (1 + val) / 2 |
| 135 | res = max( |
| 136 | 0, min(1, res) |
| 137 | ) # float precision can cause values like 1.0000000596046448 |
| 138 | return res |
| 139 | |
| 140 | |
| 141 | def get_comparator(condition: str): |
| 142 | def comparator(data: dict[str, Any]): |
| 143 | try: |
| 144 | result = simple_eval(condition, names=data) |
| 145 | return result |
| 146 | except Exception as e: |
| 147 | # PrintStyle.error(f"Error evaluating condition: {e}") |
| 148 | return False |
| 149 | |
| 150 | return comparator |