You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

248 lines
7.6 KiB
TypeScript

import { Low } from "lowdb";
import { JSONFile } from "lowdb/node";
import type { CommittedMessage } from "../types";
import type {
Conversation,
ConversationEntity,
Fact,
FactEntity,
FactTrigger,
FactTriggerEntity,
MessageEntity,
} from "./common";
import { nanoid } from "nanoid";
type DB = {
conversations: Array<Conversation>;
messages: Array</*{
id: string;
conversationId: string;
content: string;
role: "user" | "assistant" | "system" | "data";
index: number;
createdAt: string;
runningSummary?: string;
}*/ CommittedMessage>;
facts: Array<Fact>;
factTriggers: Array<FactTrigger>;
};
export const dbClient = new Low<DB>(new JSONFile("db.json"), {
conversations: [],
messages: [],
facts: [],
factTriggers: [],
});
/** Initialize the database. Sets `db.data` to the default state if the file doesn't exist. */
await dbClient.read();
/** Write the database to the file, in case it didn't exist before. */
await dbClient.write();
const conversations: ConversationEntity = {
construct: (conversation: Conversation) => conversation,
create: async (_conversation) => {
const conversation = { ..._conversation, id: nanoid() };
await dbClient.data.conversations.push(conversation);
await dbClient.write();
return conversation;
},
createMany: async (_conversations) => {
const conversations = _conversations.map((c) => ({ ...c, id: nanoid() }));
await dbClient.data.conversations.push(...conversations);
await dbClient.write();
return conversations;
},
findAll: async () => {
return dbClient.data.conversations;
},
findById: async (id) => {
return dbClient.data.conversations.find((c) => c.id === id);
},
update: async (id, data: Partial<Conversation>) => {
const conversationIndex = dbClient.data.conversations.findIndex(
(c) => c.id === id
);
if (conversationIndex === -1) throw new Error("Conversation not found");
dbClient.data.conversations[conversationIndex] = {
...dbClient.data.conversations[conversationIndex],
...data,
};
await dbClient.write();
},
delete: async (id) => {
dbClient.data.conversations.splice(
dbClient.data.conversations.findIndex((c) => c.id === id),
1
);
const deletedMessageIds = dbClient.data.messages
.filter((m) => m.conversationId === id)
.map((m) => m.id);
dbClient.data.messages = dbClient.data.messages.filter(
(m) => m.conversationId !== id
);
const deletedFactIds = dbClient.data.facts
.filter((fact) => deletedMessageIds.includes(fact.sourceMessageId))
.map((fact) => fact.id);
dbClient.data.facts = dbClient.data.facts.filter(
(fact) => !deletedFactIds.includes(fact.id)
);
dbClient.data.factTriggers = dbClient.data.factTriggers.filter(
(factTrigger) => !deletedFactIds.includes(factTrigger.sourceFactId)
);
await dbClient.write();
},
fetchMessages: async (conversationId) => {
const rows = await dbClient.data.messages.filter(
(m) => m.conversationId === conversationId
);
return rows as Array<CommittedMessage>;
},
};
const factTriggers: FactTriggerEntity = {
construct: (factTrigger: FactTrigger) => factTrigger,
create: async (_factTrigger) => {
const factTrigger = { ..._factTrigger, id: nanoid() };
await dbClient.data.factTriggers.push(factTrigger);
await dbClient.write();
return factTrigger;
},
createMany: async (_factTriggers) => {
const factTriggers = _factTriggers.map((f) => ({ ...f, id: nanoid() }));
await dbClient.data.factTriggers.push(...factTriggers);
await dbClient.write();
return factTriggers;
},
findAll: async () => {
return dbClient.data.factTriggers;
},
findById: async (id) => {
return dbClient.data.factTriggers.find(
(factTrigger) => factTrigger.id === id
);
},
update: async (id, data: Partial<FactTrigger>) => {
const factTriggerIndex = dbClient.data.factTriggers.findIndex(
(factTrigger) => factTrigger.id === id
);
if (factTriggerIndex === -1) throw new Error("Fact trigger not found");
dbClient.data.factTriggers[factTriggerIndex] = {
...dbClient.data.factTriggers[factTriggerIndex],
...data,
};
await dbClient.write();
},
delete: async (id) => {
const deletedFactTriggerIndex = dbClient.data.factTriggers.findIndex(
(factTrigger) => factTrigger.id === id
);
if (deletedFactTriggerIndex === -1)
throw new Error("Fact trigger not found");
dbClient.data.factTriggers.splice(deletedFactTriggerIndex, 1);
await dbClient.write();
},
findByFactId: async (factId: string) => {
return dbClient.data.factTriggers.filter(
(factTrigger) => factTrigger.sourceFactId === factId
);
},
};
const facts: FactEntity = {
construct: (fact: Fact) => fact,
create: async (_fact) => {
const fact = { ..._fact, id: nanoid() };
await dbClient.data.facts.push(fact);
await dbClient.write();
return fact;
},
createMany: async (_facts) => {
const facts = _facts.map((f) => ({ ...f, id: nanoid() }));
await dbClient.data.facts.push(...facts);
await dbClient.write();
return facts;
},
findAll: async () => {
return dbClient.data.facts;
},
findById: async (id) => {
return dbClient.data.facts.find((fact) => fact.id === id);
},
update: async (id, data: Partial<Fact>) => {
const factIndex = dbClient.data.facts.findIndex((fact) => fact.id === id);
if (factIndex === -1) throw new Error("Fact not found");
dbClient.data.facts[factIndex] = {
...dbClient.data.facts[factIndex],
...data,
};
await dbClient.write();
},
delete: async (id) => {
const deletedFactId = dbClient.data.facts.findIndex(
(fact) => fact.id === id
);
if (deletedFactId === -1) throw new Error("Fact not found");
dbClient.data.facts.splice(deletedFactId, 1);
await dbClient.write();
},
findByConversationId: async (conversationId) => {
const conversationMessageIds = dbClient.data.messages
.filter((m) => m.conversationId === conversationId)
.map((m) => m.id);
const rows = await dbClient.data.facts.filter((f) =>
conversationMessageIds.includes(f.sourceMessageId)
);
return rows as Array<Fact>;
},
};
const messages: MessageEntity = {
construct: (message: CommittedMessage) => message,
create: async (_message) => {
const message = { ..._message, id: nanoid() };
await dbClient.data.messages.push(message);
await dbClient.write();
return message;
},
createMany: async (_messages) => {
const messages = _messages.map((m) => ({ ...m, id: nanoid() }));
await dbClient.data.messages.push(...messages);
await dbClient.write();
return messages;
},
findAll: async () => {
return dbClient.data.messages;
},
findById: async (id) => {
return dbClient.data.messages.find((m) => m.id === id);
},
update: async (id, data: Partial<CommittedMessage>) => {
const messageIndex = dbClient.data.messages.findIndex((m) => m.id === id);
if (messageIndex === -1) throw new Error("Message not found");
dbClient.data.messages[messageIndex] = {
...dbClient.data.messages[messageIndex],
...data,
};
await dbClient.write();
},
delete: async (id) => {
dbClient.data.messages.splice(
dbClient.data.messages.findIndex((m) => m.id === id),
1
);
await dbClient.write();
},
findByConversationId: async (conversationId) => {
return dbClient.data.messages.filter(
(m) => m.conversationId === conversationId
);
},
};
export const db = {
conversations,
factTriggers,
facts,
messages,
};