Fretux/KrisHD_AI
3
0
Fork 0
Code Issues 3 Pull Requests Actions Packages Projects 1 Releases Wiki Activity

add response versioning support in messages for regeneration improvements

Browse Source
This commit is contained in:
fretuks
2026-09-22 14:48:13 +02:00
parent 7602692341
commit 72e429b1b3
15 changed files with 591 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/config.js
+1 -1
View File
@@ -58,7 +58,7 @@ export function createConfig(overrides = {}) {
chatSummaryMaxChars: int(process.env.CHAT_SUMMARY_MAX_CHARS, 2200),
chatMaxRetries: int(process.env.CHAT_MAX_RETRIES, 5),
modelApiBaseUrl: process.env.MODEL_API_BASE_URL || "https://ai.krishd.ch/api",
modelRequestTimeoutMs: int(process.env.MODEL_REQUEST_TIMEOUT_MS, 180000),
modelRequestYTimeoutMs: int(process.env.MODEL_REQUEST_TIMEOUT_MS, 180000),
modelConnectTimeoutMs: Math.max(1, int(process.env.MODEL_CONNECT_TIMEOUT_MS, 30000)),
modelInactivityTimeoutMs: Math.max(1, int(process.env.MODEL_INACTIVITY_TIMEOUT_MS, int(process.env.MODEL_REQUEST_TIMEOUT_MS, 180000))),
modelMaxGenerationMs: Math.max(1, int(process.env.MODEL_MAX_GENERATION_MS, 600000)),
src/db/repositories.js
+31
View File
@@ -871,6 +871,10 @@ export function createRepositories(db, config) {
if (message.sourceId != null) importedMessageIds.set(message.sourceId, Number(result.lastInsertRowid));
if (message.sources) statements.updateChatMessageSources.run(JSON.stringify(message.sources), result.lastInsertRowid, chatId);
if (message.thinking) statements.updateChatMessageThinking.run(message.thinking, result.lastInsertRowid, chatId);
for (const version of message.responseVersions || []) {
db.prepare("INSERT INTO chat_message_versions (message_id, content, thinking, sources, model_name, created_at) VALUES (?, ?, ?, ?, ?, COALESCE(?, CURRENT_TIMESTAMP))")
.run(result.lastInsertRowid, version.content, version.thinking, JSON.stringify(version.sources), version.modelName, version.createdAt);
}
for (const [position, attachment] of (message.attachments || []).slice(0, config.uploadMaxFiles || 5).entries()) {
// JSON backups contain references, not files. Only restore existing, owned images.
db.prepare("INSERT OR IGNORE INTO message_attachments(message_id,attachment_id,position) SELECT ?,id,? FROM attachments WHERE id=? AND username=? AND status='attached'")
@@ -1046,6 +1050,29 @@ export function createRepositories(db, config) {
username
),
updateChatPreferredModel: (chatId, username, model) => statements.updateChatPreferredModel.run(model || null, chatId, username),
listMessageVersions: (messageId) => db.prepare("SELECT * FROM chat_message_versions WHERE message_id = ? ORDER BY id").all(messageId),
saveMessageVersion: (chatId, messageId, version = null, selectedIndex = null) => db.transaction(() => {
const row = statements.getChatMessageById.get(messageId, chatId);
if (!row) throw new Error("Message not found");
let versions = db.prepare("SELECT * FROM chat_message_versions WHERE message_id = ? ORDER BY id").all(messageId);
const insert = db.prepare("INSERT INTO chat_message_versions (message_id, content, thinking, sources, model_name) VALUES (?, ?, ?, ?, ?)");
if (!versions.length) {
db.prepare("INSERT INTO chat_message_versions (message_id, content, thinking, sources, model_name, created_at) VALUES (?, ?, ?, ?, ?, ?)")
.run(messageId, row.content, row.thinking, row.sources, row.model_name, row.created_at);
} else {
// Edits keep their existing semantics: update the selected answer.
db.prepare("UPDATE chat_message_versions SET content = ?, thinking = ?, sources = ?, model_name = ? WHERE id = ?")
.run(row.content, row.thinking, row.sources, row.model_name, versions[row.retry_active_index]?.id ?? versions[0].id);
}
if (version) insert.run(messageId, version.content, version.thinking || null, JSON.stringify(version.sources ?? null), version.modelName);
versions = db.prepare("SELECT * FROM chat_message_versions WHERE message_id = ? ORDER BY id").all(messageId);
const index = version ? versions.length - 1 : selectedIndex;
if (!Number.isInteger(index) || !versions[index]) throw new Error("Version is out of range");
const active = versions[index];
db.prepare(`UPDATE chat_messages SET content = ?, thinking = ?, sources = ?, model_name = ?,
retry_variants = ?, retry_active_index = ?, retry_retries_used = ? WHERE id = ? AND chat_id = ?`)
.run(active.content, active.thinking, active.sources, active.model_name, JSON.stringify(versions.map(item => item.content)), index, versions.length - 1, messageId, chatId);
})(),
updateChatMessage: (chatId, messageId, content) => statements.updateChatMessage.run(content, messageId, chatId),
updateChatMessageRetryState: (chatId, messageId, payload) => statements.updateChatMessageWithRetryState.run(
payload.content,
@@ -1278,6 +1305,10 @@ export function createRepositories(db, config) {
persona_state: this.getPersonaChatState(chat.id, username),
messages: this.listChatMessages(chat.id).map((message) => ({
...message,
responseVersions: this.listMessageVersions(message.id).map(version => ({
content: version.content, thinking: version.thinking,
sources: parseJsonArray(version.sources), modelName: version.model_name, createdAt: version.created_at
})),
attachments: db.prepare("SELECT a.* FROM attachments a JOIN message_attachments m ON a.id=m.attachment_id WHERE m.message_id=? ORDER BY m.position").all(message.id).map(attachmentMetadata),
retryVariants: parseJsonArray(message.retry_variants)
}))
src/db/schema.js
+10
View File
@@ -236,6 +236,16 @@ export function initializeSchema(db) {
}
export function runMigrations(db, {dropLegacyChats = true} = {}) {
db.exec(`CREATE TABLE IF NOT EXISTS chat_message_versions (
id INTEGER PRIMARY KEY AUTOINCREMENT,
message_id INTEGER NOT NULL REFERENCES chat_messages(id) ON DELETE CASCADE,
content TEXT NOT NULL,
thinking TEXT,
sources TEXT,
model_name TEXT,
created_at DATETIME DEFAULT CURRENT_TIMESTAMP
);
CREATE INDEX IF NOT EXISTS chat_message_versions_message ON chat_message_versions(message_id);`);
const ensureColumn = (tableName, columns, name, sql) => {
if (!columns.some((column) => column.name === name)) {
db.prepare(sql).run();
src/routes/chats.js
+53 -3
View File
@@ -105,6 +105,14 @@ const preferredModelValidator = (body) => {
export function createChatsRouter({repositories, chatService, modelService, config, chatRateLimiters}) {
const router = express.Router();
router.use(requireLogin);
router.use(["/chats/:id/messages/by-index/:index", "/chats/:id/messages/:messageId"], (req, res, next) => {
const chatId = Number(req.params.id);
const messageId = req.params.messageId ? Number(req.params.messageId) : chatService.getChatMessageByIndex(chatId, req.params.index)?.id;
if (["PUT", "DELETE", "POST"].includes(req.method) && chatService.isRegeneratingMessage(chatId, messageId)) {
return res.status(409).json({error: "This response is being regenerated. Wait for it to finish."});
}
next();
});
router.get("/models", async (req, res) => {
try {
@@ -475,7 +483,21 @@ export function createChatsRouter({repositories, chatService, modelService, conf
}
});
router.post("/chats/:id/messages/:messageId/retry", ...chatRateLimiters, async (req, res) => {
router.post("/chats/:id/messages/:messageId/versions/:index", (req, res) => {
const chatId = Number(req.params.id), messageId = Number(req.params.messageId);
if (!repositories.getChat(chatId, req.session.user)) return res.status(404).json({error: "Chat not found"});
const message = repositories.getChatMessage(chatId, messageId);
if (!message || message.role !== "bot") return res.status(404).json({error: "Assistant message not found"});
if (chatService.isGenerating(chatId)) return res.status(409).json({error: "Wait for generation to finish"});
const index = Number(req.params.index);
const versions = repositories.listMessageVersions(messageId);
if (!Number.isInteger(index) || index < 0 || index >= versions.length) return res.status(400).json({error: "Version is out of range"});
repositories.saveMessageVersion(chatId, messageId, null, index);
repositories.updateChatContextSummary(chatId, req.session.user, null, 0);
res.json({message: chatService.formatChatMessage(repositories.getChatMessage(chatId, messageId))});
});
router.post(["/chats/:id/messages/:messageId/retry", "/chats/:id/messages/:messageId/retry/stream"], ...chatRateLimiters, async (req, res) => {
const chatId = Number(req.params.id);
const messageId = Number(req.params.messageId);
const chat = repositories.getChat(chatId, req.session.user);
@@ -484,19 +506,45 @@ export function createChatsRouter({repositories, chatService, modelService, conf
if (!targetMessage) return res.status(404).json({error: "Message not found"});
if (targetMessage.role !== "bot") return res.status(400).json({error: "Only assistant messages can be regenerated"});
const streaming = req.path.endsWith("/stream");
const controller = new AbortController();
let finished = false;
const sendEvent = (event, data) => {
if (!res.destroyed) res.write(`event: ${event}\ndata: ${JSON.stringify(data)}\n\n`);
};
if (streaming) {
res.once("close", () => { if (!finished) controller.abort(); });
res.writeHead(200, {"Content-Type": "text/event-stream", "Cache-Control": "no-cache, no-transform", Connection: "keep-alive", "X-Accel-Buffering": "no"});
}
try {
const result = await chatService.retryMessage({
user: req.session.user,
chatId,
selectedModel: String(req.body?.model || "qwen3:14b"),
targetMessage,
retryStyle: String(req.body?.style || "").trim() || null
retryStyle: String(req.body?.style || "").trim() || null,
webSearch: req.body?.webSearch === true,
reasoning: req.body?.reasoning === true,
signal: controller.signal,
...(streaming ? {
onChunk: (chunk, fullReply) => sendEvent("chunk", {chunk, fullReply}),
onThinking: (chunk, fullThinking) => sendEvent("thinking", {chunk, fullThinking}),
onStatus: status => sendEvent("status", status)
} : {})
});
if (streaming) {
sendEvent(result.error ? "error" : "done", result.error ? {error: result.error, searchState: result.searchState} : {...result, reply: result.message.content});
return;
}
if (result.error) return res.status(result.status).json({error: result.error});
return res.json(result);
} catch (error) {
const mapped = modelService.mapError(error);
if (streaming) { sendEvent("error", mapped.body); return; }
return res.status(mapped.status).json(mapped.body);
} finally {
finished = true;
if (streaming) res.end();
}
});
@@ -514,7 +562,9 @@ export function createChatsRouter({repositories, chatService, modelService, conf
chatId,
selectedModel: String(req.body?.model || "qwen3:14b"),
targetMessage,
retryStyle: String(req.body?.style || "").trim() || null
retryStyle: String(req.body?.style || "").trim() || null,
webSearch: req.body?.webSearch === true,
reasoning: req.body?.reasoning === true
});
if (result.error) return res.status(result.status).json({error: result.error});
return res.json(result);
src/services/chatService.js
+71 -34
View File
@@ -24,12 +24,22 @@ import {supportsReasoning} from "../../public/app/modelCapabilities.js";
export function createChatService(repositories, modelService, config, webSearchProvider = new SearXNGSearchProvider(config), attachmentService = null) {
const formatChatMessage = row => {
const message = formatBaseChatMessage(row);
if (message?.role === "bot") {
const versions = repositories.listMessageVersions(row.id);
message.versionCount = versions.length || 1;
message.versionIndex = versions.length ? Number(row.retry_active_index || 0) : 0;
if (versions.length) message.responseVersions = versions.map(version => ({
content: version.content, thinking: version.thinking,
sources: parseSources(version.sources), modelName: version.model_name, createdAt: version.created_at
}));
}
const attachments = row && attachmentService?.list(row.id);
return message && attachments?.length ? {...message, attachments} : message;
};
const chatGenerationQueues = new Map();
const summaryQueues = new Map();
const activeGenerations = new Map();
const regeneratingMessages = new Map();
const maxRetries = Number(config.chatMaxRetries ?? 5);
const now = typeof config.now === "function" ? config.now : () => new Date();
@@ -84,6 +94,8 @@ export function createChatService(repositories, modelService, config, webSearchP
});
const persistChatMessageRetryState = ({chatId, messageId, content, retryVariants, retryActiveIndex, retryRetriesUsed, retryPromptMessageId}) => {
const versioned = repositories.listMessageVersions(messageId).length > 0;
if (versioned) repositories.saveMessageVersion(chatId, messageId, null, retryActiveIndex);
repositories.updateChatMessageRetryState(chatId, messageId, {
content,
retryVariants: retryVariants || [content || ""],
@@ -91,6 +103,7 @@ export function createChatService(repositories, modelService, config, webSearchP
retryRetriesUsed: Number(retryRetriesUsed || 0),
retryPromptMessageId: retryPromptMessageId ?? null
});
if (versioned) repositories.saveMessageVersion(chatId, messageId, null, retryActiveIndex);
};
const getChatMessageByIndex = (chatId, index) => {
@@ -174,7 +187,7 @@ export function createChatService(repositories, modelService, config, webSearchP
return repositories.createChat(user, persona.name, persona.id, userPersonaId, scenarioPrompt, scenarioSummary);
};
const generateRoleplayOpener = async ({selectedModel = "qwen3:14b", assistantPersona, userPersona, scenarioPrompt, sceneSummary, onThinking = null}) => {
const generateRoleplayOpener = async ({selectedModel = "qwen3:14b", assistantPersona, userPersona, scenarioPrompt, sceneSummary, onThinking = null, onChunk = null, signal = null, generation = {}, reasoning = false}) => {
const baseMessages = [
{role: "system", content: buildPersonaPrompt(assistantPersona)},
...(userPersona ? [{role: "system", content: buildUserPersonaPrompt(userPersona)}] : []),
@@ -182,20 +195,24 @@ export function createChatService(repositories, modelService, config, webSearchP
{role: "user", content: buildRoleplayOpenerPrompt({assistantPersona, userPersona, scenarioPrompt, sceneSummary})}
];
let opener = await modelService.generateReply(selectedModel, baseMessages, {onThinking});
const generate = messages => onChunk
? modelService.streamReply(selectedModel, messages, onChunk, {onThinking, signal, generation, reasoning})
: modelService.generateReply(selectedModel, messages, {onThinking, signal, generation, reasoning});
let opener = await generate(baseMessages);
if (!containsUserVoiceInRoleplayOpener(opener, userPersona)) {
return opener;
}
onThinking?.("", "");
return modelService.generateReply(selectedModel, [
onChunk?.("", "");
return generate([
...baseMessages,
{role: "assistant", content: opener || ""},
{
role: "user",
content: "Rewrite this opening. Hard rule: only the assistant character may speak or act. Never write or imply user speech, thoughts, feelings, reactions, decisions, actions, or labels like 'User:'/'You:'. Stop before the user's reply."
}
], {onThinking});
]);
};
const generateValidatedRoleplayReply = async ({selectedModel, messagesPayload, userPersona, signal = null, generation = {}, onThinking = null, onChunk = null, reasoning = false}) => {
@@ -285,13 +302,20 @@ export function createChatService(repositories, modelService, config, webSearchP
return {session: activeSession, activePersona, activeUserPersona, messagesPayload};
};
const buildRetryPayload = async ({chatId, user, session, selectedModel, targetMessage, retryStyle = null, signal = null, generation = {}}) => {
const buildRetryPayload = async ({chatId, user, session, selectedModel, targetMessage, retryStyle = null, signal = null, generation = {}, webSearch = false, reasoning = false, onChunk = null, onThinking = null, onStatus = null}) => {
let thinking = "";
const captureThinking = (chunk, fullThinking) => {
if (!reasoning) return;
thinking = fullThinking;
onThinking?.(chunk, fullThinking);
};
onStatus?.({status: "accepted"});
const parentMessage = targetMessage.parent_message_id
? repositories.getChatMessage(chatId, targetMessage.parent_message_id)
: null;
const promptMessage = parentMessage?.role === "user"
? parentMessage
: repositories.getPreviousUserMessage(chatId, targetMessage.id);
: repositories.listChatBranchThroughMessage(chatId, targetMessage.id).slice(0, -1).reverse().find(entry => entry.role === "user");
const activePersona = repositories.getAssistantPersonaForChat(chatId, user);
const activeUserPersona = getUserPersonaForSession(chatId, user, session, activePersona);
@@ -299,9 +323,8 @@ export function createChatService(repositories, modelService, config, webSearchP
if (!activePersona || !session.scenario_summary) {
return {error: "No previous user prompt found for retry"};
}
let thinking = "";
const opener = await generateRoleplayOpener({
onThinking: (_chunk, fullThinking) => { thinking = fullThinking; },
onThinking: captureThinking, onChunk, signal, generation, reasoning,
selectedModel,
assistantPersona: activePersona,
userPersona: activeUserPersona,
@@ -314,25 +337,38 @@ export function createChatService(repositories, modelService, config, webSearchP
return {fullReply: opener, thinking, promptMessageId: null};
}
if (attachmentService?.list(promptMessage.id).length) onStatus?.({status: "analyzing"});
await attachmentService?.analyze(promptMessage.id, promptMessage.content, signal);
let sources = null;
let searchState = "skipped";
if (webSearch) {
const search = await retrieveWebSearch(promptMessage.content, webSearchProvider, {
signal,
referenceText: attachmentService?.searchContext(promptMessage.id) || "",
onSearching: () => onStatus?.({status: "searching"})
});
sources = search.sources;
searchState = search.state;
}
onStatus?.({status: "generating", searchState, sources});
const {messagesPayload} = buildConversationContext({
user,
chatId,
session,
message: null,
upToMessageId: promptMessage.id,
sources: targetMessage.sources ? parseSources(targetMessage.sources) : null
sources
});
const styledPayload = appendRetryStyleInstruction(messagesPayload, retryStyle);
let thinking = "";
const onThinking = (_chunk, fullThinking) => { thinking = fullThinking; };
const fullReply = activePersona
? await generateValidatedRoleplayReply({selectedModel, messagesPayload: styledPayload, userPersona: activeUserPersona, signal, generation, onThinking})
: await modelService.generateReply(selectedModel, styledPayload, {signal, generation, onThinking});
? await generateValidatedRoleplayReply({selectedModel, messagesPayload: styledPayload, userPersona: activeUserPersona, signal, generation, onThinking: captureThinking, onChunk, reasoning})
: onChunk
? await modelService.streamReply(selectedModel, styledPayload, onChunk, {signal, generation, onThinking: captureThinking, reasoning})
: await modelService.generateReply(selectedModel, styledPayload, {signal, generation, onThinking: captureThinking, reasoning});
if (!fullReply) return {error: "Failed to regenerate message"};
return {fullReply, thinking, promptMessageId: promptMessage.id};
return {fullReply, thinking, sources, searchState, promptMessageId: promptMessage.id};
};
const generateChatTurn = ({user, chatId, message, model, attachments = [], reuseChatId = null, webSearch = false, reasoning = false, onStatus = null, onChunk = null, onThinking = null, signal = null, parentMessageId = undefined}) =>
@@ -536,8 +572,13 @@ export function createChatService(repositories, modelService, config, webSearchP
async streamChatMessage({user, chatId, message, model, attachments = [], webSearch = false, reasoning = false, onStatus, onChunk, onThinking, signal = null}) {
return generateChatTurn({user, chatId, message, model, attachments, webSearch, reasoning, onStatus, onChunk, onThinking, signal});
},
async retryMessage({user, chatId, selectedModel, targetMessage, retryStyle = null}) {
isGenerating: (chatId) => chatGenerationQueues.has(chatId),
isRegeneratingMessage: (chatId, messageId) => regeneratingMessages.get(chatId) === messageId,
async retryMessage({user, chatId, selectedModel, targetMessage, retryStyle = null, webSearch = false, reasoning = false, onChunk = null, onThinking = null, onStatus = null, signal = null}) {
if (chatGenerationQueues.has(chatId)) return {error: "A response is already being generated", status: 409};
regeneratingMessages.set(chatId, targetMessage.id);
return withChatGenerationLock(chatId, async () => {
await (summaryQueues.get(chatId) || Promise.resolve());
const session = repositories.getChat(chatId, user);
if (!session) return {error: "Chat not found", status: 404};
@@ -545,35 +586,31 @@ export function createChatService(repositories, modelService, config, webSearchP
const existingMessage = formatChatMessage(freshTarget);
if (!existingMessage) return {error: "Message not found", status: 404};
if (existingMessage.role !== "bot") return {error: "Only assistant messages can be regenerated", status: 400};
const siblings = repositories.listSiblingChatMessages(chatId, "bot", freshTarget.parent_message_id);
if (siblings.length - 1 >= maxRetries) {
if (existingMessage.versionCount - 1 >= maxRetries) {
return {error: `Maximum of ${maxRetries} retries reached`, status: 400};
}
const {model, generation} = getGenerationSettings(session, selectedModel);
return withActiveGeneration(chatId, null, async (activeSignal) => {
return withActiveGeneration(chatId, signal, async (activeSignal) => {
activeSignal.throwIfAborted();
const retryResult = await buildRetryPayload({
chatId, user, session, selectedModel: model, targetMessage: freshTarget, retryStyle,
signal: activeSignal, generation
signal: activeSignal, generation, webSearch, reasoning: reasoning === true && supportsReasoning(model), onChunk, onThinking, onStatus
});
if (retryResult.error) return {error: retryResult.error, status: 400};
const retryVariants = [...siblings.map((sibling) => sibling.content), retryResult.fullReply];
const retryState = {
retryVariants,
retryActiveIndex: retryVariants.length - 1,
retryRetriesUsed: siblings.length,
retryPromptMessageId: retryResult.promptMessageId
};
const inserted = repositories.insertBranchChatMessage(
chatId, "bot", retryResult.fullReply, retryState, model, "complete", null, freshTarget.parent_message_id
);
if (freshTarget.sources) repositories.updateChatMessageSources(chatId, Number(inserted.lastInsertRowid), parseSources(freshTarget.sources));
if (retryResult.thinking) repositories.updateChatMessageThinking(chatId, Number(inserted.lastInsertRowid), retryResult.thinking);
activeSignal.throwIfAborted();
repositories.saveMessageVersion(chatId, freshTarget.id, {
content: retryResult.fullReply, thinking: retryResult.thinking,
sources: retryResult.sources, modelName: model
});
repositories.updateChatContextSummary(chatId, user, null, 0);
repositories.touchChat(chatId, user);
const message = formatChatMessage(repositories.getChatMessage(chatId, Number(inserted.lastInsertRowid)));
return {message, promptMessageId: retryResult.promptMessageId, retryStyle, branchedFromMessageId: freshTarget.id};
const message = formatChatMessage(repositories.getChatMessage(chatId, freshTarget.id));
return {message, promptMessageId: retryResult.promptMessageId, retryStyle};
});
});
}).catch(error => {
if (error instanceof WebSearchError) return {error: error.message, status: error.status, searchState: "failed"};
throw error;
}).finally(() => regeneratingMessages.delete(chatId));
},
async editAndResendMessage({user, chatId, targetMessage, content, model, webSearch = false, onChunk = null, signal = null}) {
if (!targetMessage || targetMessage.role !== "user") return {error: "Only user messages can be edited and resent", status: 400};
src/services/workspaceImportService.js
+4
View File
@@ -86,6 +86,10 @@ function validateMessage(message, chatIndex, messageIndex, errors) {
content,
model_name: optionalText(message.model_name || message.modelName),
thinking: optionalText(message.thinking),
responseVersions: asArray(message.responseVersions).filter(version => version && typeof version.content === "string").slice(0, 100).map(version => ({
content: version.content, thinking: optionalText(version.thinking), sources: parseSources(version.sources),
modelName: optionalText(version.modelName), createdAt: optionalText(version.createdAt)
})),
attachments: asArray(message.attachments).filter(item => item && typeof item.id === "string" && /^[a-f0-9-]{36}$/.test(item.id)).slice(0, 100).map(item => ({id: item.id})),
...(message.sources ? {sources: parseSources(message.sources)} : {}),
retryVariants,