/
migrations.py
202 lines (170 loc) · 4.68 KB
/
migrations.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
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.utcnow())}
)
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"},
)
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)