add callback mount

Author: yaozheng-fang

veadk/memory/short_term_memory.py

@@ -12,6 +12,7 @@
 # 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 DatabaseSessionService, InMemorySessionService
@@ -36,6 +37,18 @@
 DEFAULT_LOCAL_DATABASE_PATH = "/tmp/veadk_local_database.db"
 
 
+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):
     backend: Literal["local", "mysql", "sqlite", "redis", "database"] = "local"
     """Short term memory backend. `Local` for in-memory storage, `redis` for redis storage, `mysql` for mysql / PostgreSQL storage. `sqlite` for sqlite storage."""
@@ -46,8 +59,8 @@ class ShortTermMemory(BaseModel):
     db_url: str = ""
     """Database connection URL, e.g. `sqlite:///./test.db`. Once set, it will override the `backend` parameter."""
 
-    after_load_memory_callbacks: list[Callable] | None = None
-    """A list of callbacks to be called after loading memory from the backend. The callback function should accept `Session` as an input."""
+    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 model_post_init(self, __context: Any) -> None:
         if self.db_url:
@@ -78,6 +91,11 @@ def model_post_init(self, __context: Any) -> None:
                         **self.backend_configs
                     ).session_service
 
+        if self.after_load_memory_callback:
+            wrap_get_session_with_callbacks(
+                self.session_service, self.after_load_memory_callback
+            )
+
     async def create_session(
         self,
         app_name: str,