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
This commit is contained in:
chiguyong 2026-06-06 22:21:00 +08:00
parent f16dcb5ebe
commit 364fe6bd6d
2 changed files with 81 additions and 2 deletions

View File

@ -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()

View File

@ -279,6 +279,7 @@ def create_app(
try: try:
from agentkit.memory.episodic import EpisodicMemory from agentkit.memory.episodic import EpisodicMemory
from agentkit.memory.embedder import OpenAIEmbedder, EmbeddingCache from agentkit.memory.embedder import OpenAIEmbedder, EmbeddingCache
from agentkit.memory.models import EpisodeModel, create_episodic_session_factory
epi_conf = server_config.memory["episodic"] epi_conf = server_config.memory["episodic"]
embedder = None embedder = None
@ -293,9 +294,23 @@ def create_app(
base_url=epi_conf.get("embedder_base_url"), base_url=epi_conf.get("embedder_base_url"),
cache=cache, 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( episodic = EpisodicMemory(
session_factory=None, # Set externally when DB session is available session_factory=epi_session_factory,
episodic_model=None, # Set externally when ORM model is available episodic_model=epi_model,
embedder=embedder, embedder=embedder,
decay_rate=epi_conf.get("decay_rate", 0.01), decay_rate=epi_conf.get("decay_rate", 0.01),
alpha=epi_conf.get("alpha", 0.7), alpha=epi_conf.get("alpha", 0.7),