diff --git a/memonto/core/retain.py b/memonto/core/retain.py index 0ea3b58..16b819c 100644 --- a/memonto/core/retain.py +++ b/memonto/core/retain.py @@ -202,10 +202,12 @@ def _retain( message: str, id: str, auto_expand: bool, + auto_update: bool, ephemeral: bool, ) -> None: str_ontology = ontology.serialize(format="turtle") - + updated_memory = "" + if auto_expand: ontology = expand_ontology( ontology=ontology, @@ -213,17 +215,17 @@ def _retain( message=message, ) - # TODO: this should work with ephemeral as well but needs to be adjusted slightly - updated_memory = update_memory( - data=data, - llm=llm, - vector_store=vector_store, - triple_store=triple_store, - str_ontology=str_ontology, - message=message, - id=id, - ephemeral=ephemeral, - ) + if auto_update: + updated_memory = update_memory( + data=data, + llm=llm, + vector_store=vector_store, + triple_store=triple_store, + str_ontology=str_ontology, + message=message, + id=id, + ephemeral=ephemeral, + ) save_memory( ontology=ontology, diff --git a/memonto/memonto.py b/memonto/memonto.py index 0da0f34..a740004 100644 --- a/memonto/memonto.py +++ b/memonto/memonto.py @@ -25,6 +25,7 @@ class Memonto(BaseModel): triple_store: Optional[TripleStoreModel] = None vector_store: Optional[VectorStoreModel] = None auto_expand: Optional[bool] = False + auto_update: Optional[bool] = False ephemeral: Optional[bool] = False debug: Optional[bool] = False model_config = ConfigDict(arbitrary_types_allowed=True) @@ -63,6 +64,7 @@ def retain(self, message: str) -> None: message=message, id=self.id, auto_expand=self.auto_expand, + auto_update=self.auto_update, ephemeral=self.ephemeral, )