context-efficient-ai/database/postgres.ts

import { Pool } from "pg";
import { Kysely, PostgresDialect } from "kysely";
import type Database from "./generated/Database";
import type {
  ConversationEntity,
  FactEntity,
  FactTriggerEntity,
  MessageEntity,
} from "./common.ts";
import type { CommittedMessage } from "../types";
import { env } from "../server/env.js";

export const pool = new Pool({
  connectionString: env.POSTGRES_CONNECTION_STRING as string,
  // channelBinding: require ?
});

const dialect = new PostgresDialect({
  pool,
});

// Database interface is passed to Kysely's constructor, and from now on, Kysely
// knows your database structure.
// Dialect is passed to Kysely's constructor, and from now on, Kysely knows how
// to communicate with your database.
export const dbClient = new Kysely<Database>({
  dialect,
});

const conversations: ConversationEntity = {
  construct: (conversation) => conversation,
  create: async (conversation) => {
    const insertedRows = await dbClient
      .insertInto("conversations")
      .values(conversation)
      .returningAll()
      .execute();
    return insertedRows[0];
  },
  createMany: async (conversations) => {
    const insertedRows = await dbClient
      .insertInto("conversations")
      .values(conversations)
      .returningAll()
      .execute();
    return insertedRows;
  },
  findAll: async () => {
    const rows = await dbClient
      .selectFrom("conversations")
      .selectAll()
      .execute();
    return rows;
  },
  findById: async (id) => {
    const row = await dbClient
      .selectFrom("conversations")
      .selectAll()
      .where("id", "=", id)
      .execute();
    return row[0];
  },
  update: async (id, data) => {
    await dbClient
      .updateTable("conversations")
      .set(data)
      .where("id", "=", id)
      .execute();
  },
  delete: async (id) => {
    await dbClient.deleteFrom("conversations").where("id", "=", id).execute();
  },
  fetchMessages: async (conversationId) => {
    const rows = await dbClient
      .selectFrom("messages")
      .selectAll()
      .where("conversationId", "=", conversationId)
      .execute();
    return rows as Array<CommittedMessage>;
  },
};

const facts: FactEntity = {
  construct: (fact) => fact,
  create: async (fact) => {
    const insertedRows = await dbClient
      .insertInto("facts")
      .values(fact)
      .returningAll()
      .execute();
    return insertedRows[0];
  },
  createMany: async (facts) => {
    const insertedRows = await dbClient
      .insertInto("facts")
      .values(facts)
      .returningAll()
      .execute();
    return insertedRows;
  },
  findAll: async () => {
    const rows = await dbClient.selectFrom("facts").selectAll().execute();
    return rows;
  },
  findById: async (id) => {
    const row = await dbClient
      .selectFrom("facts")
      .selectAll()
      .where("id", "=", id)
      .execute();
    return row[0];
  },
  update: async (id, data) => {
    await dbClient
      .updateTable("facts")
      .set(data)
      .where("id", "=", id)
      .execute();
  },
  delete: async (id) => {
    await dbClient.deleteFrom("facts").where("id", "=", id).execute();
  },
  findByConversationId: async (conversationId) => {
    const rows = await dbClient
      .selectFrom("facts")
      .innerJoin("messages", "messages.id", "facts.sourceMessageId")
      .selectAll("facts")
      .where("conversationId", "=", conversationId)
      .execute();
    return rows;
  },
};

const factTriggers: FactTriggerEntity = {
  construct: (factTrigger) => factTrigger,
  create: async (factTrigger) => {
    const insertedRows = await dbClient
      .insertInto("fact_triggers")
      .values(factTrigger)
      .returningAll()
      .execute();
    return insertedRows[0];
  },
  createMany: async (factTriggers) => {
    const insertedRows = await dbClient
      .insertInto("fact_triggers")
      .values(factTriggers)
      .returningAll()
      .execute();
    return insertedRows;
  },
  findAll: async () => {
    const rows = await dbClient
      .selectFrom("fact_triggers")
      .selectAll()
      .execute();
    return rows;
  },
  findById: async (id) => {
    const row = await dbClient
      .selectFrom("fact_triggers")
      .selectAll()
      .where("id", "=", id)
      .execute();
    return row[0];
  },
  update: async (id, data) => {
    await dbClient
      .updateTable("fact_triggers")
      .set(data)
      .where("id", "=", id)
      .execute();
  },
  delete: async (id) => {
    await dbClient.deleteFrom("fact_triggers").where("id", "=", id).execute();
  },
  findByFactId: async (factId) => {
    const rows = await dbClient
      .selectFrom("fact_triggers")
      .innerJoin("facts", "facts.id", "fact_triggers.sourceFactId")
      .selectAll("fact_triggers")
      .where("sourceFactId", "=", factId)
      .execute();
    return rows;
  },
  findByConversationId: async (conversationId) => {
    const rows = await dbClient
      .selectFrom("fact_triggers")
      .innerJoin("facts", "facts.id", "fact_triggers.sourceFactId")
      .innerJoin("messages", "messages.id", "facts.sourceMessageId")
      .selectAll("fact_triggers")
      .where("messages.conversationId", "=", conversationId)
      .execute();
    return rows;
  },
};

const messages: MessageEntity = {
  construct: (message) => message,
  create: async (message) => {
    const insertedRows = await dbClient
      .insertInto("messages")
      .values({ ...message, parts: JSON.stringify(message.parts) })
      .returningAll()
      .execute();
    return insertedRows[0] as CommittedMessage;
  },
  createMany: async (messages) => {
    const insertedRows = await dbClient
      .insertInto("messages")
      .values(
        messages.map((message) => ({
          ...message,
          parts: JSON.stringify(message.parts),
        }))
      )
      .returningAll()
      .execute();
    return insertedRows as Array<CommittedMessage>;
  },
  findAll: async () => {
    const rows = await dbClient.selectFrom("messages").selectAll().execute();
    return rows as Array<CommittedMessage>;
  },
  findById: async (id) => {
    const row = await dbClient
      .selectFrom("messages")
      .selectAll()
      .where("id", "=", id)
      .execute();
    return row[0] as CommittedMessage;
  },
  update: async (id, data) => {
    await dbClient
      .updateTable("messages")
      .set(data)
      .where("id", "=", id)
      .execute();
  },
  delete: async (id) => {
    await dbClient.deleteFrom("messages").where("id", "=", id).execute();
  },
  findByConversationId: async (conversationId) => {
    const rows = (await dbClient
      .selectFrom("messages")
      .selectAll()
      .where("conversationId", "=", conversationId)
      .execute()) as Array<CommittedMessage>;
    return rows;
  },
};

export const db = {
  conversations,
  facts,
  factTriggers,
  messages,
};