mirror of
https://github.com/badlogic/pi-mono.git
synced 2026-11-05 23:37:25 +00:00
feat(sqlite): more helpers
This commit is contained in:
Cristina Poncela Cubeiro
parent
9c1b446870
commit
cf82e2e2ae
4 changed files with 276 additions and 5 deletions
|
|
@ -0,0 +1,136 @@
|
|||
import type {
|
||||
BranchSummaryEntry,
|
||||
CompactionEntry,
|
||||
CustomEntry,
|
||||
Entry,
|
||||
EntryScan,
|
||||
EntryStructure,
|
||||
MessageEntry,
|
||||
} from "@earendil-works/pi-agent-core";
|
||||
import { joinSqlFragments, type SqlQuery, sql } from "../sql.ts";
|
||||
import type { SqliteDatabase } from "../types.ts";
|
||||
|
||||
export interface EntryRow {
|
||||
id: string;
|
||||
parent_id: string | null;
|
||||
seq: number;
|
||||
type: Entry["type"];
|
||||
custom_type: string | null;
|
||||
timestamp: number;
|
||||
payload: string;
|
||||
}
|
||||
|
||||
type StoredEntryPayload<TEntry extends Entry> = Omit<
|
||||
TEntry,
|
||||
"id" | "parentId" | "seq" | "timestamp" | "type" | "customType"
|
||||
>;
|
||||
|
||||
function entryPayload(entry: Entry): StoredEntryPayload<Entry> {
|
||||
switch (entry.type) {
|
||||
case "message": {
|
||||
const payload: StoredEntryPayload<MessageEntry> = {
|
||||
message: entry.message,
|
||||
...(entry.terminate === undefined ? {} : { terminate: entry.terminate }),
|
||||
};
|
||||
return payload;
|
||||
}
|
||||
case "compaction": {
|
||||
const payload: StoredEntryPayload<CompactionEntry> = {
|
||||
summary: entry.summary,
|
||||
retainedTail: entry.retainedTail,
|
||||
tokensBefore: entry.tokensBefore,
|
||||
...(entry.details === undefined ? {} : { details: entry.details }),
|
||||
...(entry.usage === undefined ? {} : { usage: entry.usage }),
|
||||
fromHook: entry.fromHook,
|
||||
};
|
||||
return payload;
|
||||
}
|
||||
case "branch_summary": {
|
||||
const payload: StoredEntryPayload<BranchSummaryEntry> = {
|
||||
fromId: entry.fromId,
|
||||
summary: entry.summary,
|
||||
...(entry.details === undefined ? {} : { details: entry.details }),
|
||||
...(entry.usage === undefined ? {} : { usage: entry.usage }),
|
||||
fromHook: entry.fromHook,
|
||||
};
|
||||
return payload;
|
||||
}
|
||||
case "custom": {
|
||||
const payload: StoredEntryPayload<CustomEntry> = entry.data === undefined ? {} : { data: entry.data };
|
||||
return payload;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
function parsePayload<TEntry extends Entry>(row: EntryRow): StoredEntryPayload<TEntry> {
|
||||
return JSON.parse(row.payload) as StoredEntryPayload<TEntry>;
|
||||
}
|
||||
|
||||
export function insertEntryRow(db: SqliteDatabase, entry: Entry): void {
|
||||
sql`INSERT INTO entries (id, parent_id, seq, type, custom_type, timestamp, payload)
|
||||
VALUES (
|
||||
${entry.id},
|
||||
${entry.parentId},
|
||||
${entry.seq},
|
||||
${entry.type},
|
||||
${entry.type === "custom" ? entry.customType : null},
|
||||
${entry.timestamp},
|
||||
${JSON.stringify(entryPayload(entry))}
|
||||
)`.run(db);
|
||||
}
|
||||
|
||||
export function decodeEntryRow(row: EntryRow): Entry {
|
||||
const base = {
|
||||
id: row.id,
|
||||
parentId: row.parent_id,
|
||||
seq: row.seq,
|
||||
timestamp: row.timestamp,
|
||||
};
|
||||
switch (row.type) {
|
||||
case "message":
|
||||
return { ...base, type: "message", ...parsePayload<MessageEntry>(row) };
|
||||
case "compaction":
|
||||
return { ...base, type: "compaction", ...parsePayload<CompactionEntry>(row) };
|
||||
case "branch_summary":
|
||||
return { ...base, type: "branch_summary", ...parsePayload<BranchSummaryEntry>(row) };
|
||||
case "custom":
|
||||
if (row.custom_type === null) throw new Error(`Custom entry ${row.id} is missing custom_type`);
|
||||
return { ...base, type: "custom", customType: row.custom_type, ...parsePayload<CustomEntry>(row) };
|
||||
}
|
||||
}
|
||||
|
||||
export function entryStructureFromRow(row: EntryRow): EntryStructure {
|
||||
return {
|
||||
id: row.id,
|
||||
parentId: row.parent_id,
|
||||
seq: row.seq,
|
||||
timestamp: row.timestamp,
|
||||
type: row.type,
|
||||
...(row.custom_type === null ? {} : { customType: row.custom_type }),
|
||||
};
|
||||
}
|
||||
|
||||
export function readEntryRows(db: SqliteDatabase, ids: readonly string[]): EntryRow[] {
|
||||
if (ids.length === 0) return [];
|
||||
const placeholders = joinSqlFragments(
|
||||
ids.map((id) => sql`${id}`),
|
||||
", ",
|
||||
);
|
||||
return sql`SELECT id, parent_id, seq, type, custom_type, timestamp, payload
|
||||
FROM entries
|
||||
WHERE id IN (${placeholders})`.all<EntryRow>(db);
|
||||
}
|
||||
|
||||
export function scanEntryRows(db: SqliteDatabase, query: EntryScan): EntryRow[] {
|
||||
const filters: SqlQuery[] = [];
|
||||
if (query.type !== undefined) filters.push(sql`type = ${query.type}`);
|
||||
if (query.customType !== undefined) filters.push(sql`custom_type = ${query.customType}`);
|
||||
if (query.fromSeq !== undefined) filters.push(sql`seq >= ${query.fromSeq}`);
|
||||
if (query.toSeq !== undefined) filters.push(sql`seq <= ${query.toSeq}`);
|
||||
|
||||
const where = filters.length === 0 ? sql`` : sql`WHERE ${joinSqlFragments(filters, " AND ")}`;
|
||||
const order = query.order === "desc" ? sql`ORDER BY seq DESC` : sql`ORDER BY seq ASC`;
|
||||
const limit = query.limit === undefined ? sql`` : sql`LIMIT ${Math.max(0, query.limit)}`;
|
||||
return sql`SELECT id, parent_id, seq, type, custom_type, timestamp, payload
|
||||
FROM entries ${where} ${order} ${limit}`.all<EntryRow>(db);
|
||||
}
|
||||
|
|
@ -1,11 +1,66 @@
|
|||
import type { LaneState } from "@earendil-works/pi-agent-core";
|
||||
import type { LaneState, Register, RegisterNamespace, RegisterValues } from "@earendil-works/pi-agent-core";
|
||||
import { sql } from "../sql.ts";
|
||||
import type { SqliteDatabase } from "../types.ts";
|
||||
|
||||
export interface RegisterRow {
|
||||
namespace: RegisterNamespace;
|
||||
key: string;
|
||||
seq: number;
|
||||
value: string;
|
||||
}
|
||||
|
||||
export function insertInitialMainLaneRegisters(db: SqliteDatabase): void {
|
||||
const laneState = { currentOperationId: null, pendingNextRun: [] } satisfies LaneState;
|
||||
sql`INSERT INTO registers (namespace, key, seq, value)
|
||||
VALUES (${"lane.leaf"}, ${"main"}, ${1}, ${JSON.stringify(null)})`.run(db);
|
||||
sql`INSERT INTO registers (namespace, key, seq, value)
|
||||
VALUES (${"lane.state"}, ${"main"}, ${2}, ${JSON.stringify(laneState)})`.run(db);
|
||||
setRegisterRow(db, "lane.leaf", "main", 1, null);
|
||||
setRegisterRow(db, "lane.state", "main", 2, laneState);
|
||||
}
|
||||
|
||||
export function setRegisterRow<TNamespace extends RegisterNamespace>(
|
||||
db: SqliteDatabase,
|
||||
namespace: TNamespace,
|
||||
key: string,
|
||||
seq: number,
|
||||
value: RegisterValues[TNamespace],
|
||||
): void {
|
||||
sql`INSERT INTO registers (namespace, key, seq, value)
|
||||
VALUES (${namespace}, ${key}, ${seq}, ${JSON.stringify(value)})
|
||||
ON CONFLICT(namespace, key) DO UPDATE SET seq = excluded.seq, value = excluded.value`.run(db);
|
||||
}
|
||||
|
||||
export function deleteRegisterRow(db: SqliteDatabase, namespace: RegisterNamespace, key: string): void {
|
||||
sql`DELETE FROM registers WHERE namespace = ${namespace} AND key = ${key}`.run(db);
|
||||
}
|
||||
|
||||
export function decodeRegisterRow<TNamespace extends RegisterNamespace>(
|
||||
namespace: TNamespace,
|
||||
row: RegisterRow,
|
||||
): Register<TNamespace> {
|
||||
if (row.namespace !== namespace) throw new Error(`Expected register namespace ${namespace}, found ${row.namespace}`);
|
||||
return {
|
||||
namespace,
|
||||
key: row.key,
|
||||
seq: row.seq,
|
||||
value: JSON.parse(row.value) as RegisterValues[TNamespace],
|
||||
};
|
||||
}
|
||||
|
||||
export function readRegisterRow<TNamespace extends RegisterNamespace>(
|
||||
db: SqliteDatabase,
|
||||
namespace: TNamespace,
|
||||
key: string,
|
||||
): Register<TNamespace> | undefined {
|
||||
const row = sql`SELECT namespace, key, seq, value FROM registers
|
||||
WHERE namespace = ${namespace} AND key = ${key}`.get<RegisterRow>(db);
|
||||
return row === undefined ? undefined : decodeRegisterRow(namespace, row);
|
||||
}
|
||||
|
||||
export function listRegisterRows<TNamespace extends RegisterNamespace>(
|
||||
db: SqliteDatabase,
|
||||
namespace: TNamespace,
|
||||
keyPrefix = "",
|
||||
): Register<TNamespace>[] {
|
||||
const rows = sql`SELECT namespace, key, seq, value FROM registers
|
||||
WHERE namespace = ${namespace} AND key >= ${keyPrefix} AND key < ${`${keyPrefix}\uffff`}
|
||||
ORDER BY key ASC`.all<RegisterRow>(db);
|
||||
return rows.map((row) => decodeRegisterRow(namespace, row));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -0,0 +1,45 @@
|
|||
import type { SessionStats, UsageRow } from "@earendil-works/pi-agent-core";
|
||||
import type { Usage } from "@earendil-works/pi-ai";
|
||||
import { sql } from "../sql.ts";
|
||||
import type { SqliteDatabase } from "../types.ts";
|
||||
import { readSingleSessionRow } from "./session-row.ts";
|
||||
|
||||
function addUsage(left: Usage, right: Usage): Usage {
|
||||
return {
|
||||
input: left.input + right.input,
|
||||
output: left.output + right.output,
|
||||
cacheRead: left.cacheRead + right.cacheRead,
|
||||
cacheWrite: left.cacheWrite + right.cacheWrite,
|
||||
...(left.cacheWrite1h === undefined && right.cacheWrite1h === undefined
|
||||
? {}
|
||||
: { cacheWrite1h: (left.cacheWrite1h ?? 0) + (right.cacheWrite1h ?? 0) }),
|
||||
...(left.reasoning === undefined && right.reasoning === undefined
|
||||
? {}
|
||||
: { reasoning: (left.reasoning ?? 0) + (right.reasoning ?? 0) }),
|
||||
totalTokens: left.totalTokens + right.totalTokens,
|
||||
cost: {
|
||||
input: left.cost.input + right.cost.input,
|
||||
output: left.cost.output + right.cost.output,
|
||||
cacheRead: left.cost.cacheRead + right.cost.cacheRead,
|
||||
cacheWrite: left.cost.cacheWrite + right.cost.cacheWrite,
|
||||
total: left.cost.total + right.cost.total,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
export function readSessionStats(db: SqliteDatabase): SessionStats {
|
||||
const row = readSingleSessionRow(db);
|
||||
return {
|
||||
messageCount: row.message_count,
|
||||
usage: JSON.parse(row.usage_payload) as Usage,
|
||||
};
|
||||
}
|
||||
|
||||
export function incrementMessageCount(db: SqliteDatabase): void {
|
||||
sql`UPDATE session SET message_count = message_count + 1`.run(db);
|
||||
}
|
||||
|
||||
export function addUsageToSessionStats(db: SqliteDatabase, usage: UsageRow["usage"]): void {
|
||||
const current = readSessionStats(db).usage;
|
||||
sql`UPDATE session SET usage_payload = ${JSON.stringify(addUsage(current, usage))}`.run(db);
|
||||
}
|
||||
|
|
@ -0,0 +1,35 @@
|
|||
import type { UsageRow } from "@earendil-works/pi-agent-core";
|
||||
import { sql } from "../sql.ts";
|
||||
import type { SqliteDatabase } from "../types.ts";
|
||||
|
||||
export interface UsageLedgerRow {
|
||||
id: string;
|
||||
seq: number;
|
||||
entry_id: string | null;
|
||||
adjustment: number;
|
||||
usage: string;
|
||||
details: string | null;
|
||||
}
|
||||
|
||||
export function insertUsageLedgerRow(db: SqliteDatabase, row: UsageRow): void {
|
||||
sql`INSERT INTO usage_ledger (id, seq, entry_id, adjustment, usage, details)
|
||||
VALUES (
|
||||
${row.id},
|
||||
${row.seq},
|
||||
${row.entryId ?? null},
|
||||
${row.adjustment ? 1 : 0},
|
||||
${JSON.stringify(row.usage)},
|
||||
${row.details === undefined ? null : JSON.stringify(row.details)}
|
||||
)`.run(db);
|
||||
}
|
||||
|
||||
export function decodeUsageLedgerRow(row: UsageLedgerRow): UsageRow {
|
||||
return {
|
||||
id: row.id,
|
||||
seq: row.seq,
|
||||
usage: JSON.parse(row.usage) as UsageRow["usage"],
|
||||
...(row.entry_id === null ? {} : { entryId: row.entry_id }),
|
||||
adjustment: row.adjustment !== 0,
|
||||
...(row.details === null ? {} : { details: JSON.parse(row.details) as UsageRow["details"] }),
|
||||
};
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue