| 12 | |
| 13 | |
| 14 | class Memory: |
| 15 | def __init__( |
| 16 | self, |
| 17 | project_path: str, |
| 18 | db_name: str = '.sa', |
| 19 | platform: str = 'OpenAI', |
| 20 | api_key: str = None, |
| 21 | embedding_model: str = "text-embedding-3-small" |
| 22 | ): |
| 23 | """ |
| 24 | Memory: memory and external knowledge management. |
| 25 | Args: |
| 26 | project_path: the path to store the data. |
| 27 | embedding_model: the embedding model to use, default will use the embedding model from ChromaDB, |
| 28 | if the OpenAI has been set in the configuration, it will use the OpenAI embedding model |
| 29 | "text-embedding-ada-002". |
| 30 | """ |
| 31 | self.db_name = db_name |
| 32 | self.collection_name = 'memory' |
| 33 | self.client = chromadb.PersistentClient(path=os.path.join(project_path, self.db_name)) |
| 34 | self.client.get_or_create_collection( |
| 35 | self.collection_name, |
| 36 | ) |
| 37 | # use the OpenAI embedding function if the openai section is set in the configuration. |
| 38 | if platform == 'OpenAI': |
| 39 | openai_client = OpenAI(api_key=api_key or os.environ["OPENAI_API_KEY"]) |
| 40 | self.embedder = lambda x: [i.embedding for i in openai_client.embeddings.create(input=x, model=embedding_model).data] |
| 41 | else: |
| 42 | # self.embedder = embedding_functions.DefaultEmbeddingFunction() |
| 43 | self.embedder = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="all-MiniLM-L6-v2") |
| 44 | |
| 45 | def add_query( |
| 46 | self, |
| 47 | queries: List[Dict[str, str]], |
| 48 | collection: str = None, |
| 49 | idx: List[str] = None |
| 50 | ): |
| 51 | """ |
| 52 | add_query: add the queries to the memery. |
| 53 | Args: |
| 54 | queries: the queries to add to the memery. Should be in the format of |
| 55 | { |
| 56 | "query": "the query", |
| 57 | "response": "the response" |
| 58 | } |
| 59 | collection: the name of the collection to add the queries. |
| 60 | idx: the ids of the queries, should be in the same length as the queries. |
| 61 | If not provided, the ids will be generated by UUID. |
| 62 | |
| 63 | Return: A list of generated IDs. |
| 64 | """ |
| 65 | if idx: |
| 66 | ids = idx |
| 67 | else: |
| 68 | ids = [str(uuid.uuid4()) for _ in range(len(queries))] |
| 69 | |
| 70 | if not collection: |
| 71 | collection = self.collection_name |
no outgoing calls
no test coverage detected