forked from wangziqi/gongxue-base
feat: 集成 AI 对话与只读查询工具
This commit is contained in:
252
apps/server/src/ai-chat/ai-model-stream.service.ts
Normal file
252
apps/server/src/ai-chat/ai-model-stream.service.ts
Normal file
@@ -0,0 +1,252 @@
|
||||
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;
|
||||
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';
|
||||
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;
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user