veadk-python 0.2.7__py3-none-any.whl → 0.2.9__py3-none-any.whl

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/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/runner.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import asyncio
+import os
 import functools
 from types import MethodType
 from typing import Union
@@ -47,7 +47,7 @@ RunnerMessage = Union[
 ]
 
 
-def pre_run_process(self, process_func, new_message, user_id, session_id):
+async def pre_run_process(self, process_func, new_message, user_id, session_id):
     if new_message.parts:
         for part in new_message.parts:
             if (
@@ -55,13 +55,12 @@ def pre_run_process(self, process_func, new_message, user_id, session_id):
                 and part.inline_data.mime_type == "image/png"
                 and self.upload_inline_data_to_tos
             ):
-                process_func(
+                await process_func(
                     part,
                     self.app_name,
                     user_id,
                     session_id,
                 )
-    return
 
 
 def post_run_process(self):
@@ -79,7 +78,7 @@ def intercept_new_message(process_func):
             new_message: types.Content,
             **kwargs,
         ):
-            pre_run_process(self, process_func, new_message, user_id, session_id)
+            await pre_run_process(self, process_func, new_message, user_id, session_id)
 
             async for event in func(
                 user_id=user_id,
@@ -137,27 +136,21 @@ def _convert_messages(
     return _messages
 
 
-def _upload_image_to_tos(
+async def _upload_image_to_tos(
     part: genai.types.Part, app_name: str, user_id: str, session_id: str
 ) -> None:
     try:
         if part.inline_data and part.inline_data.display_name and part.inline_data.data:
             from veadk.integrations.ve_tos.ve_tos import VeTOS
 
+            filename = os.path.basename(part.inline_data.display_name)
+            object_key = f"{app_name}/{user_id}-{session_id}-{filename}"
             ve_tos = VeTOS()
-
-            object_key, tos_url = ve_tos.build_tos_url(
-                user_id=user_id,
-                app_name=app_name,
-                session_id=session_id,
-                data_path=part.inline_data.display_name,
+            tos_url = ve_tos.build_tos_signed_url(object_key=object_key)
+            await ve_tos.async_upload_bytes(
+                object_key=object_key,
+                data=part.inline_data.data,
             )
-
-            upload_task = ve_tos.upload(object_key, part.inline_data.data)
-
-            if upload_task is not None:
-                asyncio.create_task(upload_task)
-
             part.inline_data.display_name = tos_url
     except Exception as e:
         logger.error(f"Upload to TOS failed: {e}")
@@ -226,7 +219,7 @@ class Runner(ADKRunner):
         )
 
         self.run_async = MethodType(
-            intercept_new_message(_upload_image_to_tos)(self.run_async), self
+            intercept_new_message(_upload_image_to_tos)(super().run_async), self
         )
 
     async def run(