fix: short-term-mem postgresql_backend

Author: zakahan

pyproject.toml

@@ -33,6 +33,8 @@ dependencies = [
     "llama-index-llms-openai-like>=0.5.1",
     "llama-index-vector-stores-opensearch>=0.6.1",
     "llama-index-vector-stores-redis>=0.6.1",
+    "mcp-server-vikingdb-memory",
+    "psycopg2-binary>=2.9.10",
 ]
 
 [project.scripts]
@@ -83,3 +85,6 @@ exclude = [
     "veadk/integrations/ve_faas/template/*",
     "veadk/integrations/ve_faas/web_template/*"
 ]
+
+[tool.uv.sources]
+mcp-server-vikingdb-memory = { git = "https://github.com/volcengine/mcp-server", subdirectory = "server/mcp_server_vikingdb_memory" }

veadk/memory/short_term_memory.py

@@ -28,7 +28,6 @@
 from veadk.memory.short_term_memory_backends.postgresql_backend import (
     PostgreSqlSTMBackend,
 )
-from veadk.memory.short_term_memory_backends.redis_backend import RedisSTMBackend
 from veadk.memory.short_term_memory_backends.sqlite_backend import (
     SQLiteSTMBackend,
 )
@@ -50,10 +49,8 @@ def wrapper(*args, **kwargs):
 
 
 class ShortTermMemory(BaseModel):
-    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: 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."""
 
     backend_configs: dict = Field(default_factory=dict)
     """Backend specific configurations."""
@@ -90,10 +87,6 @@ def model_post_init(self, __context: Any) -> None:
                     self._session_service = SQLiteSTMBackend(
                         local_path=self.local_database_path
                     ).session_service
-                case "redis":
-                    self._session_service = RedisSTMBackend(
-                        **self.backend_configs
-                    ).session_service
                 case "postgresql":
                     self._session_service = PostgreSqlSTMBackend(
                         **self.backend_configs

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+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+psycopg2://{self.postgresql_config.user}:{self.postgresql_config.password}@{self.postgresql_config.host}:{self.postgresql_config.port}/{self.postgresql_config.database}"
 
     @cached_property
     @override