@lobehub/chat 1.11.9 → 1.12.1

package/src/database/server/schemas/lobechat/agent.ts

@@ -0,0 +1,93 @@
+/* eslint-disable sort-keys-fix/sort-keys-fix  */
+import { boolean, jsonb, pgTable, primaryKey, text, varchar } from 'drizzle-orm/pg-core';
+import { createInsertSchema } from 'drizzle-zod';
+
+import { LobeAgentChatConfig, LobeAgentTTSConfig } from '@/types/agent';
+
+import { idGenerator, randomSlug } from '../../utils/idGenerator';
+import { createdAt, updatedAt } from './_helpers';
+import { files, knowledgeBases } from './file';
+import { users } from './user';
+
+// Agent table is the main table for storing agents
+// agent is a model that represents the assistant that is created by the user
+// agent can have its own knowledge base and files
+
+export const agents = pgTable('agents', {
+  id: text('id')
+    .primaryKey()
+    .$defaultFn(() => idGenerator('agents'))
+    .notNull(),
+  slug: varchar('slug', { length: 100 })
+    .$defaultFn(() => randomSlug())
+    .unique(),
+  title: text('title'),
+  description: text('description'),
+  tags: jsonb('tags').$type<string[]>().default([]),
+  avatar: text('avatar'),
+  backgroundColor: text('background_color'),
+
+  plugins: jsonb('plugins').$type<string[]>().default([]),
+  userId: text('user_id')
+    .references(() => users.id, { onDelete: 'cascade' })
+    .notNull(),
+
+  chatConfig: jsonb('chat_config').$type<LobeAgentChatConfig>(),
+
+  fewShots: jsonb('few_shots'),
+  model: text('model'),
+  params: jsonb('params').default({}),
+  provider: text('provider'),
+  systemRole: text('system_role'),
+  tts: jsonb('tts').$type<LobeAgentTTSConfig>(),
+
+  createdAt: createdAt(),
+  updatedAt: updatedAt(),
+});
+
+export const insertAgentSchema = createInsertSchema(agents);
+
+export type NewAgent = typeof agents.$inferInsert;
+export type AgentItem = typeof agents.$inferSelect;
+
+export const agentsKnowledgeBases = pgTable(
+  'agents_knowledge_bases',
+  {
+    agentId: text('agent_id')
+      .references(() => agents.id, { onDelete: 'cascade' })
+      .notNull(),
+    knowledgeBaseId: text('knowledge_base_id')
+      .references(() => knowledgeBases.id, { onDelete: 'cascade' })
+      .notNull(),
+    userId: text('user_id')
+      .references(() => users.id, { onDelete: 'cascade' })
+      .notNull(),
+    enabled: boolean('enabled').default(true),
+    createdAt: createdAt(),
+    updatedAt: updatedAt(),
+  },
+  (t) => ({
+    pk: primaryKey({ columns: [t.agentId, t.knowledgeBaseId] }),
+  }),
+);
+
+export const agentsFiles = pgTable(
+  'agents_files',
+  {
+    fileId: text('file_id')
+      .notNull()
+      .references(() => files.id, { onDelete: 'cascade' }),
+    agentId: text('agent_id')
+      .notNull()
+      .references(() => agents.id, { onDelete: 'cascade' }),
+    enabled: boolean('enabled').default(true),
+    userId: text('user_id')
+      .references(() => users.id, { onDelete: 'cascade' })
+      .notNull(),
+    createdAt: createdAt(),
+    updatedAt: updatedAt(),
+  },
+  (t) => ({
+    pk: primaryKey({ columns: [t.fileId, t.agentId, t.userId] }),
+  }),
+);

package/src/database/server/schemas/lobechat/discover.ts

@@ -2,7 +2,7 @@
 import { integer, pgTable, primaryKey, serial, text } from 'drizzle-orm/pg-core';
 
 import { createdAt, updatedAt } from './_helpers';
