fix: short-term-mem test

Author: zakahan

tests/test_short_term_memory.py

@@ -15,7 +15,6 @@
 import asyncio
 import os
 
-import veadk.memory.short_term_memory
 from veadk.memory.short_term_memory import ShortTermMemory
 from veadk.utils.misc import formatted_timestamp
 
@@ -35,11 +34,11 @@ def test_short_term_memory():
     )
     assert session is not None
 
-    # database - local
-    veadk.memory.short_term_memory.DEFAULT_LOCAL_DATABASE_PATH = (
-        f"/tmp/tmp_for_test_{formatted_timestamp()}.db"
+    # sqlite
+    memory = ShortTermMemory(
+        backend="sqlite",
+        local_database_path=f"/tmp/tmp_for_test_{formatted_timestamp()}.db",
     )
-    memory = ShortTermMemory(backend="database")
     asyncio.run(
         memory.session_service.create_session(
             app_name="app", user_id="user", session_id="session"
@@ -51,5 +50,5 @@ def test_short_term_memory():
         )
     )
     assert session is not None
-    assert os.path.exists(veadk.memory.short_term_memory.DEFAULT_LOCAL_DATABASE_PATH)
-    os.remove(veadk.memory.short_term_memory.DEFAULT_LOCAL_DATABASE_PATH)
+    assert os.path.exists(memory.local_database_path)
+    os.remove(memory.local_database_path)

veadk/memory/short_term_memory.py

@@ -50,7 +50,9 @@ def wrapper(*args, **kwargs):
 
 
 class ShortTermMemory(BaseModel):
-    backend: Literal["local", "mysql", "sqlite", "redis", "database"] = "local"
+    backend: Literal["local", "mysql", "sqlite", "redis", "postgresql", "database"] = (
+        "local"
+    )
     """Short term memory backend. `Local` for in-memory storage, `redis` for redis storage, `mysql` for mysql / PostgreSQL storage. `sqlite` for sqlite storage."""
 
     backend_configs: dict = Field(default_factory=dict)

veadk/memory/short_term_memory_backends/postgresql_backend.py

@@ -32,7 +32,7 @@ class PostgreSqlSTMBackend(BaseShortTermMemoryBackend):
     postgresql_config: PostgreSqlConfig = Field(default_factory=PostgreSqlConfig)
 
     def model_post_init(self, context: Any) -> None:
-        self._db_url = f"postgresql://{self.postgresql_config.user}:{self.postgresql_config.password}@{self.postgresql_config.host}:{self.postgresql_config.port}/{self.postgresql_config.database}"
+        self._db_url = f"postgresql+postgresql://{self.postgresql_config.user}:{self.postgresql_config.password}@{self.postgresql_config.host}:{self.postgresql_config.port}/{self.postgresql_config.database}"
 
     @cached_property
     @override