| 1 | import { HumanMessage } from '@langchain/core/messages' |
| 2 | import { isSupportedImageMimeType, normalizeImageMimeType } from '@shared/image-mime' |
| 3 | import { resolveModelTimeoutMs } from '@shared/model-timeout' |
| 4 | import log from 'electron-log/main.js' |
| 5 | import { extractModelText } from '../../model/result' |
| 6 | import { resolveModel } from '../../model/resolve' |
| 7 | import type { ModelRuntimeConfig } from '../../model/usage' |
| 8 | |
| 9 | const combineAbortSignals = (signals: AbortSignal[]): { signal: AbortSignal; dispose: () => void } => { |
| 10 | if (signals.length === 1) return { signal: signals[0], dispose: () => undefined } |
| 11 | |
| 12 | const controller = new AbortController() |
| 13 | const abort = (signal: AbortSignal): void => controller.abort(signal.reason) |
| 14 | const listeners = signals.map((signal) => { |
| 15 | const listener = (): void => abort(signal) |
| 16 | if (signal.aborted) listener() |
| 17 | else signal.addEventListener('abort', listener, { once: true }) |
| 18 | return { signal, listener } |
| 19 | }) |
| 20 | return { |
| 21 | signal: controller.signal, |
| 22 | dispose: () => { |
| 23 | for (const { signal, listener } of listeners) { |
| 24 | signal.removeEventListener('abort', listener) |
| 25 | } |
| 26 | } |
| 27 | } |
| 28 | } |
| 29 | |
| 30 | export async function invokeVisionModelText(args: { |
| 31 | imageBase64: string |
| 32 | mimeType: string |
| 33 | prompt: string |
| 34 | provider: string |
| 35 | apiKey: string |
| 36 | model: string |
| 37 | baseUrl: string |
| 38 | maxTokens?: number |
| 39 | modelRuntime?: ModelRuntimeConfig |
| 40 | modelTimeoutMs: number |
| 41 | logTag: string |
| 42 | signal?: AbortSignal |
| 43 | }): Promise<string> { |
| 44 | const mimeType = normalizeImageMimeType(args.mimeType) |
| 45 | const imageBase64 = String(args.imageBase64 || '').trim() |
| 46 | if (!isSupportedImageMimeType(args.mimeType)) { |
| 47 | throw new Error(`不支持的图片格式：${mimeType || 'unknown'}`) |
| 48 | } |
| 49 | if (!imageBase64) { |
| 50 | throw new Error('图片数据为空') |
| 51 | } |
| 52 | |
| 53 | const imageBytes = Buffer.byteLength(imageBase64, 'base64') |
| 54 | log.info(`[${args.logTag}] invoke vision model`, { |
| 55 | provider: args.provider, |
| 56 | model: args.model, |
| 57 | mimeType, |
| 58 | imageBytes |
| 59 | }) |
| 60 | |
| 61 | const model = resolveModel( |
| 62 | args.provider, |
| 63 | args.apiKey, |
| 64 | args.model, |
| 65 | args.baseUrl, |
| 66 | 0.2, |
| 67 | args.maxTokens, |
| 68 | args.modelRuntime |
| 69 | ) |
| 70 | const imageUrl = `data:${mimeType};base64,${imageBase64}` |
| 71 | const timeoutSignal = AbortSignal.timeout(resolveModelTimeoutMs(args.modelTimeoutMs, 'document')) |
| 72 | const combinedSignal = combineAbortSignals( |
| 73 | args.signal ? [args.signal, timeoutSignal] : [timeoutSignal] |
| 74 | ) |
| 75 | try { |
| 76 | const result = await model.invoke( |
| 77 | [ |
| 78 | new HumanMessage({ |
| 79 | content: [ |
| 80 | { type: 'text', text: args.prompt }, |
| 81 | { type: 'image_url', image_url: { url: imageUrl } } |
| 82 | ] |
| 83 | }) |
| 84 | ], |
| 85 | { signal: combinedSignal.signal } |
| 86 | ) |
| 87 | return extractModelText(result) |
| 88 | } finally { |
| 89 | combinedSignal.dispose() |
| 90 | } |
| 91 | } |
| 92 |