@lobehub/chat 1.98.2 → 1.99.1

package/src/database/repositories/aiInfra/index.test.ts

@@ -242,6 +242,62 @@ describe('AiInfraRepos', () => {
         settings: mockSettings,
       });
     });
+
+    it('should return all models including disabled ones when filterEnabled is false', async () => {
+      const mockProviders = [
+        { enabled: true, id: 'openai', name: 'OpenAI', source: 'builtin' },
+        { enabled: false, id: 'anthropic', name: 'Anthropic', source: 'builtin' },
+      ] as AiProviderListItem[];
+
+      const mockAllModels = [
+        {
+          abilities: {},
+          enabled: false,
+          id: 'claude-3',
+          providerId: 'anthropic',
+          type: 'chat' as const,
+        },
+      ] as EnabledAiModel[];
+
+      vi.spyOn(repo, 'getAiProviderList').mockResolvedValue(mockProviders);
+      vi.spyOn(repo.aiModelModel, 'getAllModels').mockResolvedValue(mockAllModels);
+      vi.spyOn(repo as any, 'fetchBuiltinModels').mockResolvedValue([
+        {
+          enabled: true,
+          id: 'gpt-4',
+          type: 'chat' as const,
+        },
+        {
+          enabled: false,
+          id: 'claude-3',
+          type: 'chat' as const,
+        },
+      ]);
+
+      const result = await repo.getEnabledModels(false);
+
+      // Should include both enabled and disabled models
+      expect(result).toEqual(
+        expect.arrayContaining([
+          expect.objectContaining({
+            id: 'gpt-4',
+            enabled: true,
+            providerId: 'openai',
+          }),
+          expect.objectContaining({
+            id: 'claude-3',
+            enabled: false,
+            providerId: 'anthropic',
+          }),
+        ]),
+      );
+      // Verify we have at least the expected models (may have more from builtin models)
+      expect(result.length).toBeGreaterThanOrEqual(2);
+
+      // Verify disabled models are included
+      const disabledModels = result.filter((model) => !model.enabled);
+      expect(disabledModels.length).toBeGreaterThan(0);
+    });
   });
 
   describe('getAiProviderModelList', () => {
@@ -332,7 +388,9 @@ describe('AiInfraRepos', () => {
         openai: { apiKey: 'test-key' },
       } as unknown as Record<string, AiProviderRuntimeConfig>;
       const mockEnabledProviders = [{ id: 'openai', name: 'OpenAI' }] as EnabledProvider[];
-      const mockEnabledModels = [{ id: 'gpt-4', providerId: 'openai' }] as EnabledAiModel[];
+      const mockEnabledModels = [
+        { id: 'gpt-4', providerId: 'openai', enabled: true },
+      ] as EnabledAiModel[];
 
       vi.spyOn(repo.aiProviderModel, 'getAiProviderRuntimeConfig').mockResolvedValue(
         mockRuntimeConfig,
@@ -385,6 +443,10 @@ describe('AiInfraRepos', () => {
           }),
         ],
         enabledAiProviders: [{ id: 'openai', logo: 'logo1', name: 'OpenAI', source: 'builtin' }],
+        enabledChatAiProviders: [
+          { id: 'openai', logo: 'logo1', name: 'OpenAI', source: 'builtin' },
+        ],
+        enabledImageAiProviders: [],
         runtimeConfig: {
           openai: {
             apiKey: 'test-key',
@@ -393,6 +455,100 @@ describe('AiInfraRepos', () => {
         },
       });
     });
+
+    it('should return provider runtime state with enabledImageAiProviders', async () => {
+      const mockRuntimeConfig = {
+        fal: {
+          apiKey: 'test-fal-key',
+        },
+        openai: {
+          apiKey: 'test-openai-key',
+        },
+      } as unknown as Record<string, AiProviderRuntimeConfig>;
+
+      vi.spyOn(repo.aiProviderModel, 'getAiProviderRuntimeConfig').mockResolvedValue(
+        mockRuntimeConfig,
+      );
+
+      // Mock providers including fal for image generation
+      vi.spyOn(repo, 'getUserEnabledProviderList').mockResolvedValue([
+        { id: 'openai', logo: 'openai-logo', name: 'OpenAI', source: 'builtin' },
+        { id: 'fal', logo: 'fal-logo', name: 'Fal', source: 'builtin' },
+      ]);
+
+      // Mock models including image models from fal
+      vi.spyOn(repo, 'getEnabledModels').mockResolvedValue([
+        {
+          abilities: {},
+          enabled: true,
+          id: 'gpt-4',
+          providerId: 'openai',
+          type: 'chat',
+        },
+        {
+          abilities: {},
+          enabled: true,
+          id: 'flux/schnell',
+          providerId: 'fal',
+          type: 'image',
+        },
+        {
+          abilities: {},
+          enabled: true,
+          id: 'flux-kontext/dev',
+          providerId: 'fal',
+          type: 'image',
+        },
+      ]);
+
+      const result = await repo.getAiProviderRuntimeState();
+
+      expect(result).toEqual({
+        enabledAiModels: [
+          expect.objectContaining({
+            enabled: true,
+            id: 'gpt-4',
+            providerId: 'openai',
+            type: 'chat',
+          }),
+          expect.objectContaining({
+            enabled: true,
+            id: 'flux/schnell',
+            providerId: 'fal',
+            type: 'image',
+          }),
+          expect.objectContaining({
+            enabled: true,
+            id: 'flux-kontext/dev',
+            providerId: 'fal',
+            type: 'image',
+          }),
+        ],
+        enabledAiProviders: [
+          { id: 'openai', logo: 'openai-logo', name: 'OpenAI', source: 'builtin' },
+          { id: 'fal', logo: 'fal-logo', name: 'Fal', source: 'builtin' },
+        ],
+        enabledChatAiProviders: [
+          { id: 'openai', logo: 'openai-logo', name: 'OpenAI', source: 'builtin' },
+        ],
+        enabledImageAiProviders: [
+          expect.objectContaining({
+            id: 'fal',
+            name: 'Fal',
+          }),
+        ],
+        runtimeConfig: {
+          fal: {
+            apiKey: 'test-fal-key',
+            enabled: undefined,
+          },
+          openai: {
+            apiKey: 'test-openai-key',
+            enabled: true,
+          },
+        },
+      });
+    });
   });
 
   describe('getAiProviderDetail', () => {

package/src/database/repositories/aiInfra/index.ts

@@ -91,14 +91,14 @@ export class AiInfraRepos {
   /**
    * used in the chat page. to show the enabled models
    */
-  getEnabledModels = async () => {
-    const providers = await this.getAiProviderList();
-    const enabledProviders = providers.filter((item) => item.enabled);
-
-    const allModels = await this.aiModelModel.getAllModels();
-    const userEnabledModels = allModels.filter((item) => item.enabled);
-
-    const modelList = await pMap(
+  getEnabledModels = async (filterEnabled: boolean = true) => {
+    const [providers, allModels] = await Promise.all([
+      this.getAiProviderList(),
+      this.aiModelModel.getAllModels(),
+    ]);
+    const enabledProviders = providers.filter((item) => (filterEnabled ? item.enabled : true));
+
+    const builtinModelList = await pMap(
       enabledProviders,
       async (provider) => {
         const aiModels = await this.fetchBuiltinModels(provider.id);
@@ -114,6 +114,7 @@ export class AiInfraRepos {
               };
 
             return {
+              ...item,
               abilities: !isEmpty(user.abilities) ? user.abilities : item.abilities || {},
               config: !isEmpty(user.config) ? user.config : item.config,
               contextWindowTokens:
@@ -129,32 +130,49 @@ export class AiInfraRepos {
               type: item.type,
             };
           })
-          .filter((i) => i.enabled);
+          .filter((item) => (filterEnabled ? item.enabled : true));
       },
       { concurrency: 10 },
     );
 
-    return [...modelList.flat(), ...userEnabledModels].sort(
-      (a, b) => (a?.sort || -1) - (b?.sort || -1),
-    ) as EnabledAiModel[];
+    const enabledProviderIds = new Set(enabledProviders.map((item) => item.id));
+
+    return [
+      ...builtinModelList.flat(),
+      ...allModels.filter((item) =>
+        filterEnabled ? enabledProviderIds.has(item.providerId) && item.enabled : true,
+      ),
+    ].sort((a, b) => (a?.sort || -1) - (b?.sort || -1)) as EnabledAiModel[];
   };
 
   getAiProviderRuntimeState = async (
     decryptor?: DecryptUserKeyVaults,
   ): Promise<AiProviderRuntimeState> => {
-    const result = await this.aiProviderModel.getAiProviderRuntimeConfig(decryptor);
+    const [result, enabledAiProviders, allModels] = await Promise.all([
+      this.aiProviderModel.getAiProviderRuntimeConfig(decryptor),
+      this.getUserEnabledProviderList(),
+      this.getEnabledModels(false),
+    ]);
 
     const runtimeConfig = result;
-
     Object.entries(result).forEach(([key, value]) => {
       runtimeConfig[key] = merge(this.providerConfigs[key] || {}, value);
     });
+    const enabledAiModels = allModels.filter((model) => model.enabled);
+    const enabledChatAiProviders = enabledAiProviders.filter((provider) => {
+      return allModels.some((model) => model.providerId === provider.id && model.type === 'chat');
+    });
+    const enabledImageAiProviders = enabledAiProviders.filter((provider) => {
+      return allModels.some((model) => model.providerId === provider.id && model.type === 'image');
+    });
 
-    const enabledAiProviders = await this.getUserEnabledProviderList();
-
-    const enabledAiModels = await this.getEnabledModels();
-
-    return { enabledAiModels, enabledAiProviders, runtimeConfig };
+    return {
+      enabledAiModels,
+      enabledAiProviders,
+      enabledChatAiProviders,
+      enabledImageAiProviders,
+      runtimeConfig,
+    };
   };
 
   getAiProviderModelList = async (providerId: string) => {

package/src/database/repositories/tableViewer/index.test.ts

@@ -23,7 +23,7 @@ describe('TableViewerRepo', () => {
     it('should return all tables with counts', async () => {
       const result = await repo.getAllTables();
 
-      expect(result.length).toEqual(55);
+      expect(result.length).toEqual(58);
       expect(result[0]).toEqual({ name: 'agents', count: 0, type: 'BASE TABLE' });
     });
 

package/src/database/schemas/file.ts

@@ -14,6 +14,7 @@ import {
 import { createInsertSchema } from 'drizzle-zod';
 
 import { idGenerator } from '@/database/utils/idGenerator';
+import { FileSource } from '@/types/files';
 
 import { accessedAt, createdAt, timestamps } from './_helpers';
 import { asyncTasks } from './asyncTask';
@@ -45,13 +46,20 @@ export const files = pgTable(
     userId: text('user_id')
       .references(() => users.id, { onDelete: 'cascade' })
       .notNull(),
+    /**
+     * mime
+     */
     fileType: varchar('file_type', { length: 255 }).notNull(),
+    /**
+     * sha256
+     */
     fileHash: varchar('file_hash', { length: 64 }).references(() => globalFiles.hashId, {
       onDelete: 'no action',
     }),
     name: text('name').notNull(),
     size: integer('size').notNull(),
     url: text('url').notNull(),
+    source: text('source').$type<FileSource>(),
 
     clientId: text('client_id'),
     metadata: jsonb('metadata'),

package/src/database/schemas/generation.ts

@@ -0,0 +1,127 @@
+/* eslint-disable sort-keys-fix/sort-keys-fix  */
+import { integer, jsonb, pgTable, text, uuid, varchar } from 'drizzle-orm/pg-core';
+import { createInsertSchema } from 'drizzle-zod';
+
+import { idGenerator } from '@/database/utils/idGenerator';
+import { ImageGenerationAsset } from '@/types/generation';
+
+import { timestamps } from './_helpers';
+import { AsyncTaskSelectItem, asyncTasks } from './asyncTask';
+import { files } from './file';
+import { users } from './user';
+
+/**
+ * 生成主题表 - 用于组织和管理 AI 生成内容的主题
+ */
+export const generationTopics = pgTable('generation_topics', {
+  id: text('id')
+    .$defaultFn(() => idGenerator('generationTopics'))
+    .notNull()
+    .primaryKey(),
+
+  userId: text('user_id')
+    .references(() => users.id, { onDelete: 'cascade' })
+    .notNull(),
+
+  /** 简要描述主题内容，由 LLM 生成 */
+  title: text('title'),
+
+  /** 主题封面图片 URL */
+  coverUrl: text('cover_url'),
+
+  ...timestamps,
+});
+
+export const insertGenerationTopicSchema = createInsertSchema(generationTopics);
+
+export type NewGenerationTopic = typeof generationTopics.$inferInsert;
+export type GenerationTopicItem = typeof generationTopics.$inferSelect;
+
+/**
+ * 生成批次表 - 存储一次生成请求的配置信息
+ */
+export const generationBatches = pgTable('generation_batches', {
+  id: text('id')
+    .$defaultFn(() => idGenerator('generationBatches'))
+    .notNull()
+    .primaryKey(),
+
+  userId: text('user_id')
+    .references(() => users.id, { onDelete: 'cascade' })
+    .notNull(),
+
+  generationTopicId: text('generation_topic_id')
+    .notNull()
+    .references(() => generationTopics.id, { onDelete: 'cascade' }),
+
+  /** 服务商名称 */
+  provider: text('provider').notNull(),
+
+  /** 模型名称 */
+  model: text('model').notNull(),
+
+  /** 生成提示词 */
+  prompt: text('prompt').notNull(),
+
+  /** 图片宽度 */
+  width: integer('width'),
+
+  /** 图片高度 */
+  height: integer('height'),
+
+  /** 图片比例 */
+  ratio: varchar('ratio', { length: 64 }),
+
+  /** 存储生成批次的配置，存放不需要建立索引的公共配置 */
+  config: jsonb('config'),
+
+  ...timestamps,
+});
+
+export const insertGenerationBatchSchema = createInsertSchema(generationBatches);
+
+export type NewGenerationBatch = typeof generationBatches.$inferInsert;
+export type GenerationBatchItem = typeof generationBatches.$inferSelect;
+export type GenerationBatchWithGenerations = GenerationBatchItem & {
+  generations: GenerationWithAsyncTask[];
+};
+
+/**
+ *  存储单个 AI 生成信息
+ */
+export const generations = pgTable('generations', {
+  id: text('id')
+    .$defaultFn(() => idGenerator('generations'))
+    .notNull()
+    .primaryKey(),
+
+  userId: text('user_id')
+    .references(() => users.id, { onDelete: 'cascade' })
+    .notNull(),
+
+  generationBatchId: varchar('generation_batch_id', { length: 64 })
+    .notNull()
+    .references(() => generationBatches.id, { onDelete: 'cascade' }),
+
+  /** 关联的异步任务 ID */
+  asyncTaskId: uuid('async_task_id').references(() => asyncTasks.id, {
+    onDelete: 'set null',
+  }),
+
+  /** 关联的生成文件 ID，删除文件时连带删除生成记录 */
+  fileId: text('file_id').references(() => files.id, { onDelete: 'cascade' }),
+
+  /** 生成种子值 */
+  seed: integer('seed'),
+
+  /** 生成的资源信息，包含存储在 s3 上的 key, 图片实际宽高，缩略图 key 等 */
+  asset: jsonb('asset').$type<ImageGenerationAsset>(),
+
+  ...timestamps,
+});
+
+export const insertGenerationSchema = createInsertSchema(generations);
+
+export type NewGeneration = typeof generations.$inferInsert;
+export type GenerationItem = typeof generations.$inferSelect;
+export type GenerationWithAsyncTask = GenerationItem & { asyncTask?: AsyncTaskSelectItem };

package/src/database/schemas/index.ts

@@ -3,6 +3,7 @@ export * from './aiInfra';
 export * from './asyncTask';
 export * from './document';
 export * from './file';
+export * from './generation';
 export * from './message';
 export * from './nextauth';
 export * from './oidc';

package/src/database/schemas/relations.ts

@@ -8,6 +8,7 @@ import { agents, agentsFiles, agentsKnowledgeBases } from './agent';
 import { asyncTasks } from './asyncTask';
 import { documentChunks, documents } from './document';
 import { files, knowledgeBases } from './file';
+import { generationBatches, generationTopics, generations } from './generation';
 import { messages, messagesFiles } from './message';
 import { chunks, unstructuredChunks } from './rag';
 import { sessionGroups, sessions } from './session';
@@ -198,7 +199,10 @@ export const filesRelations = relations(files, ({ many, one }) => ({
   sessions: many(filesToSessions),
   agents: many(agentsFiles),
   documents: many(documents, { relationName: 'fileDocuments' }),
-
+  generation: one(generations, {
+    fields: [files.id],
+    references: [generations.fileId],
+  }),
   chunkingTask: one(asyncTasks, {
     fields: [files.chunkTaskId],
     references: [asyncTasks.id],
@@ -237,3 +241,43 @@ export const documentChunksRelations = relations(documentChunks, ({ one }) => ({
     references: [documents.id],
   }),
 }));
+
+// Generation 相关关系定义
+export const generationTopicsRelations = relations(generationTopics, ({ one, many }) => ({
+  user: one(users, {
+    fields: [generationTopics.userId],
+    references: [users.id],
+  }),
+  batches: many(generationBatches),
+}));
+
+export const generationBatchesRelations = relations(generationBatches, ({ one, many }) => ({
+  user: one(users, {
+    fields: [generationBatches.userId],
+    references: [users.id],
+  }),
+  topic: one(generationTopics, {
+    fields: [generationBatches.generationTopicId],
+    references: [generationTopics.id],
+  }),
+  generations: many(generations),
+}));
+
+export const generationsRelations = relations(generations, ({ one }) => ({
+  user: one(users, {
+    fields: [generations.userId],
+    references: [users.id],
+  }),
+  batch: one(generationBatches, {
+    fields: [generations.generationBatchId],
+    references: [generationBatches.id],
+  }),
+  asyncTask: one(asyncTasks, {
+    fields: [generations.asyncTaskId],
+    references: [asyncTasks.id],
+  }),
+  file: one(files, {
+    fields: [generations.fileId],
+    references: [files.id],
+  }),
+}));

package/src/database/type.ts

@@ -5,3 +5,5 @@ import * as schema from './schemas';
 export type LobeChatDatabaseSchema = typeof schema;
 
 export type LobeChatDatabase = NeonDatabase<LobeChatDatabaseSchema>;
+
+export type Transaction = Parameters<Parameters<LobeChatDatabase['transaction']>[0]>[0];

package/src/database/utils/idGenerator.ts

@@ -6,6 +6,9 @@ const prefixes = {
   agents: 'agt',
   documents: 'docs',
   files: 'file',
+  generationBatches: 'gb',
+  generationTopics: 'gt',
+  generations: 'gen',
   knowledgeBases: 'kb',
   messages: 'msg',
   plugins: 'plg',

package/src/features/Conversation/Error/ChatInvalidApiKey.tsx

@@ -0,0 +1,39 @@
+import { memo } from 'react';
+import { useTranslation } from 'react-i18next';
+
+import InvalidAPIKey from '@/components/InvalidAPIKey';
+import { useProviderName } from '@/hooks/useProviderName';
+import { useChatStore } from '@/store/chat';
+import { GlobalLLMProviderKey } from '@/types/user/settings/modelProvider';
+
+interface ChatInvalidAPIKeyProps {
+  id: string;
+  provider?: string;
+}
+const ChatInvalidAPIKey = memo<ChatInvalidAPIKeyProps>(({ id, provider }) => {
+  const { t } = useTranslation('modelProvider');
+  const { t: modelProviderErrorT } = useTranslation(['modelProvider', 'error']);
+  const [resend, deleteMessage] = useChatStore((s) => [s.regenerateMessage, s.deleteMessage]);
+  const providerName = useProviderName(provider as GlobalLLMProviderKey);
+
+  return (
+    <InvalidAPIKey
+      bedrockDescription={t('bedrock.unlock.description')}
+      description={modelProviderErrorT(`unlock.apiKey.description`, {
+        name: providerName,
+        ns: 'error',
+      })}
+      id={id}
+      onClose={() => {
+        deleteMessage(id);
+      }}
+      onRecreate={() => {
+        resend(id);
+        deleteMessage(id);
+      }}
+      provider={provider}
+    />
+  );
+});
+
+export default ChatInvalidAPIKey;

package/src/features/Conversation/Error/InvalidAccessCode.tsx

@@ -8,8 +8,8 @@ import { Flexbox } from 'react-layout-kit';
 import { useServerConfigStore } from '@/store/serverConfig';
 import { featureFlagsSelectors, serverConfigSelectors } from '@/store/serverConfig/selectors';
 
-import APIKeyForm from './APIKeyForm';
 import AccessCodeForm from './AccessCodeForm';
+import ChatInvalidAPIKey from './ChatInvalidApiKey';
 import OAuthForm from './OAuthForm';
 import { ErrorActionContainer } from './style';
 
@@ -69,7 +69,7 @@ const InvalidAccessCode = memo<InvalidAccessCodeProps>(({ id, provider }) => {
 
       <Flexbox gap={24}>
         {mode === Tab.Password && <AccessCodeForm id={id} />}
-        {showOpenAIApiKey && mode === Tab.Api && <APIKeyForm id={id} provider={provider} />}
+        {showOpenAIApiKey && mode === Tab.Api && <ChatInvalidAPIKey id={id} provider={provider} />}
         {isEnabledOAuth && mode === Tab.Oauth && <OAuthForm id={id} />}
       </Flexbox>
     </ErrorActionContainer>

package/src/features/Conversation/Error/index.tsx

@@ -10,9 +10,9 @@ import { AgentRuntimeErrorType, ILobeAgentRuntimeErrorType } from '@/libs/model-
 import { ChatErrorType, ErrorType } from '@/types/fetch';
 import { ChatMessage, ChatMessageError } from '@/types/message';
 
+import ChatInvalidAPIKey from './ChatInvalidApiKey';
 import ClerkLogin from './ClerkLogin';
 import ErrorJsonViewer from './ErrorJsonViewer';
-import InvalidAPIKey from './InvalidAPIKey';
 import InvalidAccessCode from './InvalidAccessCode';
 import { ErrorActionContainer } from './style';
 
@@ -113,13 +113,13 @@ const ErrorMessageExtra = memo<{ data: ChatMessage }>(({ data }) => {
 
     case AgentRuntimeErrorType.NoOpenAIAPIKey: {
       {
-        return <InvalidAPIKey id={data.id} provider={data.error?.body?.provider} />;
+        return <ChatInvalidAPIKey id={data.id} provider={data.error?.body?.provider} />;
       }
     }
   }
 
   if (error.type.toString().includes('Invalid')) {
-    return <InvalidAPIKey id={data.id} provider={data.error?.body?.provider} />;
+    return <ChatInvalidAPIKey id={data.id} provider={data.error?.body?.provider} />;
   }
 
   return <ErrorJsonViewer error={data.error} id={data.id} />;