From b5708913d17372a81e172a94c065f4df79844d6e Mon Sep 17 00:00:00 2001 From: 7000pctAUTO Date: Sun, 22 Mar 2026 16:18:53 +0000 Subject: [PATCH] Add memory_manager source files and tests --- src/memory_manager/db/models.py | 161 +++++++++++++++++++------------- 1 file changed, 97 insertions(+), 64 deletions(-) diff --git a/src/memory_manager/db/models.py b/src/memory_manager/db/models.py index bde12f5..842f67f 100644 --- a/src/memory_manager/db/models.py +++ b/src/memory_manager/db/models.py @@ -1,27 +1,19 @@ -"""SQLAlchemy models for memory manager.""" +"""SQLAlchemy database models for the memory manager.""" + import enum from datetime import datetime -from typing import List, Optional +from typing import Any -from sqlalchemy import ( - Column, - DateTime, - Enum, - ForeignKey, - Integer, - String, - Text, - JSON, - Table, -) -from sqlalchemy.orm import DeclarativeBase, relationship, Mapped, mapped_column +from sqlalchemy import JSON, DateTime, Enum, Index, String, Text, func, text +from sqlalchemy.ext.asyncio import create_async_engine +from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column class Base(DeclarativeBase): pass -class MemoryCategory(str, enum.Enum): +class MemoryCategory(enum.StrEnum): DECISION = "decision" FEATURE = "feature" REFACTORING = "refactoring" @@ -30,72 +22,113 @@ class MemoryCategory(str, enum.Enum): NOTE = "note" -memory_tags = Table( - "memory_tags", - Base.metadata, - Column("memory_entry_id", Integer, ForeignKey("memory_entries.id"), primary_key=True), - Column("tag", String(100), primary_key=True), -) - - class MemoryEntry(Base): __tablename__ = "memory_entries" - id: Mapped[int] = mapped_column(Integer, primary_key=True) + id: Mapped[int] = mapped_column(primary_key=True) title: Mapped[str] = mapped_column(String(255), nullable=False) content: Mapped[str] = mapped_column(Text, nullable=False) - category: Mapped[MemoryCategory] = mapped_column( - Enum(MemoryCategory), nullable=False - ) - tags: Mapped[List[str]] = mapped_column(JSON, default=list) - agent_id: Mapped[str] = mapped_column(String(100), default="unknown") - project_path: Mapped[str] = mapped_column(String(500), default=".") - created_at: Mapped[datetime] = mapped_column( - DateTime, default=datetime.utcnow - ) - updated_at: Mapped[datetime] = mapped_column( - DateTime, default=datetime.utcnow, onupdate=datetime.utcnow - ) - parent_id: Mapped[Optional[int]] = mapped_column( - Integer, ForeignKey("memory_entries.id"), nullable=True + category: Mapped[MemoryCategory] = mapped_column(Enum(MemoryCategory), nullable=False) + tags: Mapped[list[str]] = mapped_column(JSON, default=list) + agent_id: Mapped[str] = mapped_column(String(100), nullable=False) + project_path: Mapped[str] = mapped_column(String(500), nullable=False) + created_at: Mapped[datetime] = mapped_column(DateTime, default=func.now()) + updated_at: Mapped[datetime] = mapped_column(DateTime, default=func.now(), onupdate=func.now()) + + __table_args__ = ( + Index("idx_category", "category"), + Index("idx_agent_id", "agent_id"), + Index("idx_project_path", "project_path"), + Index("idx_created_at", "created_at"), ) - parent: Mapped[Optional["MemoryEntry"]] = relationship( - "MemoryEntry", remote_side=[id], back_populates="children" - ) - children: Mapped[List["MemoryEntry"]] = relationship( - "MemoryEntry", back_populates="parent" - ) + def to_dict(self) -> dict[str, Any]: + return { + "id": self.id, + "title": self.title, + "content": self.content, + "category": self.category.value, + "tags": self.tags, + "agent_id": self.agent_id, + "project_path": self.project_path, + "created_at": self.created_at.isoformat() if self.created_at else None, + "updated_at": self.updated_at.isoformat() if self.updated_at else None, + } class Commit(Base): __tablename__ = "commits" - id: Mapped[int] = mapped_column(Integer, primary_key=True) + id: Mapped[int] = mapped_column(primary_key=True) hash: Mapped[str] = mapped_column(String(40), unique=True, nullable=False) message: Mapped[str] = mapped_column(Text, nullable=False) - agent_id: Mapped[str] = mapped_column(String(100), default="unknown") - project_path: Mapped[str] = mapped_column(String(500), default=".") - created_at: Mapped[datetime] = mapped_column( - DateTime, default=datetime.utcnow + agent_id: Mapped[str] = mapped_column(String(100), nullable=False) + project_path: Mapped[str] = mapped_column(String(500), nullable=False) + snapshot: Mapped[list[dict[str, Any]]] = mapped_column(JSON, nullable=False) + created_at: Mapped[datetime] = mapped_column(DateTime, default=func.now()) + + __table_args__ = ( + Index("idx_commit_hash", "hash"), + Index("idx_commit_agent_id", "agent_id"), + Index("idx_commit_created_at", "created_at"), ) - entries: Mapped[List["CommitEntry"]] = relationship( - "CommitEntry", back_populates="commit", cascade="all, delete-orphan" - ) + def to_dict(self) -> dict[str, Any]: + return { + "id": self.id, + "hash": self.hash, + "message": self.message, + "agent_id": self.agent_id, + "project_path": self.project_path, + "snapshot": self.snapshot, + "created_at": self.created_at.isoformat() if self.created_at else None, + } -class CommitEntry(Base): - __tablename__ = "commit_entries" +class Tag(Base): + __tablename__ = "tags" - id: Mapped[int] = mapped_column(Integer, primary_key=True) - commit_id: Mapped[int] = mapped_column( - Integer, ForeignKey("commits.id"), nullable=False - ) - memory_entry_id: Mapped[int] = mapped_column( - Integer, ForeignKey("memory_entries.id"), nullable=False - ) - entry_snapshot: Mapped[str] = mapped_column(JSON, nullable=False) + id: Mapped[int] = mapped_column(primary_key=True) + name: Mapped[str] = mapped_column(String(100), unique=True, nullable=False) + created_at: Mapped[datetime] = mapped_column(DateTime, default=func.now()) - commit: Mapped["Commit"] = relationship("Commit", back_populates="entries") - memory_entry: Mapped["MemoryEntry"] = relationship("MemoryEntry") + def to_dict(self) -> dict[str, Any]: + return { + "id": self.id, + "name": self.name, + "created_at": self.created_at.isoformat() if self.created_at else None, + } + + +async def init_db(db_path: str): + engine = create_async_engine(f"sqlite+aiosqlite:///{db_path}", echo=False) + + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + + await conn.execute(text( + "CREATE VIRTUAL TABLE IF NOT EXISTS memory_entries_fts USING fts5(" + "title, content, tags, category, agent_id, project_path, content='memory_entries', content_rowid='id')" + )) + + await conn.execute(text( + "CREATE TRIGGER IF NOT EXISTS memory_entries_ai AFTER INSERT ON memory_entries BEGIN " + "INSERT INTO memory_entries_fts(rowid, title, content, tags, category, agent_id, project_path) " + "VALUES (new.id, new.title, new.content, new.tags, new.category, new.agent_id, new.project_path); END" + )) + + await conn.execute(text( + "CREATE TRIGGER IF NOT EXISTS memory_entries_ad AFTER DELETE ON memory_entries BEGIN " + "INSERT INTO memory_entries_fts(memory_entries_fts, rowid, title, content, tags, category, agent_id, project_path) " + "VALUES ('delete', old.id, old.title, old.content, old.tags, old.category, old.agent_id, old.project_path); END" + )) + + await conn.execute(text( + "CREATE TRIGGER IF NOT EXISTS memory_entries_au AFTER UPDATE ON memory_entries BEGIN " + "INSERT INTO memory_entries_fts(memory_entries_fts, rowid, title, content, tags, category, agent_id, project_path) " + "VALUES ('delete', old.id, old.title, old.content, old.tags, old.category, old.agent_id, old.project_path); " + "INSERT INTO memory_entries_fts(rowid, title, content, tags, category, agent_id, project_path) " + "VALUES (new.id, new.title, new.content, new.tags, new.category, new.agent_id, new.project_path); END" + )) + + return engine