/llm/migrations.py

import datetime
from typing import Callable, List

MIGRATIONS: List[Callable] = []
migration = MIGRATIONS.append


def migrate(db):
    ensure_migrations_table(db)
    already_applied = {r["name"] for r in db["_llm_migrations"].rows}
    for fn in MIGRATIONS:
        name = fn.__name__
        if name not in already_applied:
            fn(db)
            db["_llm_migrations"].insert(
                {
                    "name": name,
                    "applied_at": str(datetime.datetime.now(datetime.timezone.utc)),
                }
            )
            already_applied.add(name)


def ensure_migrations_table(db):
    if not db["_llm_migrations"].exists():
        db["_llm_migrations"].create(
            {
                "name": str,
                "applied_at": str,
            },
            pk="name",
        )


@migration
def m001_initial(db):
    # Ensure the original table design exists, so other migrations can run
    if db["log"].exists():
        # It needs to have the chat_id column
        if "chat_id" not in db["log"].columns_dict:
            db["log"].add_column("chat_id")
        return
    db["log"].create(
        {
            "provider": str,
            "system": str,
            "prompt": str,
            "chat_id": str,
            "response": str,
            "model": str,
            "timestamp": str,
        }
    )


@migration
def m002_id_primary_key(db):
    db["log"].transform(pk="id")


@migration
def m003_chat_id_foreign_key(db):
    db["log"].transform(types={"chat_id": int})
    db["log"].add_foreign_key("chat_id", "log", "id")


@migration
def m004_column_order(db):
    db["log"].transform(
        column_order=(
            "id",
            "model",
            "timestamp",
            "prompt",
            "system",
            "response",
            "chat_id",
        )
    )


@migration
def m004_drop_provider(db):
    db["log"].transform(drop=("provider",))


@migration
def m005_debug(db):
    db["log"].add_column("debug", str)
    db["log"].add_column("duration_ms", int)


@migration
def m006_new_logs_table(db):
    columns = db["log"].columns_dict
    for column, type in (
        ("options_json", str),
        ("prompt_json", str),
        ("response_json", str),
        ("reply_to_id", int),
    ):
        # It's possible people running development code like myself
        # might have accidentally created these columns already
        if column not in columns:
            db["log"].add_column(column, type)

    # Use .transform() to rename options and timestamp_utc, and set new order
    db["log"].transform(
        column_order=(
            "id",
            "model",
            "prompt",
            "system",
            "prompt_json",
            "options_json",
            "response",
            "response_json",
            "reply_to_id",
            "chat_id",
            "duration_ms",
            "timestamp_utc",
        ),
        rename={
            "timestamp": "timestamp_utc",
            "options": "options_json",
        },
    )


@migration
def m007_finish_logs_table(db):
    db["log"].transform(
        drop={"debug"},
        rename={"timestamp_utc": "datetime_utc"},
        drop_foreign_keys=("chat_id",),
    )
    with db.conn:
        db.execute("alter table log rename to logs")


@migration
def m008_reply_to_id_foreign_key(db):
    db["logs"].add_foreign_key("reply_to_id", "logs", "id")


@migration
def m008_fix_column_order_in_logs(db):
    # reply_to_id ended up at the end after foreign key added
    db["logs"].transform(
        column_order=(
            "id",
            "model",
            "prompt",
            "system",
            "prompt_json",
            "options_json",
            "response",
            "response_json",
            "reply_to_id",
            "chat_id",
            "duration_ms",
            "timestamp_utc",
        ),
    )


@migration
def m009_delete_logs_table_if_empty(db):
    # We moved to a new table design, but we don't delete the table
    # if someone has put data in it
    if not db["logs"].count:
        db["logs"].drop()


@migration
def m010_create_new_log_tables(db):
    db["conversations"].create(
        {
            "id": str,
            "name": str,
            "model": str,
        },
        pk="id",
    )
    db["responses"].create(
        {
            "id": str,
            "model": str,
            "prompt": str,
            "system": str,
            "prompt_json": str,
            "options_json": str,
            "response": str,
            "response_json": str,
            "conversation_id": str,
            "duration_ms": int,
            "datetime_utc": str,
        },
        pk="id",
        foreign_keys=(("conversation_id", "conversations", "id"),),
    )


@migration
def m011_fts_for_responses(db):
    db["responses"].enable_fts(["prompt", "response"], create_triggers=True)


