Memory: memory and external knowledge management. Args: project_path: the path to store the data. embedding_model: the embedding model to use, default will use the embedding model from ChromaDB, if the OpenAI has been set in the configuration, it
(
self,
project_path: str,
db_name: str = '.sa',
platform: str = 'OpenAI',
api_key: str = None,
embedding_model: str = "text-embedding-3-small"
)
| 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, |
nothing calls this directly
no outgoing calls
no test coverage detected