WIP

apps/server/package.json

@@ -37,6 +37,7 @@
     "@fastify/cookie": "^11.0.2",
     "@fastify/multipart": "^9.0.3",
     "@fastify/static": "^8.2.0",
+    "@langchain/textsplitters": "^0.1.0",
     "@nestjs/bullmq": "^11.0.2",
     "@nestjs/common": "^11.1.3",
     "@nestjs/config": "^4.0.2",
@@ -78,6 +79,7 @@
     "passport-jwt": "^4.0.1",
     "pg": "^8.16.0",
     "pg-tsquery": "^8.4.2",
+    "pgvector": "^0.2.1",
     "postmark": "^4.0.5",
     "react": "^18.3.1",
     "reflect-metadata": "^0.2.2",

apps/server/src/collaboration/extensions/persistence.extension.ts

@@ -35,6 +35,7 @@ export class PersistenceExtension implements Extension {
     @InjectKysely() private readonly db: KyselyDB,
     private eventEmitter: EventEmitter2,
     @InjectQueue(QueueName.GENERAL_QUEUE) private generalQueue: Queue,
+    @InjectQueue(QueueName.AI_QUEUE) private aiQueue: Queue,
   ) {}
 
   async onLoadDocument(data: onLoadDocumentPayload) {
@@ -181,10 +182,22 @@ export class PersistenceExtension implements Extension {
     }
 
     this.contributors.get(documentName).add(userId);
+
+    console.log('embedd me')
+    const pageId = getPageId(documentName);
+
+    await this.aiQueue.add(QueueJob.GENERATE_PAGE_EMBEDDINGS, {
+      pageId: pageId,
+    });
   }
 
   async afterUnloadDocument(data: afterUnloadDocumentPayload) {
     const documentName = data.documentName;
+    const pageId = getPageId(documentName);
+
     this.contributors.delete(documentName);
+
+    // should only queue embed after unload// should delay so we dont embed always
+
   }
 }

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

@@ -63,6 +63,7 @@ export class PageRepo {
     pageId: string,
     opts?: {
       includeContent?: boolean;
+      includeText?: boolean;
       includeYdoc?: boolean;
       includeSpace?: boolean;
       includeCreator?: boolean;
@@ -79,6 +80,7 @@ export class PageRepo {
       .selectFrom('pages')
       .select(this.baseFields)
       .$if(opts?.includeContent, (qb) => qb.select('content'))
+      .$if(opts?.includeText, (qb) => qb.select('textContent'))
       .$if(opts?.includeYdoc, (qb) => qb.select('ydoc'))
       .$if(opts?.includeHasChildren, (qb) =>
         qb.select((eb) => this.withHasChildren(eb)),

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

@@ -0,0 +1,45 @@
+import {
+  Attachments,
+  AuthAccounts,
+  AuthProviders,
+  Backlinks,
+  Billing,
+  Comments,
+  FileTasks,
+  Groups,
+  GroupUsers,
+  PageHistory,
+  Pages,
+  Shares,
+  SpaceMembers,
+  Spaces,
+  UserMfa,
+  Users,
+  UserTokens,
+  WorkspaceInvitations,
+  Workspaces,
+} from '@docmost/db/types/db';
+import { Embeddings } from '@docmost/db/types/embeddings.types';
+
+export interface DbInterface {
+  attachments: Attachments;
+  authAccounts: AuthAccounts;
+  authProviders: AuthProviders;
+  backlinks: Backlinks;
+  billing: Billing;
+  comments: Comments;
+  fileTasks: FileTasks;
+  groups: Groups;
+  groupUsers: GroupUsers;
+  pageEmbeddings: Embeddings;
+  pageHistory: PageHistory;
+  pages: Pages;
+  shares: Shares;
+  spaceMembers: SpaceMembers;
+  spaces: Spaces;
+  userMfa: UserMfa;
+  users: Users;
+  userTokens: UserTokens;
+  workspaceInvitations: WorkspaceInvitations;
+  workspaces: Workspaces;
+}

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

@@ -0,0 +1,17 @@
+import { Json, Timestamp, Generated } from '@docmost/db/types/db';
+
+// page_embeddings type
+export interface Embeddings {
+  id: Generated<string>;
+  pageId: string;
+  spaceId: string;
+  workspaceId: string;
+  embedding: number[] | Buffer | string;
+  chunkIndex: Generated<number>;
+  chunkStart: Generated<number>;
+  chunkLength: Generated<number>;
+  metadata: Generated<Json>;
+  createdAt: Generated<Timestamp>;
+  updatedAt: Generated<Timestamp>;
+  deletedAt: Timestamp | null;
+}

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

@@ -20,6 +20,7 @@ import {
   FileTasks,
   UserMfa as _UserMFA,
 } from './db';
+import { Embeddings } from '@docmost/db/types/embeddings.types';
 
 // Workspace
 export type Workspace = Selectable<Workspaces>;
@@ -119,3 +120,8 @@ export type UpdatableFileTask = Updateable<Omit<FileTasks, 'id'>>;
 export type UserMFA = Selectable<_UserMFA>;
 export type InsertableUserMFA = Insertable<_UserMFA>;
 export type UpdatableUserMFA = Updateable<Omit<_UserMFA, 'id'>>;
+
+// Page Embedding
+export type PageEmbedding = Selectable<Embeddings>;
+export type InsertablePageEmbedding = Insertable<Embeddings>;
+export type UpdatablePageEmbedding = Updateable<Omit<Embeddings, 'id'>>;

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

@@ -1,5 +1,5 @@
-import { DB } from './db';
 import { Kysely, Transaction } from 'kysely';
+import { DbInterface } from '@docmost/db/types/db.interface';
 
-export type KyselyDB = Kysely<DB>;
-export type KyselyTransaction = Transaction<DB>;
+export type KyselyDB = Kysely<DbInterface>;
+export type KyselyTransaction = Transaction<DbInterface>;

apps/server/src/integrations/queue/constants/queue.constants.ts

@@ -4,6 +4,7 @@ export enum QueueName {
   GENERAL_QUEUE = '{general-queue}',
   BILLING_QUEUE = '{billing-queue}',
   FILE_TASK_QUEUE = '{file-task-queue}',
+  AI_QUEUE = '{ai-queue}',
 }
 
 export enum QueueJob {
@@ -23,4 +24,7 @@ export enum QueueJob {
 
   IMPORT_TASK = 'import-task',
   EXPORT_TASK = 'export-task',
+
+  GENERATE_PAGE_EMBEDDINGS = 'generate-page-embeddings',
+  DELETE_PAGE_EMBEDDINGS = 'delete-page-embeddings',
 }

apps/server/src/integrations/queue/queue.module.ts

@@ -57,6 +57,14 @@ import { BacklinksProcessor } from './processors/backlinks.processor';
         attempts: 1,
       },
     }),
+    BullModule.registerQueue({
+      name: QueueName.AI_QUEUE,
+      defaultJobOptions: {
+        removeOnComplete: true,
+        removeOnFail: true,
+        attempts: 1,
+      },
+    }),
   ],
   exports: [BullModule],
   providers: [BacklinksProcessor],