From a01bd4abfc9525fc98765cacd96a4074d181b9bd Mon Sep 17 00:00:00 2001 From: Ariel Memory Date: Fri, 7 Aug 2026 23:21:37 +0200 Subject: [PATCH] feat: implement Alembic migrations for memory.db --- alembic.ini | 149 +++++++ alembic/README | 1 + alembic/env.py | 87 ++++ alembic/script.py.mako | 28 ++ .../versions/a38d67fcd99e_init_v8_schema.py | 275 ++++++++++++ shared/migrations.py | 393 ++---------------- 6 files changed, 581 insertions(+), 352 deletions(-) create mode 100644 alembic.ini create mode 100644 alembic/README create mode 100644 alembic/env.py create mode 100644 alembic/script.py.mako create mode 100644 alembic/versions/a38d67fcd99e_init_v8_schema.py diff --git a/alembic.ini b/alembic.ini new file mode 100644 index 00000000..8cc5ddd5 --- /dev/null +++ b/alembic.ini @@ -0,0 +1,149 @@ +# A generic, single database configuration. + +[alembic] +# path to migration scripts. +# this is typically a path given in POSIX (e.g. forward slashes) +# format, relative to the token %(here)s which refers to the location of this +# ini file +script_location = %(here)s/alembic + +# template used to generate migration file names; The default value is %%(rev)s_%%(slug)s +# Uncomment the line below if you want the files to be prepended with date and time +# see https://alembic.sqlalchemy.org/en/latest/tutorial.html#editing-the-ini-file +# for all available tokens +# file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s +# Or organize into date-based subdirectories (requires recursive_version_locations = true) +# file_template = %%(year)d/%%(month).2d/%%(day).2d_%%(hour).2d%%(minute).2d_%%(second).2d_%%(rev)s_%%(slug)s + +# sys.path path, will be prepended to sys.path if present. +# defaults to the current working directory. for multiple paths, the path separator +# is defined by "path_separator" below. +prepend_sys_path = . + + +# timezone to use when rendering the date within the migration file +# as well as the filename. +# If specified, requires the tzdata library which can be installed by adding +# `alembic[tz]` to the pip requirements. +# string value is passed to ZoneInfo() +# leave blank for localtime +# timezone = + +# max length of characters to apply to the "slug" field +# truncate_slug_length = 40 + +# set to 'true' to run the environment during +# the 'revision' command, regardless of autogenerate +# revision_environment = false + +# set to 'true' to allow .pyc and .pyo files without +# a source .py file to be detected as revisions in the +# versions/ directory +# sourceless = false + +# version location specification; This defaults +# to /versions. When using multiple version +# directories, initial revisions must be specified with --version-path. +# The path separator used here should be the separator specified by "path_separator" +# below. +# version_locations = %(here)s/bar:%(here)s/bat:%(here)s/alembic/versions + +# path_separator; This indicates what character is used to split lists of file +# paths, including version_locations and prepend_sys_path within configparser +# files such as alembic.ini. +# The default rendered in new alembic.ini files is "os", which uses os.pathsep +# to provide os-dependent path splitting. +# +# Note that in order to support legacy alembic.ini files, this default does NOT +# take place if path_separator is not present in alembic.ini. If this +# option is omitted entirely, fallback logic is as follows: +# +# 1. Parsing of the version_locations option falls back to using the legacy +# "version_path_separator" key, which if absent then falls back to the legacy +# behavior of splitting on spaces and/or commas. +# 2. Parsing of the prepend_sys_path option falls back to the legacy +# behavior of splitting on spaces, commas, or colons. +# +# Valid values for path_separator are: +# +# path_separator = : +# path_separator = ; +# path_separator = space +# path_separator = newline +# +# Use os.pathsep. Default configuration used for new projects. +path_separator = os + +# set to 'true' to search source files recursively +# in each "version_locations" directory +# new in Alembic version 1.10 +# recursive_version_locations = false + +# the output encoding used when revision files +# are written from script.py.mako +# output_encoding = utf-8 + +# database URL. This is consumed by the user-maintained env.py script only. +# other means of configuring database URLs may be customized within the env.py +# file. +sqlalchemy.url = sqlite:///memory.db + + +[post_write_hooks] +# post_write_hooks defines scripts or Python functions that are run +# on newly generated revision scripts. See the documentation for further +# detail and examples + +# format using "black" - use the console_scripts runner, against the "black" entrypoint +# hooks = black +# black.type = console_scripts +# black.entrypoint = black +# black.options = -l 79 REVISION_SCRIPT_FILENAME + +# lint with attempts to fix using "ruff" - use the module runner, against the "ruff" module +# hooks = ruff +# ruff.type = module +# ruff.module = ruff +# ruff.options = check --fix REVISION_SCRIPT_FILENAME + +# Alternatively, use the exec runner to execute a binary found on your PATH +# hooks = ruff +# ruff.type = exec +# ruff.executable = ruff +# ruff.options = check --fix REVISION_SCRIPT_FILENAME + +# Logging configuration. This is also consumed by the user-maintained +# env.py script only. +[loggers] +keys = root,sqlalchemy,alembic + +[handlers] +keys = console + +[formatters] +keys = generic + +[logger_root] +level = WARNING +handlers = console +qualname = + +[logger_sqlalchemy] +level = WARNING +handlers = +qualname = sqlalchemy.engine + +[logger_alembic] +level = INFO +handlers = +qualname = alembic + +[handler_console] +class = StreamHandler +args = (sys.stderr,) +level = NOTSET +formatter = generic + +[formatter_generic] +format = %(levelname)-5.5s [%(name)s] %(message)s +datefmt = %H:%M:%S diff --git a/alembic/README b/alembic/README new file mode 100644 index 00000000..98e4f9c4 --- /dev/null +++ b/alembic/README @@ -0,0 +1 @@ +Generic single-database configuration. \ No newline at end of file diff --git a/alembic/env.py b/alembic/env.py new file mode 100644 index 00000000..86714e7e --- /dev/null +++ b/alembic/env.py @@ -0,0 +1,87 @@ +import os +from pathlib import Path +from logging.config import fileConfig + +from sqlalchemy import engine_from_config +from sqlalchemy import pool + +from alembic import context + +# this is the Alembic Config object, which provides +# access to the values within the .ini file in use. +config = context.config + +def get_url(): + data_dir = os.environ.get("MCP_MEMORY_DATA_DIR", str(Path.home() / ".mcp-ariel-memory")) + db_path = Path(data_dir) / "memory.db" + return f"sqlite:///{db_path}" + +# Interpret the config file for Python logging. +# This line sets up loggers basically. +if config.config_file_name is not None: + fileConfig(config.config_file_name) + +# add your model's MetaData object here +# for 'autogenerate' support +# from myapp import mymodel +# target_metadata = mymodel.Base.metadata +target_metadata = None + +# other values from the config, defined by the needs of env.py, +# can be acquired: +# my_important_option = config.get_main_option("my_important_option") +# ... etc. + + +def run_migrations_offline() -> None: + """Run migrations in 'offline' mode. + + This configures the context with just a URL + and not an Engine, though an Engine is acceptable + here as well. By skipping the Engine creation + we don't even need a DBAPI to be available. + + Calls to context.execute() here emit the given string to the + script output. + + """ + url = get_url() + context.configure( + url=url, + target_metadata=target_metadata, + literal_binds=True, + dialect_opts={"paramstyle": "named"}, + ) + + with context.begin_transaction(): + context.run_migrations() + + +def run_migrations_online() -> None: + """Run migrations in 'online' mode. + + In this scenario we need to create an Engine + and associate a connection with the context. + + """ + configuration = config.get_section(config.config_ini_section, {}) + configuration["sqlalchemy.url"] = get_url() + connectable = engine_from_config( + configuration, + prefix="sqlalchemy.", + poolclass=pool.NullPool, + ) + + with connectable.connect() as connection: + context.configure( + connection=connection, target_metadata=target_metadata + ) + + with context.begin_transaction(): + context.run_migrations() + + +if context.is_offline_mode(): + run_migrations_offline() +else: + run_migrations_online() diff --git a/alembic/script.py.mako b/alembic/script.py.mako new file mode 100644 index 00000000..11016301 --- /dev/null +++ b/alembic/script.py.mako @@ -0,0 +1,28 @@ +"""${message} + +Revision ID: ${up_revision} +Revises: ${down_revision | comma,n} +Create Date: ${create_date} + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +${imports if imports else ""} + +# revision identifiers, used by Alembic. +revision: str = ${repr(up_revision)} +down_revision: Union[str, Sequence[str], None] = ${repr(down_revision)} +branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)} +depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)} + + +def upgrade() -> None: + """Upgrade schema.""" + ${upgrades if upgrades else "pass"} + + +def downgrade() -> None: + """Downgrade schema.""" + ${downgrades if downgrades else "pass"} diff --git a/alembic/versions/a38d67fcd99e_init_v8_schema.py b/alembic/versions/a38d67fcd99e_init_v8_schema.py new file mode 100644 index 00000000..ea1905d4 --- /dev/null +++ b/alembic/versions/a38d67fcd99e_init_v8_schema.py @@ -0,0 +1,275 @@ +"""init_v8_schema + +Revision ID: a38d67fcd99e +Revises: +Create Date: 2026-08-07 23:05:00.000000 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = 'a38d67fcd99e' +down_revision: Union[str, None] = None +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # 1. Core tables + op.execute(""" + CREATE TABLE IF NOT EXISTS core_memory ( + entry_id INTEGER PRIMARY KEY AUTOINCREMENT, + user_id TEXT NOT NULL, key TEXT NOT NULL, value TEXT NOT NULL, + importance REAL DEFAULT 0.5, is_conflict INTEGER DEFAULT 0, + conflict_group_id TEXT, created_at REAL NOT NULL, updated_at REAL NOT NULL, + memory_kind TEXT, expires_at REAL, source TEXT DEFAULT 'manual', metadata TEXT + ) + """) + op.execute("CREATE INDEX IF NOT EXISTS idx_core_user ON core_memory(user_id)") + op.execute("CREATE UNIQUE INDEX IF NOT EXISTS idx_core_user_key ON core_memory(user_id, key)") + op.execute("CREATE INDEX IF NOT EXISTS idx_core_memory_kind ON core_memory(user_id, memory_kind)") + + op.execute(""" + CREATE TABLE IF NOT EXISTS sessions ( + session_id TEXT PRIMARY KEY, user_id TEXT NOT NULL, summary TEXT, + state_deltas TEXT, topics TEXT, message_count INTEGER DEFAULT 0, + started_at REAL NOT NULL, ended_at REAL + ) + """) + op.execute("CREATE INDEX IF NOT EXISTS idx_sessions_user ON sessions(user_id)") + + op.execute(""" + CREATE TABLE IF NOT EXISTS episodes ( + episode_id INTEGER PRIMARY KEY AUTOINCREMENT, + user_id TEXT NOT NULL, summary TEXT NOT NULL, + emotional_weight REAL DEFAULT 0.5, tags TEXT, created_at REAL NOT NULL, + memory_kind TEXT + ) + """) + op.execute("CREATE INDEX IF NOT EXISTS idx_episodes_user ON episodes(user_id)") + op.execute("CREATE INDEX IF NOT EXISTS idx_episodes_kind ON episodes(user_id, memory_kind)") + + # 2. Support tables + op.execute(""" + CREATE TABLE IF NOT EXISTS staging_memories ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + user_id TEXT NOT NULL DEFAULT 'default', session_id TEXT NOT NULL, + event_id TEXT, content TEXT NOT NULL, importance REAL DEFAULT 0.5, + metadata TEXT DEFAULT '{}', created_at DATETIME DEFAULT CURRENT_TIMESTAMP + ) + """) + + op.execute(""" + CREATE TABLE IF NOT EXISTS archived_memories ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + user_id TEXT NOT NULL DEFAULT 'default', original_id INTEGER, + content TEXT NOT NULL, memory_type TEXT, importance REAL, + archive_reason TEXT NOT NULL, archived_at DATETIME DEFAULT CURRENT_TIMESTAMP + ) + """) + + op.execute(""" + CREATE TABLE IF NOT EXISTS audit_log ( + log_id INTEGER PRIMARY KEY AUTOINCREMENT, + user_id TEXT NOT NULL, action TEXT NOT NULL, layer TEXT, + target_id TEXT, details TEXT, timestamp REAL NOT NULL + ) + """) + + op.execute(""" + CREATE TABLE IF NOT EXISTS rate_limits ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + user_id TEXT NOT NULL, timestamp REAL NOT NULL + ) + """) + + op.execute(""" + CREATE TABLE IF NOT EXISTS embedding_cache ( + text_hash TEXT PRIMARY KEY, embedding BLOB NOT NULL, + model_name TEXT NOT NULL, created_at DATETIME DEFAULT CURRENT_TIMESTAMP + ) + """) + + # 3. RAG + op.execute(""" + CREATE TABLE IF NOT EXISTS rag_pages ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + layer TEXT NOT NULL DEFAULT 'user', user_id TEXT NOT NULL DEFAULT 'default', + title TEXT NOT NULL, path TEXT, content TEXT NOT NULL, + sha256_hash TEXT, wiki_type TEXT, + created_at REAL DEFAULT (strftime('%s','now')), + updated_at REAL DEFAULT (strftime('%s','now')) + ) + """) + op.execute("CREATE INDEX IF NOT EXISTS idx_rag_user ON rag_pages(user_id)") + + op.execute(""" + CREATE TABLE IF NOT EXISTS rag_chunks ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + page_id INTEGER NOT NULL, chunk_index INTEGER NOT NULL, + content TEXT NOT NULL, bin_embedding BLOB, memory_kind TEXT + ) + """) + op.execute("CREATE INDEX IF NOT EXISTS idx_rag_chunks_bin ON rag_chunks(page_id, id) WHERE bin_embedding IS NOT NULL") + op.execute("CREATE INDEX IF NOT EXISTS idx_rag_chunks_page_idx ON rag_chunks(page_id, chunk_index)") + + op.execute(""" + CREATE TABLE IF NOT EXISTS rag_relations ( + source_id INTEGER NOT NULL, target_id INTEGER NOT NULL, + relation_type TEXT NOT NULL DEFAULT 'elaborates', + weight REAL DEFAULT 0.8, + PRIMARY KEY (source_id, target_id, relation_type) + ) + """) + + # 4. Graph + op.execute(""" + CREATE TABLE IF NOT EXISTS epi_nodes ( + node_id INTEGER PRIMARY KEY AUTOINCREMENT, + layer TEXT NOT NULL DEFAULT 'user', + user_id TEXT NOT NULL, content TEXT NOT NULL, + node_type TEXT NOT NULL, tags TEXT, + confidence REAL DEFAULT 0.5, created_at REAL NOT NULL + ) + """) + + op.execute(""" + CREATE TABLE IF NOT EXISTS epi_edges ( + source_id INTEGER NOT NULL, target_id INTEGER NOT NULL, + relation TEXT NOT NULL, weight REAL DEFAULT 0.8, + created_at REAL NOT NULL, + PRIMARY KEY (source_id, target_id, relation) + ) + """) + + op.execute(""" + CREATE TABLE IF NOT EXISTS epi_tags ( + node_id INTEGER NOT NULL, + tag TEXT NOT NULL, + PRIMARY KEY (node_id, tag) + ) + """) + op.execute("CREATE INDEX IF NOT EXISTS idx_epi_tags_tag ON epi_tags(tag)") + + op.execute(""" + CREATE TABLE IF NOT EXISTS temporal_events ( + event_id INTEGER PRIMARY KEY AUTOINCREMENT, + user_id TEXT NOT NULL, event_type TEXT NOT NULL, + content TEXT NOT NULL, timestamp REAL NOT NULL, + importance REAL DEFAULT 0.5, metadata TEXT + ) + """) + + op.execute(""" + CREATE TABLE IF NOT EXISTS temporal_links ( + from_event INTEGER NOT NULL, to_event INTEGER NOT NULL, + link_type TEXT NOT NULL DEFAULT 'follows', + strength REAL DEFAULT 0.5, + PRIMARY KEY (from_event, to_event, link_type) + ) + """) + + # 5. Wiki + op.execute(""" + CREATE TABLE IF NOT EXISTS user_wiki ( + entry_id INTEGER PRIMARY KEY AUTOINCREMENT, + user_id TEXT NOT NULL, wiki_type TEXT NOT NULL, + title TEXT NOT NULL, content TEXT NOT NULL, + tags TEXT, importance REAL DEFAULT 0.5, + source TEXT DEFAULT 'manual', + created_at REAL NOT NULL, updated_at REAL NOT NULL + ) + """) + + op.execute(""" + CREATE TABLE IF NOT EXISTS agent_wiki ( + entry_id INTEGER PRIMARY KEY AUTOINCREMENT, + user_id TEXT NOT NULL, wiki_type TEXT NOT NULL, + title TEXT NOT NULL, content TEXT NOT NULL, + tags TEXT, importance REAL DEFAULT 0.5, + source TEXT DEFAULT 'manual', + created_at REAL NOT NULL, updated_at REAL NOT NULL + ) + """) + + op.execute(""" + CREATE TABLE IF NOT EXISTS wiki_index ( + entry_id INTEGER PRIMARY KEY AUTOINCREMENT, + layer TEXT NOT NULL, wiki_type TEXT NOT NULL, + title TEXT NOT NULL, file_path TEXT NOT NULL, + tags TEXT, importance REAL DEFAULT 0.5, + content TEXT DEFAULT '', content_hash TEXT, + created_at REAL NOT NULL, updated_at REAL NOT NULL + ) + """) + op.execute("CREATE UNIQUE INDEX IF NOT EXISTS idx_wiki_path ON wiki_index(file_path)") + + # 6. Registry + op.execute(""" + CREATE TABLE IF NOT EXISTS memory_kind_registry ( + kind TEXT PRIMARY KEY, + display_name TEXT NOT NULL, + default_importance REAL NOT NULL, + decay_rate REAL NOT NULL, + never_archive INTEGER NOT NULL DEFAULT 0, + requires_expires_at INTEGER NOT NULL DEFAULT 0, + boost_on_keywords TEXT NOT NULL DEFAULT '', + description TEXT + ) + """) + + op.execute(""" + INSERT OR IGNORE INTO memory_kind_registry VALUES + ('instruction','Instruction',0.9,0.0,1,0,'обязательно,важно,critical,never forget,rule,инструкция','Правило/инструкция, не подлежит забыванию'), + ('fact','Fact',0.5,0.01,0,0,'факт,fact,имя,возраст,день рождения','Атомарный факт'), + ('decision','Decision',0.7,0.005,0,0,'решение,decided,chose,decision','Принятое решение'), + ('goal','Goal',0.8,0.005,0,1,'цель,goal,plan,к концу','Цель с дедлайном'), + ('preference','Preference',0.7,0.003,0,0,'предпочитаю,prefer,like,нравится,не люблю','Предпочтение'), + ('commitment','Commitment',0.85,0.0,1,1,'обещаю,обязуюсь,commit,promise,согласен','Обязательство'), + ('relationship','Relationship',0.6,0.002,0,0,'знаком,друг,коллега,knows,friend','Связь'), + ('observation','Observation',0.4,0.02,0,0,'видел,заметил,noticed,observed','Наблюдение'), + ('rule','Rule',0.85,0.0,1,0,'запрещено,нельзя,do not,forbidden,rule','Жёсткое правило'), + ('todo','Todo',0.6,0.005,0,1,'todo,сделать,do later,remind','Задача с дедлайном'), + ('question','Open Question',0.5,0.05,0,0,'вопрос,?,уточнить,ask later','Открытый вопрос'), + ('hypothesis','Hypothesis',0.45,0.03,0,0,'возможно,наверное,probably,hypothesis','Гипотеза'), + ('context','Context',0.3,0.05,0,0,'контекст,background,context','Фоновый контекст') + """) + + # 7. Audit & Conflicts + op.execute(""" + CREATE TABLE IF NOT EXISTS importance_audit ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + user_id TEXT NOT NULL, + chunk_id INTEGER, + source TEXT NOT NULL, + old_importance REAL, + new_importance REAL, + signal_breakdown TEXT, + reason TEXT, + rescored_at REAL NOT NULL + ) + """) + op.execute("CREATE INDEX IF NOT EXISTS idx_importance_audit_user ON importance_audit(user_id, rescored_at DESC)") + + op.execute(""" + CREATE TABLE IF NOT EXISTS memory_conflicts ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + user_id TEXT NOT NULL, content TEXT NOT NULL, + is_conflict INTEGER DEFAULT 0, conflict_group_id TEXT, + created_at DATETIME DEFAULT CURRENT_TIMESTAMP + ) + """) + + # 8. FTS5 (must be done statement by statement) + op.execute("CREATE VIRTUAL TABLE IF NOT EXISTS rag_fts USING fts5(title, content, wiki_type, content=rag_pages, content_rowid=id)") + op.execute("CREATE VIRTUAL TABLE IF NOT EXISTS user_wiki_fts USING fts5(title, content, wiki_type, tags, content=user_wiki, content_rowid=entry_id)") + op.execute("CREATE VIRTUAL TABLE IF NOT EXISTS agent_wiki_fts USING fts5(title, content, wiki_type, tags, content=agent_wiki, content_rowid=entry_id)") + op.execute("CREATE VIRTUAL TABLE IF NOT EXISTS wiki_fts USING fts5(title, content, wiki_type, tags, content=wiki_index, content_rowid=entry_id)") + + +def downgrade() -> None: + pass # Baseline migration doesn't drop anything for safety diff --git a/shared/migrations.py b/shared/migrations.py index ae8323e8..d126f255 100644 --- a/shared/migrations.py +++ b/shared/migrations.py @@ -1,377 +1,66 @@ """ -DB Migrations — async, unified memory.db -All tables in one file. wiki/graph/audit can be split out later under load. +DB Migrations — async, unified memory.db using Alembic """ import logging -import sqlite3 -import time -from collections.abc import Callable +import os +import asyncio +from pathlib import Path from typing import Any +from alembic import command as alembic_command +from alembic.config import Config as AlembicConfig from shared.connection import AsyncConnectionManager, connection_manager -import contextlib logger = logging.getLogger(__name__) DB_NAME = "memory.db" -class Migration: - def __init__(self, version: int, name: str, up: Callable): - self.version = version - self.name = name - self.up = up - - -def _get_migrations() -> list[Migration]: - migrations = [] - - async def v1_init(conn): - """All tables in a single memory.db.""" - await conn.executescript(""" - -- === L2-L4 Core === - CREATE TABLE IF NOT EXISTS core_memory ( - entry_id INTEGER PRIMARY KEY AUTOINCREMENT, - user_id TEXT NOT NULL, key TEXT NOT NULL, value TEXT NOT NULL, - importance REAL DEFAULT 0.5, is_conflict INTEGER DEFAULT 0, - conflict_group_id TEXT, created_at REAL NOT NULL, updated_at REAL NOT NULL - ); - CREATE INDEX IF NOT EXISTS idx_core_user ON core_memory(user_id); - CREATE UNIQUE INDEX IF NOT EXISTS idx_core_user_key ON core_memory(user_id, key); - - CREATE TABLE IF NOT EXISTS sessions ( - session_id TEXT PRIMARY KEY, user_id TEXT NOT NULL, summary TEXT, - state_deltas TEXT, topics TEXT, message_count INTEGER DEFAULT 0, - started_at REAL NOT NULL, ended_at REAL - ); - CREATE INDEX IF NOT EXISTS idx_sessions_user ON sessions(user_id); - - CREATE TABLE IF NOT EXISTS episodes ( - episode_id INTEGER PRIMARY KEY AUTOINCREMENT, - user_id TEXT NOT NULL, summary TEXT NOT NULL, - emotional_weight REAL DEFAULT 0.5, tags TEXT, created_at REAL NOT NULL - ); - CREATE INDEX IF NOT EXISTS idx_episodes_user ON episodes(user_id); - - -- === Staging + Archived === - CREATE TABLE IF NOT EXISTS staging_memories ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - user_id TEXT NOT NULL DEFAULT 'default', session_id TEXT NOT NULL, - event_id TEXT, content TEXT NOT NULL, importance REAL DEFAULT 0.5, - metadata TEXT DEFAULT '{}', created_at DATETIME DEFAULT CURRENT_TIMESTAMP - ); - - CREATE TABLE IF NOT EXISTS archived_memories ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - user_id TEXT NOT NULL DEFAULT 'default', original_id INTEGER, - content TEXT NOT NULL, memory_type TEXT, importance REAL, - archive_reason TEXT NOT NULL, archived_at DATETIME DEFAULT CURRENT_TIMESTAMP - ); - - -- === Audit === - CREATE TABLE IF NOT EXISTS audit_log ( - log_id INTEGER PRIMARY KEY AUTOINCREMENT, - user_id TEXT NOT NULL, action TEXT NOT NULL, layer TEXT, - target_id TEXT, details TEXT, timestamp REAL NOT NULL - ); - - -- === Rate Limit === - CREATE TABLE IF NOT EXISTS rate_limits ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - user_id TEXT NOT NULL, timestamp REAL NOT NULL - ); - - -- === Embeddings === - CREATE TABLE IF NOT EXISTS embedding_cache ( - text_hash TEXT PRIMARY KEY, embedding BLOB NOT NULL, - model_name TEXT NOT NULL, created_at DATETIME DEFAULT CURRENT_TIMESTAMP - ); - - -- === RAG === - CREATE TABLE IF NOT EXISTS rag_pages ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - layer TEXT NOT NULL DEFAULT 'user', user_id TEXT NOT NULL DEFAULT 'default', - title TEXT NOT NULL, path TEXT, content TEXT NOT NULL, - sha256_hash TEXT, wiki_type TEXT, - created_at REAL DEFAULT (strftime('%s','now')), - updated_at REAL DEFAULT (strftime('%s','now')) - ); - CREATE TABLE IF NOT EXISTS rag_chunks ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - page_id INTEGER NOT NULL, chunk_index INTEGER NOT NULL, - content TEXT NOT NULL, embedding BLOB - ); - CREATE TABLE IF NOT EXISTS rag_relations ( - source_id INTEGER NOT NULL, target_id INTEGER NOT NULL, - relation_type TEXT NOT NULL DEFAULT 'elaborates', - weight REAL DEFAULT 0.8, - PRIMARY KEY (source_id, target_id, relation_type) - ); - CREATE INDEX IF NOT EXISTS idx_rag_user ON rag_pages(user_id); - - -- === Graph === - CREATE TABLE IF NOT EXISTS epi_nodes ( - node_id INTEGER PRIMARY KEY AUTOINCREMENT, - layer TEXT NOT NULL DEFAULT 'user', - user_id TEXT NOT NULL, content TEXT NOT NULL, - node_type TEXT NOT NULL, tags TEXT, - confidence REAL DEFAULT 0.5, created_at REAL NOT NULL - ); - CREATE TABLE IF NOT EXISTS epi_edges ( - source_id INTEGER NOT NULL, target_id INTEGER NOT NULL, - relation TEXT NOT NULL, weight REAL DEFAULT 0.8, - created_at REAL NOT NULL, - PRIMARY KEY (source_id, target_id, relation) - ); - CREATE TABLE IF NOT EXISTS temporal_events ( - event_id INTEGER PRIMARY KEY AUTOINCREMENT, - user_id TEXT NOT NULL, event_type TEXT NOT NULL, - content TEXT NOT NULL, timestamp REAL NOT NULL, - importance REAL DEFAULT 0.5, metadata TEXT - ); - CREATE TABLE IF NOT EXISTS temporal_links ( - from_event INTEGER NOT NULL, to_event INTEGER NOT NULL, - link_type TEXT NOT NULL DEFAULT 'follows', - strength REAL DEFAULT 0.5, - PRIMARY KEY (from_event, to_event, link_type) - ); - - -- === Wiki === - CREATE TABLE IF NOT EXISTS user_wiki ( - entry_id INTEGER PRIMARY KEY AUTOINCREMENT, - user_id TEXT NOT NULL, wiki_type TEXT NOT NULL, - title TEXT NOT NULL, content TEXT NOT NULL, - tags TEXT, importance REAL DEFAULT 0.5, - source TEXT DEFAULT 'manual', - created_at REAL NOT NULL, updated_at REAL NOT NULL - ); - CREATE TABLE IF NOT EXISTS agent_wiki ( - entry_id INTEGER PRIMARY KEY AUTOINCREMENT, - user_id TEXT NOT NULL, wiki_type TEXT NOT NULL, - title TEXT NOT NULL, content TEXT NOT NULL, - tags TEXT, importance REAL DEFAULT 0.5, - source TEXT DEFAULT 'manual', - created_at REAL NOT NULL, updated_at REAL NOT NULL - ); - - -- === FileWiki === - CREATE TABLE IF NOT EXISTS wiki_index ( - entry_id INTEGER PRIMARY KEY AUTOINCREMENT, - layer TEXT NOT NULL, wiki_type TEXT NOT NULL, - title TEXT NOT NULL, file_path TEXT NOT NULL, - tags TEXT, importance REAL DEFAULT 0.5, - content TEXT DEFAULT '', content_hash TEXT, - created_at REAL NOT NULL, updated_at REAL NOT NULL - ); - CREATE UNIQUE INDEX IF NOT EXISTS idx_wiki_path ON wiki_index(file_path); - - -- === FTS5 indexes === - CREATE VIRTUAL TABLE IF NOT EXISTS rag_fts USING fts5( - title, content, wiki_type, content=rag_pages, content_rowid=id - ); - CREATE VIRTUAL TABLE IF NOT EXISTS user_wiki_fts USING fts5( - title, content, wiki_type, tags, content=user_wiki, content_rowid=entry_id - ); - CREATE VIRTUAL TABLE IF NOT EXISTS agent_wiki_fts USING fts5( - title, content, wiki_type, tags, content=agent_wiki, content_rowid=entry_id - ); - CREATE VIRTUAL TABLE IF NOT EXISTS wiki_fts USING fts5( - title, content, wiki_type, tags, content=wiki_index, content_rowid=entry_id - ); - - -- === Conflict tracking === - CREATE TABLE IF NOT EXISTS memory_conflicts ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - user_id TEXT NOT NULL, content TEXT NOT NULL, - is_conflict INTEGER DEFAULT 0, conflict_group_id TEXT, - created_at DATETIME DEFAULT CURRENT_TIMESTAMP - ); - - -- === Migration log === - CREATE TABLE IF NOT EXISTS migration_log ( - version INTEGER PRIMARY KEY, name TEXT NOT NULL, applied_at REAL NOT NULL - ); - """) - - migrations.append(Migration(1, "init_unified_schema", v1_init)) - - async def v2_binary_embeddings(conn): - """Add binary embeddings column for MIB search.""" - try: - await conn.execute("ALTER TABLE rag_chunks ADD COLUMN bin_embedding BLOB") - except sqlite3.OperationalError: - pass # Column already exists - - await conn.execute(""" - CREATE INDEX IF NOT EXISTS idx_rag_chunks_bin - ON rag_chunks(page_id, id) - WHERE bin_embedding IS NOT NULL - """) - - migrations.append(Migration(2, "binary_embeddings", v2_binary_embeddings)) - - async def v3_epi_tags(conn): - """Add epi_tags table for fast tag lookups.""" - try: - await conn.execute(""" - CREATE TABLE IF NOT EXISTS epi_tags ( - node_id INTEGER NOT NULL, - tag TEXT NOT NULL, - PRIMARY KEY (node_id, tag) - ) - """) - await conn.execute("CREATE INDEX IF NOT EXISTS idx_epi_tags_tag ON epi_tags(tag)") - except sqlite3.OperationalError: - pass - - migrations.append(Migration(3, "epi_tags", v3_epi_tags)) - - async def v4_rag_chunks_index(conn): - """Add index on rag_chunks(page_id, chunk_index) for JOINs.""" - try: - await conn.execute("CREATE INDEX IF NOT EXISTS idx_rag_chunks_page_idx ON rag_chunks(page_id, chunk_index)") - except sqlite3.OperationalError: - pass # Index already exists - - migrations.append(Migration(4, "rag_chunks_index", v4_rag_chunks_index)) - - async def v5_typed_memory(conn): - """Add memory_kind column and memory_kind_registry table.""" - try: - # Registry table - await conn.execute(""" - CREATE TABLE IF NOT EXISTS memory_kind_registry ( - kind TEXT PRIMARY KEY, - display_name TEXT NOT NULL, - default_importance REAL NOT NULL, - decay_rate REAL NOT NULL, - never_archive INTEGER NOT NULL DEFAULT 0, - requires_expires_at INTEGER NOT NULL DEFAULT 0, - boost_on_keywords TEXT NOT NULL DEFAULT '', - description TEXT - ) - """) - # Seed 13 types - await conn.execute(""" - INSERT OR IGNORE INTO memory_kind_registry VALUES - ('instruction','Instruction',0.9,0.0,1,0,'обязательно,важно,critical,never forget,rule,инструкция','Правило/инструкция, не подлежит забыванию'), - ('fact','Fact',0.5,0.01,0,0,'факт,fact,имя,возраст,день рождения','Атомарный факт'), - ('decision','Decision',0.7,0.005,0,0,'решение,decided,chose,decision','Принятое решение'), - ('goal','Goal',0.8,0.005,0,1,'цель,goal,plan,к концу','Цель с дедлайном'), - ('preference','Preference',0.7,0.003,0,0,'предпочитаю,prefer,like,нравится,не люблю','Предпочтение'), - ('commitment','Commitment',0.85,0.0,1,1,'обещаю,обязуюсь,commit,promise,согласен','Обязательство'), - ('relationship','Relationship',0.6,0.002,0,0,'знаком,друг,коллега,knows,friend','Связь'), - ('observation','Observation',0.4,0.02,0,0,'видел,заметил,noticed,observed','Наблюдение'), - ('rule','Rule',0.85,0.0,1,0,'запрещено,нельзя,do not,forbidden,rule','Жёсткое правило'), - ('todo','Todo',0.6,0.005,0,1,'todo,сделать,do later,remind','Задача с дедлайном'), - ('question','Open Question',0.5,0.05,0,0,'вопрос,?,уточнить,ask later','Открытый вопрос'), - ('hypothesis','Hypothesis',0.45,0.03,0,0,'возможно,наверное,probably,hypothesis','Гипотеза'), - ('context','Context',0.3,0.05,0,0,'контекст,background,context','Фоновый контекст') - """) - except sqlite3.OperationalError: - pass - # Add memory_kind to core_memory - with contextlib.suppress(sqlite3.OperationalError): - await conn.execute("ALTER TABLE core_memory ADD COLUMN memory_kind TEXT") - # Add expires_at, source, metadata to core_memory - with contextlib.suppress(sqlite3.OperationalError): - await conn.execute("ALTER TABLE core_memory ADD COLUMN expires_at REAL") - with contextlib.suppress(sqlite3.OperationalError): - await conn.execute("ALTER TABLE core_memory ADD COLUMN source TEXT DEFAULT 'manual'") - with contextlib.suppress(sqlite3.OperationalError): - await conn.execute("ALTER TABLE core_memory ADD COLUMN metadata TEXT") - await conn.execute("CREATE INDEX IF NOT EXISTS idx_core_memory_kind ON core_memory(user_id, memory_kind)") - # Add memory_kind to episodes - with contextlib.suppress(sqlite3.OperationalError): - await conn.execute("ALTER TABLE episodes ADD COLUMN memory_kind TEXT") - await conn.execute("CREATE INDEX IF NOT EXISTS idx_episodes_kind ON episodes(user_id, memory_kind)") - # Add memory_kind to rag_chunks (for type boost in search) - with contextlib.suppress(sqlite3.OperationalError): - await conn.execute("ALTER TABLE rag_chunks ADD COLUMN memory_kind TEXT") - - migrations.append(Migration(5, "typed_memory", v5_typed_memory)) - - async def v6_core_memory_columns(conn): - """Add expires_at, source, metadata columns to core_memory.""" - with contextlib.suppress(sqlite3.OperationalError): - await conn.execute("ALTER TABLE core_memory ADD COLUMN expires_at REAL") - with contextlib.suppress(sqlite3.OperationalError): - await conn.execute("ALTER TABLE core_memory ADD COLUMN source TEXT DEFAULT 'manual'") - with contextlib.suppress(sqlite3.OperationalError): - await conn.execute("ALTER TABLE core_memory ADD COLUMN metadata TEXT") - - migrations.append(Migration(6, "core_memory_columns", v6_core_memory_columns)) - - async def v7_drop_float_embeddings(conn): - """Drop float embedding column to save disk space (keep binary only).""" - try: - # Check if we should drop (config-driven via keep_float_blobs) - await conn.execute("ALTER TABLE rag_chunks DROP COLUMN embedding") - logger.info("Dropped float embedding column from rag_chunks") - except sqlite3.OperationalError: - pass # Column may already be dropped - - migrations.append(Migration(7, "drop_float_embeddings", v7_drop_float_embeddings)) - - async def v8_importance_audit(conn): - """Create importance_audit table for scheduler logging.""" - try: - await conn.execute(""" - CREATE TABLE IF NOT EXISTS importance_audit ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - user_id TEXT NOT NULL, - chunk_id INTEGER, - source TEXT NOT NULL, - old_importance REAL, - new_importance REAL, - signal_breakdown TEXT, - reason TEXT, - rescored_at REAL NOT NULL - ) - """) - await conn.execute("CREATE INDEX IF NOT EXISTS idx_importance_audit_user ON importance_audit(user_id, rescored_at DESC)") - except sqlite3.OperationalError: - pass - - migrations.append(Migration(8, "importance_audit", v8_importance_audit)) - - return migrations - - class MigrationManager: def __init__(self, cm: AsyncConnectionManager | None = None): self._cm = cm or connection_manager - self._migrations = _get_migrations() + self._repo_root = Path(__file__).parent.parent + self._alembic_ini = self._repo_root / "alembic.ini" + + def _get_alembic_config(self) -> AlembicConfig: + cfg = AlembicConfig(str(self._alembic_ini)) + # Ensure alembic uses the correct directory for versions + cfg.set_main_option("script_location", str(self._repo_root / "alembic")) + return cfg - async def get_current_version(self) -> int: + async def get_current_version(self) -> str | None: conn = await self._cm.get(DB_NAME) try: - row = await (await conn.execute("SELECT MAX(version) as v FROM migration_log")).fetchone() - return row["v"] if row and row["v"] else 0 - except sqlite3.OperationalError: - return 0 + row = await (await conn.execute("SELECT version_num FROM alembic_version")).fetchone() + return row["version_num"] if row else None + except Exception: + return None async def migrate(self) -> dict[str, Any]: + """Run all pending migrations using Alembic.""" current = await self.get_current_version() - applied = [] - for migration in self._migrations: - if migration.version <= current: - continue - logger.info("Applying migration v%d: %s" % (migration.version, migration.name)) - conn = await self._cm.get(DB_NAME) - await migration.up(conn) - await conn.execute( - "INSERT INTO migration_log (version, name, applied_at) VALUES (?, ?, ?)", - (migration.version, migration.name, time.time()), - ) - await conn.commit() - applied.append({"version": migration.version, "name": migration.name}) - return {"current_version": current, "applied": applied, "new_version": await self.get_current_version()} - - async def get_pending(self) -> list[dict[str, Any]]: - current = await self.get_current_version() - return [{"version": m.version, "name": m.name} for m in self._migrations if m.version > current] + + # Run Alembic upgrade in a thread to avoid blocking async loop + # (Alembic/SQLAlchemy sync nature) + def run_upgrade(): + cfg = self._get_alembic_config() + alembic_command.upgrade(cfg, "head") + + logger.info("Starting DB migration via Alembic...") + await asyncio.to_thread(run_upgrade) + + new_version = await self.get_current_version() + + return { + "current_version": current, + "new_version": new_version, + "status": "up_to_date" if new_version else "initialized" + } + + async def get_pending(self) -> list[str]: + # Simple check: if current != head + return ["Update to head"] migration_manager = MigrationManager()