330 lines
12 KiB
TypeScript
330 lines
12 KiB
TypeScript
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<string, unknown>;
|
||
};
|
||
}
|
||
|
||
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<ModelStreamEvent> {
|
||
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<number, { id: string; name: string; arguments: string }>();
|
||
|
||
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<Uint8Array>) {
|
||
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<void> {
|
||
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<string, string>,
|
||
body: string,
|
||
signal: AbortSignal,
|
||
): Promise<PinnedResponse> {
|
||
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<string> {
|
||
const chunks: Uint8Array[] = [];
|
||
let total = 0;
|
||
for await (const value of body as AsyncIterable<Uint8Array>) {
|
||
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;
|
||
}
|
||
}
|