veadk-python 0.2.6__py3-none-any.whl → 0.2.8__py3-none-any.whl

veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py

@@ -0,0 +1,148 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import re
+import time
+import uuid
+from typing import Any
+
+from pydantic import Field
+from typing_extensions import override
+
+import veadk.config  # noqa E401
+from veadk.config import getenv
+from veadk.memory.long_term_memory_backends.base_backend import (
+    BaseLongTermMemoryBackend,
+)
+from veadk.utils.logger import get_logger
+
+try:
+    from mcp_server_vikingdb_memory.common.memory_client import VikingDBMemoryService
+except ImportError:
+    raise ImportError(
+        "Please install VeADK extensions\npip install veadk-python[extensions]"
+    )
+
+logger = get_logger(__name__)
+
+
+class VikingDBLTMBackend(BaseLongTermMemoryBackend):
+    volcengine_access_key: str = Field(
+        default_factory=lambda: getenv("VOLCENGINE_ACCESS_KEY")
+    )
+
+    volcengine_secret_key: str = Field(
+        default_factory=lambda: getenv("VOLCENGINE_SECRET_KEY")
+    )
+
+    region: str = "cn-beijing"
+    """VikingDB memory region"""
+
+    def precheck_index_naming(self):
+        if not (
+            isinstance(self.index, str)
+            and 1 <= len(self.index) <= 128
+            and re.fullmatch(r"^[a-zA-Z][a-zA-Z0-9_]*$", self.index)
+        ):
+            raise ValueError(
+                "The index name does not conform to the rules: it must start with an English letter, contain only letters, numbers, and underscores, and have a length of 1-128."
+            )
+
+    def model_post_init(self, __context: Any) -> None:
+        self._client = VikingDBMemoryService(
+            ak=self.volcengine_access_key,
+            sk=self.volcengine_secret_key,
+            region=self.region,
+        )
+
+        # check whether collection exist, if not, create it
+        if not self._collection_exist():
+            self._create_collection()
+
+    def _collection_exist(self) -> bool:
+        try:
+            self._client.get_collection(collection_name=self.index)
+            return True
+        except Exception:
+            return False
+
+    def _create_collection(self) -> None:
+        response = self._client.create_collection(
+            collection_name=self.index,
+            description="Created by Volcengine Agent Development Kit VeADK",
+            builtin_event_types=["sys_event_v1"],
+        )
+        return response
+
+    @override
+    def save_memory(self, event_strings: list[str], **kwargs) -> bool:
+        user_id = kwargs.get("user_id")
+        if user_id is None:
+            raise ValueError("user_id is required")
+        session_id = str(uuid.uuid1())
+        messages = []
+        for raw_events in event_strings:
+            event = json.loads(raw_events)
+            content = event["parts"][0]["text"]
+            role = (
+                "user" if event["role"] == "user" else "assistant"
+            )  # field 'role': viking memory only allow 'assistant','system','user',
+            messages.append({"role": role, "content": content})
+        metadata = {
+            "default_user_id": user_id,
+            "default_assistant_id": "assistant",
+            "time": int(time.time() * 1000),
+        }
+        response = self._client.add_messages(
+            collection_name=self.index,
+            messages=messages,
+            metadata=metadata,
+            session_id=session_id,
+        )
+
+        if not response.get("code") == 0:
+            raise ValueError(f"Save VikingDB memory error: {response}")
+
+        return True
+
+    @override
+    def search_memory(self, query: str, top_k: int, **kwargs) -> list[str]:
+        user_id = kwargs.get("user_id")
+        if user_id is None:
+            raise ValueError("user_id is required")
+        filter = {
+            "user_id": user_id,
+            "memory_type": ["sys_event_v1"],
+        }
+        response = self._client.search_memory(
+            collection_name=self.index, query=query, filter=filter, limit=top_k
+        )
+
+        if not response.get("code") == 0:
+            raise ValueError(f"Search VikingDB memory error: {response}")
+
+        result = response.get("data", {}).get("result_list", [])
+        if result:
+            return [
+                json.dumps(
+                    {
+                        "role": "user",
+                        "parts": [{"text": r.get("memory_info").get("summary")}],
+                    },
+                    ensure_ascii=False,
+                )
+                for r in result
+            ]
+        return []