-import { agents } from './chat';
+import { agents } from './agent';
 import { users } from './user';
 
 export const tags = pgTable('tags', {

package/src/database/server/schemas/lobechat/file.ts

@@ -1,10 +1,50 @@
 /* eslint-disable sort-keys-fix/sort-keys-fix  */
-import { integer, jsonb, pgTable, text, varchar } from 'drizzle-orm/pg-core';
+import {
+  boolean,
+  integer,
+  jsonb,
+  pgTable,
+  primaryKey,
+  text,
+  uuid,
+  varchar,
+} from 'drizzle-orm/pg-core';
+import { createInsertSchema } from 'drizzle-zod';
+
+import { chunks } from '@/database/server/schemas/lobechat/rag';
 
 import { idGenerator } from '../../utils/idGenerator';
 import { createdAt, updatedAt } from './_helpers';
 import { users } from './user';
 
+export const asyncTasks = pgTable('async_tasks', {
+  id: uuid('id').defaultRandom().primaryKey(),
+  type: text('type'),
+  status: text('status'),
+  error: jsonb('error'),
+  userId: text('user_id')
+    .references(() => users.id, { onDelete: 'cascade' })
+    .notNull(),
+  duration: integer('duration'),
+  createdAt: createdAt(),
+  updatedAt: updatedAt(),
+});
+
+export type NewAsyncTaskItem = typeof asyncTasks.$inferInsert;
+export type AsyncTaskSelectItem = typeof asyncTasks.$inferSelect;
+
+export const globalFiles = pgTable('global_files', {
+  hashId: varchar('hash_id', { length: 64 }).primaryKey(),
+  fileType: varchar('file_type', { length: 255 }).notNull(),
+  size: integer('size').notNull(),
+  url: text('url').notNull(),
+  metadata: jsonb('metadata'),
+  createdAt: createdAt(),
+});
+
+export type NewGlobalFile = typeof globalFiles.$inferInsert;
+export type GlobalFileItem = typeof globalFiles.$inferSelect;
+
 export const files = pgTable('files', {
   id: text('id')
     .$defaultFn(() => idGenerator('files'))
@@ -14,11 +54,18 @@ export const files = pgTable('files', {
     .references(() => users.id, { onDelete: 'cascade' })
     .notNull(),
   fileType: varchar('file_type', { length: 255 }).notNull(),
+  fileHash: varchar('file_hash', { length: 64 }).references(() => globalFiles.hashId, {
+    onDelete: 'no action',
+  }),
   name: text('name').notNull(),
   size: integer('size').notNull(),
   url: text('url').notNull(),
 
   metadata: jsonb('metadata'),
+  chunkTaskId: uuid('chunk_task_id').references(() => asyncTasks.id, { onDelete: 'set null' }),
+  embeddingTaskId: uuid('embedding_task_id').references(() => asyncTasks.id, {
+    onDelete: 'set null',
+  }),
 
   createdAt: createdAt(),
   updatedAt: updatedAt(),
@@ -26,3 +73,73 @@ export const files = pgTable('files', {
 
 export type NewFile = typeof files.$inferInsert;
 export type FileItem = typeof files.$inferSelect;
+
+export const fileChunks = pgTable(
+  'file_chunks',
+  {
+    fileId: varchar('file_id').references(() => files.id, { onDelete: 'cascade' }),
+    chunkId: uuid('chunk_id').references(() => chunks.id, { onDelete: 'cascade' }),
+    createdAt: createdAt(),
+  },
+  (t) => ({
+    pk: primaryKey({ columns: [t.fileId, t.chunkId] }),
+  }),
+);
+
+export type NewFileChunkItem = typeof fileChunks.$inferInsert;
+
+export const knowledgeBases = pgTable('knowledge_bases', {
+  id: text('id')
+    .$defaultFn(() => idGenerator('knowledgeBases'))
+    .primaryKey(),
+
+  name: text('name').notNull(),
+  description: text('description'),
+  avatar: text('avatar'),
+
+  // different types of knowledge bases need to be distinguished
+  type: text('type'),
+  userId: text('user_id')
+    .references(() => users.id, { onDelete: 'cascade' })
+    .notNull(),
+
+  isPublic: boolean('is_public').default(false),
+
+  settings: jsonb('settings'),
+
+  createdAt: createdAt(),
+  updatedAt: updatedAt(),
+});
+
+export const insertKnowledgeBasesSchema = createInsertSchema(knowledgeBases);
+
+export type NewKnowledgeBase = typeof knowledgeBases.$inferInsert;
+export type KnowledgeBaseItem = typeof knowledgeBases.$inferSelect;
+
+export const knowledgeBaseFiles = pgTable(
+  'knowledge_base_files',
+  {
+    knowledgeBaseId: text('knowledge_base_id')
+      .references(() => knowledgeBases.id, { onDelete: 'cascade' })
+      .notNull(),
+
+    fileId: text('file_id')
+      .references(() => files.id, { onDelete: 'cascade' })
+      .notNull(),
+
+    // userId: text('user_id')
+    //   .references(() => users.id, { onDelete: 'cascade' })
+    //   .notNull(),
+
+    createdAt: createdAt(),
+  },
+  (t) => ({
+    pk: primaryKey({
+      columns: [
+        t.knowledgeBaseId,
+        t.fileId,
+        // t.userId
+      ],
+    }),
+  }),
+);

package/src/database/server/schemas/lobechat/index.ts

@@ -1,6 +1,10 @@
-export * from './chat';
+export * from './agent';
 export * from './discover';
 export * from './file';
+export * from './message';
 export * from './nextauth';
+export * from './rag';
 export * from './relations';
+export * from './session';
+export * from './topic';
 export * from './user';

package/src/database/server/schemas/lobechat/message.ts

@@ -0,0 +1,169 @@
+/* eslint-disable sort-keys-fix/sort-keys-fix  */
+import {
+  boolean,
+  index,
+  jsonb,
+  numeric,
+  pgTable,
+  primaryKey,
+  text,
+  uniqueIndex,
+  uuid,
+} from 'drizzle-orm/pg-core';
+import { createSelectSchema } from 'drizzle-zod';
+
+import { idGenerator } from '../../utils/idGenerator';
+import { createdAt, updatedAt } from './_helpers';
+import { agents } from './agent';
+import { files } from './file';
+import { chunks, embeddings } from './rag';
+import { sessions } from './session';
+import { topics } from './topic';
+import { users } from './user';
+
+// @ts-ignore
+export const messages = pgTable(
+  'messages',
+  {
+    id: text('id')
+      .$defaultFn(() => idGenerator('messages'))
+      .primaryKey(),
+
+    role: text('role', { enum: ['user', 'system', 'assistant', 'tool'] }).notNull(),
+    content: text('content'),
+
+    model: text('model'),
+    provider: text('provider'),
+
+    favorite: boolean('favorite').default(false),
+    error: jsonb('error'),
+
+    tools: jsonb('tools'),
+
+    traceId: text('trace_id'),
+    observationId: text('observation_id'),
+
+    clientId: text('client_id'),
+
+    // foreign keys
+    userId: text('user_id')
+      .references(() => users.id, { onDelete: 'cascade' })
+      .notNull(),
+    sessionId: text('session_id').references(() => sessions.id, { onDelete: 'cascade' }),
+    topicId: text('topic_id').references(() => topics.id, { onDelete: 'cascade' }),
+    // @ts-ignore
+    parentId: text('parent_id').references(() => messages.id, { onDelete: 'set null' }),
+    quotaId: text('quota_id').references(() => messages.id, { onDelete: 'set null' }),
+
+    // used for group chat
+    agentId: text('agent_id').references(() => agents.id, { onDelete: 'set null' }),
+
+    createdAt: createdAt(),
+    updatedAt: updatedAt(),
+  },
+  (table) => ({
+    createdAtIdx: index('messages_created_at_idx').on(table.createdAt),
+    messageClientIdUnique: uniqueIndex('message_client_id_user_unique').on(
+      table.clientId,
+      table.userId,
+    ),
+  }),
+);
+
+export type NewMessage = typeof messages.$inferInsert;
+export type MessageItem = typeof messages.$inferSelect;
+
+// if the message container a plugin
+export const messagePlugins = pgTable('message_plugins', {
+  id: text('id')
+    .references(() => messages.id, { onDelete: 'cascade' })
+    .primaryKey(),
+
+  toolCallId: text('tool_call_id'),
+  type: text('type', {
+    enum: ['default', 'markdown', 'standalone', 'builtin'],
+  }).default('default'),
+
+  apiName: text('api_name'),
+  arguments: text('arguments'),
+  identifier: text('identifier'),
+  state: jsonb('state'),
+  error: jsonb('error'),
+});
+
+export type MessagePluginItem = typeof messagePlugins.$inferSelect;
+export const updateMessagePluginSchema = createSelectSchema(messagePlugins);
+
+export const messageTTS = pgTable('message_tts', {
+  id: text('id')
+    .references(() => messages.id, { onDelete: 'cascade' })
+    .primaryKey(),
+  contentMd5: text('content_md5'),
+  fileId: text('file_id').references(() => files.id, { onDelete: 'cascade' }),
+  voice: text('voice'),
+});
+
+export const messageTranslates = pgTable('message_translates', {
+  id: text('id')
+    .references(() => messages.id, { onDelete: 'cascade' })
+    .primaryKey(),
+  content: text('content'),
+  from: text('from'),
+  to: text('to'),
+});
+
+// if the message contains a file
+// save the file id and message id
+export const messagesFiles = pgTable(
+  'messages_files',
+  {
+    fileId: text('file_id')
+      .notNull()
+      .references(() => files.id, { onDelete: 'cascade' }),
+    messageId: text('message_id')
+      .notNull()
+      .references(() => messages.id, { onDelete: 'cascade' }),
+  },
+  (t) => ({
+    pk: primaryKey({ columns: [t.fileId, t.messageId] }),
+  }),
+);
+
+export const messageQueries = pgTable('message_queries', {
+  id: uuid('id').defaultRandom().primaryKey(),
+  messageId: text('message_id')
+    .references(() => messages.id, { onDelete: 'cascade' })
+    .notNull(),
+  rewriteQuery: text('rewrite_query'),
+  userQuery: text('user_query'),
+  embeddingsId: uuid('embeddings_id').references(() => embeddings.id, { onDelete: 'set null' }),
+});
+
+export type NewMessageQuery = typeof messageQueries.$inferInsert;
+
+export const messageQueryChunks = pgTable(
+  'message_query_chunks',
+  {
+    messageId: text('id').references(() => messages.id, { onDelete: 'cascade' }),
+    queryId: uuid('query_id').references(() => messageQueries.id, { onDelete: 'cascade' }),
+    chunkId: uuid('chunk_id').references(() => chunks.id, { onDelete: 'cascade' }),
+    similarity: numeric('similarity', { precision: 6, scale: 5 }),
+  },
+  (t) => ({
+    pk: primaryKey({ columns: [t.chunkId, t.messageId, t.queryId] }),
+  }),
+);
+export type NewMessageFileChunk = typeof messageQueryChunks.$inferInsert;
+
+// convert message content to the chunks
+// then we can use message as the RAG source
+export const messageChunks = pgTable(
+  'message_chunks',
+  {
+    messageId: text('message_id').references(() => messages.id, { onDelete: 'cascade' }),
+    chunkId: uuid('chunk_id').references(() => chunks.id, { onDelete: 'cascade' }),
+  },
+  (t) => ({
+    pk: primaryKey({ columns: [t.chunkId, t.messageId] }),
+  }),
+);

package/src/database/server/schemas/lobechat/rag.ts

@@ -0,0 +1,53 @@
+/* eslint-disable sort-keys-fix/sort-keys-fix  */
+import { integer, jsonb, pgTable, text, uuid, varchar, vector } from 'drizzle-orm/pg-core';
+
+import { createdAt, updatedAt } from './_helpers';
+import { files } from './file';
+import { users } from './user';
+
+export const chunks = pgTable('chunks', {
+  id: uuid('id').defaultRandom().primaryKey(),
+  text: text('text'),
+  abstract: text('abstract'),
+  metadata: jsonb('metadata'),
+  index: integer('index'),
+  type: varchar('type'),
+
+  createdAt: createdAt(),
+  updatedAt: updatedAt(),
+
+  userId: text('user_id').references(() => users.id, { onDelete: 'cascade' }),
+});
+
+export type NewChunkItem = typeof chunks.$inferInsert & { fileId?: string };
+
+export const unstructuredChunks = pgTable('unstructured_chunks', {
+  id: uuid('id').defaultRandom().primaryKey(),
+  text: text('text'),
+  metadata: jsonb('metadata'),
+  index: integer('index'),
+  type: varchar('type'),
+
+  createdAt: createdAt(),
+  updatedAt: updatedAt(),
+
+  parentId: varchar('parent_id'),
+  compositeId: uuid('composite_id').references(() => chunks.id, { onDelete: 'cascade' }),
+  userId: text('user_id').references(() => users.id, { onDelete: 'cascade' }),
+  fileId: varchar('file_id').references(() => files.id, { onDelete: 'cascade' }),
+});
+
+export type NewUnstructuredChunkItem = typeof unstructuredChunks.$inferInsert;
+
+export const embeddings = pgTable('embeddings', {
+  id: uuid('id').defaultRandom().primaryKey(),
+  chunkId: uuid('chunk_id')
+    .references(() => chunks.id, { onDelete: 'cascade' })
+    .unique(),
+  embeddings: vector('embeddings', { dimensions: 1024 }),
+  model: text('model'),
+  userId: text('user_id').references(() => users.id, { onDelete: 'cascade' }),
+});
+
+export type NewEmbeddingsItem = typeof embeddings.$inferInsert;
+export type EmbeddingsSelectItem = typeof embeddings.$inferSelect;

package/src/database/server/schemas/lobechat/relations.ts

@@ -1,25 +1,44 @@
 /* eslint-disable sort-keys-fix/sort-keys-fix  */
 import { relations } from 'drizzle-orm';
+import { pgTable, primaryKey, text } from 'drizzle-orm/pg-core';
 
-import {
-  agents,
-  agentsToSessions,
-  filesToAgents,
-  filesToMessages,
-  filesToSessions,
-  messages,
-  sessionGroups,
-  sessions,
-  topics,
-} from './chat';
+import { agents, agentsFiles, agentsKnowledgeBases } from './agent';
 import { agentsTags, plugins, pluginsTags, tags } from './discover';
-import { files } from './file';
-
-export const filesRelations = relations(files, ({ many }) => ({
-  filesToMessages: many(filesToMessages),
-  filesToSessions: many(filesToSessions),
-  filesToAgents: many(filesToAgents),
-}));
+import { asyncTasks, files, knowledgeBases } from './file';
+import { messages, messagesFiles } from './message';
+import { unstructuredChunks } from './rag';
+import { sessionGroups, sessions } from './session';
+import { topics } from './topic';
+
+export const agentsToSessions = pgTable(
+  'agents_to_sessions',
+  {
+    agentId: text('agent_id')
+      .notNull()
+      .references(() => agents.id, { onDelete: 'cascade' }),
+    sessionId: text('session_id')
+      .notNull()
+      .references(() => sessions.id, { onDelete: 'cascade' }),
+  },
+  (t) => ({
+    pk: primaryKey({ columns: [t.agentId, t.sessionId] }),
+  }),
+);
+
+export const filesToSessions = pgTable(
+  'files_to_sessions',
+  {
+    fileId: text('file_id')
+      .notNull()
+      .references(() => files.id, { onDelete: 'cascade' }),
+    sessionId: text('session_id')
+      .notNull()
+      .references(() => sessions.id, { onDelete: 'cascade' }),
+  },
+  (t) => ({
+    pk: primaryKey({ columns: [t.fileId, t.sessionId] }),
+  }),
+);
 
 export const topicRelations = relations(topics, ({ one }) => ({
   session: one(sessions, {
@@ -49,7 +68,7 @@ export const tagsRelations = relations(tags, ({ many }) => ({
 }));
 
 export const messagesRelations = relations(messages, ({ many, one }) => ({
-  filesToMessages: many(filesToMessages),
+  filesToMessages: many(messagesFiles),
 
   session: one(sessions, {
     fields: [messages.sessionId],
@@ -69,7 +88,8 @@ export const messagesRelations = relations(messages, ({ many, one }) => ({
 
 export const agentsRelations = relations(agents, ({ many }) => ({
   agentsToSessions: many(agentsToSessions),
-  filesToAgents: many(filesToAgents),
+  knowledgeBases: many(agentsKnowledgeBases),
+  files: many(agentsFiles),
   agentsTags: many(agentsTags),
 }));
 
@@ -84,37 +104,15 @@ export const agentsToSessionsRelations = relations(agentsToSessions, ({ one }) =
   }),
 }));
 
-export const filesToAgentsRelations = relations(filesToAgents, ({ one }) => ({
+export const agentsKnowledgeBasesRelations = relations(agentsKnowledgeBases, ({ one }) => ({
+  knowledgeBase: one(knowledgeBases, {
+    fields: [agentsKnowledgeBases.knowledgeBaseId],
+    references: [knowledgeBases.id],
+  }),
   agent: one(agents, {
-    fields: [filesToAgents.agentId],
+    fields: [agentsKnowledgeBases.agentId],
     references: [agents.id],
   }),
-  file: one(files, {
-    fields: [filesToAgents.fileId],
-    references: [files.id],
-  }),
-}));
-
-export const filesToMessagesRelations = relations(filesToMessages, ({ one }) => ({
-  file: one(files, {
-    fields: [filesToMessages.fileId],
-    references: [files.id],
-  }),
-  message: one(messages, {
-    fields: [filesToMessages.messageId],
-    references: [messages.id],
-  }),
-}));
-
-export const filesToSessionsRelations = relations(filesToSessions, ({ one }) => ({
-  file: one(files, {
-    fields: [filesToSessions.fileId],
-    references: [files.id],
-  }),
-  session: one(sessions, {
-    fields: [filesToSessions.sessionId],
-    references: [sessions.id],
-  }),
 }));
 
 export const agentsTagsRelations = relations(agentsTags, ({ one }) => ({
@@ -136,3 +134,25 @@ export const sessionsRelations = relations(sessions, ({ many, one }) => ({
     references: [sessionGroups.id],
   }),
 }));
+
+export const chunksRelations = relations(unstructuredChunks, ({ one }) => ({
+  file: one(files, {
+    fields: [unstructuredChunks.fileId],
+    references: [files.id],
+  }),
+}));
+
+export const filesRelations = relations(files, ({ many, one }) => ({
+  messages: many(messagesFiles),
+  sessions: many(filesToSessions),
+  agents: many(agentsFiles),
+
+  chunkingTask: one(asyncTasks, {
+    fields: [files.chunkTaskId],
+    references: [asyncTasks.id],
+  }),
+  embeddingTask: one(asyncTasks, {
+    fields: [files.embeddingTaskId],
+    references: [asyncTasks.id],
+  }),
+}));