33 lines
965 B
Python
33 lines
965 B
Python
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
|
|
from backend.database.migrations.runner import Migration
|
|
|
|
|
|
def extend_llm_audit(connection: sqlite3.Connection) -> None:
|
|
columns = {
|
|
str(row["name"])
|
|
for row in connection.execute("PRAGMA table_info(llm_usage)")
|
|
}
|
|
additions = (
|
|
("role", "TEXT NOT NULL DEFAULT ''"),
|
|
("prompt_version", "TEXT NOT NULL DEFAULT ''"),
|
|
("error_code", "TEXT NOT NULL DEFAULT ''"),
|
|
("input_tokens", "INTEGER NOT NULL DEFAULT 0"),
|
|
("output_tokens", "INTEGER NOT NULL DEFAULT 0"),
|
|
)
|
|
for name, declaration in additions:
|
|
if name not in columns:
|
|
connection.execute(
|
|
f"ALTER TABLE llm_usage ADD COLUMN {name} {declaration}"
|
|
)
|
|
|
|
|
|
MIGRATION = Migration(
|
|
version="0003",
|
|
name="extend_llm_audit",
|
|
action=extend_llm_audit,
|
|
signature="llm-audit:v1:role,prompt-version,error-code,input-tokens,output-tokens",
|
|
)
|