veadk/memory/short_term_memory.py

@@ -12,95 +12,103 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import os
-from typing import Literal
+from functools import wraps
+from typing import Any, Callable, Literal
+
+from google.adk.sessions import (
+    BaseSessionService,
+    DatabaseSessionService,
+    InMemorySessionService,
+)
+from pydantic import BaseModel, Field, PrivateAttr
+
+from veadk.memory.short_term_memory_backends.mysql_backend import (
+    MysqlSTMBackend,
+)
+from veadk.memory.short_term_memory_backends.postgresql_backend import (
+    PostgreSqlSTMBackend,
+)
+from veadk.memory.short_term_memory_backends.sqlite_backend import (
+    SQLiteSTMBackend,
+)
+from veadk.utils.logger import get_logger
 
-from google.adk.sessions import DatabaseSessionService, InMemorySessionService
+logger = get_logger(__name__)
 
-from veadk.config import getenv
-from veadk.utils.logger import get_logger
 
-from .short_term_memory_processor import ShortTermMemoryProcessor
+def wrap_get_session_with_callbacks(obj, callback_fn: Callable):
+    get_session_fn = getattr(obj, "get_session")
 
-logger = get_logger(__name__)
+    @wraps(get_session_fn)
+    def wrapper(*args, **kwargs):
+        result = get_session_fn(*args, **kwargs)
+        callback_fn(result, *args, **kwargs)
+        return result
 
-DEFAULT_LOCAL_DATABASE_PATH = "/tmp/veadk_local_database.db"
+    setattr(obj, "get_session", wrapper)
 
 
-class ShortTermMemory:
-    """
-    Short term memory class.
+class ShortTermMemory(BaseModel):
+    backend: Literal["local", "mysql", "sqlite", "postgresql", "database"] = "local"
+    """Short term memory backend. `Local` for in-memory storage, `mysql` for mysql / PostgreSQL storage. `sqlite` for sqlite storage."""
 
-    This class is used to store short term memory.
-    """
+    backend_configs: dict = Field(default_factory=dict)
+    """Backend specific configurations."""
 
-    def __init__(
-        self,
-        backend: Literal["local", "database", "mysql"] = "local",
-        db_url: str = "",
-        enable_memory_optimization: bool = False,
-    ):
-        self.backend = backend
-        self.db_url = db_url
-
-        if self.backend == "mysql":
-            host = getenv("DATABASE_MYSQL_HOST")
-            user = getenv("DATABASE_MYSQL_USER")
-            password = getenv("DATABASE_MYSQL_PASSWORD")
-            database = getenv("DATABASE_MYSQL_DATABASE")
-            db_url = f"mysql+pymysql://{user}:{password}@{host}/{database}"
-
-            self.db_url = db_url
-            self.backend = "database"
-
-        if self.backend == "local":
-            logger.warning(
-                f"Short term memory backend: {self.backend}, the history will be lost after application shutdown."
-            )
-            self.session_service = InMemorySessionService()
-        elif self.backend == "database":
-            if self.db_url == "" or self.db_url is None:
-                logger.warning("The `db_url` is an empty or None string.")
-                self._use_default_database()
-            else:
-                try:
-                    self.session_service = DatabaseSessionService(db_url=self.db_url)
-                    logger.info("Connected to database with db_url.")
-                except Exception as e:
-                    logger.error(f"Failed to connect to database, error: {e}.")
-                    self._use_default_database()
-        else:
-            raise ValueError(f"Unknown short term memory backend: {self.backend}")
+    db_url: str = ""
+    """Database connection URL, e.g. `sqlite:///./test.db`. Once set, it will override the `backend` parameter."""
 
