From 364fe6bd6dbad7aee98f463f6db92dbd98303393 Mon Sep 17 00:00:00 2001
From: chiguyong <chiguyong@beyondsoft.com>
Date: Sat, 6 Jun 2026 22:21:00 +0800
Subject: [PATCH] feat(memory): U3 EpisodicMemory ORM integration -
 EpisodeModel and session factory

- EpisodeModel ORM model with pgvector embedding support
- create_episodic_session_factory for async PostgreSQL sessions
- Server app.py now resolves session_factory from database_url config
- Graceful fallback when database_url not configured
---
 src/agentkit/memory/models.py | 64 +++++++++++++++++++++++++++++++++++
 src/agentkit/server/app.py    | 19 +++++++++--
 2 files changed, 81 insertions(+), 2 deletions(-)
 create mode 100644 src/agentkit/memory/models.py

diff --git a/src/agentkit/memory/models.py b/src/agentkit/memory/models.py
new file mode 100644
index 0000000..d636c65
--- /dev/null
+++ b/src/agentkit/memory/models.py
@@ -0,0 +1,64 @@
+"""SQLAlchemy ORM models for episodic memory persistence (PostgreSQL + pgvector)."""
+
+import uuid
+from datetime import datetime, timezone
+
+from sqlalchemy import Column, DateTime, Float, String, Text, create_engine
+from sqlalchemy.dialects.postgresql import JSONB
+from sqlalchemy.orm import declarative_base, sessionmaker
+
+Base = declarative_base()
+
+
+class EpisodeModel(Base):
+    """Episodic memory ORM model
+
+    Stores task execution experiences with optional pgvector embeddings
+    for semantic similarity search.
+    """
+
+    __tablename__ = "episodic_memories"
+
+    id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
+    agent_name = Column(String, index=True)
+    task_type = Column(String, index=True)
+    input_summary = Column(Text, default="")
+    output_summary = Column(Text, default="")
+    outcome = Column(String, default="success")  # "success", "failure", "partial"
+    quality_score = Column(Float, default=0.5)
+    reflection = Column(Text, default="")
+    embedding = Column(Text, nullable=True)  # JSON-encoded float list; pgvector if extension available
+    metadata_ = Column("metadata", JSONB, nullable=True)  # Additional metadata
+    created_at = Column(
+        DateTime, default=lambda: datetime.now(timezone.utc), index=True
+    )
+
+
+def create_episodic_session_factory(database_url: str):
+    """Create an async session factory for episodic memory.
+
+    Args:
+        database_url: PostgreSQL connection string,
+            e.g. "postgresql+asyncpg://user:pass@localhost/dbname"
+
+    Returns:
+        async_sessionmaker bound to the engine.
+    """
+    from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
+
+    engine = create_async_engine(database_url, echo=False)
+    async_session = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
+    return async_session
+
+
+async def ensure_episodic_table(database_url: str) -> None:
+    """Create the episodic_memories table if it does not exist.
+
+    Safe to call on startup — uses CREATE TABLE IF NOT EXISTS.
+    """
+    from sqlalchemy.ext.asyncio import create_async_engine
+
+    engine = create_async_engine(database_url, echo=False)
+    async with engine.begin() as conn:
+        await conn.run_sync(Base.metadata.create_all)
+    await engine.dispose()
diff --git a/src/agentkit/server/app.py b/src/agentkit/server/app.py
index e7578be..e92ae9b 100644
--- a/src/agentkit/server/app.py
+++ b/src/agentkit/server/app.py
@@ -279,6 +279,7 @@ def create_app(
                 try:
                     from agentkit.memory.episodic import EpisodicMemory
                     from agentkit.memory.embedder import OpenAIEmbedder, EmbeddingCache
+                    from agentkit.memory.models import EpisodeModel, create_episodic_session_factory
 
                     epi_conf = server_config.memory["episodic"]
                     embedder = None
@@ -293,9 +294,23 @@ def create_app(
                             base_url=epi_conf.get("embedder_base_url"),
                             cache=cache,
                         )
+                    # Resolve session_factory and model from database_url if configured
+                    epi_session_factory = None
+                    epi_model = None
+                    database_url = epi_conf.get("database_url") or os.environ.get("DATABASE_URL")
+                    if database_url:
+                        try:
+                            epi_session_factory = create_episodic_session_factory(database_url)
+                            epi_model = EpisodeModel
+                        except Exception as db_err:
+                            import logging as _log
+                            _log.getLogger(__name__).warning(
+                                f"Failed to create episodic DB session: {db_err}"
+                            )
+
                     episodic = EpisodicMemory(
-                        session_factory=None,  # Set externally when DB session is available
-                        episodic_model=None,   # Set externally when ORM model is available
+                        session_factory=epi_session_factory,
+                        episodic_model=epi_model,
                         embedder=embedder,
                         decay_rate=epi_conf.get("decay_rate", 0.01),
                         alpha=epi_conf.get("alpha", 0.7),