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:
parent
f16dcb5ebe
commit
364fe6bd6d
|
|
@ -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()
|
||||
|
|
@ -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),
|
||||
|
|
|
|||
Loading…
Reference in New Issue