veadk-python 0.2.10__py3-none-any.whl → 0.2.12__py3-none-any.whl

veadk/memory/long_term_memory.py

@@ -72,10 +72,6 @@ def _get_backend_cls(backend: str) -> type[BaseLongTermMemoryBackend]:
     raise ValueError(f"Unsupported long term memory backend: {backend}")
 
 
-def build_long_term_memory_index(app_name: str, user_id: str):
-    return f"{app_name}_{user_id}"
-
-
 class LongTermMemory(BaseMemoryService, BaseModel):
     backend: Union[
         Literal["local", "opensearch", "redis", "viking", "viking_mem", "mem0"],
@@ -89,54 +85,48 @@ class LongTermMemory(BaseMemoryService, BaseModel):
     top_k: int = 5
     """Number of top similar documents to retrieve during search."""
 
+    index: str = ""
+
     app_name: str = ""
 
     user_id: str = ""
+    """Deprecated attribute"""
 
     def model_post_init(self, __context: Any) -> None:
-        if self.backend == "viking_mem":
-            logger.warning(
-                "The `viking_mem` backend is deprecated, please use `viking` instead."
-            )
-            self.backend = "viking"
-
-        self._backend = None
-
         # Once user define a backend instance, use it directly
         if isinstance(self.backend, BaseLongTermMemoryBackend):
             self._backend = self.backend
+            self.index = self._backend.index
             logger.info(
-                f"Initialized long term memory with provided backend instance {self._backend.__class__.__name__}"
+                f"Initialized long term memory with provided backend instance {self._backend.__class__.__name__}, index={self.index}"
             )
             return
 
+        # Once user define backend config, use it directly
         if self.backend_config:
-            logger.warning(
-                f"Initialized long term memory backend {self.backend} with config. We will ignore `app_name` and `user_id` if provided."
-            )
             self._backend = _get_backend_cls(self.backend)(**self.backend_config)
-            _index = self.backend_config.get("index", None)
-            if _index:
-                self._index = _index
-                logger.info(f"Long term memory index set to {self._index}.")
-            else:
-                logger.warning(
-                    "Cannot find index via backend_config, please set `index` parameter."
-                )
             return
 
-        if self.app_name and self.user_id:
-            self._index = build_long_term_memory_index(
-                app_name=self.app_name, user_id=self.user_id
-            )
-            logger.info(f"Long term memory index set to {self._index}.")
-            self._backend = _get_backend_cls(self.backend)(
-                index=self._index, **self.backend_config if self.backend_config else {}
+        # Check index
+        self.index = self.index or self.app_name
+        if not self.index:
+            logger.warning(
+                "Attribute `index` or `app_name` not provided, use `default_app` instead."
             )
-        else:
+            self.index = "default_app"
+
+        # Forward compliance
+        if self.backend == "viking_mem":
             logger.warning(
-                "Neither `backend_instance`, `backend_config`, nor (`app_name`/`user_id`) is provided, the long term memory storage will initialize when adding a session."
+                "The `viking_mem` backend is deprecated, change to `viking` instead."
             )
+            self.backend = "viking"
+
+        self._backend = _get_backend_cls(self.backend)(index=self.index)
+
+        logger.info(
+            f"Initialized long term memory with provided backend instance {self._backend.__class__.__name__}, index={self.index}"
+        )
 
     def _filter_and_convert_events(self, events: list[Event]) -> list[str]:
         final_events = []
@@ -164,75 +154,32 @@ class LongTermMemory(BaseMemoryService, BaseModel):
         self,
         session: Session,
     ):
-        app_name = session.app_name
         user_id = session.user_id
-
-        if not self._backend and isinstance(self.backend, str):
-            self._index = build_long_term_memory_index(app_name, user_id)
-            self._backend = _get_backend_cls(self.backend)(
-                index=self._index, **self.backend_config if self.backend_config else {}
-            )
-            logger.info(
-                f"Initialize long term memory backend now, index is {self._index}"
-            )
-
-        if not self._index and self._index != build_long_term_memory_index(
-            app_name, user_id
-        ):
-            logger.warning(
-                f"The `app_name` or `user_id` is different from the initialized one, skip add session to memory. Initialized index: {self._index}, current built index: {build_long_term_memory_index(app_name, user_id)}"
-            )
-            return
         event_strings = self._filter_and_convert_events(session.events)
 
         logger.info(
-            f"Adding {len(event_strings)} events to long term memory: index={self._index}"
+            f"Adding {len(event_strings)} events to long term memory: index={self.index}"
+        )
+        self._backend.save_memory(user_id=user_id, event_strings=event_strings)
+        logger.info(
+            f"Added {len(event_strings)} events to long term memory: index={self.index}, user_id={user_id}"
         )
-
-        if self._backend:
-            self._backend.save_memory(event_strings=event_strings, user_id=user_id)
-
-            logger.info(
-                f"Added {len(event_strings)} events to long term memory: index={self._index}"
-            )
-        else:
-            logger.error(
-                "Long term memory backend initialize failed, cannot add session to memory."
-            )
 
     @override
-    async def search_memory(self, *, app_name: str, user_id: str, query: str):
-        # prevent model invoke `load_memory` before add session to this memory
-        if not self._backend and isinstance(self.backend, str):
-            self._index = build_long_term_memory_index(app_name, user_id)
-            self._backend = _get_backend_cls(self.backend)(
-                index=self._index, **self.backend_config if self.backend_config else {}
-            )
-            logger.info(
-                f"Initialize long term memory backend now, index is {self._index}"
-            )
+    async def search_memory(
+        self, *, app_name: str, user_id: str, query: str
+    ) -> SearchMemoryResponse:
+        logger.info(f"Search memory with query={query}")
 
-        if not self._index and self._index != build_long_term_memory_index(
-            app_name, user_id
-        ):
-            logger.warning(
-                f"The `app_name` or `user_id` is different from the initialized one. Initialized index: {self._index}, current built index: {build_long_term_memory_index(app_name, user_id)}. Search memory return empty list."
+        memory_chunks = []
+        try:
+            memory_chunks = self._backend.search_memory(
+                query=query, top_k=self.top_k, user_id=user_id
             )
-            return SearchMemoryResponse(memories=[])
-
-        if not self._backend:
+        except Exception as e:
             logger.error(
-                "Long term memory backend is not initialized, cannot search memory."
+                f"Exception orrcus during memory search: {e}. Return empty memory chunks"
             )
-            return SearchMemoryResponse(memories=[])
-
-        logger.info(
-            f"Searching long term memory: query={query} index={self._index} top_k={self.top_k}"
-        )
-
-        memory_chunks = self._backend.search_memory(
-            query=query, top_k=self.top_k, user_id=user_id
-        )
 
         memory_events = []
         for memory in memory_chunks:
@@ -260,6 +207,6 @@ class LongTermMemory(BaseMemoryService, BaseModel):
             )
 
         logger.info(
-            f"Return {len(memory_events)} memory events for query: {query} index={self._index}"
+            f"Return {len(memory_events)} memory events for query: {query} index={self.index} user_id={user_id}"
         )
         return SearchMemoryResponse(memories=memory_events)

veadk/memory/long_term_memory_backends/base_backend.py

@@ -25,9 +25,11 @@ class BaseLongTermMemoryBackend(ABC, BaseModel):
         """Check the index name is valid or not"""
 
     @abstractmethod
-    def save_memory(self, event_strings: list[str], **kwargs) -> bool:
+    def save_memory(self, user_id: str, event_strings: list[str], **kwargs) -> bool:
         """Save memory to long term memory backend"""
 
     @abstractmethod
-    def search_memory(self, query: str, top_k: int, **kwargs) -> list[str]:
+    def search_memory(
+        self, user_id: str, query: str, top_k: int, **kwargs
+    ) -> list[str]:
         """Retrieve memory from long term memory backend"""

veadk/memory/long_term_memory_backends/in_memory_backend.py

@@ -29,10 +29,6 @@ class InMemoryLTMBackend(BaseLongTermMemoryBackend):
     embedding_config: EmbeddingModelConfig = Field(default_factory=EmbeddingModelConfig)
     """Embedding model configs"""
 
-    def precheck_index_naming(self):
-        # no checking
-        pass
-
     def model_post_init(self, __context: Any) -> None:
         self._embed_model = OpenAILikeEmbedding(
             model_name=self.embedding_config.name,
@@ -41,8 +37,12 @@ class InMemoryLTMBackend(BaseLongTermMemoryBackend):
         )
         self._vector_index = VectorStoreIndex([], embed_model=self._embed_model)
 
+    def precheck_index_naming(self):
+        # no checking
+        pass
+
     @override
-    def save_memory(self, event_strings: list[str], **kwargs) -> bool:
+    def save_memory(self, user_id: str, event_strings: list[str], **kwargs) -> bool:
         for event_string in event_strings:
             document = Document(text=event_string)
             nodes = self._split_documents([document])
@@ -50,7 +50,9 @@ class InMemoryLTMBackend(BaseLongTermMemoryBackend):
         return True
 
     @override
-    def search_memory(self, query: str, top_k: int, **kwargs) -> list[str]:
+    def search_memory(
+        self, user_id: str, query: str, top_k: int, **kwargs
+    ) -> list[str]:
         _retriever = self._vector_index.as_retriever(similarity_top_k=top_k)
         retrieved_nodes = _retriever.retrieve(query)
         return [node.text for node in retrieved_nodes]

veadk/memory/long_term_memory_backends/mem0_backend.py

@@ -13,12 +13,11 @@
 # limitations under the License.
 
 from typing import Any
-from typing_extensions import override
+
 from pydantic import Field
+from typing_extensions import override
 
 from veadk.configs.database_configs import Mem0Config
-
-
 from veadk.memory.long_term_memory_backends.base_backend import (
     BaseLongTermMemoryBackend,
 )
@@ -66,7 +65,9 @@ class Mem0LTMBackend(BaseLongTermMemoryBackend):
         pass
 
     @override
-    def save_memory(self, event_strings: list[str], **kwargs) -> bool:
+    def save_memory(
+        self, event_strings: list[str], user_id: str = "default_user", **kwargs
+    ) -> bool:
         """Save memory to Mem0
 
         Args:
@@ -76,8 +77,6 @@ class Mem0LTMBackend(BaseLongTermMemoryBackend):
         Returns:
             bool: True if saved successfully, False otherwise
         """
-        user_id = kwargs.get("user_id", "default_user")
-
         try:
             logger.info(
                 f"Saving {len(event_strings)} events to Mem0 for user: {user_id}"
@@ -100,7 +99,9 @@ class Mem0LTMBackend(BaseLongTermMemoryBackend):
             return False
 
     @override
-    def search_memory(self, query: str, top_k: int, **kwargs) -> list[str]:
+    def search_memory(
+        self, query: str, top_k: int, user_id: str = "default_user", **kwargs
+    ) -> list[str]:
         """Search memory from Mem0
 
         Args:
@@ -111,7 +112,6 @@ class Mem0LTMBackend(BaseLongTermMemoryBackend):
         Returns:
             list[str]: List of memory strings
         """
-        user_id = kwargs.get("user_id", "default_user")
 
         try:
             logger.info(

veadk/memory/long_term_memory_backends/opensearch_backend.py

@@ -14,11 +14,7 @@
 
 import re
 
-from llama_index.core import (
-    Document,
-    StorageContext,
-    VectorStoreIndex,
-)
+from llama_index.core import Document, VectorStoreIndex
 from llama_index.core.schema import BaseNode
 from llama_index.embeddings.openai_like import OpenAILikeEmbedding
 from pydantic import Field
@@ -31,6 +27,7 @@ from veadk.knowledgebase.backends.utils import get_llama_index_splitter
 from veadk.memory.long_term_memory_backends.base_backend import (
     BaseLongTermMemoryBackend,
 )
+from veadk.utils.logger import get_logger
 
 try:
     from llama_index.vector_stores.opensearch import (
@@ -42,6 +39,8 @@ except ImportError:
         "Please install VeADK extensions\npip install veadk-python[extensions]"
     )
 
+logger = get_logger(__name__)
+
 
 class OpensearchLTMBackend(BaseLongTermMemoryBackend):
     opensearch_config: OpensearchConfig = Field(default_factory=OpensearchConfig)
@@ -52,19 +51,30 @@ class OpensearchLTMBackend(BaseLongTermMemoryBackend):
     )
     """Embedding model configs"""
 
-    def precheck_index_naming(self):
+    def model_post_init(self, __context: Any) -> None:
+        self._embed_model = OpenAILikeEmbedding(
+            model_name=self.embedding_config.name,
+            api_key=self.embedding_config.api_key,
+            api_base=self.embedding_config.api_base,
+        )
+
+    def precheck_index_naming(self, index: str):
         if not (
-            isinstance(self.index, str)
-            and not self.index.startswith(("_", "-"))
-            and self.index.islower()
-            and re.match(r"^[a-z0-9_\-.]+$", self.index)
+            isinstance(index, str)
+            and not index.startswith(("_", "-"))
+            and index.islower()
+            and re.match(r"^[a-z0-9_\-.]+$", index)
         ):
             raise ValueError(
-                "The index name does not conform to the naming rules of OpenSearch"
+                f"The index name {index} does not conform to the naming rules of OpenSearch"
             )
 
-    def model_post_init(self, __context: Any) -> None:
-        self._opensearch_client = OpensearchVectorClient(
+    def _create_vector_index(self, index: str) -> VectorStoreIndex:
+        logger.info(f"Create OpenSearch vector index with index={index}")
+
+        self.precheck_index_naming(index)
+
+        opensearch_client = OpensearchVectorClient(
             endpoint=self.opensearch_config.host,
             port=self.opensearch_config.port,
             http_auth=(
@@ -74,39 +84,33 @@ class OpensearchLTMBackend(BaseLongTermMemoryBackend):
             use_ssl=True,
             verify_certs=False,
             dim=self.embedding_config.dim,
-            index=self.index,  # collection name
+            index=index,
         )
-
-        self._vector_store = OpensearchVectorStore(client=self._opensearch_client)
-
-        self._storage_context = StorageContext.from_defaults(
-            vector_store=self._vector_store
-        )
-
-        self._embed_model = OpenAILikeEmbedding(
-            model_name=self.embedding_config.name,
-            api_key=self.embedding_config.api_key,
-            api_base=self.embedding_config.api_base,
-        )
-
-        self._vector_index = VectorStoreIndex.from_documents(
-            documents=[],
-            storage_context=self._storage_context,
-            embed_model=self._embed_model,
+        vector_store = OpensearchVectorStore(client=opensearch_client)
+        return VectorStoreIndex.from_vector_store(
+            vector_store=vector_store, embed_model=self._embed_model
         )
-        self._retriever = self._vector_index.as_retriever()
 
     @override
-    def save_memory(self, event_strings: list[str], **kwargs) -> bool:
+    def save_memory(self, user_id: str, event_strings: list[str], **kwargs) -> bool:
+        index = f"{self.index}_{user_id}"
+        vector_index = self._create_vector_index(index)
+
         for event_string in event_strings:
             document = Document(text=event_string)
             nodes = self._split_documents([document])
-            self._vector_index.insert_nodes(nodes)
+            vector_index.insert_nodes(nodes)
         return True
 
     @override
-    def search_memory(self, query: str, top_k: int, **kwargs) -> list[str]:
-        _retriever = self._vector_index.as_retriever(similarity_top_k=top_k)
+    def search_memory(
+        self, user_id: str, query: str, top_k: int, **kwargs
+    ) -> list[str]:
+        index = f"{self.index}_{user_id}"
+
+        vector_index = self._create_vector_index(index)
+
+        _retriever = vector_index.as_retriever(similarity_top_k=top_k)
         retrieved_nodes = _retriever.retrieve(query)
         return [node.text for node in retrieved_nodes]
 

veadk/memory/long_term_memory_backends/redis_backend.py

@@ -12,11 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from llama_index.core import (
-    Document,
-    StorageContext,
-    VectorStoreIndex,
-)
+from llama_index.core import Document, VectorStoreIndex
 from llama_index.core.schema import BaseNode
 from llama_index.embeddings.openai_like import OpenAILikeEmbedding
 from pydantic import Field
@@ -29,21 +25,22 @@ from veadk.knowledgebase.backends.utils import get_llama_index_splitter
 from veadk.memory.long_term_memory_backends.base_backend import (
     BaseLongTermMemoryBackend,
 )
+from veadk.utils.logger import get_logger
 
 try:
     from llama_index.vector_stores.redis import RedisVectorStore
-    from llama_index.vector_stores.redis.schema import (
-        RedisIndexInfo,
-        RedisVectorStoreSchema,
-    )
     from redis import Redis
-    from redisvl.schema.fields import BaseVectorFieldAttributes
+    from redisvl.schema import IndexSchema
+
 except ImportError:
     raise ImportError(
         "Please install VeADK extensions\npip install veadk-python[extensions]"
     )
 
 
+logger = get_logger(__name__)
+
+
 class RedisLTMBackend(BaseLongTermMemoryBackend):
     redis_config: RedisConfig = Field(default_factory=RedisConfig)
     """Redis client configs"""
@@ -53,67 +50,83 @@ class RedisLTMBackend(BaseLongTermMemoryBackend):
     )
     """Embedding model configs"""
 
-    def precheck_index_naming(self):
+    def model_post_init(self, __context: Any) -> None:
+        self._embed_model = OpenAILikeEmbedding(
+            model_name=self.embedding_config.name,
+            api_key=self.embedding_config.api_key,
+            api_base=self.embedding_config.api_base,
+        )
+
+    def precheck_index_naming(self, index: str):
         # no checking
         pass
 
-    def model_post_init(self, __context: Any) -> None:
+    def _create_vector_index(self, index: str) -> VectorStoreIndex:
+        logger.info(f"Create Redis vector index with index={index}")
+
+        self.precheck_index_naming(index)
+
         # We will use `from_url` to init Redis client once the
         # AK/SK -> STS token is ready.
         # self._redis_client = Redis.from_url(url=...)
-
-        self._redis_client = Redis(
+        redis_client = Redis(
             host=self.redis_config.host,
             port=self.redis_config.port,
             db=self.redis_config.db,
             password=self.redis_config.password,
         )
 
-        self._embed_model = OpenAILikeEmbedding(
-            model_name=self.embedding_config.name,
-            api_key=self.embedding_config.api_key,
-            api_base=self.embedding_config.api_base,
+        # Create an index for each user
+        # Should be Optimized in the future
+        schema = IndexSchema.from_dict(
+            {
+                "index": {"name": index, "prefix": index, "key_separator": "_"},
+                "fields": [
+                    {"name": "id", "type": "tag", "attrs": {"sortable": False}},
+                    {"name": "doc_id", "type": "tag", "attrs": {"sortable": False}},
+                    {"name": "text", "type": "text", "attrs": {"weight": 1.0}},
+                    {
+                        "name": "vector",
+                        "type": "vector",
+                        "attrs": {
+                            "dims": self.embedding_config.dim,
+                            "algorithm": "flat",
+                            "distance_metric": "cosine",
+                        },
+                    },
+                ],
+            }
         )
+        vector_store = RedisVectorStore(schema=schema, redis_client=redis_client)
 
-        self._schema = RedisVectorStoreSchema(
-            index=RedisIndexInfo(name=self.index),
-        )
-        if "vector" in self._schema.fields:
-            vector_field = self._schema.fields["vector"]
-            if (
-                vector_field
-                and vector_field.attrs
-                and isinstance(vector_field.attrs, BaseVectorFieldAttributes)
-            ):
-                vector_field.attrs.dims = self.embedding_config.dim
-        self._vector_store = RedisVectorStore(
-            schema=self._schema,
-            redis_client=self._redis_client,
-            overwrite=True,
-            collection_name=self.index,
-        )
-
-        self._storage_context = StorageContext.from_defaults(
-            vector_store=self._vector_store
+        logger.info(
+            f"Create vector store done, index_name={vector_store.index_name} prefix={vector_store.schema.index.prefix}"
         )
 
-        self._vector_index = VectorStoreIndex.from_documents(
-            documents=[],
-            storage_context=self._storage_context,
-            embed_model=self._embed_model,
+        return VectorStoreIndex.from_vector_store(
+            vector_store=vector_store, embed_model=self._embed_model
         )
 
     @override
-    def save_memory(self, event_strings: list[str], **kwargs) -> bool:
+    def save_memory(self, user_id: str, event_strings: list[str], **kwargs) -> bool:
+        index = f"veadk-ltm/{self.index}/{user_id}"
+        vector_index = self._create_vector_index(index)
+
         for event_string in event_strings:
             document = Document(text=event_string)
             nodes = self._split_documents([document])
-            self._vector_index.insert_nodes(nodes)
+            vector_index.insert_nodes(nodes)
+
         return True
 
     @override
-    def search_memory(self, query: str, top_k: int, **kwargs) -> list[str]:
-        _retriever = self._vector_index.as_retriever(similarity_top_k=top_k)
+    def search_memory(
+        self, user_id: str, query: str, top_k: int, **kwargs
+    ) -> list[str]:
+        index = f"veadk-ltm/{self.index}/{user_id}"
+        vector_index = self._create_vector_index(index)
+
+        _retriever = vector_index.as_retriever(similarity_top_k=top_k)
         retrieved_nodes = _retriever.retrieve(query)
         return [node.text for node in retrieved_nodes]