|
| 1 | +import { Database } from "bun:sqlite" |
| 2 | +import { mkdir, rm } from "node:fs/promises" |
| 3 | +import path from "node:path" |
| 4 | +import { progress } from "./progress" |
| 5 | +import type { Options, Target } from "./types" |
| 6 | + |
| 7 | +export async function createPartialSnapshot(source: string, destination: string, options: Options, targets: Target[]) { |
| 8 | + await mkdir(path.dirname(destination), { recursive: true }) |
| 9 | + await rm(destination, { force: true }) |
| 10 | + const input = new Database(source, { readonly: true }) |
| 11 | + const schema = input |
| 12 | + .query( |
| 13 | + `SELECT type, name, sql FROM sqlite_schema |
| 14 | + WHERE sql IS NOT NULL AND name NOT LIKE 'sqlite_%' |
| 15 | + ORDER BY CASE type WHEN 'table' THEN 0 WHEN 'index' THEN 1 ELSE 2 END, name`, |
| 16 | + ) |
| 17 | + .all() as { type: string; name: string; sql: string }[] |
| 18 | + input.close() |
| 19 | + |
| 20 | + const output = new Database(destination, { create: true }) |
| 21 | + output.run("PRAGMA foreign_keys = OFF") |
| 22 | + schema.filter((item) => item.type === "table").forEach((item) => output.run(item.sql)) |
| 23 | + output.run("ATTACH DATABASE ? AS source", source) |
| 24 | + const selected = [...new Set(targets.map((target) => target.id))] |
| 25 | + const placeholders = selected.map(() => "?").join(",") |
| 26 | + |
| 27 | + for (const table of schema.filter((item) => item.type === "table").map((item) => item.name)) { |
| 28 | + progress("copying partial snapshot table", { table }) |
| 29 | + if (table === "event") continue |
| 30 | + if (table === "message") { |
| 31 | + output.run( |
| 32 | + `INSERT INTO main.message SELECT * FROM source.message |
| 33 | + WHERE (time_created >= ? AND time_created < ? AND session_id IN ( |
| 34 | + SELECT id FROM source.session WHERE parent_id IS NULL |
| 35 | + )) OR session_id IN (${placeholders})`, |
| 36 | + options.windowStart, |
| 37 | + options.windowEnd, |
| 38 | + ...selected, |
| 39 | + ) |
| 40 | + continue |
| 41 | + } |
| 42 | + if (table === "part") { |
| 43 | + output.run("INSERT INTO main.part SELECT * FROM source.part WHERE message_id IN (SELECT id FROM main.message)") |
| 44 | + continue |
| 45 | + } |
| 46 | + if (["session_context_epoch", "session_input", "session_message", "session_share", "todo"].includes(table)) { |
| 47 | + output.run( |
| 48 | + `INSERT INTO main."${table}" SELECT * FROM source."${table}" WHERE session_id IN (${placeholders})`, |
| 49 | + ...selected, |
| 50 | + ) |
| 51 | + continue |
| 52 | + } |
| 53 | + output.run(`INSERT INTO main."${table}" SELECT * FROM source."${table}"`) |
| 54 | + } |
| 55 | + output.run("DETACH DATABASE source") |
| 56 | + schema.filter((item) => item.type !== "table").forEach((item) => output.run(item.sql)) |
| 57 | + output.close() |
| 58 | +} |
| 59 | + |
| 60 | +export async function fingerprint(file: string) { |
| 61 | + const input = Bun.file(file) |
| 62 | + const hasher = new Bun.CryptoHasher("sha256") |
| 63 | + for await (const chunk of input.stream()) hasher.update(chunk) |
| 64 | + return { bytes: input.size, sha256: hasher.digest("hex") } |
| 65 | +} |
| 66 | + |
| 67 | +export function loadCorpus(options: Options) { |
| 68 | + const database = new Database(options.database, { readonly: true }) |
| 69 | + database.run("PRAGMA query_only = ON") |
| 70 | + const sessions = database |
| 71 | + .query( |
| 72 | + `SELECT id, project_id AS projectID, directory, title |
| 73 | + FROM session AS candidate |
| 74 | + WHERE parent_id IS NULL |
| 75 | + AND EXISTS ( |
| 76 | + SELECT 1 FROM message |
| 77 | + WHERE session_id = candidate.id AND time_created >= ? AND time_created < ? |
| 78 | + )`, |
| 79 | + ) |
| 80 | + .all(options.windowStart, options.windowEnd) as { id: string; projectID: string; directory: string; title: string }[] |
| 81 | + const messageRows = database.query( |
| 82 | + `SELECT id, data FROM message |
| 83 | + WHERE session_id = ? AND time_created >= ? AND time_created < ? |
| 84 | + ORDER BY time_created, id`, |
| 85 | + ) |
| 86 | + const partRows = database.query(`SELECT data FROM part WHERE message_id = ? ORDER BY id`) |
| 87 | + const ranked = sessions |
| 88 | + .map((session) => { |
| 89 | + const messages = messageRows.all(session.id, options.windowStart, options.windowEnd) as { |
| 90 | + id: string |
| 91 | + data: string |
| 92 | + }[] |
| 93 | + const parts = messages.flatMap((message) => partRows.all(message.id) as { data: string }[]) |
| 94 | + return { |
| 95 | + ...session, |
| 96 | + bytes: |
| 97 | + messages.reduce((sum, message) => sum + Buffer.byteLength(message.data), 0) + |
| 98 | + parts.reduce((sum, part) => sum + Buffer.byteLength(part.data), 0), |
| 99 | + messages: messages.length, |
| 100 | + parts: parts.length, |
| 101 | + userTurns: messages.filter((message) => JSON.parse(message.data).role === "user").length, |
| 102 | + } |
| 103 | + }) |
| 104 | + .filter((session) => session.messages > 0) |
| 105 | + .sort((a, b) => a.bytes - b.bytes || a.id.localeCompare(b.id)) |
| 106 | + if (ranked.length === 0) throw new Error("No sessions found in the profile window") |
| 107 | + const select = (label: Target["label"], percentile: number) => ({ |
| 108 | + label, |
| 109 | + ...ranked[Math.max(0, Math.ceil(ranked.length * percentile) - 1)]!, |
| 110 | + }) |
| 111 | + const targets = [select("p50", 0.5), select("p95", 0.95), select("max", 1)] satisfies Target[] |
| 112 | + const typingText = loadTypingText(database, partRows, messageRows, targets[2]!, options) |
| 113 | + const projectIDs = [...new Set(ranked.map((session) => session.projectID))] |
| 114 | + database.close() |
| 115 | + return { targets, typingText, projectIDs } |
| 116 | +} |
| 117 | + |
| 118 | +function loadTypingText( |
| 119 | + database: Database, |
| 120 | + partRows: ReturnType<Database["query"]>, |
| 121 | + messageRows: ReturnType<Database["query"]>, |
| 122 | + target: Target, |
| 123 | + options: Options, |
| 124 | +) { |
| 125 | + const messages = messageRows.all(target.id, options.windowStart, options.windowEnd) as { id: string; data: string }[] |
| 126 | + const text = messages |
| 127 | + .filter((message) => JSON.parse(message.data).role === "user") |
| 128 | + .flatMap((message) => |
| 129 | + (partRows.all(message.id) as { data: string }[]).flatMap((part) => { |
| 130 | + const data = JSON.parse(part.data) |
| 131 | + return data.type === "text" && typeof data.text === "string" ? [data.text] : [] |
| 132 | + }), |
| 133 | + ) |
| 134 | + .sort((a, b) => b.length - a.length)[0] |
| 135 | + if (!text) throw new Error("No real user prompt found for composer profiling") |
| 136 | + return text |
| 137 | +} |
0 commit comments