forked from wangziqi/gongxue-base
259 lines
8.9 KiB
TypeScript
259 lines
8.9 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 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;
|
||
|
||
// 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 timeout = AbortSignal.timeout(config.timeoutMs);
|
||
const combinedSignal = AbortSignal.any([signal, timeout]);
|
||
let response: PinnedResponse;
|
||
|
||
try {
|
||
response = await this.pinnedPost(
|
||
`${config.baseUrl.replace(/\/$/, '')}/chat/completions`,
|
||
{
|
||
Authorization: `Bearer ${config.apiKey}`,
|
||
'Content-Type': 'application/json',
|
||
Accept: 'text/event-stream',
|
||
},
|
||
JSON.stringify({
|
||
model: config.defaultModel,
|
||
messages,
|
||
stream: true,
|
||
...(tools.length ? { tools, tool_choice: 'auto' } : {}),
|
||
}),
|
||
combinedSignal,
|
||
);
|
||
} catch (error) {
|
||
if (timeout.aborted && !signal.aborted) throw new RequestTimeoutException('AI 服务响应超时');
|
||
throw error;
|
||
}
|
||
|
||
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 (timeout.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 safeUpstreamMessage(status: number, body: string): string {
|
||
if (status === 401 || status === 403) return 'AI 服务认证失败';
|
||
if (status === 429) 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;
|
||
}
|
||
}
|