-        if enable_memory_optimization and backend == "database":
-            self.processor = ShortTermMemoryProcessor()
-            intercept_get_session = self.processor.patch()
-            self.session_service.get_session = intercept_get_session(
-                self.session_service.get_session
-            )
+    local_database_path: str = "/tmp/veadk_local_database.db"
+    """Local database path, only used when `backend` is `sqlite`. Default to `/tmp/veadk_local_database.db`."""
+
+    after_load_memory_callback: Callable | None = None
+    """A callback to be called after loading memory from the backend. The callback function should accept `Session` as an input."""
 
-    def _use_default_database(self):
-        self.db_url = DEFAULT_LOCAL_DATABASE_PATH
-        logger.info(f"Using default local database {self.db_url}")
-        if not os.path.exists(self.db_url):
-            self.create_local_sqlite3_db(self.db_url)
-        self.session_service = DatabaseSessionService(db_url="sqlite:///" + self.db_url)
+    _session_service: BaseSessionService = PrivateAttr()
 
-    def create_local_sqlite3_db(self, path: str):
-        import sqlite3
+    def model_post_init(self, __context: Any) -> None:
+        if self.db_url:
+            logger.info("The `db_url` is set, ignore `backend` option.")
+            self._session_service = DatabaseSessionService(db_url=self.db_url)
+        else:
+            if self.backend == "database":
+                logger.warning(
+                    "Backend `database` is deprecated, use `sqlite` to create short term memory."
+                )
+                self.backend = "sqlite"
+            match self.backend:
+                case "local":
+                    self._session_service = InMemorySessionService()
+                case "mysql":
+                    self._session_service = MysqlSTMBackend(
+                        **self.backend_configs
+                    ).session_service
+                case "sqlite":
+                    self._session_service = SQLiteSTMBackend(
+                        local_path=self.local_database_path
+                    ).session_service
+                case "postgresql":
+                    self._session_service = PostgreSqlSTMBackend(
+                        **self.backend_configs
+                    ).session_service
+
+        if self.after_load_memory_callback:
+            wrap_get_session_with_callbacks(
+                self._session_service, self.after_load_memory_callback
+            )
 
-        conn = sqlite3.connect(path)
-        conn.close()
-        logger.debug(f"Create local sqlite3 database {path} done.")
+    @property
+    def session_service(self) -> BaseSessionService:
+        return self._session_service
 
     async def create_session(
         self,
         app_name: str,
         user_id: str,
         session_id: str,
-    ):
-        if isinstance(self.session_service, DatabaseSessionService):
-            list_sessions_response = await self.session_service.list_sessions(
+    ) -> None:
+        if isinstance(self._session_service, DatabaseSessionService):
+            list_sessions_response = await self._session_service.list_sessions(
                 app_name=app_name, user_id=user_id
             )
 
