import { BadGatewayException, Injectable, RequestTimeoutException } from '@nestjs/common'; import { lookup } from 'node:dns'; import * as http from 'node:http'; import * as https from 'node:https'; import { isIP } from 'node:net'; import type { AiRuntimeConfig } from '../ai-config/dto/ai-config.dto'; import { AiProvider } from '../ai-config/ai-config.entity'; import type { ModelMessage, ModelStreamEvent } from './ai-chat.types'; interface ChatTool { type: 'function'; function: { name: string; description: string; parameters: Record; }; } interface StreamChoiceDelta { content?: string | null; reasoning_content?: string | null; tool_calls?: Array<{ index?: number; id?: string; function?: { name?: string; arguments?: string }; }>; } const MAX_UPSTREAM_EVENT_BYTES = 1024 * 1024; const MAX_UPSTREAM_RETRIES = 3; const UPSTREAM_RETRY_DELAYS_MS = [500, 1000, 2000]; const RETRYABLE_STATUS_CODES = new Set([408, 425, 429, 500, 502, 503, 504]); const RETRYABLE_TRANSPORT_CODES = new Set([ 'ECONNRESET', 'ECONNREFUSED', 'ETIMEDOUT', 'ENOTFOUND', 'EAI_AGAIN', 'EPIPE', ]); // Known public provider hosts — trusted even if CDN resolves to private-range IPs const DNS_TRUSTED_HOSTS = new Set(['api.openai.com', 'api.deepseek.com']); const PRIVATE_IPV4_RANGES = [ /^127\./, /^10\./, /^172\.(1[6-9]|2\d|3[01])\./, /^192\.168\./, /^169\.254\./, /^0\./, /^100\.(6[4-9]|[7-9]\d|1[01]\d|12[0-7])\./, ]; interface PinnedResponse { status: number; contentType: string; body: http.IncomingMessage; } @Injectable() export class AiModelStreamService { async *stream( config: AiRuntimeConfig, messages: ModelMessage[], tools: ChatTool[], signal: AbortSignal, ): AsyncGenerator { const requestBody = JSON.stringify({ model: config.defaultModel, messages, stream: true, ...(tools.length ? { tools, tool_choice: 'auto' } : {}), // reasoning_effort 仅对支持该参数的 OpenAI 兼容服务生效; // DeepSeek 官方接口不接受该参数,避免请求被拒。 ...(config.reasoningEffort && config.reasoningEffort !== 'none' && config.provider !== AiProvider.DEEPSEEK ? { reasoning_effort: config.reasoningEffort } : {}), }); const url = `${config.baseUrl.replace(/\/$/, '')}/chat/completions`; const headers = { Authorization: `Bearer ${config.apiKey}`, 'Content-Type': 'application/json', Accept: 'text/event-stream', }; let response: PinnedResponse | null = null; let activeTimeout: AbortSignal | undefined; for (let attempt = 0; attempt <= MAX_UPSTREAM_RETRIES; attempt += 1) { activeTimeout = AbortSignal.timeout(config.timeoutMs); const combinedSignal = AbortSignal.any([signal, activeTimeout]); try { response = await this.pinnedPost(url, headers, requestBody, combinedSignal); } catch (error) { if (activeTimeout.aborted && !signal.aborted) { throw new RequestTimeoutException('AI 服务响应超时'); } if ( attempt < MAX_UPSTREAM_RETRIES && !signal.aborted && this.isRetryableTransportError(error) ) { const delayMs = UPSTREAM_RETRY_DELAYS_MS[attempt]; yield { type: 'retrying', attempt: attempt + 1, maxRetries: MAX_UPSTREAM_RETRIES, delayMs, reason: error instanceof Error ? error.message : '网络连接失败', }; await this.sleep(delayMs); continue; } throw error; } if (response.status >= 200 && response.status < 300) break; if (attempt < MAX_UPSTREAM_RETRIES && RETRYABLE_STATUS_CODES.has(response.status)) { response.body.resume?.(); const delayMs = UPSTREAM_RETRY_DELAYS_MS[attempt]; yield { type: 'retrying', attempt: attempt + 1, maxRetries: MAX_UPSTREAM_RETRIES, delayMs, reason: `上游返回 ${response.status}`, }; await this.sleep(delayMs); continue; } const body = await this.readLimitedBody(response.body); throw new BadGatewayException(this.safeUpstreamMessage(response.status, body)); } if (!response) throw new BadGatewayException('AI 服务暂时不可用'); if (response.status < 200 || response.status >= 300) { const body = await this.readLimitedBody(response.body); throw new BadGatewayException(this.safeUpstreamMessage(response.status, body)); } const contentType = response.contentType.toLowerCase(); if (!contentType.includes('text/event-stream')) { throw new BadGatewayException('AI 服务返回了无效的响应格式'); } const decoder = new TextDecoder(); let buffer = ''; const calls = new Map(); const consumeEvent = (event: string): ModelStreamEvent[] => { const output: ModelStreamEvent[] = []; const data = event .split(/\r?\n/) .filter((line) => line.startsWith('data:')) .map((line) => line.slice(5).trimStart()) .join('\n'); if (!data || data === '[DONE]') return output; const parsed = this.parseEvent(data); const delta = parsed.choices?.[0]?.delta; if (!delta) return output; if (delta.reasoning_content) output.push({ type: 'reasoning', delta: delta.reasoning_content }); if (delta.content) output.push({ type: 'content', delta: delta.content }); for (const part of delta.tool_calls ?? []) { const index = part.index ?? 0; const current = calls.get(index) ?? { id: '', name: '', arguments: '' }; if (part.id) current.id += part.id; if (part.function?.name) current.name += part.function.name; if (part.function?.arguments) current.arguments += part.function.arguments; calls.set(index, current); } return output; }; try { for await (const chunk of response.body as AsyncIterable) { buffer += decoder.decode(chunk, { stream: true }); if (buffer.length > MAX_UPSTREAM_EVENT_BYTES) { throw new BadGatewayException('AI 服务返回的单个事件过大'); } const events = buffer.split(/\r?\n\r?\n/); buffer = events.pop() ?? ''; for (const event of events) for (const parsed of consumeEvent(event)) yield parsed; } buffer += decoder.decode(); if (buffer.trim()) for (const parsed of consumeEvent(buffer)) yield parsed; } catch (error) { if (activeTimeout?.aborted && !signal.aborted) { throw new RequestTimeoutException('AI 服务响应超时'); } throw error; } yield { type: 'complete', toolCalls: [...calls.entries()] .sort(([a], [b]) => a - b) .map(([, call], index) => ({ id: call.id || `call_${index}`, name: call.name, arguments: call.arguments || '{}', })), }; } private parseEvent(data: string): { choices?: Array<{ delta?: StreamChoiceDelta }> } { try { const value: unknown = JSON.parse(data); if (!value || typeof value !== 'object') throw new Error('invalid'); return value; } catch { throw new BadGatewayException('AI 服务返回了无效的流式数据'); } } private isRetryableTransportError(error: unknown): boolean { if (!(error instanceof Error)) return false; const code = (error as NodeJS.ErrnoException).code; if (code && RETRYABLE_TRANSPORT_CODES.has(code)) return true; return /socket hang up|ECONNRESET|ETIMEDOUT|EAI_AGAIN/i.test(error.message); } private sleep(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); } private safeUpstreamMessage(status: number, body: string): string { if (status === 401 || status === 403) return 'AI 服务认证失败'; if (status === 429) return 'AI 服务请求过于频繁'; if (status === 503) return 'AI 服务繁忙,请稍后重试'; if (status >= 500) return 'AI 服务暂时不可用'; const message = this.extractErrorMessage(body); return message ? `AI 服务请求失败:${message}` : `AI 服务请求失败(${status})`; } private extractErrorMessage(body: string): string | null { try { const parsed = JSON.parse(body) as { error?: { message?: unknown } }; const message = parsed.error?.message; return typeof message === 'string' ? message.slice(0, 200) : null; } catch { return null; } } private pinnedPost( url: string, headers: Record, body: string, signal: AbortSignal, ): Promise { return new Promise((resolve, reject) => { const parsed = new URL(url); const isHttps = parsed.protocol === 'https:'; const port = parsed.port ? Number(parsed.port) : isHttps ? 443 : 80; lookup(parsed.hostname, { all: true, family: 0 }, (dnsError, addresses) => { if (dnsError || !addresses?.length) return reject(new Error('DNS 解析失败')); const allowPrivate = process.env.AI_ALLOW_PRIVATE_BASE_URL === 'true' || DNS_TRUSTED_HOSTS.has(parsed.hostname); if (!allowPrivate && addresses.some(({ address }) => this.isPrivateAddress(address))) { return reject(new Error('域名解析到内网地址')); } const target = addresses[0]; const transport = isHttps ? https : http; const request = transport.request( { hostname: target.address, port, path: parsed.pathname + parsed.search, method: 'POST', headers: { ...headers, Host: parsed.hostname, 'Content-Length': Buffer.byteLength(body).toString(), }, servername: isHttps ? parsed.hostname : undefined, rejectUnauthorized: isHttps, family: target.family === 6 ? 6 : 4, signal, }, (response) => { const status = response.statusCode ?? 500; if (status >= 300 && status < 400) { response.resume(); response.destroy(); reject(new Error('禁止重定向')); return; } resolve({ status, contentType: String(response.headers['content-type'] ?? ''), body: response, }); }, ); request.once('error', reject); request.end(body); }); }); } private async readLimitedBody(body: http.IncomingMessage): Promise { const chunks: Uint8Array[] = []; let total = 0; for await (const value of body as AsyncIterable) { total += value.length; if (total > MAX_UPSTREAM_EVENT_BYTES) { body.destroy(); return ''; } chunks.push(value); } return Buffer.concat(chunks).toString('utf8'); } private isPrivateAddress(rawAddress: string): boolean { const address = rawAddress.toLowerCase(); if (isIP(address) === 4) return PRIVATE_IPV4_RANGES.some((range) => range.test(address)); if (isIP(address) !== 6) return true; if (address === '::1' || address === '::') return true; if (address.startsWith('fc') || address.startsWith('fd')) return true; if (/^fe[89ab]/.test(address)) return true; if (address.startsWith('::ffff:') && isIP(address.slice(7)) === 4) { return PRIVATE_IPV4_RANGES.some((range) => range.test(address.slice(7))); } return false; } }