@migration
def m012_attachments_tables(db):
    db["attachments"].create(
        {
            "id": str,
            "type": str,
            "path": str,
            "url": str,
            "content": bytes,
        },
        pk="id",
    )
    db["prompt_attachments"].create(
        {
            "response_id": str,
            "attachment_id": str,
            "order": int,
        },
        foreign_keys=(
            ("response_id", "responses", "id"),
            ("attachment_id", "attachments", "id"),
        ),
        pk=("response_id", "attachment_id"),
    )


@migration
def m013_usage(db):
    db["responses"].add_column("input_tokens", int)
    db["responses"].add_column("output_tokens", int)
    db["responses"].add_column("token_details", str)


@migration
def m014_schemas(db):
    db["schemas"].create(
        {
            "id": str,
            "content": str,
        },
        pk="id",
    )
    db["responses"].add_column("schema_id", str, fk="schemas", fk_col="id")
    # Clean up SQL create table indentation
    db["responses"].transform()
    # These changes may have dropped the FTS configuration, fix that
    db["responses"].enable_fts(
        ["prompt", "response"], create_triggers=True, replace=True
    )


@migration
def m015_fragments_tables(db):
    db["fragments"].create(
        {
            "id": int,
            "hash": str,
            "content": str,
            "datetime_utc": str,
            "source": str,
        },
        pk="id",
    )
    db["fragments"].create_index(["hash"], unique=True)
    db["fragment_aliases"].create(
        {
            "alias": str,
            "fragment_id": int,
        },
        foreign_keys=(("fragment_id", "fragments", "id"),),
        pk="alias",
    )
    db["prompt_fragments"].create(
        {
            "response_id": str,
            "fragment_id": int,
            "order": int,
        },
        foreign_keys=(
            ("response_id", "responses", "id"),
            ("fragment_id", "fragments", "id"),
        ),
        pk=("response_id", "fragment_id"),
    )
    db["system_fragments"].create(
        {
            "response_id": str,
            "fragment_id": int,
            "order": int,
        },
        foreign_keys=(
            ("response_id", "responses", "id"),
            ("fragment_id", "fragments", "id"),
        ),
        pk=("response_id", "fragment_id"),
    )


@migration
def m016_fragments_table_pks(db):
    # The same fragment can be attached to a response multiple times
    # https://github.com/simonw/llm/issues/863#issuecomment-2781720064
    db["prompt_fragments"].transform(pk=("response_id", "fragment_id", "order"))
    db["system_fragments"].transform(pk=("response_id", "fragment_id", "order"))


@migration
def m017_tools_tables(db):
    db["tools"].create(
        {
            "id": int,
            "hash": str,
            "name": str,
            "description": str,
            "input_schema": str,
        },
        pk="id",
    )
    db["tools"].create_index(["hash"], unique=True)
    # Many-to-many relationship between tools and responses
    db["tool_responses"].create(
        {
            "tool_id": int,
            "response_id": str,
        },
        foreign_keys=(
            ("tool_id", "tools", "id"),
            ("response_id", "responses", "id"),
        ),
        pk=("tool_id", "response_id"),
    )
    # tool_calls and tool_results are one-to-many against responses
    db["tool_calls"].create(
        {
            "id": int,
            "response_id": str,
            "tool_id": int,
            "name": str,
            "arguments": str,
            "tool_call_id": str,
        },
        pk="id",
        foreign_keys=(
            ("response_id", "responses", "id"),
            ("tool_id", "tools", "id"),
        ),
    )
    db["tool_results"].create(
        {
            "id": int,
            "response_id": str,
            "tool_id": int,
            "name": str,
            "output": str,
            "tool_call_id": str,
        },
        pk="id",
        foreign_keys=(
            ("response_id", "responses", "id"),
            ("tool_id", "tools", "id"),
        ),
    )


@migration
def m017_tools_plugin(db):
    db["tools"].add_column("plugin")


@migration
def m018_tool_instances(db):
    # Used to track instances of Toolbox classes that may be
    # used multiple times by different tools
    db["tool_instances"].create(
        {
            "id": int,
            "plugin": str,
            "name": str,
            "arguments": str,
        },
        pk="id",
    )
    # We record which instance was used only on the results
    db["tool_results"].add_column("instance_id", fk="tool_instances")


@migration
def m019_resolved_model(db):
    # For models like gemini-1.5-flash-latest where we wish to record
    # the resolved model name in addition to the alias
    db["responses"].add_column("resolved_model", str)


