feat(ee): ai chat (#2098)

* feat: ai chat

* feat: ai chat

* sync

* cleanup

* view space button

apps/server/src/database/migrations/20260409T132415-ai-chat.ts

@@ -0,0 +1,118 @@
+import { type Kysely, sql } from 'kysely';
+
+export async function up(db: Kysely<any>): Promise<void> {
+  await db.schema
+    .createTable('ai_chats')
+    .ifNotExists()
+    .addColumn('id', 'uuid', (col) =>
+      col.primaryKey().defaultTo(sql`gen_uuid_v7()`),
+    )
+    .addColumn('workspace_id', 'uuid', (col) =>
+      col.references('workspaces.id').onDelete('cascade').notNull(),
+    )
+    .addColumn('creator_id', 'uuid', (col) =>
+      col.references('users.id').notNull(),
+    )
+    .addColumn('title', 'varchar', (col) => col)
+    .addColumn('created_at', 'timestamptz', (col) =>
+      col.notNull().defaultTo(sql`now()`),
+    )
+    .addColumn('updated_at', 'timestamptz', (col) =>
+      col.notNull().defaultTo(sql`now()`),
+    )
+    .addColumn('deleted_at', 'timestamptz', (col) => col)
+    .execute();
+
+  await db.schema
+    .createIndex('idx_ai_chats_workspace_creator')
+    .ifNotExists()
+    .on('ai_chats')
+    .columns(['workspace_id', 'creator_id', 'id'])
+    .execute();
+
+  await db.schema
+    .createTable('ai_chat_messages')
+    .ifNotExists()
+    .addColumn('id', 'uuid', (col) =>
+      col.primaryKey().defaultTo(sql`gen_uuid_v7()`),
+    )
+    .addColumn('chat_id', 'uuid', (col) =>
+      col.references('ai_chats.id').onDelete('cascade').notNull(),
+    )
+    .addColumn('workspace_id', 'uuid', (col) =>
+      col.references('workspaces.id').onDelete('cascade').notNull(),
+    )
+    .addColumn('user_id', 'uuid', (col) =>
+      col.references('users.id').onDelete('set null'),
+    )
+    .addColumn('role', 'varchar', (col) => col.notNull())
+    .addColumn('content', 'text', (col) => col)
+    .addColumn('tool_calls', 'jsonb', (col) => col)
+    .addColumn('metadata', 'jsonb', (col) => col)
+    .addColumn('tsv', sql`tsvector`, (col) => col)
+    .addColumn('created_at', 'timestamptz', (col) =>
+      col.notNull().defaultTo(sql`now()`),
+    )
+    .addColumn('updated_at', 'timestamptz', (col) =>
+      col.notNull().defaultTo(sql`now()`),
+    )
+    .addColumn('deleted_at', 'timestamptz', (col) => col)
+    .execute();
+
+  await db.schema
+    .createIndex('idx_ai_chat_messages_chat_id')
+    .ifNotExists()
+    .on('ai_chat_messages')
+    .columns(['chat_id', 'id'])
+    .execute();
+
+  await db.schema
+    .createIndex('idx_ai_chat_messages_tsv')
+    .ifNotExists()
+    .on('ai_chat_messages')
+    .using('GIN')
+    .column('tsv')
+    .execute();
+
+  //ts-vector
+  await sql`
+    CREATE OR REPLACE FUNCTION ai_chat_messages_tsvector_trigger() RETURNS trigger AS $$
+    BEGIN
+      NEW.tsv := to_tsvector('english', f_unaccent(substring(coalesce(NEW.content, ''), 1, 100000)));
+      RETURN NEW;
+    END;
+    $$ LANGUAGE plpgsql;
+  `.execute(db);
+
+  await sql`
+    CREATE OR REPLACE TRIGGER ai_chat_messages_tsvector_update
+      BEFORE INSERT OR UPDATE ON ai_chat_messages
+      FOR EACH ROW EXECUTE FUNCTION ai_chat_messages_tsvector_trigger();
+  `.execute(db);
+
+  await db.schema
+    .alterTable('attachments')
+    .addColumn('ai_chat_id', 'uuid', (col) => col)
+    .execute();
+
+  await db.schema
+    .createIndex('idx_attachments_ai_chat_id')
+    .ifNotExists()
+    .on('attachments')
+    .column('ai_chat_id')
+    .execute();
+}
+
+export async function down(db: Kysely<any>): Promise<void> {
+  await db.schema.dropIndex('idx_attachments_ai_chat_id').execute();
+  await db.schema.alterTable('attachments').dropColumn('ai_chat_id').execute();
+
+  await sql`DROP TRIGGER IF EXISTS ai_chat_messages_tsvector_update ON ai_chat_messages`.execute(
+    db,
+  );
+  await sql`DROP FUNCTION IF EXISTS ai_chat_messages_tsvector_trigger`.execute(
+    db,
+  );
+  await db.schema.dropTable('ai_chat_messages').execute();
+  await db.schema.dropTable('ai_chats').execute();
+}

apps/server/src/database/repos/attachment/attachment.repo.ts

@@ -7,6 +7,7 @@ import {
   InsertableAttachment,
   UpdatableAttachment,
 } from '@docmost/db/types/entity.types';
+import { AttachmentType } from '../../../core/attachment/attachment.constants';
 
 @Injectable()
 export class AttachmentRepo {
@@ -23,6 +24,7 @@ export class AttachmentRepo {
     'creatorId',
     'pageId',
     'spaceId',
+    'aiChatId',
     'workspaceId',
     'createdAt',
     'updatedAt',
@@ -44,6 +46,21 @@ export class AttachmentRepo {
       .executeTakeFirst();
   }
 
+  async findByIdWithContent(
+    attachmentId: string,
+    opts?: {
+      trx?: KyselyTransaction;
+    },
+  ): Promise<Attachment> {
+    const db = dbOrTx(this.db, opts?.trx);
+
+    return db
+      .selectFrom('attachments')
+      .select([...this.baseFields, 'textContent'])
+      .where('id', '=', attachmentId)
+      .executeTakeFirst();
+  }
+
   async insertAttachment(
     insertableAttachment: InsertableAttachment,
     trx?: KyselyTransaction,
@@ -72,6 +89,21 @@ export class AttachmentRepo {
       .execute();
   }
 
+  async findByAiChatId(
+    aiChatId: string,
+    opts?: {
+      trx?: KyselyTransaction;
+    },
+  ): Promise<Attachment[]> {
+    const db = dbOrTx(this.db, opts?.trx);
+
+    return db
+      .selectFrom('attachments')
+      .select(this.baseFields)
+      .where('aiChatId', '=', aiChatId)
+      .execute();
+  }
+
   updateAttachmentsByPageId(
     updatableAttachment: UpdatableAttachment,
     pageIds: string[],
@@ -97,6 +129,25 @@ export class AttachmentRepo {
       .executeTakeFirst();
   }
 
+  async claimAttachmentsForChat(
+    attachmentIds: string[],
+    aiChatId: string,
+    creatorId: string,
+    workspaceId: string,
+  ): Promise<void> {
+    if (attachmentIds.length === 0) return;
+
+    await this.db
+      .updateTable('attachments')
+      .set({ aiChatId })
+      .where('id', 'in', attachmentIds)
+      .where('creatorId', '=', creatorId)
+      .where('workspaceId', '=', workspaceId)
+      .where('type', '=', AttachmentType.Chat)
+      .where('aiChatId', 'is', null)
+      .execute();
+  }
+
   async deleteAttachmentById(attachmentId: string): Promise<void> {
     await this.db
       .deleteFrom('attachments')

apps/server/src/database/types/db.d.ts

@@ -43,6 +43,7 @@ export interface ApiKeys {
 }
 
 export interface Attachments {
+  aiChatId: string | null;
   createdAt: Generated<Timestamp>;
   creatorId: string;
   deletedAt: Timestamp | null;
@@ -429,6 +430,31 @@ export interface PagePermissions {
   updatedAt: Generated<Timestamp>;
 }
 
+export interface AiChats {
+  id: Generated<string>;
+  workspaceId: string;
+  creatorId: string;
+  title: string | null;
+  createdAt: Generated<Timestamp>;
+  updatedAt: Generated<Timestamp>;
+  deletedAt: Timestamp | null;
+}
+
+export interface AiChatMessages {
+  id: Generated<string>;
+  chatId: string;
+  workspaceId: string;
+  userId: string | null;
+  role: string;
+  content: string | null;
+  toolCalls: Json | null;
+  metadata: Json | null;
+  tsv: string | null;
+  createdAt: Generated<Timestamp>;
+  updatedAt: Generated<Timestamp>;
+  deletedAt: Timestamp | null;
+}
+
 export interface UserSessions {
   id: Generated<string>;
   userId: string;
@@ -445,6 +471,8 @@ export interface UserSessions {
 }
 
 export interface DB {
+  aiChats: AiChats;
+  aiChatMessages: AiChatMessages;
   apiKeys: ApiKeys;
   attachments: Attachments;
   audit: Audit;

apps/server/src/database/types/entity.types.ts

@@ -1,5 +1,7 @@
 import { Insertable, Selectable, Updateable } from 'kysely';
 import {
+  AiChats,
+  AiChatMessages,
   Attachments,
   Comments,
   Groups,
@@ -29,6 +31,21 @@ import {
 } from './db';
 import { PageEmbeddings } from '@docmost/db/types/embeddings.types';
 
+// AI Chat
+export type AiChat = Selectable<AiChats>;
+export type InsertableAiChat = Insertable<AiChats>;
+export type UpdatableAiChat = Updateable<Omit<AiChats, 'id'>>;
+
+// AI Chat Message
+// `tsv` is an internal tsvector column maintained by a trigger for
+// full-text search. It is omitted from the public type so it never leaks
+// into HTTP responses or the chat history fed to the language model.
+export type AiChatMessage = Omit<Selectable<AiChatMessages>, 'tsv'>;
+export type InsertableAiChatMessage = Omit<
+  Insertable<AiChatMessages>,
+  'tsv'
+>;
+
 // Workspace
 export type Workspace = Selectable<Workspaces>;
 export type InsertableWorkspace = Insertable<Workspaces>;