236 lines
10 KiB
TypeScript
236 lines
10 KiB
TypeScript
import {randomUUID} from "node:crypto";
|
|
import {mkdirSync} from "node:fs";
|
|
import path from "node:path";
|
|
import {Database} from "bun:sqlite";
|
|
import {normalizeGenerationSettings, type GenerationSettings} from "./generation-settings";
|
|
import type {ChatIdentity} from "./key-vault";
|
|
import type {Conversation, ConversationSummary, ResponseMetadata, StoredChatMessage} from "./conversation-types";
|
|
|
|
const databasePath = process.env.CHAT_DATABASE_PATH || "/data/chat.db";
|
|
let database: Database | undefined;
|
|
|
|
function getDatabase() {
|
|
if (database) return database;
|
|
mkdirSync(path.dirname(databasePath), {recursive: true});
|
|
const opened = new Database(databasePath, {create: true, strict: true});
|
|
opened.run(`
|
|
PRAGMA journal_mode = WAL;
|
|
PRAGMA foreign_keys = ON;
|
|
PRAGMA busy_timeout = 5000;
|
|
CREATE TABLE IF NOT EXISTS chat_conversation (
|
|
id TEXT PRIMARY KEY,
|
|
owner_issuer TEXT NOT NULL,
|
|
owner_sub TEXT NOT NULL,
|
|
title TEXT NOT NULL,
|
|
provider_id TEXT NOT NULL,
|
|
model TEXT NOT NULL,
|
|
settings_json TEXT NOT NULL DEFAULT '{}',
|
|
created_at TEXT NOT NULL,
|
|
updated_at TEXT NOT NULL
|
|
);
|
|
CREATE INDEX IF NOT EXISTS chat_conversation_owner_updated
|
|
ON chat_conversation (owner_issuer, owner_sub, updated_at DESC);
|
|
CREATE TABLE IF NOT EXISTS chat_message (
|
|
conversation_id TEXT NOT NULL REFERENCES chat_conversation(id) ON DELETE CASCADE,
|
|
id TEXT NOT NULL,
|
|
ordinal INTEGER NOT NULL,
|
|
role TEXT NOT NULL CHECK (role IN ('system', 'user', 'assistant')),
|
|
parts_json TEXT NOT NULL,
|
|
created_at TEXT NOT NULL,
|
|
PRIMARY KEY (conversation_id, id),
|
|
UNIQUE (conversation_id, ordinal)
|
|
);
|
|
`);
|
|
const columns = opened.query("PRAGMA table_info(chat_conversation)").all() as Array<{name: string}>;
|
|
if (!columns.some((column) => column.name === "settings_json")) {
|
|
opened.run("ALTER TABLE chat_conversation ADD COLUMN settings_json TEXT NOT NULL DEFAULT '{}'");
|
|
}
|
|
database = opened;
|
|
return opened;
|
|
}
|
|
|
|
type ConversationRow = {
|
|
id: string;
|
|
title: string;
|
|
provider_id: string;
|
|
model: string;
|
|
settings_json: string;
|
|
message_count: number;
|
|
created_at: string;
|
|
updated_at: string;
|
|
};
|
|
|
|
type MessageRow = {id: string; role: StoredChatMessage["role"]; parts_json: string};
|
|
|
|
function now() {
|
|
return new Date().toISOString();
|
|
}
|
|
|
|
function requiredString(value: unknown, field: string, maximum: number) {
|
|
if (typeof value !== "string" || !value.trim()) throw new Error(`${field} is required`);
|
|
return value.trim().slice(0, maximum);
|
|
}
|
|
|
|
function conversationSummary(row: ConversationRow): ConversationSummary {
|
|
return {
|
|
id: row.id,
|
|
title: row.title,
|
|
providerId: row.provider_id,
|
|
model: row.model,
|
|
messageCount: Number(row.message_count),
|
|
createdAt: row.created_at,
|
|
updatedAt: row.updated_at
|
|
};
|
|
}
|
|
|
|
function normalizeMessages(value: unknown): StoredChatMessage[] {
|
|
if (!Array.isArray(value) || value.length > 500) throw new Error("messages must contain at most 500 entries");
|
|
return value.map((message, index) => {
|
|
if (!message || typeof message !== "object" || Array.isArray(message)) throw new Error(`messages[${index}] is invalid`);
|
|
const record = message as Record<string, unknown>;
|
|
const role = record.role;
|
|
if (!["system", "user", "assistant"].includes(String(role))) throw new Error(`messages[${index}].role is invalid`);
|
|
if (!Array.isArray(record.parts)) throw new Error(`messages[${index}].parts is required`);
|
|
const partsJson = JSON.stringify(record.parts);
|
|
if (Buffer.byteLength(partsJson) > 1024 * 1024) throw new Error(`messages[${index}] is too large`);
|
|
const metadataJson = JSON.stringify(record.metadata || {});
|
|
if (Buffer.byteLength(metadataJson) > 64 * 1024) throw new Error(`messages[${index}].metadata is too large`);
|
|
return {
|
|
id: typeof record.id === "string" && record.id.trim() ? record.id.trim().slice(0, 160) : randomUUID(),
|
|
role: role as StoredChatMessage["role"],
|
|
parts: JSON.parse(partsJson) as StoredChatMessage["parts"],
|
|
...(record.metadata && typeof record.metadata === "object" && !Array.isArray(record.metadata) ? {metadata: JSON.parse(metadataJson) as StoredChatMessage["metadata"]} : {})
|
|
};
|
|
});
|
|
}
|
|
|
|
function inferredTitle(messages: StoredChatMessage[]) {
|
|
const userMessage = messages.find((message) => message.role === "user");
|
|
const text = userMessage?.parts
|
|
.filter((part) => part.type === "text" && typeof part.text === "string")
|
|
.map((part) => String(part.text))
|
|
.join(" ")
|
|
.replace(/\s+/g, " ")
|
|
.trim();
|
|
return text ? text.slice(0, 60) : "新对话";
|
|
}
|
|
|
|
function ownedConversation(identity: ChatIdentity, id: string) {
|
|
return getDatabase().query(`
|
|
SELECT c.id, c.title, c.provider_id, c.model, c.settings_json, c.created_at, c.updated_at,
|
|
(SELECT COUNT(*) FROM chat_message m WHERE m.conversation_id = c.id) AS message_count
|
|
FROM chat_conversation c
|
|
WHERE c.id = ? AND c.owner_issuer = ? AND c.owner_sub = ?
|
|
`).get(id, identity.issuer, identity.sub) as ConversationRow | undefined;
|
|
}
|
|
|
|
export function listConversations(identity: ChatIdentity): ConversationSummary[] {
|
|
return (getDatabase().query(`
|
|
SELECT c.id, c.title, c.provider_id, c.model, c.settings_json, c.created_at, c.updated_at,
|
|
(SELECT COUNT(*) FROM chat_message m WHERE m.conversation_id = c.id) AS message_count
|
|
FROM chat_conversation c
|
|
WHERE c.owner_issuer = ? AND c.owner_sub = ?
|
|
ORDER BY c.updated_at DESC
|
|
LIMIT 100
|
|
`).all(identity.issuer, identity.sub) as ConversationRow[]).map(conversationSummary);
|
|
}
|
|
|
|
export function createConversation(identity: ChatIdentity, input: {providerId: unknown; model: unknown; generationSettings?: unknown}): Conversation {
|
|
const providerId = requiredString(input.providerId, "providerId", 80);
|
|
const model = requiredString(input.model, "model", 300);
|
|
const generationSettings = normalizeGenerationSettings(input.generationSettings);
|
|
const timestamp = now();
|
|
const id = randomUUID();
|
|
getDatabase().run("BEGIN IMMEDIATE");
|
|
try {
|
|
getDatabase().query(`
|
|
DELETE FROM chat_conversation
|
|
WHERE owner_issuer = ? AND owner_sub = ?
|
|
AND NOT EXISTS (SELECT 1 FROM chat_message WHERE conversation_id = chat_conversation.id)
|
|
`).run(identity.issuer, identity.sub);
|
|
getDatabase().query(`
|
|
INSERT INTO chat_conversation (id, owner_issuer, owner_sub, title, provider_id, model, settings_json, created_at, updated_at)
|
|
VALUES (?, ?, ?, '新对话', ?, ?, ?, ?, ?)
|
|
`).run(id, identity.issuer, identity.sub, providerId, model, JSON.stringify(generationSettings), timestamp, timestamp);
|
|
getDatabase().run("COMMIT");
|
|
} catch (error) {
|
|
getDatabase().run("ROLLBACK");
|
|
throw error;
|
|
}
|
|
return {...conversationSummary({id, title: "新对话", provider_id: providerId, model, settings_json: JSON.stringify(generationSettings), message_count: 0, created_at: timestamp, updated_at: timestamp}), generationSettings, messages: []};
|
|
}
|
|
|
|
export function getConversation(identity: ChatIdentity, id: string): Conversation | null {
|
|
const row = ownedConversation(identity, id);
|
|
if (!row) return null;
|
|
const messages = (getDatabase().query(`
|
|
SELECT id, role, parts_json FROM chat_message WHERE conversation_id = ? ORDER BY ordinal
|
|
`).all(id) as MessageRow[]).map((message) => {
|
|
const parsed = JSON.parse(message.parts_json) as StoredChatMessage["parts"];
|
|
const metadataPart = parsed.find((part) => part.type === "data-response-metadata" && part.data && typeof part.data === "object");
|
|
return {
|
|
id: message.id,
|
|
role: message.role,
|
|
parts: parsed.filter((part) => part.type !== "data-response-metadata"),
|
|
...(metadataPart ? {metadata: {custom: {response: metadataPart.data as ResponseMetadata}}} : {})
|
|
};
|
|
});
|
|
return {...conversationSummary(row), generationSettings: normalizeGenerationSettings(JSON.parse(row.settings_json || "{}")), messages};
|
|
}
|
|
|
|
export function updateConversationSettings(identity: ChatIdentity, id: string, input: {providerId: unknown; model: unknown; generationSettings?: unknown}) {
|
|
if (!ownedConversation(identity, id)) return false;
|
|
const providerId = requiredString(input.providerId, "providerId", 80);
|
|
const model = requiredString(input.model, "model", 300);
|
|
const generationSettings = normalizeGenerationSettings(input.generationSettings);
|
|
getDatabase().query(`
|
|
UPDATE chat_conversation SET provider_id = ?, model = ?, settings_json = ?, updated_at = ?
|
|
WHERE id = ? AND owner_issuer = ? AND owner_sub = ?
|
|
`).run(providerId, model, JSON.stringify(generationSettings), now(), id, identity.issuer, identity.sub);
|
|
return true;
|
|
}
|
|
|
|
export function saveConversationMessages(identity: ChatIdentity, id: string, input: {
|
|
providerId: unknown;
|
|
model: unknown;
|
|
messages: unknown;
|
|
}) {
|
|
const existing = ownedConversation(identity, id);
|
|
if (!existing) return null;
|
|
const providerId = requiredString(input.providerId, "providerId", 80);
|
|
const model = requiredString(input.model, "model", 300);
|
|
const messages = normalizeMessages(input.messages);
|
|
const timestamp = now();
|
|
const title = existing.title === "新对话" ? inferredTitle(messages) : existing.title;
|
|
getDatabase().run("BEGIN IMMEDIATE");
|
|
try {
|
|
getDatabase().query("DELETE FROM chat_message WHERE conversation_id = ?").run(id);
|
|
const insert = getDatabase().query(`
|
|
INSERT INTO chat_message (conversation_id, id, ordinal, role, parts_json, created_at)
|
|
VALUES (?, ?, ?, ?, ?, ?)
|
|
`);
|
|
messages.forEach((message, ordinal) => {
|
|
const parts = message.metadata?.custom?.response
|
|
? [...message.parts, {type: "data-response-metadata", data: message.metadata.custom.response}]
|
|
: message.parts;
|
|
insert.run(id, message.id, ordinal, message.role, JSON.stringify(parts), timestamp);
|
|
});
|
|
getDatabase().query(`
|
|
UPDATE chat_conversation SET title = ?, provider_id = ?, model = ?, updated_at = ?
|
|
WHERE id = ? AND owner_issuer = ? AND owner_sub = ?
|
|
`).run(title, providerId, model, timestamp, id, identity.issuer, identity.sub);
|
|
getDatabase().run("COMMIT");
|
|
} catch (error) {
|
|
getDatabase().run("ROLLBACK");
|
|
throw error;
|
|
}
|
|
return getConversation(identity, id);
|
|
}
|
|
|
|
export function deleteConversation(identity: ChatIdentity, id: string) {
|
|
const result = getDatabase().query(`
|
|
DELETE FROM chat_conversation WHERE id = ? AND owner_issuer = ? AND owner_sub = ?
|
|
`).run(id, identity.issuer, identity.sub);
|
|
return result.changes > 0;
|
|
}
|