Enable FTS indexes on responses prompt and response columns, refs #109

This commit is contained in:
Simon Willison 2023-07-15 16:03:35 -07:00
parent ed8cd776a4
commit 75e40e7077
3 changed files with 17 additions and 4 deletions

View file

@ -96,7 +96,7 @@ def cleanup_sql(sql):
return first_line + '(\n ' + ',\n '.join(columns) + '\n);'
cog.out("```sql\n")
for table in ("conversations", "responses"):
for table in ("conversations", "responses", "responses_fts"):
schema = db[table].schema
cog.out(format(cleanup_sql(schema)))
cog.out("\n")
@ -121,5 +121,11 @@ CREATE TABLE [responses] (
[duration_ms] INTEGER,
[datetime_utc] TEXT
);
CREATE VIRTUAL TABLE [responses_fts] USING FTS5 (
[prompt],
[response],
content=[responses]
);
```
<!-- [[[end]]] -->
`responses_fts` configures [SQLite full-text search](https://www.sqlite.org/fts5.html) against the `prompt` and `response` columns in the `responses` table.

View file

@ -195,3 +195,8 @@ def m010_create_new_log_tables(db):
pk="id",
foreign_keys=(("conversation_id", "conversations", "id"),),
)
@migration
def m011_fts_for_responses(db):
db["responses"].enable_fts(["prompt", "response"], create_triggers=True)

View file

@ -21,7 +21,9 @@ EXPECTED = {
def test_migrate_blank():
db = sqlite_utils.Database(memory=True)
migrate(db)
assert set(db.table_names()) == {"_llm_migrations", "conversations", "responses"}
assert set(db.table_names()).issuperset(
{"_llm_migrations", "conversations", "responses", "responses_fts"}
)
assert db["responses"].columns_dict == EXPECTED
foreign_keys = db["responses"].foreign_keys
@ -65,7 +67,7 @@ def test_migrate_from_original_schema(has_record):
}
)
migrate(db)
expected_tables = {"_llm_migrations", "conversations", "responses"}
expected_tables = {"_llm_migrations", "conversations", "responses", "responses_fts"}
if has_record:
expected_tables.add("logs")
assert set(db.table_names()) == expected_tables
assert set(db.table_names()).issuperset(expected_tables)