@@ -109,12 +117,12 @@ class ShortTermMemory:
             )
 
         if (
-            await self.session_service.get_session(
+            await self._session_service.get_session(
                 app_name=app_name, user_id=user_id, session_id=session_id
             )
             is None
         ):
             # create a new session for this running
-            await self.session_service.create_session(
+            await self._session_service.create_session(
                 app_name=app_name, user_id=user_id, session_id=session_id
             )

veadk/memory/short_term_memory_backends/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

veadk/memory/short_term_memory_backends/base_backend.py

@@ -0,0 +1,31 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from abc import ABC, abstractmethod
+from functools import cached_property
+
+from google.adk.sessions import BaseSessionService
+from pydantic import BaseModel
+
+
+class BaseShortTermMemoryBackend(ABC, BaseModel):
+    """
+    Base class for short term memory backend.
+    """
+
+    @cached_property
+    @abstractmethod
+    def session_service(self) -> BaseSessionService:
+        """Return the session service instance."""

veadk/memory/short_term_memory_backends/mysql_backend.py

@@ -0,0 +1,41 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from functools import cached_property
+from typing import Any
+
+from google.adk.sessions import (
+    BaseSessionService,
+    DatabaseSessionService,
+)
+from pydantic import Field
+from typing_extensions import override
+
+import veadk.config  # noqa E401
+from veadk.configs.database_configs import MysqlConfig
+from veadk.memory.short_term_memory_backends.base_backend import (
+    BaseShortTermMemoryBackend,
+)
+
+
+class MysqlSTMBackend(BaseShortTermMemoryBackend):
+    mysql_config: MysqlConfig = Field(default_factory=MysqlConfig)
+
+    def model_post_init(self, context: Any) -> None:
+        self._db_url = f"mysql+pymysql://{self.mysql_config.user}:{self.mysql_config.password}@{self.mysql_config.host}/{self.mysql_config.database}"
+
+    @cached_property
+    @override
+    def session_service(self) -> BaseSessionService:
+        return DatabaseSessionService(db_url=self._db_url)

veadk/memory/short_term_memory_backends/postgresql_backend.py

@@ -0,0 +1,41 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from functools import cached_property
+from typing import Any
+
+from google.adk.sessions import (
+    BaseSessionService,
+    DatabaseSessionService,
+)
+from pydantic import Field
+from typing_extensions import override
+
+import veadk.config  # noqa E401
+from veadk.configs.database_configs import PostgreSqlConfig
+from veadk.memory.short_term_memory_backends.base_backend import (
+    BaseShortTermMemoryBackend,
+)
+
+
+class PostgreSqlSTMBackend(BaseShortTermMemoryBackend):
+    postgresql_config: PostgreSqlConfig = Field(default_factory=PostgreSqlConfig)
+
+    def model_post_init(self, context: Any) -> None:
+        self._db_url = f"postgresql+psycopg2://{self.postgresql_config.user}:{self.postgresql_config.password}@{self.postgresql_config.host}:{self.postgresql_config.port}/{self.postgresql_config.database}"
+
+    @cached_property
+    @override
+    def session_service(self) -> BaseSessionService:
+        return DatabaseSessionService(db_url=self._db_url)

veadk/memory/short_term_memory_backends/sqlite_backend.py

@@ -0,0 +1,48 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import sqlite3
+from functools import cached_property
+from typing import Any
+
+from google.adk.sessions import (
+    BaseSessionService,
+    DatabaseSessionService,
+)
+from typing_extensions import override
+
+from veadk.memory.short_term_memory_backends.base_backend import (
+    BaseShortTermMemoryBackend,
+)
+
+
+class SQLiteSTMBackend(BaseShortTermMemoryBackend):
+    local_path: str
+
+    def model_post_init(self, context: Any) -> None:
+        # if the DB file not exists, create it
+        if not self._db_exists():
+            conn = sqlite3.connect(self.local_path)
+            conn.close()
+
+        self._db_url = f"sqlite:///{self.local_path}"
+
+    @cached_property
+    @override
+    def session_service(self) -> BaseSessionService:
+        return DatabaseSessionService(db_url=self._db_url)
+
+    def _db_exists(self) -> bool:
+        return os.path.exists(self.local_path)

veadk/memory/short_term_memory_processor.py

@@ -20,7 +20,12 @@ from google.adk.sessions import Session
 from google.genai.types import Content, Part
 from litellm import completion
 
-from veadk.config import getenv
+from veadk.config import settings
+from veadk.consts import (
+    DEFAULT_MODEL_AGENT_API_BASE,
+    DEFAULT_MODEL_AGENT_NAME,
+    DEFAULT_MODEL_AGENT_PROVIDER,
+)
 from veadk.prompts.prompt_memory_processor import render_prompt
 from veadk.utils.logger import get_logger
 
@@ -62,9 +67,9 @@ class ShortTermMemoryProcessor:
         prompt = render_prompt(messages=messages)
 
         res = completion(
-            model=getenv("MODEL_AGENT_PROVIDER") + "/" + getenv("MODEL_AGENT_NAME"),
-            base_url=getenv("MODEL_AGENT_API_BASE"),
-            api_key=getenv("MODEL_AGENT_API_KEY"),
+            model=DEFAULT_MODEL_AGENT_PROVIDER + "/" + DEFAULT_MODEL_AGENT_NAME,
+            base_url=DEFAULT_MODEL_AGENT_API_BASE,
+            api_key=settings.model.api_key,
             messages=[
                 {
                     "role": "user",