veadk-python 0.2.27__py3-none-any.whl

veadk/memory/long_term_memory_backends/vikingdb_memory_backend.py

@@ -0,0 +1,189 @@
+# 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 os
+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.auth.veauth.utils import get_credential_from_vefaas_iam
+from veadk.integrations.ve_viking_db_memory.ve_viking_db_memory import (
+    VikingDBMemoryClient,
+)
+from veadk.memory.long_term_memory_backends.base_backend import (
+    BaseLongTermMemoryBackend,
+)
+from veadk.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class VikingDBLTMBackend(BaseLongTermMemoryBackend):
+    volcengine_access_key: str | None = Field(
+        default_factory=lambda: os.getenv("VOLCENGINE_ACCESS_KEY")
+    )
+
+    volcengine_secret_key: str | None = Field(
+        default_factory=lambda: os.getenv("VOLCENGINE_SECRET_KEY")
+    )
+
+    session_token: str = ""
+
+    region: str = "cn-beijing"
+    """VikingDB memory region"""
+
+    memory_type: list[str] = Field(default_factory=list)
+
+    def model_post_init(self, __context: Any) -> None:
+        # We get memory type from:
+        # 1. user input
+        # 2. environment variable
+        # 3. default value
+        if not self.memory_type:
+            env_memory_type = os.getenv("DATABASE_VIKINGMEM_MEMORY_TYPE")
+            if env_memory_type:
+                # "event_1, event_2" -> ["event_1", "event_2"]
+                self.memory_type = [x.strip() for x in env_memory_type.split(",")]
+            else:
+                # self.memory_type = ["sys_event_v1", "event_v1"]
+                self.memory_type = ["sys_event_v1"]
+
+        logger.info(f"Using memory type: {self.memory_type}")
+
+        # check whether collection exist, if not, create it
+        if not self._collection_exist():
+            self._create_collection()
+
+    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 _collection_exist(self) -> bool:
+        try:
+            client = self._get_client()
+            client.get_collection(collection_name=self.index)
+            logger.info(f"Collection {self.index} exist.")
+            return True
+        except Exception:
+            logger.info(f"Collection {self.index} not exist.")
+            return False
+
+    def _create_collection(self) -> None:
+        logger.info(
+            f"Create collection with collection_name={self.index}, builtin_event_types={self.memory_type}"
+        )
+        client = self._get_client()
+        response = client.create_collection(
+            collection_name=self.index,
+            description="Created by Volcengine Agent Development Kit VeADK",
+            builtin_event_types=self.memory_type,
+        )
+        logger.debug(f"Create collection with response {response}")
+        return response
+
+    def _get_client(self) -> VikingDBMemoryClient:
+        if not (self.volcengine_access_key and self.volcengine_secret_key):
+            cred = get_credential_from_vefaas_iam()
+            self.volcengine_access_key = cred.access_key_id
+            self.volcengine_secret_key = cred.secret_access_key
+            self.session_token = cred.session_token
+
+        return VikingDBMemoryClient(
+            ak=self.volcengine_access_key,
+            sk=self.volcengine_secret_key,
+            sts_token=self.session_token,
+            region=self.region,
+        )
+
+    @override
+    def save_memory(self, user_id: str, event_strings: list[str], **kwargs) -> bool:
+        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),
+        }
+
+        logger.debug(
+            f"Request for add {len(messages)} memory to VikingDB: collection_name={self.index}, metadata={metadata}, session_id={session_id}"
+        )
+
+        client = self._get_client()
+        response = client.add_messages(
+            collection_name=self.index,
+            messages=messages,
+            metadata=metadata,
+            session_id=session_id,
+        )
+
+        logger.debug(f"Response from add memory to VikingDB: {response}")
+
+        if not response.get("code") == 0:
+            raise ValueError(f"Save VikingDB memory error: {response}")
+
+        return True
+
+    @override
+    def search_memory(
+        self, user_id: str, query: str, top_k: int, **kwargs
+    ) -> list[str]:
+        filter = {"user_id": user_id, "memory_type": self.memory_type}
+
+        logger.debug(
+            f"Request for search memory in VikingDB: filter={filter}, collection_name={self.index}, query={query}, limit={top_k}"
+        )
+
+        client = self._get_client()
+        response = client.search_memory(
+            collection_name=self.index, query=query, filter=filter, limit=top_k
+        )
+
+        logger.debug(f"Response from search memory in VikingDB: {response}")
+
+        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

