feat: add db_kwargs for sql connection config (#338)

Author: zakahan

veadk/memory/short_term_memory.py

@@ -126,6 +126,8 @@ class ShortTermMemory(BaseModel):
 
     backend_configs: dict = Field(default_factory=dict)
 
+    db_kwargs: dict = Field(default_factory=dict)
+
     db_url: str = ""
 
     local_database_path: str = "/tmp/veadk_local_database.db"
@@ -143,7 +145,9 @@ def model_post_init(self, __context: Any) -> None:
                     "Please encode `username` or `password` with `urllib.parse.quote_plus`. "
                     "Examples: p@ssword→p%40ssword."
                 )
-            self._session_service = DatabaseSessionService(db_url=self.db_url)
+            self._session_service = DatabaseSessionService(
+                db_url=self.db_url, **self.db_kwargs
+            )
         else:
             if self.backend == "database":
                 logger.warning(
@@ -155,15 +159,15 @@ def model_post_init(self, __context: Any) -> None:
                     self._session_service = InMemorySessionService()
                 case "mysql":
                     self._session_service = MysqlSTMBackend(
-                        **self.backend_configs
+                        db_kwargs=self.db_kwargs, **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
+                        db_kwargs=self.db_kwargs, **self.backend_configs
                     ).session_service
 
         if self.after_load_memory_callback:

veadk/memory/short_term_memory_backends/mysql_backend.py

@@ -34,6 +34,7 @@
 
 class MysqlSTMBackend(BaseShortTermMemoryBackend):
     mysql_config: MysqlConfig = Field(default_factory=MysqlConfig)
+    db_kwargs: dict = Field(default_factory=dict)
 
     def model_post_init(self, context: Any) -> None:
         encoded_username = quote_plus(self.mysql_config.user)
@@ -46,4 +47,4 @@ def model_post_init(self, context: Any) -> None:
     @cached_property
     @override
     def session_service(self) -> BaseSessionService:
-        return DatabaseSessionService(db_url=self._db_url)
+        return DatabaseSessionService(db_url=self._db_url, **self.db_kwargs)

veadk/memory/short_term_memory_backends/postgresql_backend.py

@@ -34,6 +34,7 @@
 
 class PostgreSqlSTMBackend(BaseShortTermMemoryBackend):
     postgresql_config: PostgreSqlConfig = Field(default_factory=PostgreSqlConfig)
+    db_kwargs: dict = Field(default_factory=dict)
 
     def model_post_init(self, context: Any) -> None:
         encoded_username = quote_plus(self.postgresql_config.user)
@@ -46,4 +47,4 @@ def model_post_init(self, context: Any) -> None:
     @cached_property
     @override
     def session_service(self) -> BaseSessionService:
-        return DatabaseSessionService(db_url=self._db_url)
+        return DatabaseSessionService(db_url=self._db_url, **self.db_kwargs)