@migration
def m020_tool_results_attachments(db):
    db["tool_results_attachments"].create(
        {
            "tool_result_id": int,
            "attachment_id": str,
            "order": int,
        },
        foreign_keys=(
            ("tool_result_id", "tool_results", "id"),
            ("attachment_id", "attachments", "id"),
        ),
        pk=("tool_result_id", "attachment_id"),
    )


@migration
def m021_tool_results_exception(db):
    db["tool_results"].add_column("exception", str)


@migration
def m022_parts_table(db):
    db["parts"].create(
        {
            "id": int,
            "response_id": str,
            "direction": str,  # "input" or "output"
            "role": str,  # "user", "assistant", "system", "tool"
            "part_type": str,  # "text", "reasoning", "tool_call", "tool_result", "attachment"
            "order": int,
            "content": str,  # Text content for text/reasoning parts
            "content_json": str,  # JSON for structured data
            "tool_call_id": str,
            "server_executed": int,  # 1 for server-side tool calls/results
        },
        pk="id",
        foreign_keys=[("response_id", "responses", "id")],
    )
    db.execute(
        'CREATE UNIQUE INDEX "idx_parts_response_order" ON parts (response_id, direction, "order")'
    )


@migration
def m023_messages_table(db):
    # DAG-shaped message store. See plans/dag-schema.md.
    #
    # A message is identified by (parent_id, content_hash). parent_id is
    # NOT NULL; chain roots point at the self-referencing sentinel row
    # "root" so the unique (parent_id, content_hash) index works uniformly
    # at every chain position — see the "NULL-parent uniqueness footgun"
    # section of plans/dag-schema.md for why.
    db["messages"].create(
        {
            "id": str,  # ULID, or "root" for the sentinel.
            "parent_id": str,
            "content_hash": str,
            "role": str,
            "provider_metadata_json": str,
            "created_at": str,
        },
        pk="id",
        foreign_keys=[("parent_id", "messages", "id")],
        not_null={"parent_id", "content_hash", "role", "created_at"},
    )
    db.execute(
        'CREATE UNIQUE INDEX "idx_messages_parent_hash_unique" '
        "ON messages(parent_id, content_hash)"
    )
    db.execute(
        'CREATE INDEX "idx_messages_parent_id" ON messages(parent_id)'
    )
    db.execute(
        'CREATE INDEX "idx_messages_content_hash" ON messages(content_hash)'
    )

    # Sentinel root row — self-references so parent_id can be NOT NULL.
    db["messages"].insert(
        {
            "id": "root",
            "parent_id": "root",
            "content_hash": "",
            "role": "root",
            "provider_metadata_json": None,
            "created_at": str(
                datetime.datetime.now(datetime.timezone.utc)
            ),
        }
    )

    db["message_parts"].create(
        {
            "id": str,  # ULID
            "message_id": str,
            "order": int,
            "part_type": str,
            "content": str,
            "content_json": str,
            "tool_call_id": str,
            "server_executed": int,
        },
        pk="id",
        foreign_keys=[("message_id", "messages", "id")],
    )
    db.execute(
        'CREATE UNIQUE INDEX "idx_message_parts_message_order" '
        'ON message_parts (message_id, "order")'
    )

    # conversations gains a head pointer. Walking parent_id from
    # head_message_id reconstructs the full conversation history.
    db["conversations"].add_column(
        "head_message_id", str, fk="messages", fk_col="id"
    )

    # One row per LLM call. Messages are shared across calls (dedup);
    # this table records per-call metadata (model, timing, usage).
    db["calls"].create(
        {
            "id": str,  # ULID
            "conversation_id": str,
            "head_input_message_id": str,
            "head_output_message_id": str,
            "model": str,
            "resolved_model": str,
            "started_at": str,
            "duration_ms": int,
            "input_tokens": int,
            "output_tokens": int,
            "token_details_json": str,
            "prompt_json": str,
            "response_json": str,
            "error": str,
        },
        pk="id",
        foreign_keys=[
            ("conversation_id", "conversations", "id"),
            ("head_input_message_id", "messages", "id"),
            ("head_output_message_id", "messages", "id"),
        ],
        not_null={"model", "started_at"},
    )
    db.execute(
        'CREATE INDEX "idx_calls_conversation" ON calls(conversation_id)'
    )
    db.execute(
        'CREATE INDEX "idx_calls_started_at" ON calls(started_at)'
    )
    db.execute(
        'CREATE INDEX "idx_calls_head_input" ON calls(head_input_message_id)'
    )
    db.execute(
        'CREATE INDEX "idx_calls_head_output" ON calls(head_output_message_id)'
    )