@@ -0,0 +1,252 @@
+# 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 wraps
+from typing import Any, Callable, Literal
+
+from google.adk.sessions import (
+    BaseSessionService,
+    DatabaseSessionService,
+    InMemorySessionService,
+    Session,
+)
+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
+
+logger = get_logger(__name__)
+
+
+def wrap_get_session_with_callbacks(obj, callback_fn: Callable):
+    get_session_fn = getattr(obj, "get_session")
+
+    @wraps(get_session_fn)
+    def wrapper(*args, **kwargs):
+        result = get_session_fn(*args, **kwargs)
+        callback_fn(result, *args, **kwargs)
+        return result
+
+    setattr(obj, "get_session", wrapper)
+
+
+class ShortTermMemory(BaseModel):
+    """Short term memory for agent execution.
+
+    The short term memory represents the context of the agent model. All content in the short term memory will be sent to agent model directly, including the system prompt, historical user prompt, and historical model responses.
+
+    Attributes:
+        backend (Literal["local", "mysql", "sqlite", "postgresql", "database"]):
+            The backend of short term memory:
+            - `local` for in-memory storage
+            - `mysql` for mysql / PostgreSQL storage
+            - `sqlite` for locally sqlite storage
+        backend_configs (dict): Configuration dict for init short term memory backend.
+        db_url (str):
+            Database connection url for init short term memory backend.
+            For example, `sqlite:///./test.db`. Once set, it will override the `backend` parameter.
+        local_database_path (str):
+            Local database path, only used when `backend` is `sqlite`.
+            Default to `/tmp/veadk_local_database.db`.
+        after_load_memory_callback (Callable | None):
+            A callback to be called after loading memory from the backend. The callback function should accept `Session` as an input.
+
+    Examples:
+        ### In-memory simple memory
+
+        You can initialize a short term memory with in-memory storage:
+
+        ```python
+        from veadk import Agent, Runner
+        from veadk.memory.short_term_memory import ShortTermMemory
+        import asyncio
+
+        session_id = "veadk_playground_session"
+
+        agent = Agent()
+        short_term_memory = ShortTermMemory(backend="local")
+
+        runner = Runner(
+            agent=agent, short_term_memory=short_term_memory)
+
+        # This invocation will be stored in short-term memory
+        response = asyncio.run(runner.run(
+            messages="My name is VeADK", session_id=session_id
+        ))
+        print(response)
+
+        # The history invocation can be fetched by model
+        response = asyncio.run(runner.run(
+            messages="Do you remember my name?", session_id=session_id # keep the same `session_id`
+        ))
+        print(response)
+        ```
+
+        ### Memory with a Database URL
+
+        Also you can use a databasae connection URL to initialize a short-term memory:
+
+        ```python
+        from veadk.memory.short_term_memory import ShortTermMemory
+
+        short_term_memory = ShortTermMemory(db_url="...")
+        ```
+
+        ### Memory with SQLite
+
+        Once you want to start the short term memory with a local SQLite, you can specify the backend to `sqlite`. It will create a local database in `local_database_path`:
+
+        ```python
+        from veadk.memory.short_term_memory import ShortTermMemory
+
+        short_term_memory = ShortTermMemory(backend="sqlite", local_database_path="")
+        ```
+    """
+
+    backend: Literal["local", "mysql", "sqlite", "postgresql", "database"] = "local"
+
+    backend_configs: dict = Field(default_factory=dict)
+
+    db_url: str = ""
+
+    local_database_path: str = "/tmp/veadk_local_database.db"
+
+    after_load_memory_callback: Callable | None = None
+
+    _session_service: BaseSessionService = PrivateAttr()
+
+    def model_post_init(self, __context: Any) -> None:
+        if self.db_url:
+            logger.info("The `db_url` is set, ignore `backend` option.")
+            if self.db_url.count("@") > 1 or self.db_url.count(":") > 2:
+                logger.warning(
+                    "Multiple `@` or `:` symbols detected in the database URL. "
+                    "Please encode `username` or `password` with `urllib.parse.quote_plus`. "
+                    "Examples: p@ssword→p%40ssword."
+                )
+            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
+            )
+
+    @property
+    def session_service(self) -> BaseSessionService:
+        return self._session_service
+
+    async def create_session(
+        self,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+    ) -> Session | None:
+        """Create or retrieve a user session.
+
+        Short term memory can attempt to create a new session for a given application and user. If a session with the same `session_id` already exists, it will be returned instead of creating a new one.
+
+        If the underlying session service is backed by a database (`DatabaseSessionService`), the method first lists all existing sessions for the given `app_name` and `user_id` and logs the number of sessions found. It then checks whether a session with the specified `session_id` already exists:
+        - If it exists → returns the existing session.
+        - If it does not exist → creates and returns a new session.
+
+        Args:
+            app_name (str): The name of the application associated with the session.
+            user_id (str): The unique identifier of the user.
+            session_id (str): The unique identifier of the session to be created or retrieved.
+
+        Returns:
+            Session | None: The retrieved or newly created `Session` object, or `None` if the session creation failed.
+
+        Examples:
+            Create a new session manually:
+
+            ```python
+            import asyncio
+
+            from veadk.memory import ShortTermMemory
+
+            app_name = "app_name"
+            user_id = "user_id"
+            session_id = "session_id"
+
+            short_term_memory = ShortTermMemory()
+
+            session = asyncio.run(
+                short_term_memory.create_session(
+                    app_name=app_name, user_id=user_id, session_id=session_id
+                )
+            )
+
+            print(session)
+
+            session = asyncio.run(
+                short_term_memory.session_service.get_session(
+                    app_name=app_name, user_id=user_id, session_id=session_id
+                )
+            )
+
+            print(session)
+            ```
+        """
+        if isinstance(self._session_service, DatabaseSessionService):
+            list_sessions_response = await self._session_service.list_sessions(
+                app_name=app_name, user_id=user_id
+            )
+
+            logger.debug(
+                f"Loaded {len(list_sessions_response.sessions)} sessions from db {self.db_url}."
+            )
+
+        session = await self._session_service.get_session(
+            app_name=app_name, user_id=user_id, session_id=session_id
+        )
+
+        if session:
+            logger.info(
+                f"Session {session_id} already exists with app_name={app_name} user_id={user_id}."
+            )
+            return session
+        else:
+            return 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,49 @@
+# 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 import version as adk_version
+from google.adk.sessions import (
+    BaseSessionService,
+    DatabaseSessionService,
+)
+from packaging.version import parse as parse_version
+from pydantic import Field
+from typing_extensions import override
+from urllib.parse import quote_plus
+
+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:
+        encoded_username = quote_plus(self.mysql_config.user)
+        encoded_password = quote_plus(self.mysql_config.password)
+        if parse_version(adk_version.__version__) < parse_version("1.19.0"):
+            self._db_url = f"mysql+pymysql://{encoded_username}:{encoded_password}@{self.mysql_config.host}/{self.mysql_config.database}"
+        else:
+            self._db_url = f"mysql+aiomysql://{encoded_username}:{encoded_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,49 @@
+# 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 urllib.parse import quote_plus
+
+from google.adk import version as adk_version
+from google.adk.sessions import (
+    BaseSessionService,
+    DatabaseSessionService,
+)
+from packaging.version import parse as parse_version
+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:
+        encoded_username = quote_plus(self.postgresql_config.user)
+        encoded_password = quote_plus(self.postgresql_config.password)
+        if parse_version(adk_version.__version__) < parse_version("1.19.0"):
+            self._db_url = f"postgresql://{encoded_username}:{encoded_password}@{self.postgresql_config.host}:{self.postgresql_config.port}/{self.postgresql_config.database}"
+        else:
+            self._db_url = f"postgresql+asyncpg://{encoded_username}:{encoded_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,55 @@
+# 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 import version as adk_version
+from google.adk.sessions import (
+    BaseSessionService,
+    DatabaseSessionService,
+)
+from packaging.version import parse as parse_version
+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:
+        self.local_path = os.path.abspath(self.local_path)
+        # if the DB file not exists, create it
+        if not self._db_exists():
+            os.makedirs(os.path.dirname(self.local_path), exist_ok=True)
+            conn = sqlite3.connect(self.local_path)
+            conn.close()
+
+        if parse_version(adk_version.__version__) < parse_version("1.19.0"):
+            self._db_url = f"sqlite:///{self.local_path}"
+        else:
+            self._db_url = f"sqlite+aiosqlite:///{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)