From 4f7bc5acd2d38f8e6be3ce44c2d80c8f915546a5 Mon Sep 17 00:00:00 2001 From: YuTengjing Date: Tue, 11 Nov 2025 19:39:36 +0800 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix:=20add=20SSRF=20protection?= =?UTF-8?q?=20=20(#10152)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .env.example | 11 +++ .../environment-variables/basic.mdx | 52 ++++++++++- .../environment-variables/basic.zh-CN.mdx | 53 ++++++++++- package.json | 1 + packages/fetch-sse/package.json | 29 ++++++ .../src}/__tests__/fetchSSE.test.ts | 8 +- .../src}/__tests__/parseError.test.ts | 11 ++- .../src/fetch => fetch-sse/src}/fetchSSE.ts | 4 +- .../src/fetch => fetch-sse/src}/headers.ts | 0 .../src/fetch => fetch-sse/src}/index.ts | 0 .../src/fetch => fetch-sse/src}/parseError.ts | 6 +- .../src/fetch => fetch-sse/src}/request.ts | 0 .../core/contextBuilders/anthropic.test.ts | 28 +++--- .../src/core/contextBuilders/anthropic.ts | 2 +- .../src/core/contextBuilders/google.test.ts | 2 +- .../src/core/contextBuilders/google.ts | 9 +- .../src/core/contextBuilders/openai.test.ts | 6 +- .../src/core/contextBuilders/openai.ts | 2 +- .../createImage.test.ts | 2 +- .../openaiCompatibleFactory/createImage.ts | 2 +- .../openaiCompatibleFactory/index.test.ts | 9 +- .../streams/openai/responsesStream.test.ts | 2 +- .../src/helpers/mergeChatMethodOptions.ts | 3 +- .../src/providers/aihubmix/index.test.ts | 2 +- .../anthropic/generateObject.test.ts | 2 +- .../src/providers/anthropic/index.test.ts | 2 +- .../src/providers/baichuan/index.test.ts | 2 +- .../src/providers/bedrock/index.test.ts | 2 +- .../src/providers/bfl/createImage.test.ts | 8 +- .../src/providers/bfl/createImage.ts | 2 +- .../src/providers/cloudflare/index.test.ts | 2 +- .../src/providers/cohere/index.test.ts | 2 +- .../src/providers/google/createImage.test.ts | 4 +- .../src/providers/google/createImage.ts | 2 +- .../providers/google/generateObject.test.ts | 2 +- .../src/providers/google/index.test.ts | 5 +- .../src/providers/groq/index.test.ts | 2 +- .../src/providers/hunyuan/index.test.ts | 2 +- .../src/providers/minimax/createImage.test.ts | 2 +- .../src/providers/mistral/index.test.ts | 2 +- .../src/providers/moonshot/index.test.ts | 2 +- .../src/providers/novita/index.test.ts | 2 +- .../src/providers/ollama/index.test.ts | 75 ++++++++------- .../src/providers/ollama/index.ts | 38 ++++++-- .../src/providers/openrouter/index.test.ts | 2 +- .../src/providers/perplexity/index.test.ts | 2 +- .../src/providers/ppio/index.test.ts | 2 +- .../src/providers/qwen/createImage.test.ts | 2 +- .../src/providers/search1api/index.test.ts | 2 +- .../src/providers/siliconcloud/createImage.ts | 2 +- .../src/providers/taichu/index.test.ts | 2 +- .../src/providers/wenxin/index.test.ts | 2 +- .../src/providers/zhipu/index.test.ts | 2 +- .../src/utils/errorResponse.test.ts | 2 +- .../src/utils/imageToBase64.test.ts | 91 ------------------- .../model-runtime/src/utils/imageToBase64.ts | 62 ------------- packages/ssrf-safe-fetch/index.browser.ts | 14 +++ packages/ssrf-safe-fetch/package.json | 9 +- packages/utils/src/imageToBase64.ts | 27 ++++-- packages/utils/src/index.ts | 2 +- .../AgentTTS/SelectWithTTSPreview.tsx | 2 +- src/features/AgentSetting/store/action.ts | 2 +- .../ChatInput/ActionBar/STT/browser.tsx | 2 +- .../ChatInput/ActionBar/STT/openai.tsx | 2 +- .../components/Extras/TTS/InitPlayer.tsx | 2 +- .../genServerAiProviderConfig.test.ts | 10 +- .../globalConfig/genServerAiProviderConfig.ts | 2 +- src/services/chat/chat.test.ts | 10 +- src/services/chat/clientModelRuntime.test.ts | 2 +- src/services/chat/index.ts | 12 +-- src/services/chat/types.ts | 3 +- src/services/models.ts | 3 +- .../utils}/electron/desktopRemoteRPCFetch.ts | 2 +- .../__snapshots__/parseModels.test.ts.snap | 0 .../utils/server}/parseModels.test.ts | 0 .../src => src/utils/server}/parseModels.ts | 3 +- vitest.config.mts | 2 + 77 files changed, 367 insertions(+), 321 deletions(-) create mode 100644 packages/fetch-sse/package.json rename packages/{utils/src/fetch => fetch-sse/src}/__tests__/fetchSSE.test.ts (98%) rename packages/{utils/src/fetch => fetch-sse/src}/__tests__/parseError.test.ts (93%) rename packages/{utils/src/fetch => fetch-sse/src}/fetchSSE.ts (99%) rename packages/{utils/src/fetch => fetch-sse/src}/headers.ts (100%) rename packages/{utils/src/fetch => fetch-sse/src}/index.ts (100%) rename packages/{utils/src/fetch => fetch-sse/src}/parseError.ts (69%) rename packages/{utils/src/fetch => fetch-sse/src}/request.ts (100%) delete mode 100644 packages/model-runtime/src/utils/imageToBase64.test.ts delete mode 100644 packages/model-runtime/src/utils/imageToBase64.ts create mode 100644 packages/ssrf-safe-fetch/index.browser.ts rename {packages/utils/src => src/utils}/electron/desktopRemoteRPCFetch.ts (97%) rename {packages/utils/src => src/utils/server}/__snapshots__/parseModels.test.ts.snap (100%) rename {packages/utils/src => src/utils/server}/parseModels.test.ts (100%) rename {packages/utils/src => src/utils/server}/parseModels.ts (99%) diff --git a/.env.example b/.env.example index 3500cddf13..7675bc50c9 100644 --- a/.env.example +++ b/.env.example @@ -13,6 +13,17 @@ # Default is '0' (enabled) # ENABLED_CSP=1 +# SSRF Protection Settings +# Set to '1' to allow connections to private IP addresses (disable SSRF protection) +# WARNING: Only enable this in trusted environments +# Default is '0' (SSRF protection enabled) +# SSRF_ALLOW_PRIVATE_IP_ADDRESS=0 + +# Whitelist of allowed private IP addresses (comma-separated) +# Only takes effect when SSRF_ALLOW_PRIVATE_IP_ADDRESS is '0' +# Example: Allow specific internal servers while keeping SSRF protection +# SSRF_ALLOW_IP_ADDRESS_LIST=192.168.1.100,10.0.0.50 + ######################################## ########## AI Provider Service ######### ######################################## diff --git a/docs/self-hosting/environment-variables/basic.mdx b/docs/self-hosting/environment-variables/basic.mdx index e837551f3f..ef69cdfa12 100644 --- a/docs/self-hosting/environment-variables/basic.mdx +++ b/docs/self-hosting/environment-variables/basic.mdx @@ -127,16 +127,62 @@ For specific content, please refer to the [Feature Flags](/docs/self-hosting/adv ### `SSRF_ALLOW_PRIVATE_IP_ADDRESS` - Type: Optional -- Description: Allow to connect private IP address. In a trusted environment, it can be set to true to turn off SSRF protection. +- Description: Controls whether to allow connections to private IP addresses. Set to `1` to disable SSRF protection and allow all private IP addresses. In a trusted environment (e.g., internal network), this can be enabled to allow access to internal resources. - Default: `0` - Example: `1` or `0` + + **Security Notice**: Enabling this option will disable SSRF protection and allow connections to private + IP addresses (127.0.0.0/8, 10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16, etc.). Only enable this in + trusted environments where you need to access internal network resources. + + +**Use Cases**: + +LobeChat performs SSRF security checks in the following scenarios: + +1. **Image/Video URL to Base64 Conversion**: When processing media messages (e.g., vision models, multimodal models), LobeChat converts image and video URLs to base64 format. This check prevents malicious users from accessing internal network resources. + + Examples: + + - Image: A user sends an image message with URL `http://192.168.1.100/admin/secrets.png` + - Video: A user sends a video message with URL `http://10.0.0.50/internal/meeting.mp4` + + Without SSRF protection, these requests could expose internal network resources. + +2. **Web Crawler**: When using web crawling features to fetch external content. + +3. **Proxy Requests**: When proxying external API requests. + +**Configuration Examples**: + +```bash +# Scenario 1: Public deployment (recommended) +# Block all private IP addresses for security +SSRF_ALLOW_PRIVATE_IP_ADDRESS=0 + +# Scenario 2: Internal deployment +# Allow all private IP addresses to access internal image servers +SSRF_ALLOW_PRIVATE_IP_ADDRESS=1 + +# Scenario 3: Hybrid deployment (most common) +# Block private IPs by default, but allow specific trusted internal servers +SSRF_ALLOW_PRIVATE_IP_ADDRESS=0 +SSRF_ALLOW_IP_ADDRESS_LIST=192.168.1.100,10.0.0.50 +``` + ### `SSRF_ALLOW_IP_ADDRESS_LIST` - Type: Optional -- Description: Allow private IP address list, multiple IP addresses are separated by commas. Only when `SSRF_ALLOW_PRIVATE_IP_ADDRESS` is `0`, it takes effect. +- Description: Whitelist of allowed IP addresses, separated by commas. Only takes effect when `SSRF_ALLOW_PRIVATE_IP_ADDRESS` is `0`. Use this to allow specific internal IP addresses while keeping SSRF protection enabled for other private IPs. - Default: - -- Example: `198.18.1.62,224.0.0.3` +- Example: `192.168.1.100,10.0.0.50,172.16.0.10` + +**Common Use Cases**: + +- Allow access to internal image storage server: `192.168.1.100` +- Allow access to internal API gateway: `10.0.0.50` +- Allow access to internal documentation server: `172.16.0.10` ### `ENABLE_AUTH_PROTECTION` diff --git a/docs/self-hosting/environment-variables/basic.zh-CN.mdx b/docs/self-hosting/environment-variables/basic.zh-CN.mdx index 8e92183c1e..c181811160 100644 --- a/docs/self-hosting/environment-variables/basic.zh-CN.mdx +++ b/docs/self-hosting/environment-variables/basic.zh-CN.mdx @@ -123,16 +123,61 @@ LobeChat 在部署时提供了一些额外的配置项，你可以使用环境 ### `SSRF_ALLOW_PRIVATE_IP_ADDRESS` - 类型：可选 -- 描述：是否允许连接私有 IP 地址。在可信环境中可以设置为 true 来关闭 SSRF 防护。 +- 描述：控制是否允许连接私有 IP 地址。设置为 `1` 时将关闭 SSRF 防护并允许所有私有 IP 地址。在可信环境（如内网部署）中，可以启用此选项以访问内部资源。 - 默认值：`0` -- 示例：`1` or `0` +- 示例：`1` 或 `0` + + + **安全提示**：启用此选项将关闭 SSRF 防护，允许连接私有 IP 地址段（127.0.0.0/8、10.0.0.0/8、172.16.0.0/12、192.168.0.0/16 + 等）。仅在需要访问内网资源的可信环境中启用。 + + +**应用场景**： + +LobeChat 会在以下场景执行 SSRF 安全检查： + +1. **图片 / 视频 URL 转 Base64**：在处理媒体消息时（例如视觉模型、多模态模型），LobeChat 会将图片和视频 URL 转换为 base64 格式。此检查可防止恶意用户通过媒体 URL 访问内网资源。 + + 举例： + + - 图片：用户发送图片消息，URL 为 `http://192.168.1.100/admin/secrets.png` + - 视频：用户发送视频消息，URL 为 `http://10.0.0.50/internal/meeting.mp4` + + 若无 SSRF 防护，这些请求可能导致内网资源泄露。 + +2. **网页爬取**：使用网页爬取功能获取外部内容时。 + +3. **代理请求**：代理外部 API 请求时。 + +**配置示例**： + +```bash +# 场景 1：公网部署（推荐） +# 阻止所有私有 IP 访问，保证安全 +SSRF_ALLOW_PRIVATE_IP_ADDRESS=0 + +# 场景 2：内网部署 +# 允许所有私有 IP，可访问内网图片服务器等资源 +SSRF_ALLOW_PRIVATE_IP_ADDRESS=1 + +# 场景 3：混合部署（最常见） +# 默认阻止私有 IP，但允许特定可信的内网服务器 +SSRF_ALLOW_PRIVATE_IP_ADDRESS=0 +SSRF_ALLOW_IP_ADDRESS_LIST=192.168.1.100,10.0.0.50 +``` ### `SSRF_ALLOW_IP_ADDRESS_LIST` - 类型：可选 -- 说明：允许的私有 IP 地址列表，多个 IP 地址用逗号分隔。仅在 `SSRF_ALLOW_PRIVATE_IP_ADDRESS` 为 `0` 时生效。 +- 描述：允许访问的 IP 地址白名单，多个 IP 地址用逗号分隔。仅在 `SSRF_ALLOW_PRIVATE_IP_ADDRESS` 为 `0` 时生效。使用此选项可以在保持 SSRF 防护的同时，允许访问特定的内网 IP 地址。 - 默认值：- -- 示例：`198.18.1.62,224.0.0.3` +- 示例：`192.168.1.100,10.0.0.50,172.16.0.10` + +**常见使用场景**： + +- 允许访问内网图片存储服务器：`192.168.1.100` +- 允许访问内网 API 网关：`10.0.0.50` +- 允许访问内网文档服务器：`172.16.0.10` ### `ENABLE_AUTH_PROTECTION` diff --git a/package.json b/package.json index c7f02cbf01..34d4ab2a22 100644 --- a/package.json +++ b/package.json @@ -154,6 +154,7 @@ "@lobechat/database": "workspace:*", "@lobechat/electron-client-ipc": "workspace:*", "@lobechat/electron-server-ipc": "workspace:*", + "@lobechat/fetch-sse": "workspace:*", "@lobechat/file-loaders": "workspace:*", "@lobechat/model-runtime": "workspace:*", "@lobechat/observability-otel": "workspace:*", diff --git a/packages/fetch-sse/package.json b/packages/fetch-sse/package.json new file mode 100644 index 0000000000..305b778623 --- /dev/null +++ b/packages/fetch-sse/package.json @@ -0,0 +1,29 @@ +{ + "name": "@lobechat/fetch-sse", + "version": "1.0.0", + "private": true, + "description": "SSE fetch utilities with streaming support", + "exports": { + ".": { + "types": "./src/index.ts", + "default": "./src/index.ts" + }, + "./parseError": { + "types": "./src/parseError.ts", + "default": "./src/parseError.ts" + } + }, + "main": "./src/index.ts", + "types": "./src/index.ts", + "scripts": { + "test": "vitest", + "test:coverage": "vitest --coverage --silent='passed-only'" + }, + "dependencies": { + "@lobechat/const": "workspace:*", + "@lobechat/model-runtime": "workspace:*", + "@lobechat/types": "workspace:*", + "@lobechat/utils": "workspace:*", + "i18next": "^24.2.1" + } +} diff --git a/packages/utils/src/fetch/__tests__/fetchSSE.test.ts b/packages/fetch-sse/src/__tests__/fetchSSE.test.ts similarity index 98% rename from packages/utils/src/fetch/__tests__/fetchSSE.test.ts rename to packages/fetch-sse/src/__tests__/fetchSSE.test.ts index b68c25af30..9a2111451a 100644 --- a/packages/utils/src/fetch/__tests__/fetchSSE.test.ts +++ b/packages/fetch-sse/src/__tests__/fetchSSE.test.ts @@ -1,10 +1,10 @@ import { MESSAGE_CANCEL_FLAT } from '@lobechat/const'; import { ChatMessageError } from '@lobechat/types'; +import { FetchEventSourceInit } from '@lobechat/utils/client/fetchEventSource/index'; +import { fetchEventSource } from '@lobechat/utils/client/fetchEventSource/index'; +import { sleep } from '@lobechat/utils/sleep'; import { afterEach, describe, expect, it, vi } from 'vitest'; -import { FetchEventSourceInit } from '../../client/fetchEventSource'; -import { fetchEventSource } from '../../client/fetchEventSource'; -import { sleep } from '../../sleep'; import { fetchSSE } from '../fetchSSE'; // 模拟 i18next @@ -12,7 +12,7 @@ vi.mock('i18next', () => ({ t: vi.fn((key) => `translated_${key}`), })); -vi.mock('../../client/fetchEventSource', () => ({ +vi.mock('@lobechat/utils/client/fetchEventSource/index', () => ({ fetchEventSource: vi.fn(), })); diff --git a/packages/utils/src/fetch/__tests__/parseError.test.ts b/packages/fetch-sse/src/__tests__/parseError.test.ts similarity index 93% rename from packages/utils/src/fetch/__tests__/parseError.test.ts rename to packages/fetch-sse/src/__tests__/parseError.test.ts index 23c239b3ae..16a2a1e7c8 100644 --- a/packages/utils/src/fetch/__tests__/parseError.test.ts +++ b/packages/fetch-sse/src/__tests__/parseError.test.ts @@ -1,14 +1,14 @@ import { ErrorResponse } from '@lobechat/types'; -import { afterEach, describe, expect, it, vi } from 'vitest'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { getMessageError } from '../parseError'; -// 模拟 i18next +// Mock i18next vi.mock('i18next', () => ({ t: vi.fn((key) => `translated_${key}`), })); -// 模拟 Response +// Mock Response const createMockResponse = (body: any, ok: boolean, status: number = 200) => ({ ok, status, @@ -38,11 +38,14 @@ const createMockResponse = (body: any, ok: boolean, status: number = 200) => ({ }, }); -// 在每次测试后清理所有模拟 afterEach(() => { vi.restoreAllMocks(); }); +beforeEach(() => { + vi.clearAllMocks(); +}); + describe('getMessageError', () => { it('should handle business error correctly', async () => { const mockErrorResponse: ErrorResponse = { diff --git a/packages/utils/src/fetch/fetchSSE.ts b/packages/fetch-sse/src/fetchSSE.ts similarity index 99% rename from packages/utils/src/fetch/fetchSSE.ts rename to packages/fetch-sse/src/fetchSSE.ts index f8a5f3aea4..a3e3bf50c0 100644 --- a/packages/utils/src/fetch/fetchSSE.ts +++ b/packages/fetch-sse/src/fetchSSE.ts @@ -12,9 +12,9 @@ import { ResponseAnimation, ResponseAnimationStyle, } from '@lobechat/types'; +import { fetchEventSource } from '@lobechat/utils/client/fetchEventSource/index'; +import { nanoid } from '@lobechat/utils/uuid'; -import { fetchEventSource } from '../client/fetchEventSource'; -import { nanoid } from '../uuid'; import { getMessageError } from './parseError'; type SSEFinishType = 'done' | 'error' | 'abort'; diff --git a/packages/utils/src/fetch/headers.ts b/packages/fetch-sse/src/headers.ts similarity index 100% rename from packages/utils/src/fetch/headers.ts rename to packages/fetch-sse/src/headers.ts diff --git a/packages/utils/src/fetch/index.ts b/packages/fetch-sse/src/index.ts similarity index 100% rename from packages/utils/src/fetch/index.ts rename to packages/fetch-sse/src/index.ts diff --git a/packages/utils/src/fetch/parseError.ts b/packages/fetch-sse/src/parseError.ts similarity index 69% rename from packages/utils/src/fetch/parseError.ts rename to packages/fetch-sse/src/parseError.ts index ff47d4df3c..356cec003a 100644 --- a/packages/utils/src/fetch/parseError.ts +++ b/packages/fetch-sse/src/parseError.ts @@ -1,7 +1,7 @@ import { ChatMessageError, ErrorResponse, ErrorType } from '@lobechat/types'; import { t } from 'i18next'; -export const getMessageError = async (response: Response) => { +export const getMessageError = async (response: Response): Promise => { let chatMessageError: ChatMessageError; // try to get the biz error @@ -9,13 +9,13 @@ export const getMessageError = async (response: Response) => { const data = (await response.json()) as ErrorResponse; chatMessageError = { body: data.body, - message: t(`response.${data.errorType}` as any, { ns: 'error' }), + message: t(`response.${data.errorType}`, { ns: 'error' }), type: data.errorType, }; } catch { // if not return, then it's a common error chatMessageError = { - message: t(`response.${response.status}` as any, { ns: 'error' }), + message: t(`response.${response.status}`, { ns: 'error' }), type: response.status as ErrorType, }; } diff --git a/packages/utils/src/fetch/request.ts b/packages/fetch-sse/src/request.ts similarity index 100% rename from packages/utils/src/fetch/request.ts rename to packages/fetch-sse/src/request.ts diff --git a/packages/model-runtime/src/core/contextBuilders/anthropic.test.ts b/packages/model-runtime/src/core/contextBuilders/anthropic.test.ts index fc4c936565..3dee2c417e 100644 --- a/packages/model-runtime/src/core/contextBuilders/anthropic.test.ts +++ b/packages/model-runtime/src/core/contextBuilders/anthropic.test.ts @@ -1,8 +1,8 @@ +import { imageUrlToBase64 } from '@lobechat/utils'; import { OpenAI } from 'openai'; -import { describe, expect, it, vi } from 'vitest'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; import { OpenAIChatMessage, UserMessageContentPart } from '../../types/chat'; -import { imageUrlToBase64 } from '../../utils/imageToBase64'; import { parseDataUri } from '../../utils/uriParser'; import { buildAnthropicBlock, @@ -12,16 +12,22 @@ import { } from './anthropic'; // Mock the parseDataUri function since it's an implementation detail -vi.mock('../../utils/uriParser', () => ({ - parseDataUri: vi.fn().mockReturnValue({ - mimeType: 'image/jpeg', - base64: 'base64EncodedString', - type: 'base64', - }), +vi.mock('../../utils/uriParser'); +vi.mock('@lobechat/utils', () => ({ + imageUrlToBase64: vi.fn(), })); -vi.mock('../../utils/imageToBase64'); describe('anthropicHelpers', () => { + beforeEach(() => { + vi.resetAllMocks(); + // Set default mock implementation for parseDataUri + vi.mocked(parseDataUri).mockReturnValue({ + mimeType: 'image/jpeg', + base64: 'base64EncodedString', + type: 'base64', + }); + }); + describe('buildAnthropicBlock', () => { it('should return the content as is for text type', async () => { const content: UserMessageContentPart = { type: 'text', text: 'Hello!' }; @@ -52,7 +58,7 @@ describe('anthropicHelpers', () => { base64: null, type: 'url', }); - vi.mocked(imageUrlToBase64).mockResolvedValue({ + vi.mocked(imageUrlToBase64).mockResolvedValueOnce({ base64: 'convertedBase64String', mimeType: 'image/jpg', }); @@ -82,7 +88,7 @@ describe('anthropicHelpers', () => { base64: null, type: 'url', }); - vi.mocked(imageUrlToBase64).mockResolvedValue({ + vi.mocked(imageUrlToBase64).mockResolvedValueOnce({ base64: 'convertedBase64String', mimeType: 'image/png', }); diff --git a/packages/model-runtime/src/core/contextBuilders/anthropic.ts b/packages/model-runtime/src/core/contextBuilders/anthropic.ts index d8e97ebb5e..efe900812a 100644 --- a/packages/model-runtime/src/core/contextBuilders/anthropic.ts +++ b/packages/model-runtime/src/core/contextBuilders/anthropic.ts @@ -1,8 +1,8 @@ import Anthropic from '@anthropic-ai/sdk'; +import { imageUrlToBase64 } from '@lobechat/utils'; import OpenAI from 'openai'; import { OpenAIChatMessage, UserMessageContentPart } from '../../types'; -import { imageUrlToBase64 } from '../../utils/imageToBase64'; import { parseDataUri } from '../../utils/uriParser'; export const buildAnthropicBlock = async ( diff --git a/packages/model-runtime/src/core/contextBuilders/google.test.ts b/packages/model-runtime/src/core/contextBuilders/google.test.ts index 4e65422044..145d800990 100644 --- a/packages/model-runtime/src/core/contextBuilders/google.test.ts +++ b/packages/model-runtime/src/core/contextBuilders/google.test.ts @@ -1,9 +1,9 @@ // @vitest-environment node import { Type as SchemaType } from '@google/genai'; +import * as imageToBase64Module from '@lobechat/utils'; import { describe, expect, it, vi } from 'vitest'; import { ChatCompletionTool, OpenAIChatMessage, UserMessageContentPart } from '../../types'; -import * as imageToBase64Module from '../../utils/imageToBase64'; import { parseDataUri } from '../../utils/uriParser'; import { buildGoogleMessage, diff --git a/packages/model-runtime/src/core/contextBuilders/google.ts b/packages/model-runtime/src/core/contextBuilders/google.ts index 768e0814ed..9e885baa0f 100644 --- a/packages/model-runtime/src/core/contextBuilders/google.ts +++ b/packages/model-runtime/src/core/contextBuilders/google.ts @@ -5,9 +5,9 @@ import { Part, Type as SchemaType, } from '@google/genai'; +import { imageUrlToBase64 } from '@lobechat/utils'; import { ChatCompletionTool, OpenAIChatMessage, UserMessageContentPart } from '../../types'; -import { imageUrlToBase64 } from '../../utils/imageToBase64'; import { safeParseJSON } from '../../utils/safeParseJSON'; import { parseDataUri } from '../../utils/uriParser'; @@ -64,12 +64,9 @@ export const buildGooglePart = async ( } if (type === 'url') { - // For video URLs, we need to fetch and convert to base64 + // Use imageUrlToBase64 for SSRF protection (works for any binary data including videos) // Note: This might need size/duration limits for practical use - const response = await fetch(content.video_url.url); - const arrayBuffer = await response.arrayBuffer(); - const base64 = Buffer.from(arrayBuffer).toString('base64'); - const mimeType = response.headers.get('content-type') || 'video/mp4'; + const { base64, mimeType } = await imageUrlToBase64(content.video_url.url); return { inlineData: { data: base64, mimeType }, diff --git a/packages/model-runtime/src/core/contextBuilders/openai.test.ts b/packages/model-runtime/src/core/contextBuilders/openai.test.ts index 49c0a24b9e..89c17a1182 100644 --- a/packages/model-runtime/src/core/contextBuilders/openai.test.ts +++ b/packages/model-runtime/src/core/contextBuilders/openai.test.ts @@ -1,8 +1,8 @@ +import { imageUrlToBase64 } from '@lobechat/utils'; import OpenAI from 'openai'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { OpenAIChatMessage } from '../../types'; -import { imageUrlToBase64 } from '../../utils/imageToBase64'; import { parseDataUri } from '../../utils/uriParser'; import { convertImageUrlToFile, @@ -12,7 +12,9 @@ import { } from './openai'; // 模拟依赖 -vi.mock('../../utils/imageToBase64'); +vi.mock('@lobechat/utils', () => ({ + imageUrlToBase64: vi.fn(), +})); vi.mock('../../utils/uriParser'); describe('convertMessageContent', () => { diff --git a/packages/model-runtime/src/core/contextBuilders/openai.ts b/packages/model-runtime/src/core/contextBuilders/openai.ts index a89e09ef75..868f70a931 100644 --- a/packages/model-runtime/src/core/contextBuilders/openai.ts +++ b/packages/model-runtime/src/core/contextBuilders/openai.ts @@ -1,8 +1,8 @@ +import { imageUrlToBase64 } from '@lobechat/utils'; import OpenAI, { toFile } from 'openai'; import { disableStreamModels, systemToUserModels } from '../../const/models'; import { ChatStreamPayload, OpenAIChatMessage } from '../../types'; -import { imageUrlToBase64 } from '../../utils/imageToBase64'; import { parseDataUri } from '../../utils/uriParser'; export const convertMessageContent = async ( diff --git a/packages/model-runtime/src/core/openaiCompatibleFactory/createImage.test.ts b/packages/model-runtime/src/core/openaiCompatibleFactory/createImage.test.ts index e7293a6db9..96fe900e45 100644 --- a/packages/model-runtime/src/core/openaiCompatibleFactory/createImage.test.ts +++ b/packages/model-runtime/src/core/openaiCompatibleFactory/createImage.test.ts @@ -1,9 +1,9 @@ // @vitest-environment node +import * as imageToBase64Module from '@lobechat/utils'; import OpenAI from 'openai'; import { beforeEach, describe, expect, it, vi } from 'vitest'; import { CreateImagePayload } from '../../types/image'; -import * as imageToBase64Module from '../../utils/imageToBase64'; import * as uriParserModule from '../../utils/uriParser'; import { createOpenAICompatibleImage } from './createImage'; diff --git a/packages/model-runtime/src/core/openaiCompatibleFactory/createImage.ts b/packages/model-runtime/src/core/openaiCompatibleFactory/createImage.ts index b1d3822787..63a53095f5 100644 --- a/packages/model-runtime/src/core/openaiCompatibleFactory/createImage.ts +++ b/packages/model-runtime/src/core/openaiCompatibleFactory/createImage.ts @@ -1,3 +1,4 @@ +import { imageUrlToBase64 } from '@lobechat/utils'; import { cleanObject } from '@lobechat/utils/object'; import createDebug from 'debug'; import { RuntimeImageGenParamsValue } from 'model-bank'; @@ -5,7 +6,6 @@ import OpenAI from 'openai'; import { CreateImagePayload, CreateImageResponse } from '../../types/image'; import { getModelPricing } from '../../utils/getModelPricing'; -import { imageUrlToBase64 } from '../../utils/imageToBase64'; import { parseDataUri } from '../../utils/uriParser'; import { convertImageUrlToFile } from '../contextBuilders/openai'; import { convertOpenAIImageUsage } from '../usageConverters/openai'; diff --git a/packages/model-runtime/src/core/openaiCompatibleFactory/index.test.ts b/packages/model-runtime/src/core/openaiCompatibleFactory/index.test.ts index e4b8d30360..19a70f9adc 100644 --- a/packages/model-runtime/src/core/openaiCompatibleFactory/index.test.ts +++ b/packages/model-runtime/src/core/openaiCompatibleFactory/index.test.ts @@ -1,15 +1,12 @@ // @vitest-environment node -import { - AgentRuntimeErrorType, - ChatStreamCallbacks, - ChatStreamPayload, - LobeOpenAICompatibleRuntime, -} from '@lobechat/model-runtime'; import { ModelProvider } from 'model-bank'; import OpenAI from 'openai'; import type { Stream } from 'openai/streaming'; import { Mock, afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { LobeOpenAICompatibleRuntime } from '../../core/BaseAI'; +import { ChatStreamCallbacks, ChatStreamPayload } from '../../types/chat'; +import { AgentRuntimeErrorType } from '../../types/error'; import * as debugStreamModule from '../../utils/debugStream'; import * as openaiHelpers from '../contextBuilders/openai'; import { createOpenAICompatibleRuntime } from './index'; diff --git a/packages/model-runtime/src/core/streams/openai/responsesStream.test.ts b/packages/model-runtime/src/core/streams/openai/responsesStream.test.ts index d30bfb261e..761358f4b6 100644 --- a/packages/model-runtime/src/core/streams/openai/responsesStream.test.ts +++ b/packages/model-runtime/src/core/streams/openai/responsesStream.test.ts @@ -1,6 +1,6 @@ -import { AgentRuntimeErrorType } from '@lobechat/model-runtime'; import { describe, expect, it, vi } from 'vitest'; +import { AgentRuntimeErrorType } from '../../../types/error'; import { FIRST_CHUNK_ERROR_KEY } from '../protocol'; import { createReadableStream, readStreamChunk } from '../utils'; import { OpenAIResponsesStream } from './responsesStream'; diff --git a/packages/model-runtime/src/helpers/mergeChatMethodOptions.ts b/packages/model-runtime/src/helpers/mergeChatMethodOptions.ts index 36794df3d2..c1588e33f7 100644 --- a/packages/model-runtime/src/helpers/mergeChatMethodOptions.ts +++ b/packages/model-runtime/src/helpers/mergeChatMethodOptions.ts @@ -1,6 +1,7 @@ -import { ChatMethodOptions } from '@lobechat/model-runtime'; import debug from 'debug'; +import { ChatMethodOptions } from '../types/chat'; + const log = debug('model-runtime:helpers:mergeChatMethodOptions'); export const mergeMultipleChatMethodOptions = (options: ChatMethodOptions[]): ChatMethodOptions => { diff --git a/packages/model-runtime/src/providers/aihubmix/index.test.ts b/packages/model-runtime/src/providers/aihubmix/index.test.ts index 7a5c5de007..29323f0cbc 100644 --- a/packages/model-runtime/src/providers/aihubmix/index.test.ts +++ b/packages/model-runtime/src/providers/aihubmix/index.test.ts @@ -1,4 +1,4 @@ -// @vitest-environment edge-runtime +// @vitest-environment node import { ModelProvider } from 'model-bank'; import { beforeEach, describe, expect, it, vi } from 'vitest'; diff --git a/packages/model-runtime/src/providers/anthropic/generateObject.test.ts b/packages/model-runtime/src/providers/anthropic/generateObject.test.ts index e8ce9c22e8..b70f8a295e 100644 --- a/packages/model-runtime/src/providers/anthropic/generateObject.test.ts +++ b/packages/model-runtime/src/providers/anthropic/generateObject.test.ts @@ -1,4 +1,4 @@ -// @vitest-environment edge-runtime +// @vitest-environment node import { describe, expect, it, vi } from 'vitest'; import { createAnthropicGenerateObject } from './generateObject'; diff --git a/packages/model-runtime/src/providers/anthropic/index.test.ts b/packages/model-runtime/src/providers/anthropic/index.test.ts index 47945d7e4c..0a063b1e5e 100644 --- a/packages/model-runtime/src/providers/anthropic/index.test.ts +++ b/packages/model-runtime/src/providers/anthropic/index.test.ts @@ -1,8 +1,8 @@ // @vitest-environment node -import { ChatCompletionTool, ChatStreamPayload } from '@lobechat/model-runtime'; import { Mock, afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import * as anthropicHelpers from '../../core/contextBuilders/anthropic'; +import { ChatCompletionTool, ChatStreamPayload } from '../../types/chat'; import * as debugStreamModule from '../../utils/debugStream'; import { LobeAnthropicAI } from './index'; diff --git a/packages/model-runtime/src/providers/baichuan/index.test.ts b/packages/model-runtime/src/providers/baichuan/index.test.ts index 9a3956093f..d00f7c4dd1 100644 --- a/packages/model-runtime/src/providers/baichuan/index.test.ts +++ b/packages/model-runtime/src/providers/baichuan/index.test.ts @@ -1,8 +1,8 @@ // @vitest-environment node -import { LobeOpenAICompatibleRuntime } from '@lobechat/model-runtime'; import { ModelProvider } from 'model-bank'; import { Mock, afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { LobeOpenAICompatibleRuntime } from '../../core/BaseAI'; import { testProvider } from '../../providerTestUtils'; import { LobeBaichuanAI, params } from './index'; diff --git a/packages/model-runtime/src/providers/bedrock/index.test.ts b/packages/model-runtime/src/providers/bedrock/index.test.ts index f7b2edf66f..259ce3f287 100644 --- a/packages/model-runtime/src/providers/bedrock/index.test.ts +++ b/packages/model-runtime/src/providers/bedrock/index.test.ts @@ -3,10 +3,10 @@ import { InvokeModelCommand, InvokeModelWithResponseStreamCommand, } from '@aws-sdk/client-bedrock-runtime'; -import { AgentRuntimeErrorType } from '@lobechat/model-runtime'; import { ModelProvider } from 'model-bank'; import { Mock, afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { AgentRuntimeErrorType } from '../../types/error'; import * as debugStreamModule from '../../utils/debugStream'; import { LobeBedrockAI, experimental_buildLlama2Prompt } from './index'; diff --git a/packages/model-runtime/src/providers/bfl/createImage.test.ts b/packages/model-runtime/src/providers/bfl/createImage.test.ts index 72820c9e64..a32236adc9 100644 --- a/packages/model-runtime/src/providers/bfl/createImage.test.ts +++ b/packages/model-runtime/src/providers/bfl/createImage.test.ts @@ -6,7 +6,7 @@ import { createBflImage } from './createImage'; import { BflStatusResponse } from './types'; // Mock external dependencies -vi.mock('../../utils/imageToBase64', () => ({ +vi.mock('@lobechat/utils', () => ({ imageUrlToBase64: vi.fn(), })); @@ -187,7 +187,7 @@ describe('createBflImage', () => { it('should convert single imageUrl to image_prompt base64', async () => { // Arrange const { parseDataUri } = await import('../../utils/uriParser'); - const { imageUrlToBase64 } = await import('../../utils/imageToBase64'); + const { imageUrlToBase64 } = await import('@lobechat/utils'); const { asyncifyPolling } = await import('../../utils/asyncifyPolling'); const mockParseDataUri = vi.mocked(parseDataUri); @@ -290,7 +290,7 @@ describe('createBflImage', () => { it('should convert multiple imageUrls for Kontext models', async () => { // Arrange const { parseDataUri } = await import('../../utils/uriParser'); - const { imageUrlToBase64 } = await import('../../utils/imageToBase64'); + const { imageUrlToBase64 } = await import('@lobechat/utils'); const { asyncifyPolling } = await import('../../utils/asyncifyPolling'); const mockParseDataUri = vi.mocked(parseDataUri); @@ -350,7 +350,7 @@ describe('createBflImage', () => { it('should limit imageUrls to maximum 4 images', async () => { // Arrange const { parseDataUri } = await import('../../utils/uriParser'); - const { imageUrlToBase64 } = await import('../../utils/imageToBase64'); + const { imageUrlToBase64 } = await import('@lobechat/utils'); const { asyncifyPolling } = await import('../../utils/asyncifyPolling'); const mockParseDataUri = vi.mocked(parseDataUri); diff --git a/packages/model-runtime/src/providers/bfl/createImage.ts b/packages/model-runtime/src/providers/bfl/createImage.ts index f8d8a1d693..a106f027a2 100644 --- a/packages/model-runtime/src/providers/bfl/createImage.ts +++ b/packages/model-runtime/src/providers/bfl/createImage.ts @@ -1,3 +1,4 @@ +import { imageUrlToBase64 } from '@lobechat/utils'; import createDebug from 'debug'; import { RuntimeImageGenParamsValue } from 'model-bank'; @@ -5,7 +6,6 @@ import { AgentRuntimeErrorType } from '../../types/error'; import { CreateImagePayload, CreateImageResponse } from '../../types/image'; import { type TaskResult, asyncifyPolling } from '../../utils/asyncifyPolling'; import { AgentRuntimeError } from '../../utils/createError'; -import { imageUrlToBase64 } from '../../utils/imageToBase64'; import { parseDataUri } from '../../utils/uriParser'; import { BFL_ENDPOINTS, diff --git a/packages/model-runtime/src/providers/cloudflare/index.test.ts b/packages/model-runtime/src/providers/cloudflare/index.test.ts index de040ec916..043faa346e 100644 --- a/packages/model-runtime/src/providers/cloudflare/index.test.ts +++ b/packages/model-runtime/src/providers/cloudflare/index.test.ts @@ -1,7 +1,7 @@ // @vitest-environment node -import { ChatCompletionTool } from '@lobechat/model-runtime'; import { Mock, afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { ChatCompletionTool } from '../../types/chat'; import * as debugStreamModule from '../../utils/debugStream'; import { LobeCloudflareAI } from './index'; diff --git a/packages/model-runtime/src/providers/cohere/index.test.ts b/packages/model-runtime/src/providers/cohere/index.test.ts index 9a7b9c7b83..fde6650734 100644 --- a/packages/model-runtime/src/providers/cohere/index.test.ts +++ b/packages/model-runtime/src/providers/cohere/index.test.ts @@ -1,8 +1,8 @@ // @vitest-environment node -import { LobeOpenAICompatibleRuntime } from '@lobechat/model-runtime'; import { ModelProvider } from 'model-bank'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { LobeOpenAICompatibleRuntime } from '../../core/BaseAI'; import { testProvider } from '../../providerTestUtils'; import { LobeCohereAI, params } from './index'; diff --git a/packages/model-runtime/src/providers/google/createImage.test.ts b/packages/model-runtime/src/providers/google/createImage.test.ts index ca63eab03c..a9ee024aaa 100644 --- a/packages/model-runtime/src/providers/google/createImage.test.ts +++ b/packages/model-runtime/src/providers/google/createImage.test.ts @@ -1,9 +1,9 @@ -// @vitest-environment edge-runtime +// @vitest-environment node import { GoogleGenAI } from '@google/genai'; +import * as imageToBase64Module from '@lobechat/utils'; import { beforeEach, describe, expect, it, vi } from 'vitest'; import { CreateImagePayload } from '../../types/image'; -import * as imageToBase64Module from '../../utils/imageToBase64'; import { createGoogleImage } from './createImage'; const provider = 'google'; diff --git a/packages/model-runtime/src/providers/google/createImage.ts b/packages/model-runtime/src/providers/google/createImage.ts index 9fbea6dc91..e340822290 100644 --- a/packages/model-runtime/src/providers/google/createImage.ts +++ b/packages/model-runtime/src/providers/google/createImage.ts @@ -1,11 +1,11 @@ import { Content, GenerateContentConfig, GoogleGenAI, Part } from '@google/genai'; +import { imageUrlToBase64 } from '@lobechat/utils'; import { convertGoogleAIUsage } from '../../core/usageConverters/google-ai'; import { CreateImagePayload, CreateImageResponse } from '../../types/image'; import { AgentRuntimeError } from '../../utils/createError'; import { getModelPricing } from '../../utils/getModelPricing'; import { parseGoogleErrorMessage } from '../../utils/googleErrorParser'; -import { imageUrlToBase64 } from '../../utils/imageToBase64'; import { parseDataUri } from '../../utils/uriParser'; // Maximum number of images allowed for processing diff --git a/packages/model-runtime/src/providers/google/generateObject.test.ts b/packages/model-runtime/src/providers/google/generateObject.test.ts index 1e4bfe209c..d6a3c7cf7f 100644 --- a/packages/model-runtime/src/providers/google/generateObject.test.ts +++ b/packages/model-runtime/src/providers/google/generateObject.test.ts @@ -1,4 +1,4 @@ -// @vitest-environment edge-runtime +// @vitest-environment node import { Type as SchemaType } from '@google/genai'; import { describe, expect, it, vi } from 'vitest'; diff --git a/packages/model-runtime/src/providers/google/index.test.ts b/packages/model-runtime/src/providers/google/index.test.ts index 2e8f307281..d0a88e213b 100644 --- a/packages/model-runtime/src/providers/google/index.test.ts +++ b/packages/model-runtime/src/providers/google/index.test.ts @@ -1,14 +1,11 @@ -// @vitest-environment edge-runtime +// @vitest-environment node import { GenerateContentResponse, Tool } from '@google/genai'; -import { OpenAIChatMessage } from '@lobechat/model-runtime'; -import { ChatStreamPayload } from '@lobechat/types'; import OpenAI from 'openai'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { LOBE_ERROR_KEY } from '../../core/streams'; import { AgentRuntimeErrorType } from '../../types/error'; import * as debugStreamModule from '../../utils/debugStream'; -import * as imageToBase64Module from '../../utils/imageToBase64'; import { LobeGoogleAI, resolveModelThinkingBudget } from './index'; const provider = 'google'; diff --git a/packages/model-runtime/src/providers/groq/index.test.ts b/packages/model-runtime/src/providers/groq/index.test.ts index b359db7c0b..aab9946886 100644 --- a/packages/model-runtime/src/providers/groq/index.test.ts +++ b/packages/model-runtime/src/providers/groq/index.test.ts @@ -1,7 +1,7 @@ // @vitest-environment node -import { LobeOpenAICompatibleRuntime } from '@lobechat/model-runtime'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { LobeOpenAICompatibleRuntime } from '../../core/BaseAI'; import { testProvider } from '../../providerTestUtils'; import { AgentRuntimeErrorType } from '../../types/error'; import { LobeGroq, params } from './index'; diff --git a/packages/model-runtime/src/providers/hunyuan/index.test.ts b/packages/model-runtime/src/providers/hunyuan/index.test.ts index d0db00c739..dfa4b48268 100644 --- a/packages/model-runtime/src/providers/hunyuan/index.test.ts +++ b/packages/model-runtime/src/providers/hunyuan/index.test.ts @@ -1,8 +1,8 @@ // @vitest-environment node -import { LobeOpenAICompatibleRuntime } from '@lobechat/model-runtime'; import { ModelProvider } from 'model-bank'; import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { LobeOpenAICompatibleRuntime } from '../../core/BaseAI'; import { testProvider } from '../../providerTestUtils'; import { LobeHunyuanAI, params } from './index'; diff --git a/packages/model-runtime/src/providers/minimax/createImage.test.ts b/packages/model-runtime/src/providers/minimax/createImage.test.ts index 6dbdf9f763..b3e28311ee 100644 --- a/packages/model-runtime/src/providers/minimax/createImage.test.ts +++ b/packages/model-runtime/src/providers/minimax/createImage.test.ts @@ -1,4 +1,4 @@ -// @vitest-environment edge-runtime +// @vitest-environment node import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { CreateImageOptions } from '../../core/openaiCompatibleFactory'; diff --git a/packages/model-runtime/src/providers/mistral/index.test.ts b/packages/model-runtime/src/providers/mistral/index.test.ts index 8b2d220f44..71ce84f35c 100644 --- a/packages/model-runtime/src/providers/mistral/index.test.ts +++ b/packages/model-runtime/src/providers/mistral/index.test.ts @@ -1,7 +1,7 @@ // @vitest-environment node -import { LobeOpenAICompatibleRuntime } from '@lobechat/model-runtime'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { LobeOpenAICompatibleRuntime } from '../../core/BaseAI'; import { testProvider } from '../../providerTestUtils'; import { LobeMistralAI, params } from './index'; diff --git a/packages/model-runtime/src/providers/moonshot/index.test.ts b/packages/model-runtime/src/providers/moonshot/index.test.ts index 1f163abde7..c1eb67a9ec 100644 --- a/packages/model-runtime/src/providers/moonshot/index.test.ts +++ b/packages/model-runtime/src/providers/moonshot/index.test.ts @@ -1,7 +1,7 @@ // @vitest-environment node -import { LobeOpenAICompatibleRuntime } from '@lobechat/model-runtime'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { LobeOpenAICompatibleRuntime } from '../../core/BaseAI'; import { testProvider } from '../../providerTestUtils'; import { LobeMoonshotAI, params } from './index'; diff --git a/packages/model-runtime/src/providers/novita/index.test.ts b/packages/model-runtime/src/providers/novita/index.test.ts index fd848541b6..32a7a96743 100644 --- a/packages/model-runtime/src/providers/novita/index.test.ts +++ b/packages/model-runtime/src/providers/novita/index.test.ts @@ -1,8 +1,8 @@ // @vitest-environment node -import { LobeOpenAICompatibleRuntime } from '@lobechat/model-runtime'; import { ModelProvider } from 'model-bank'; import { Mock, afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { LobeOpenAICompatibleRuntime } from '../../core/BaseAI'; import { testProvider } from '../../providerTestUtils'; import models from './fixtures/models.json'; import { LobeNovitaAI } from './index'; diff --git a/packages/model-runtime/src/providers/ollama/index.test.ts b/packages/model-runtime/src/providers/ollama/index.test.ts index 9cb13f82b6..0e97fc6a74 100644 --- a/packages/model-runtime/src/providers/ollama/index.test.ts +++ b/packages/model-runtime/src/providers/ollama/index.test.ts @@ -1,4 +1,5 @@ // @vitest-environment node +import { imageUrlToBase64 } from '@lobechat/utils'; import { ModelProvider } from 'model-bank'; import { Ollama } from 'ollama/browser'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; @@ -9,6 +10,13 @@ import * as debugStreamModule from '../../utils/debugStream'; import { LobeOllamaAI, params } from './index'; vi.mock('ollama/browser'); +vi.mock('@lobechat/utils', async () => { + const actual = await vi.importActual('@lobechat/utils'); + return { + ...actual, + imageUrlToBase64: vi.fn(), + }; +}); // Mock the console.error to avoid polluting test output vi.spyOn(console, 'error').mockImplementation(() => {}); @@ -462,13 +470,13 @@ describe('LobeOllamaAI', () => { }); describe('buildOllamaMessages', () => { - it('should convert OpenAIChatMessage array to OllamaMessage array', () => { + it('should convert OpenAIChatMessage array to OllamaMessage array', async () => { const messages = [ { content: 'Hello', role: 'user' }, { content: 'Hi there!', role: 'assistant' }, ]; - const ollamaMessages = ollamaAI['buildOllamaMessages'](messages as any); + const ollamaMessages = await ollamaAI['buildOllamaMessages'](messages as any); expect(ollamaMessages).toEqual([ { content: 'Hello', role: 'user' }, @@ -476,15 +484,15 @@ describe('LobeOllamaAI', () => { ]); }); - it('should handle empty message array', () => { + it('should handle empty message array', async () => { const messages: any[] = []; - const ollamaMessages = ollamaAI['buildOllamaMessages'](messages); + const ollamaMessages = await ollamaAI['buildOllamaMessages'](messages); expect(ollamaMessages).toEqual([]); }); - it('should handle multiple messages with different roles', () => { + it('should handle multiple messages with different roles', async () => { const messages = [ { content: 'Hello', role: 'system' }, { content: 'Hi', role: 'user' }, @@ -492,7 +500,7 @@ describe('LobeOllamaAI', () => { { content: 'How are you?', role: 'user' }, ]; - const ollamaMessages = ollamaAI['buildOllamaMessages'](messages as any); + const ollamaMessages = await ollamaAI['buildOllamaMessages'](messages as any); expect(ollamaMessages).toHaveLength(4); expect(ollamaMessages[0].role).toBe('system'); @@ -503,26 +511,26 @@ describe('LobeOllamaAI', () => { }); describe('convertContentToOllamaMessage', () => { - it('should convert string content to OllamaMessage', () => { + it('should convert string content to OllamaMessage', async () => { const message = { content: 'Hello', role: 'user' }; - const ollamaMessage = ollamaAI['convertContentToOllamaMessage'](message as any); + const ollamaMessage = await ollamaAI['convertContentToOllamaMessage'](message as any); expect(ollamaMessage).toEqual({ content: 'Hello', role: 'user' }); }); - it('should convert text content to OllamaMessage', () => { + it('should convert text content to OllamaMessage', async () => { const message = { content: [{ type: 'text', text: 'Hello' }], role: 'user', }; - const ollamaMessage = ollamaAI['convertContentToOllamaMessage'](message as any); + const ollamaMessage = await ollamaAI['convertContentToOllamaMessage'](message as any); expect(ollamaMessage).toEqual({ content: 'Hello', role: 'user' }); }); - it('should convert image_url content to OllamaMessage with images', () => { + it('should convert image_url content to OllamaMessage with images', async () => { const message = { content: [ { @@ -533,7 +541,7 @@ describe('LobeOllamaAI', () => { role: 'user', }; - const ollamaMessage = ollamaAI['convertContentToOllamaMessage'](message as any); + const ollamaMessage = await ollamaAI['convertContentToOllamaMessage'](message as any); expect(ollamaMessage).toEqual({ content: '', @@ -542,7 +550,7 @@ describe('LobeOllamaAI', () => { }); }); - it('should ignore invalid image_url content', () => { + it('should ignore invalid image_url content', async () => { const message = { content: [ { @@ -553,7 +561,7 @@ describe('LobeOllamaAI', () => { role: 'user', }; - const ollamaMessage = ollamaAI['convertContentToOllamaMessage'](message as any); + const ollamaMessage = await ollamaAI['convertContentToOllamaMessage'](message as any); expect(ollamaMessage).toEqual({ content: '', @@ -561,7 +569,7 @@ describe('LobeOllamaAI', () => { }); }); - it('should handle mixed text and image content', () => { + it('should handle mixed text and image content', async () => { const message = { content: [ { type: 'text', text: 'First text' }, @@ -578,7 +586,7 @@ describe('LobeOllamaAI', () => { role: 'user', }; - const ollamaMessage = ollamaAI['convertContentToOllamaMessage'](message as any); + const ollamaMessage = await ollamaAI['convertContentToOllamaMessage'](message as any); expect(ollamaMessage).toEqual({ content: 'Second text', // Should keep latest text @@ -587,13 +595,13 @@ describe('LobeOllamaAI', () => { }); }); - it('should handle content with empty text', () => { + it('should handle content with empty text', async () => { const message = { content: [{ type: 'text', text: '' }], role: 'user', }; - const ollamaMessage = ollamaAI['convertContentToOllamaMessage'](message as any); + const ollamaMessage = await ollamaAI['convertContentToOllamaMessage'](message as any); expect(ollamaMessage).toEqual({ content: '', @@ -601,7 +609,7 @@ describe('LobeOllamaAI', () => { }); }); - it('should handle content with only images (no text)', () => { + it('should handle content with only images (no text)', async () => { const message = { content: [ { @@ -612,7 +620,7 @@ describe('LobeOllamaAI', () => { role: 'user', }; - const ollamaMessage = ollamaAI['convertContentToOllamaMessage'](message as any); + const ollamaMessage = await ollamaAI['convertContentToOllamaMessage'](message as any); expect(ollamaMessage).toEqual({ content: '', @@ -621,7 +629,7 @@ describe('LobeOllamaAI', () => { }); }); - it('should handle multiple images without text', () => { + it('should handle multiple images without text', async () => { const message = { content: [ { @@ -640,7 +648,7 @@ describe('LobeOllamaAI', () => { role: 'user', }; - const ollamaMessage = ollamaAI['convertContentToOllamaMessage'](message as any); + const ollamaMessage = await ollamaAI['convertContentToOllamaMessage'](message as any); expect(ollamaMessage).toEqual({ content: '', @@ -649,7 +657,10 @@ describe('LobeOllamaAI', () => { }); }); - it('should ignore images with invalid data URIs', () => { + it('should ignore images with invalid data URIs', async () => { + // Mock imageUrlToBase64 to simulate conversion failure for external URLs + vi.mocked(imageUrlToBase64).mockRejectedValue(new Error('Network error')); + const message = { content: [ { type: 'text', text: 'Hello' }, @@ -665,7 +676,7 @@ describe('LobeOllamaAI', () => { role: 'user', }; - const ollamaMessage = ollamaAI['convertContentToOllamaMessage'](message as any); + const ollamaMessage = await ollamaAI['convertContentToOllamaMessage'](message as any); expect(ollamaMessage).toEqual({ content: 'Hello', @@ -674,7 +685,7 @@ describe('LobeOllamaAI', () => { }); }); - it('should handle complex interleaved content', () => { + it('should handle complex interleaved content', async () => { const message = { content: [ { type: 'text', text: 'Text 1' }, @@ -692,7 +703,7 @@ describe('LobeOllamaAI', () => { role: 'user', }; - const ollamaMessage = ollamaAI['convertContentToOllamaMessage'](message as any); + const ollamaMessage = await ollamaAI['convertContentToOllamaMessage'](message as any); expect(ollamaMessage).toEqual({ content: 'Text 3', // Should keep latest text @@ -701,7 +712,7 @@ describe('LobeOllamaAI', () => { }); }); - it('should handle assistant role with images', () => { + it('should handle assistant role with images', async () => { const message = { content: [ { type: 'text', text: 'Here is the image' }, @@ -713,7 +724,7 @@ describe('LobeOllamaAI', () => { role: 'assistant', }; - const ollamaMessage = ollamaAI['convertContentToOllamaMessage'](message as any); + const ollamaMessage = await ollamaAI['convertContentToOllamaMessage'](message as any); expect(ollamaMessage).toEqual({ content: 'Here is the image', @@ -722,13 +733,13 @@ describe('LobeOllamaAI', () => { }); }); - it('should handle system role with text', () => { + it('should handle system role with text', async () => { const message = { content: [{ type: 'text', text: 'You are a helpful assistant' }], role: 'system', }; - const ollamaMessage = ollamaAI['convertContentToOllamaMessage'](message as any); + const ollamaMessage = await ollamaAI['convertContentToOllamaMessage'](message as any); expect(ollamaMessage).toEqual({ content: 'You are a helpful assistant', @@ -736,13 +747,13 @@ describe('LobeOllamaAI', () => { }); }); - it('should handle empty content array', () => { + it('should handle empty content array', async () => { const message = { content: [], role: 'user', }; - const ollamaMessage = ollamaAI['convertContentToOllamaMessage'](message as any); + const ollamaMessage = await ollamaAI['convertContentToOllamaMessage'](message as any); expect(ollamaMessage).toEqual({ content: '', diff --git a/packages/model-runtime/src/providers/ollama/index.ts b/packages/model-runtime/src/providers/ollama/index.ts index 6754db9cdf..04aaebfafb 100644 --- a/packages/model-runtime/src/providers/ollama/index.ts +++ b/packages/model-runtime/src/providers/ollama/index.ts @@ -1,4 +1,5 @@ import type { ChatModelCard } from '@lobechat/types'; +import { imageUrlToBase64 } from '@lobechat/utils'; import { ModelProvider } from 'model-bank'; import { Ollama, Tool } from 'ollama/browser'; import { ClientOptions } from 'openai'; @@ -61,7 +62,7 @@ export class LobeOllamaAI implements LobeRuntimeAI { options?.signal?.addEventListener('abort', abort); const response = await this.client.chat({ - messages: this.buildOllamaMessages(payload.messages), + messages: await this.buildOllamaMessages(payload.messages), model: payload.model, options: { frequency_penalty: payload.frequency_penalty, @@ -169,11 +170,13 @@ export class LobeOllamaAI implements LobeRuntimeAI { } }; - private buildOllamaMessages(messages: OpenAIChatMessage[]) { - return messages.map((message) => this.convertContentToOllamaMessage(message)); + private async buildOllamaMessages(messages: OpenAIChatMessage[]) { + return Promise.all(messages.map((message) => this.convertContentToOllamaMessage(message))); } - private convertContentToOllamaMessage = (message: OpenAIChatMessage): OllamaMessage => { + private convertContentToOllamaMessage = async ( + message: OpenAIChatMessage, + ): Promise => { if (typeof message.content === 'string') { return { content: message.content, role: message.role }; } @@ -183,6 +186,9 @@ export class LobeOllamaAI implements LobeRuntimeAI { role: message.role, }; + // Collect image processing tasks for parallel execution + const imagePromises: Array | string> = []; + for (const content of message.content) { switch (content.type) { case 'text': { @@ -191,16 +197,34 @@ export class LobeOllamaAI implements LobeRuntimeAI { break; } case 'image_url': { - const { base64 } = parseDataUri(content.image_url.url); + const { base64, type } = parseDataUri(content.image_url.url); + + // If already base64 format, use it directly if (base64) { - ollamaMessage.images ??= []; - ollamaMessage.images.push(base64); + imagePromises.push(base64); + } + // If it's a URL, add async conversion task with error handling + else if (type === 'url') { + imagePromises.push( + imageUrlToBase64(content.image_url.url) + .then((result) => result.base64) + .catch(() => null), // Silently ignore failed conversions + ); } break; } } } + // Process all images in parallel and filter out failed conversions + if (imagePromises.length > 0) { + const results = await Promise.all(imagePromises); + const validImages = results.filter((img): img is string => img !== null); + if (validImages.length > 0) { + ollamaMessage.images = validImages; + } + } + return ollamaMessage; }; diff --git a/packages/model-runtime/src/providers/openrouter/index.test.ts b/packages/model-runtime/src/providers/openrouter/index.test.ts index 88058f11b1..9dc64bea59 100644 --- a/packages/model-runtime/src/providers/openrouter/index.test.ts +++ b/packages/model-runtime/src/providers/openrouter/index.test.ts @@ -1,7 +1,7 @@ // @vitest-environment node -import { LobeOpenAICompatibleRuntime } from '@lobechat/model-runtime'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { LobeOpenAICompatibleRuntime } from '../../core/BaseAI'; import { testProvider } from '../../providerTestUtils'; import { LobeOpenRouterAI, params } from './index'; diff --git a/packages/model-runtime/src/providers/perplexity/index.test.ts b/packages/model-runtime/src/providers/perplexity/index.test.ts index e4d3d061a4..2c0b7aed24 100644 --- a/packages/model-runtime/src/providers/perplexity/index.test.ts +++ b/packages/model-runtime/src/providers/perplexity/index.test.ts @@ -1,8 +1,8 @@ // @vitest-environment node -import { LobeOpenAICompatibleRuntime } from '@lobechat/model-runtime'; import { ModelProvider } from 'model-bank'; import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { LobeOpenAICompatibleRuntime } from '../../core/BaseAI'; import { testProvider } from '../../providerTestUtils'; import { LobePerplexityAI } from './index'; diff --git a/packages/model-runtime/src/providers/ppio/index.test.ts b/packages/model-runtime/src/providers/ppio/index.test.ts index b02f045be3..57b894cec1 100644 --- a/packages/model-runtime/src/providers/ppio/index.test.ts +++ b/packages/model-runtime/src/providers/ppio/index.test.ts @@ -1,8 +1,8 @@ // @vitest-environment node -import { LobeOpenAICompatibleRuntime } from '@lobechat/model-runtime'; import { ModelProvider } from 'model-bank'; import { Mock, afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { LobeOpenAICompatibleRuntime } from '../../core/BaseAI'; import { testProvider } from '../../providerTestUtils'; import models from './fixtures/models.json'; import { LobePPIOAI } from './index'; diff --git a/packages/model-runtime/src/providers/qwen/createImage.test.ts b/packages/model-runtime/src/providers/qwen/createImage.test.ts index 9f13c29d78..a4d3086baa 100644 --- a/packages/model-runtime/src/providers/qwen/createImage.test.ts +++ b/packages/model-runtime/src/providers/qwen/createImage.test.ts @@ -1,4 +1,4 @@ -// @vitest-environment edge-runtime +// @vitest-environment node import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { CreateImageOptions } from '../../core/openaiCompatibleFactory'; diff --git a/packages/model-runtime/src/providers/search1api/index.test.ts b/packages/model-runtime/src/providers/search1api/index.test.ts index 46b8b7414a..11e6b61714 100644 --- a/packages/model-runtime/src/providers/search1api/index.test.ts +++ b/packages/model-runtime/src/providers/search1api/index.test.ts @@ -1,7 +1,7 @@ // @vitest-environment node -import { LobeOpenAICompatibleRuntime } from '@lobechat/model-runtime'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { LobeOpenAICompatibleRuntime } from '../../core/BaseAI'; import { testProvider } from '../../providerTestUtils'; import { LobeSearch1API, params } from './index'; diff --git a/packages/model-runtime/src/providers/siliconcloud/createImage.ts b/packages/model-runtime/src/providers/siliconcloud/createImage.ts index 87bc84a147..11ad5769e0 100644 --- a/packages/model-runtime/src/providers/siliconcloud/createImage.ts +++ b/packages/model-runtime/src/providers/siliconcloud/createImage.ts @@ -1,3 +1,4 @@ +import { imageUrlToBase64 } from '@lobechat/utils'; import createDebug from 'debug'; import { RuntimeImageGenParamsValue } from 'model-bank'; @@ -5,7 +6,6 @@ import { CreateImageOptions } from '../../core/openaiCompatibleFactory'; import { CreateImagePayload, CreateImageResponse } from '../../types'; import { AgentRuntimeErrorType } from '../../types/error'; import { AgentRuntimeError } from '../../utils/createError'; -import { imageUrlToBase64 } from '../../utils/imageToBase64'; import { parseDataUri } from '../../utils/uriParser'; const log = createDebug('lobe-image:siliconcloud'); diff --git a/packages/model-runtime/src/providers/taichu/index.test.ts b/packages/model-runtime/src/providers/taichu/index.test.ts index 3510d982c4..48b7d10508 100644 --- a/packages/model-runtime/src/providers/taichu/index.test.ts +++ b/packages/model-runtime/src/providers/taichu/index.test.ts @@ -1,9 +1,9 @@ // @vitest-environment node -import { LobeOpenAICompatibleRuntime } from '@lobechat/model-runtime'; import { ModelProvider } from 'model-bank'; import OpenAI from 'openai'; import { Mock, afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { LobeOpenAICompatibleRuntime } from '../../core/BaseAI'; import { testProvider } from '../../providerTestUtils'; import * as debugStreamModule from '../../utils/debugStream'; import { LobeTaichuAI } from './index'; diff --git a/packages/model-runtime/src/providers/wenxin/index.test.ts b/packages/model-runtime/src/providers/wenxin/index.test.ts index da59cdd421..fee9b14fd5 100644 --- a/packages/model-runtime/src/providers/wenxin/index.test.ts +++ b/packages/model-runtime/src/providers/wenxin/index.test.ts @@ -1,8 +1,8 @@ // @vitest-environment node -import { LobeOpenAICompatibleRuntime } from '@lobechat/model-runtime'; import { ModelProvider } from 'model-bank'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { LobeOpenAICompatibleRuntime } from '../../core/BaseAI'; import { testProvider } from '../../providerTestUtils'; import { LobeWenxinAI, params } from './index'; diff --git a/packages/model-runtime/src/providers/zhipu/index.test.ts b/packages/model-runtime/src/providers/zhipu/index.test.ts index 3fac0c153d..15b6794957 100644 --- a/packages/model-runtime/src/providers/zhipu/index.test.ts +++ b/packages/model-runtime/src/providers/zhipu/index.test.ts @@ -1,7 +1,7 @@ // @vitest-environment node -import { LobeOpenAICompatibleRuntime } from '@lobechat/model-runtime'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { LobeOpenAICompatibleRuntime } from '../../core/BaseAI'; import { testProvider } from '../../providerTestUtils'; import { LobeZhipuAI, params } from './index'; diff --git a/packages/model-runtime/src/utils/errorResponse.test.ts b/packages/model-runtime/src/utils/errorResponse.test.ts index 5abe11f840..cfbd3d8e5b 100644 --- a/packages/model-runtime/src/utils/errorResponse.test.ts +++ b/packages/model-runtime/src/utils/errorResponse.test.ts @@ -1,7 +1,7 @@ -import { AgentRuntimeErrorType } from '@lobechat/model-runtime'; import { ChatErrorType } from '@lobechat/types'; import { describe, expect, it, vi } from 'vitest'; +import { AgentRuntimeErrorType } from '../types/error'; import { createErrorResponse } from './errorResponse'; describe('createErrorResponse', () => { diff --git a/packages/model-runtime/src/utils/imageToBase64.test.ts b/packages/model-runtime/src/utils/imageToBase64.test.ts deleted file mode 100644 index d1fc8de27d..0000000000 --- a/packages/model-runtime/src/utils/imageToBase64.test.ts +++ /dev/null @@ -1,91 +0,0 @@ -import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; - -import { imageToBase64, imageUrlToBase64 } from './imageToBase64'; - -describe('imageToBase64', () => { - let mockImage: HTMLImageElement; - let mockCanvas: HTMLCanvasElement; - let mockContext: CanvasRenderingContext2D; - - beforeEach(() => { - mockImage = { - width: 200, - height: 100, - } as HTMLImageElement; - - mockContext = { - drawImage: vi.fn(), - } as unknown as CanvasRenderingContext2D; - - mockCanvas = { - width: 0, - height: 0, - getContext: vi.fn().mockReturnValue(mockContext), - toDataURL: vi.fn().mockReturnValue('data:image/webp;base64,mockBase64Data'), - } as unknown as HTMLCanvasElement; - - vi.spyOn(document, 'createElement').mockReturnValue(mockCanvas); - }); - - afterEach(() => { - vi.restoreAllMocks(); - }); - - it('should convert image to base64 with correct size and type', () => { - const result = imageToBase64({ img: mockImage, size: 100, type: 'image/jpeg' }); - - expect(document.createElement).toHaveBeenCalledWith('canvas'); - expect(mockCanvas.width).toBe(100); - expect(mockCanvas.height).toBe(100); - expect(mockCanvas.getContext).toHaveBeenCalledWith('2d'); - expect(mockContext.drawImage).toHaveBeenCalledWith(mockImage, 50, 0, 100, 100, 0, 0, 100, 100); - expect(mockCanvas.toDataURL).toHaveBeenCalledWith('image/jpeg'); - expect(result).toBe('data:image/webp;base64,mockBase64Data'); - }); - - it('should use default type when not specified', () => { - imageToBase64({ img: mockImage, size: 100 }); - expect(mockCanvas.toDataURL).toHaveBeenCalledWith('image/webp'); - }); - - it('should handle taller images correctly', () => { - mockImage.width = 100; - mockImage.height = 200; - imageToBase64({ img: mockImage, size: 100 }); - expect(mockContext.drawImage).toHaveBeenCalledWith(mockImage, 0, 50, 100, 100, 0, 0, 100, 100); - }); -}); - -describe('imageUrlToBase64', () => { - const mockFetch = vi.fn(); - const mockArrayBuffer = new ArrayBuffer(8); - - beforeEach(() => { - global.fetch = mockFetch; - global.btoa = vi.fn().mockReturnValue('mockBase64String'); - }); - - afterEach(() => { - vi.restoreAllMocks(); - }); - - it('should convert image URL to base64 string', async () => { - mockFetch.mockResolvedValue({ - arrayBuffer: () => Promise.resolve(mockArrayBuffer), - blob: () => Promise.resolve(new Blob([mockArrayBuffer], { type: 'image/jpg' })), - }); - - const result = await imageUrlToBase64('https://example.com/image.jpg'); - - expect(mockFetch).toHaveBeenCalledWith('https://example.com/image.jpg'); - expect(global.btoa).toHaveBeenCalled(); - expect(result).toEqual({ base64: 'mockBase64String', mimeType: 'image/jpg' }); - }); - - it('should throw an error when fetch fails', async () => { - const mockError = new Error('Fetch failed'); - mockFetch.mockRejectedValue(mockError); - - await expect(imageUrlToBase64('https://example.com/image.jpg')).rejects.toThrow('Fetch failed'); - }); -}); diff --git a/packages/model-runtime/src/utils/imageToBase64.ts b/packages/model-runtime/src/utils/imageToBase64.ts deleted file mode 100644 index 80eaabe20d..0000000000 --- a/packages/model-runtime/src/utils/imageToBase64.ts +++ /dev/null @@ -1,62 +0,0 @@ -export const imageToBase64 = ({ - size, - img, - type = 'image/webp', -}: { - img: HTMLImageElement; - size: number; - type?: string; -}) => { - const canvas = document.createElement('canvas'); - const ctx = canvas.getContext('2d') as CanvasRenderingContext2D; - let startX = 0; - let startY = 0; - - if (img.width > img.height) { - startX = (img.width - img.height) / 2; - } else { - startY = (img.height - img.width) / 2; - } - - canvas.width = size; - canvas.height = size; - - ctx.drawImage( - img, - startX, - startY, - Math.min(img.width, img.height), - Math.min(img.width, img.height), - 0, - 0, - size, - size, - ); - - return canvas.toDataURL(type); -}; - -export const imageUrlToBase64 = async ( - imageUrl: string, -): Promise<{ base64: string; mimeType: string }> => { - try { - const res = await fetch(imageUrl); - const blob = await res.blob(); - const arrayBuffer = await blob.arrayBuffer(); - - const base64 = - typeof btoa === 'function' - ? btoa( - new Uint8Array(arrayBuffer).reduce( - (data, byte) => data + String.fromCharCode(byte), - '', - ), - ) - : Buffer.from(arrayBuffer).toString('base64'); - - return { base64, mimeType: blob.type }; - } catch (error) { - console.error('Error converting image to base64:', error); - throw error; - } -}; diff --git a/packages/ssrf-safe-fetch/index.browser.ts b/packages/ssrf-safe-fetch/index.browser.ts new file mode 100644 index 0000000000..e4e1d6d5bf --- /dev/null +++ b/packages/ssrf-safe-fetch/index.browser.ts @@ -0,0 +1,14 @@ +/** + * Browser version of SSRF-safe fetch + * In browser environments, we simply use the native fetch API + * as SSRF attacks are not applicable in client-side code + */ + +/** + * Browser-safe fetch implementation + * Uses native fetch API in browser environments + */ +// eslint-disable-next-line no-undef +export const ssrfSafeFetch = async (url: string, options?: RequestInit): Promise => { + return fetch(url, options); +}; diff --git a/packages/ssrf-safe-fetch/package.json b/packages/ssrf-safe-fetch/package.json index ec98738fc3..21e7576198 100644 --- a/packages/ssrf-safe-fetch/package.json +++ b/packages/ssrf-safe-fetch/package.json @@ -2,7 +2,14 @@ "name": "ssrf-safe-fetch", "version": "1.0.0", "private": true, - "description": "", + "description": "SSRF-safe fetch implementation with browser/node conditional exports", + "exports": { + ".": { + "browser": "./index.browser.ts", + "node": "./index.ts", + "default": "./index.ts" + } + }, "main": "index.ts", "scripts": { "test": "vitest run" diff --git a/packages/utils/src/imageToBase64.ts b/packages/utils/src/imageToBase64.ts index 80eaabe20d..44cf7d0221 100644 --- a/packages/utils/src/imageToBase64.ts +++ b/packages/utils/src/imageToBase64.ts @@ -36,23 +36,30 @@ export const imageToBase64 = ({ return canvas.toDataURL(type); }; +/** + * Convert image URL to base64 + * Uses SSRF-safe fetch on server-side to prevent SSRF attacks + */ export const imageUrlToBase64 = async ( imageUrl: string, ): Promise<{ base64: string; mimeType: string }> => { try { - const res = await fetch(imageUrl); + const isServer = typeof window === 'undefined'; + + // Use SSRF-safe fetch on server-side to prevent SSRF attacks + const res = isServer + ? await import('ssrf-safe-fetch').then((m) => m.ssrfSafeFetch(imageUrl)) + : await fetch(imageUrl); + const blob = await res.blob(); const arrayBuffer = await blob.arrayBuffer(); - const base64 = - typeof btoa === 'function' - ? btoa( - new Uint8Array(arrayBuffer).reduce( - (data, byte) => data + String.fromCharCode(byte), - '', - ), - ) - : Buffer.from(arrayBuffer).toString('base64'); + // Client-side uses btoa, server-side uses Buffer + const base64 = isServer + ? Buffer.from(arrayBuffer).toString('base64') + : btoa( + new Uint8Array(arrayBuffer).reduce((data, byte) => data + String.fromCharCode(byte), ''), + ); return { base64, mimeType: blob.type }; } catch (error) { diff --git a/packages/utils/src/index.ts b/packages/utils/src/index.ts index e3929a00dc..7c8e6c0491 100644 --- a/packages/utils/src/index.ts +++ b/packages/utils/src/index.ts @@ -4,9 +4,9 @@ export * from './detectChinese'; export * from './format'; export * from './imageToBase64'; export * from './keyboard'; +export * from './merge'; export * from './number'; export * from './object'; -export * from './parseModels'; export * from './pricing'; export * from './safeParseJSON'; export * from './sleep'; diff --git a/src/features/AgentSetting/AgentTTS/SelectWithTTSPreview.tsx b/src/features/AgentSetting/AgentTTS/SelectWithTTSPreview.tsx index b8e1769ea3..b43ac15140 100644 --- a/src/features/AgentSetting/AgentTTS/SelectWithTTSPreview.tsx +++ b/src/features/AgentSetting/AgentTTS/SelectWithTTSPreview.tsx @@ -1,3 +1,4 @@ +import { getMessageError } from '@lobechat/fetch-sse'; import { ChatMessageError } from '@lobechat/types'; import { AudioPlayer } from '@lobehub/tts/react'; import { Alert, Button, Highlighter, Select, SelectProps } from '@lobehub/ui'; @@ -9,7 +10,6 @@ import { Flexbox } from 'react-layout-kit'; import { useTTS } from '@/hooks/useTTS'; import { TTSServer } from '@/types/agent'; -import { getMessageError } from '@/utils/fetch'; interface SelectWithTTSPreviewProps extends SelectProps { server: TTSServer; diff --git a/src/features/AgentSetting/store/action.ts b/src/features/AgentSetting/store/action.ts index 120b926c2b..64267a62bb 100644 --- a/src/features/AgentSetting/store/action.ts +++ b/src/features/AgentSetting/store/action.ts @@ -1,3 +1,4 @@ +import { MessageTextChunk } from '@lobechat/fetch-sse'; import { chainPickEmoji, chainSummaryAgentName, @@ -16,7 +17,6 @@ import { systemAgentSelectors } from '@/store/user/slices/settings/selectors'; import { LobeAgentChatConfig, LobeAgentConfig } from '@/types/agent'; import { MetaData } from '@/types/meta'; import { SystemAgentItem } from '@/types/user/settings'; -import { MessageTextChunk } from '@/utils/fetch'; import { merge } from '@/utils/merge'; import { setNamespace } from '@/utils/storeDebug'; diff --git a/src/features/ChatInput/ActionBar/STT/browser.tsx b/src/features/ChatInput/ActionBar/STT/browser.tsx index ee68479578..5b919a8cce 100644 --- a/src/features/ChatInput/ActionBar/STT/browser.tsx +++ b/src/features/ChatInput/ActionBar/STT/browser.tsx @@ -1,3 +1,4 @@ +import { getMessageError } from '@lobechat/fetch-sse'; import { ChatMessageError } from '@lobechat/types'; import { SpeechRecognitionOptions, useSpeechRecognition } from '@lobehub/tts/react'; import isEqual from 'fast-deep-equal'; @@ -13,7 +14,6 @@ import { useGlobalStore } from '@/store/global'; import { globalGeneralSelectors } from '@/store/global/selectors'; import { useUserStore } from '@/store/user'; import { settingsSelectors } from '@/store/user/selectors'; -import { getMessageError } from '@/utils/fetch'; import CommonSTT from './common'; diff --git a/src/features/ChatInput/ActionBar/STT/openai.tsx b/src/features/ChatInput/ActionBar/STT/openai.tsx index 4349b1b7fb..8579954841 100644 --- a/src/features/ChatInput/ActionBar/STT/openai.tsx +++ b/src/features/ChatInput/ActionBar/STT/openai.tsx @@ -1,3 +1,4 @@ +import { getMessageError } from '@lobechat/fetch-sse'; import { ChatMessageError } from '@lobechat/types'; import { getRecordMineType } from '@lobehub/tts'; import { OpenAISTTOptions, useOpenAISTT } from '@lobehub/tts/react'; @@ -16,7 +17,6 @@ import { useGlobalStore } from '@/store/global'; import { globalGeneralSelectors } from '@/store/global/selectors'; import { useUserStore } from '@/store/user'; import { settingsSelectors } from '@/store/user/selectors'; -import { getMessageError } from '@/utils/fetch'; import CommonSTT from './common'; diff --git a/src/features/Conversation/components/Extras/TTS/InitPlayer.tsx b/src/features/Conversation/components/Extras/TTS/InitPlayer.tsx index 7261ee6178..111ac065d2 100644 --- a/src/features/Conversation/components/Extras/TTS/InitPlayer.tsx +++ b/src/features/Conversation/components/Extras/TTS/InitPlayer.tsx @@ -1,3 +1,4 @@ +import { getMessageError } from '@lobechat/fetch-sse'; import { ChatMessageError, ChatTTS } from '@lobechat/types'; import { memo, useCallback, useEffect, useState } from 'react'; import { useTranslation } from 'react-i18next'; @@ -5,7 +6,6 @@ import { useTranslation } from 'react-i18next'; import { useTTS } from '@/hooks/useTTS'; import { useChatStore } from '@/store/chat'; import { useFileStore } from '@/store/file'; -import { getMessageError } from '@/utils/fetch'; import Player from './Player'; diff --git a/src/server/globalConfig/genServerAiProviderConfig.test.ts b/src/server/globalConfig/genServerAiProviderConfig.test.ts index 060086a2e1..3398dced6a 100644 --- a/src/server/globalConfig/genServerAiProviderConfig.test.ts +++ b/src/server/globalConfig/genServerAiProviderConfig.test.ts @@ -19,7 +19,7 @@ vi.mock('@/envs/llm', () => ({ })), })); -vi.mock('@/utils/parseModels', () => ({ +vi.mock('@/utils/server/parseModels', () => ({ extractEnabledModels: vi.fn(async (providerId: string, modelString?: string) => { if (!modelString) return undefined; return [`${providerId}-model-1`, `${providerId}-model-2`]; @@ -98,7 +98,7 @@ describe('genServerAiProvidersConfig', () => { it('should use environment variables for model lists', async () => { process.env.OPENAI_MODEL_LIST = '+gpt-4,+gpt-3.5-turbo'; - const { extractEnabledModels } = vi.mocked(await import('@/utils/parseModels')); + const { extractEnabledModels } = vi.mocked(await import('@/utils/server/parseModels')); extractEnabledModels.mockResolvedValue(['gpt-4', 'gpt-3.5-turbo']); const result = await genServerAiProvidersConfig({}); @@ -116,7 +116,7 @@ describe('genServerAiProvidersConfig', () => { process.env.CUSTOM_OPENAI_MODELS = '+custom-model'; - const { extractEnabledModels } = vi.mocked(await import('@/utils/parseModels')); + const { extractEnabledModels } = vi.mocked(await import('@/utils/server/parseModels')); await genServerAiProvidersConfig(specificConfig); @@ -133,7 +133,7 @@ describe('genServerAiProvidersConfig', () => { process.env.OPENAI_MODEL_LIST = '+gpt-4->deployment1'; const { extractEnabledModels, transformToAiModelList } = vi.mocked( - await import('@/utils/parseModels'), + await import('@/utils/server/parseModels'), ); await genServerAiProvidersConfig(specificConfig); @@ -206,7 +206,7 @@ describe('genServerAiProvidersConfig Error Handling', () => { getLLMConfig: vi.fn(() => ({})), })); - vi.doMock('@/utils/parseModels', () => ({ + vi.doMock('@/utils/server/parseModels', () => ({ extractEnabledModels: vi.fn(async () => undefined), transformToAiModelList: vi.fn(async () => []), })); diff --git a/src/server/globalConfig/genServerAiProviderConfig.ts b/src/server/globalConfig/genServerAiProviderConfig.ts index 9d30e366db..53dd1c6ad3 100644 --- a/src/server/globalConfig/genServerAiProviderConfig.ts +++ b/src/server/globalConfig/genServerAiProviderConfig.ts @@ -1,9 +1,9 @@ import { ProviderConfig } from '@lobechat/types'; -import { extractEnabledModels, transformToAiModelList } from '@lobechat/utils'; import { AiFullModelCard, ModelProvider } from 'model-bank'; import * as AiModels from 'model-bank'; import { getLLMConfig } from '@/envs/llm'; +import { extractEnabledModels, transformToAiModelList } from '@/utils/server/parseModels'; interface ProviderSpecificConfig { enabled?: boolean; diff --git a/src/services/chat/chat.test.ts b/src/services/chat/chat.test.ts index ae19f107e9..5aee25326a 100644 --- a/src/services/chat/chat.test.ts +++ b/src/services/chat/chat.test.ts @@ -27,7 +27,7 @@ vi.stubGlobal( ); // Mock image processing utilities -vi.mock('@/utils/fetch', async (importOriginal) => { +vi.mock('@lobechat/fetch-sse', async (importOriginal) => { const module = await importOriginal(); return { ...(module as any), getMessageError: vi.fn() }; @@ -988,7 +988,7 @@ describe('ChatService', () => { beforeEach(async () => { // Setup common fetchSSE mock for getChatCompletion tests - const { fetchSSE } = await import('@/utils/fetch'); + const { fetchSSE } = await import('@lobechat/fetch-sse'); mockFetchSSE = vi.fn().mockResolvedValue(new Response('mock response')); vi.mocked(fetchSSE).mockImplementation(mockFetchSSE); }); @@ -1049,7 +1049,7 @@ describe('ChatService', () => { it('should return InvalidAccessCode error when enableFetchOnClient is true and auth is enabled but user is not signed in', async () => { // Mock fetchSSE to call onErrorHandle with the error - const { fetchSSE } = await import('@/utils/fetch'); + const { fetchSSE } = await import('@lobechat/fetch-sse'); const mockFetchSSEWithError = vi.fn().mockImplementation((url, options) => { // Simulate the error being caught and passed to onErrorHandle @@ -1211,8 +1211,8 @@ vi.mock('../_auth', async (importOriginal) => { describe('ChatService private methods', () => { describe('getChatCompletion', () => { it('should merge responseAnimation styles correctly', async () => { - const { fetchSSE } = await import('@/utils/fetch'); - vi.mock('@/utils/fetch', async (importOriginal) => { + const { fetchSSE } = await import('@lobechat/fetch-sse'); + vi.mock('@lobechat/fetch-sse', async (importOriginal) => { const module = await importOriginal(); return { ...(module as any), diff --git a/src/services/chat/clientModelRuntime.test.ts b/src/services/chat/clientModelRuntime.test.ts index e90bc2d549..39748d7210 100644 --- a/src/services/chat/clientModelRuntime.test.ts +++ b/src/services/chat/clientModelRuntime.test.ts @@ -38,7 +38,7 @@ vi.stubGlobal( vi.fn(() => Promise.resolve(new Response(JSON.stringify({ some: 'data' })))), ); -vi.mock('@/utils/fetch', async (importOriginal) => { +vi.mock('@lobechat/fetch-sse', async (importOriginal) => { const module = await importOriginal(); return { ...(module as any), getMessageError: vi.fn() }; diff --git a/src/services/chat/index.ts b/src/services/chat/index.ts index f7f825c3c5..5ad7d258f1 100644 --- a/src/services/chat/index.ts +++ b/src/services/chat/index.ts @@ -1,3 +1,9 @@ +import { + FetchSSEOptions, + fetchSSE, + getMessageError, + standardizeAnimationStyle, +} from '@lobechat/fetch-sse'; import { AgentRuntimeError, ChatCompletionErrorPayload } from '@lobechat/model-runtime'; import { ChatErrorType, TracePayload, TraceTagMap, UIChatMessage } from '@lobechat/types'; import { PluginRequestPayload, createHeadersWithPluginSettings } from '@lobehub/chat-plugin-sdk'; @@ -25,12 +31,6 @@ import { import type { ChatStreamPayload, OpenAIChatMessage } from '@/types/openai/chat'; import { fetchWithInvokeStream } from '@/utils/electron/desktopRemoteRPCFetch'; import { createErrorResponse } from '@/utils/errorResponse'; -import { - FetchSSEOptions, - fetchSSE, - getMessageError, - standardizeAnimationStyle, -} from '@/utils/fetch'; import { createTraceHeader, getTraceId } from '@/utils/trace'; import { createHeaderWithAuth } from '../_auth'; diff --git a/src/services/chat/types.ts b/src/services/chat/types.ts index b12834024e..00391f3e6d 100644 --- a/src/services/chat/types.ts +++ b/src/services/chat/types.ts @@ -1,7 +1,6 @@ +import { FetchSSEOptions } from '@lobechat/fetch-sse'; import { TracePayload } from '@lobechat/types'; -import { FetchSSEOptions } from '@/utils/fetch'; - export interface FetchOptions extends FetchSSEOptions { historySummary?: string; signal?: AbortSignal | undefined; diff --git a/src/services/models.ts b/src/services/models.ts index 631375be37..62292c2bce 100644 --- a/src/services/models.ts +++ b/src/services/models.ts @@ -1,7 +1,8 @@ +import { getMessageError } from '@lobechat/fetch-sse'; + import { createHeaderWithAuth } from '@/services/_auth'; import { aiProviderSelectors, getAiInfraStoreState } from '@/store/aiInfra'; import { ChatModelCard } from '@/types/llm'; -import { getMessageError } from '@/utils/fetch'; import { API_ENDPOINTS } from './_url'; import { initializeWithClientStore } from './chat/clientModelRuntime'; diff --git a/packages/utils/src/electron/desktopRemoteRPCFetch.ts b/src/utils/electron/desktopRemoteRPCFetch.ts similarity index 97% rename from packages/utils/src/electron/desktopRemoteRPCFetch.ts rename to src/utils/electron/desktopRemoteRPCFetch.ts index 3121f49f9c..e2868fa37d 100644 --- a/packages/utils/src/electron/desktopRemoteRPCFetch.ts +++ b/src/utils/electron/desktopRemoteRPCFetch.ts @@ -1,10 +1,10 @@ import { isDesktop } from '@lobechat/const'; import { ProxyTRPCRequestParams, dispatch, streamInvoke } from '@lobechat/electron-client-ipc'; +import { getRequestBody, headersToRecord } from '@lobechat/fetch-sse'; import debug from 'debug'; import { getElectronStoreState } from '@/store/electron'; import { electronSyncSelectors } from '@/store/electron/selectors'; -import { getRequestBody, headersToRecord } from '@/utils/fetch'; const log = debug('utils:desktopRemoteRPCFetch'); diff --git a/packages/utils/src/__snapshots__/parseModels.test.ts.snap b/src/utils/server/__snapshots__/parseModels.test.ts.snap similarity index 100% rename from packages/utils/src/__snapshots__/parseModels.test.ts.snap rename to src/utils/server/__snapshots__/parseModels.test.ts.snap diff --git a/packages/utils/src/parseModels.test.ts b/src/utils/server/parseModels.test.ts similarity index 100% rename from packages/utils/src/parseModels.test.ts rename to src/utils/server/parseModels.test.ts diff --git a/packages/utils/src/parseModels.ts b/src/utils/server/parseModels.ts similarity index 99% rename from packages/utils/src/parseModels.ts rename to src/utils/server/parseModels.ts index b13285977d..50f0338660 100644 --- a/packages/utils/src/parseModels.ts +++ b/src/utils/server/parseModels.ts @@ -1,9 +1,8 @@ import { getModelPropertyWithFallback } from '@lobechat/model-runtime'; +import { merge } from '@lobechat/utils'; import { produce } from 'immer'; import { AiFullModelCard, AiModelType } from 'model-bank'; -import { merge } from './merge'; - /** * Parse model string to add or remove models. */ diff --git a/vitest.config.mts b/vitest.config.mts index 57c3271aa9..1f06a9b388 100644 --- a/vitest.config.mts +++ b/vitest.config.mts @@ -16,6 +16,8 @@ export default defineConfig({ // TODO: after refactor the errorResponse, we can remove it '@/utils/errorResponse': resolve(__dirname, './src/utils/errorResponse'), '@/utils/unzipFile': resolve(__dirname, './src/utils/unzipFile'), + '@/utils/server': resolve(__dirname, './src/utils/server'), + '@/utils/electron': resolve(__dirname, './src/utils/electron'), '@/utils': resolve(__dirname, './packages/utils/src'), '@/types': resolve(__dirname, './packages/types/src'), '@/const': resolve(__dirname, './packages/const/src'),