104 lines
3.7 KiB
Python
104 lines
3.7 KiB
Python
"""SQLAlchemy ORM models for episodic memory persistence (PostgreSQL + pgvector)."""
|
|
|
|
import uuid
|
|
from datetime import datetime, timezone
|
|
|
|
from sqlalchemy import Column, DateTime, Float, String, Text
|
|
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
|
|
|
|
|
|
class ExperienceModel(Base):
|
|
"""Task experience ORM model for RiskGuardLearner / ExperienceStore.
|
|
|
|
Stores task execution outcomes (success/failure/partial) with optional
|
|
pgvector embeddings for semantic similarity search.
|
|
"""
|
|
|
|
__tablename__ = "task_experiences"
|
|
|
|
id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
|
|
task_type = Column(String, index=True)
|
|
goal = Column(Text, default="")
|
|
steps_summary = Column(Text, default="")
|
|
outcome = Column(String, default="success") # "success", "failure", "partial"
|
|
duration_seconds = Column(Float, default=0.0)
|
|
success_rate = Column(Float, default=1.0)
|
|
failure_reasons = Column(JSONB, default=list) # list[str]
|
|
optimization_tips = Column(JSONB, default=list) # list[str]
|
|
embedding = Column(Text, nullable=True) # JSON-encoded float list
|
|
created_at = Column(DateTime, default=lambda: datetime.now(timezone.utc), index=True)
|
|
|
|
|
|
def create_experience_session_factory(database_url: str):
|
|
"""Create an async session factory for task experiences.
|
|
|
|
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()
|