322 lines
12 KiB
TypeScript
322 lines
12 KiB
TypeScript
import {
|
|
BadRequestException,
|
|
Injectable,
|
|
NotFoundException,
|
|
} from '@nestjs/common';
|
|
import { InjectRepository } from '@nestjs/typeorm';
|
|
import ExcelJS from 'exceljs';
|
|
import { createReadStream } from 'node:fs';
|
|
import { mkdir, readFile, unlink, writeFile } from 'node:fs/promises';
|
|
import { basename, isAbsolute, join, relative, resolve } from 'node:path';
|
|
import { randomUUID } from 'node:crypto';
|
|
import { PDFParse } from 'pdf-parse';
|
|
import { In, Repository } from 'typeorm';
|
|
import { AiAttachment } from './entities';
|
|
|
|
const MAX_FILE_BYTES = 10 * 1024 * 1024;
|
|
const MAX_EXTRACTED_CHARS = 48 * 1024;
|
|
const MAX_MODEL_IMAGE_BYTES = 20 * 1024 * 1024;
|
|
const ACCEPTED_MIME_TYPES = new Set([
|
|
'image/jpeg',
|
|
'image/png',
|
|
'image/webp',
|
|
'application/pdf',
|
|
'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
|
|
'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
|
|
]);
|
|
|
|
interface MammothResult {
|
|
value: string;
|
|
}
|
|
|
|
interface MammothModule {
|
|
extractRawText(input: { buffer: Buffer }): Promise<MammothResult>;
|
|
}
|
|
|
|
export interface AiAttachmentModelPart {
|
|
attachment: AiAttachment;
|
|
text?: string;
|
|
imageDataUrl?: string;
|
|
}
|
|
|
|
@Injectable()
|
|
export class AiAttachmentService {
|
|
private readonly storageRoot =
|
|
resolve(process.env.AI_ATTACHMENT_DIR || join(process.cwd(), 'data', 'ai-attachments'));
|
|
|
|
constructor(
|
|
@InjectRepository(AiAttachment)
|
|
private readonly attachments: Repository<AiAttachment>,
|
|
) {}
|
|
|
|
async upload(userId: number, file: Express.Multer.File): Promise<AiAttachment> {
|
|
if (!file?.buffer?.length) throw new BadRequestException('请选择附件');
|
|
if (file.size > MAX_FILE_BYTES) throw new BadRequestException('单个附件不能超过 10MB');
|
|
|
|
const mimeType = this.detectMimeType(file.buffer, file.mimetype);
|
|
if (!ACCEPTED_MIME_TYPES.has(mimeType)) {
|
|
throw new BadRequestException('仅支持图片、PDF、Word 和 Excel 文件');
|
|
}
|
|
this.assertDeclaredType(file.mimetype, mimeType);
|
|
this.assertFileExtension(file.originalname, mimeType);
|
|
|
|
await mkdir(this.storageRoot, { recursive: true });
|
|
const extension = this.extensionForMime(mimeType);
|
|
const storageKey = `${userId}/${randomUUID()}.${extension}`;
|
|
const absolutePath = this.resolveStoragePath(storageKey);
|
|
await mkdir(join(this.storageRoot, String(userId)), { recursive: true });
|
|
await writeFile(absolutePath, file.buffer, { flag: 'wx' });
|
|
|
|
let entity: AiAttachment;
|
|
try {
|
|
entity = await this.attachments.save(
|
|
this.attachments.create({
|
|
userId,
|
|
originalName: basename(file.originalname).slice(0, 255),
|
|
mimeType,
|
|
size: file.size,
|
|
storageKey,
|
|
processingStatus: 'processing',
|
|
extractedText: null,
|
|
processingError: null,
|
|
imageWidth: null,
|
|
imageHeight: null,
|
|
}),
|
|
);
|
|
} catch (error) {
|
|
await unlink(absolutePath).catch(() => undefined);
|
|
throw error;
|
|
}
|
|
|
|
try {
|
|
entity.extractedText = await this.extractText(file.buffer, mimeType);
|
|
entity.processingStatus = 'ready';
|
|
} catch {
|
|
entity.processingStatus = 'failed';
|
|
entity.processingError = '文件内容解析失败';
|
|
}
|
|
entity = await this.attachments.save(entity);
|
|
return entity;
|
|
}
|
|
|
|
async removeUnbound(userId: number, id: number): Promise<void> {
|
|
const attachment = await this.requireOwned(userId, id, true);
|
|
if (attachment.messages?.length) throw new BadRequestException('已发送的附件不能单独删除');
|
|
await this.attachments.remove(attachment);
|
|
await unlink(this.resolveStoragePath(attachment.storageKey)).catch(() => undefined);
|
|
}
|
|
|
|
async removeOrphans(userId: number, ids: number[]): Promise<void> {
|
|
const uniqueIds = [...new Set(ids)].filter((id) => Number.isInteger(id) && id > 0);
|
|
if (!uniqueIds.length) return;
|
|
const attachments = await this.attachments.find({
|
|
where: { id: In(uniqueIds), userId },
|
|
relations: { messages: true },
|
|
});
|
|
for (const attachment of attachments) {
|
|
if (attachment.messages?.length) continue;
|
|
await this.attachments.remove(attachment);
|
|
await unlink(this.resolveStoragePath(attachment.storageKey)).catch(() => undefined);
|
|
}
|
|
}
|
|
|
|
async open(userId: number, id: number): Promise<{
|
|
attachment: AiAttachment;
|
|
stream: ReturnType<typeof createReadStream>;
|
|
}> {
|
|
const attachment = await this.requireOwned(userId, id);
|
|
return {
|
|
attachment,
|
|
stream: createReadStream(this.resolveStoragePath(attachment.storageKey)),
|
|
};
|
|
}
|
|
|
|
async requireReadyOwned(userId: number, ids: number[]): Promise<AiAttachment[]> {
|
|
const uniqueIds = [...new Set(ids)];
|
|
if (uniqueIds.length > 5) throw new BadRequestException('每条消息最多添加 5 个附件');
|
|
if (!uniqueIds.length) return [];
|
|
const attachments = await this.attachments.findByIds(uniqueIds);
|
|
if (attachments.length !== uniqueIds.length || attachments.some((item) => item.userId !== userId)) {
|
|
throw new BadRequestException('附件不存在或无权访问');
|
|
}
|
|
if (attachments.some((item) => item.processingStatus !== 'ready')) {
|
|
throw new BadRequestException('附件仍在处理或处理失败');
|
|
}
|
|
return uniqueIds.map((id) => attachments.find((item) => item.id === id)!);
|
|
}
|
|
|
|
async toModelParts(
|
|
attachments: AiAttachment[],
|
|
supportsVision: boolean,
|
|
): Promise<AiAttachmentModelPart[]> {
|
|
const imageAttachments = attachments.filter((attachment) => attachment.mimeType.startsWith('image/'));
|
|
if (imageAttachments.length && !supportsVision) {
|
|
throw new BadRequestException('当前模型未启用图片理解能力');
|
|
}
|
|
const imageBytes = imageAttachments.reduce((total, attachment) => total + attachment.size, 0);
|
|
if (imageBytes > MAX_MODEL_IMAGE_BYTES) {
|
|
throw new BadRequestException('单次消息图片总大小不能超过 20MB');
|
|
}
|
|
const parts: AiAttachmentModelPart[] = [];
|
|
for (const attachment of attachments) {
|
|
if (attachment.mimeType.startsWith('image/')) {
|
|
const buffer = await readFile(this.resolveStoragePath(attachment.storageKey));
|
|
parts.push({
|
|
attachment,
|
|
imageDataUrl: `data:${attachment.mimeType};base64,${buffer.toString('base64')}`,
|
|
});
|
|
} else {
|
|
parts.push({
|
|
attachment,
|
|
text: attachment.extractedText?.slice(0, MAX_EXTRACTED_CHARS) || '',
|
|
});
|
|
}
|
|
}
|
|
return parts;
|
|
}
|
|
|
|
serialize(attachment: AiAttachment): Record<string, unknown> {
|
|
return {
|
|
id: attachment.id,
|
|
name: attachment.originalName,
|
|
mimeType: attachment.mimeType,
|
|
size: attachment.size,
|
|
status: attachment.processingStatus,
|
|
error: attachment.processingError,
|
|
url: `/api/ai/chat/attachments/${attachment.id}`,
|
|
createdAt: attachment.createdAt,
|
|
};
|
|
}
|
|
|
|
private async requireOwned(
|
|
userId: number,
|
|
id: number,
|
|
includeMessages = false,
|
|
): Promise<AiAttachment> {
|
|
const attachment = await this.attachments.findOne({
|
|
where: { id, userId },
|
|
...(includeMessages ? { relations: { messages: true } } : {}),
|
|
});
|
|
if (!attachment) throw new NotFoundException('附件不存在');
|
|
return attachment;
|
|
}
|
|
|
|
private async extractText(buffer: Buffer, mimeType: string): Promise<string | null> {
|
|
if (mimeType.startsWith('image/')) return null;
|
|
if (mimeType === 'application/pdf') {
|
|
const parser = new PDFParse({ data: buffer });
|
|
try {
|
|
const result = await parser.getText();
|
|
return this.normalizeExtractedText(result.text);
|
|
} finally {
|
|
await parser.destroy();
|
|
}
|
|
}
|
|
if (mimeType.includes('wordprocessingml')) {
|
|
const mammoth = (await import('mammoth')) as unknown as MammothModule;
|
|
const result = await mammoth.extractRawText({ buffer });
|
|
return this.normalizeExtractedText(result.value);
|
|
}
|
|
if (mimeType.includes('spreadsheetml')) {
|
|
const workbook = new ExcelJS.Workbook();
|
|
await workbook.xlsx.load(buffer as unknown as ExcelJS.Buffer);
|
|
const lines: string[] = [];
|
|
workbook.eachSheet((sheet) => {
|
|
lines.push(`# ${sheet.name}`);
|
|
sheet.eachRow((row) => {
|
|
const values = Array.isArray(row.values) ? row.values.slice(1) : [];
|
|
lines.push(values.map((value) => this.stringifyCellValue(value)).join('\t'));
|
|
});
|
|
});
|
|
return this.normalizeExtractedText(lines.join('\n'));
|
|
}
|
|
return null;
|
|
}
|
|
|
|
private normalizeExtractedText(value: string): string {
|
|
return value.split('\u0000').join('').replace(/\r\n/g, '\n').trim().slice(0, MAX_EXTRACTED_CHARS);
|
|
}
|
|
|
|
private stringifyCellValue(value: unknown): string {
|
|
if (value === null || value === undefined) return '';
|
|
if (value instanceof Date) return value.toISOString();
|
|
if (typeof value === 'string' || typeof value === 'number' || typeof value === 'boolean') {
|
|
return String(value);
|
|
}
|
|
try {
|
|
return JSON.stringify(value) || '';
|
|
} catch {
|
|
return '';
|
|
}
|
|
}
|
|
|
|
private assertDeclaredType(declared: string, detected: string): void {
|
|
if (!declared || declared === 'application/octet-stream') return;
|
|
if (declared !== detected) throw new BadRequestException('附件类型与文件内容不一致');
|
|
}
|
|
|
|
private assertFileExtension(filename: string, mimeType: string): void {
|
|
const extension = basename(filename).toLowerCase().split('.').pop();
|
|
const expected: Record<string, string[]> = {
|
|
'image/jpeg': ['jpg', 'jpeg'],
|
|
'image/png': ['png'],
|
|
'image/webp': ['webp'],
|
|
'application/pdf': ['pdf'],
|
|
'application/vnd.openxmlformats-officedocument.wordprocessingml.document': ['docx'],
|
|
'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet': ['xlsx'],
|
|
};
|
|
if (!extension || !expected[mimeType]?.includes(extension)) {
|
|
throw new BadRequestException('附件扩展名与文件内容不一致');
|
|
}
|
|
}
|
|
|
|
private detectMimeType(buffer: Buffer, declaredMimeType: string): string {
|
|
if (buffer.subarray(0, 3).equals(Buffer.from([0xff, 0xd8, 0xff]))) return 'image/jpeg';
|
|
if (buffer.subarray(0, 8).equals(Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]))) {
|
|
return 'image/png';
|
|
}
|
|
if (
|
|
buffer.subarray(0, 4).toString('ascii') === 'RIFF' &&
|
|
buffer.subarray(8, 12).toString('ascii') === 'WEBP'
|
|
) {
|
|
return 'image/webp';
|
|
}
|
|
if (buffer.subarray(0, 5).toString('ascii') === '%PDF-') return 'application/pdf';
|
|
const isZip =
|
|
buffer.subarray(0, 4).equals(Buffer.from([0x50, 0x4b, 0x03, 0x04])) ||
|
|
buffer.subarray(0, 4).equals(Buffer.from([0x50, 0x4b, 0x05, 0x06])) ||
|
|
buffer.subarray(0, 4).equals(Buffer.from([0x50, 0x4b, 0x07, 0x08]));
|
|
if (
|
|
isZip &&
|
|
(declaredMimeType.includes('wordprocessingml') ||
|
|
declaredMimeType.includes('spreadsheetml'))
|
|
) {
|
|
return declaredMimeType;
|
|
}
|
|
return 'application/octet-stream';
|
|
}
|
|
|
|
private extensionForMime(mimeType: string): string {
|
|
const extensions: Record<string, string> = {
|
|
'image/jpeg': 'jpg',
|
|
'image/png': 'png',
|
|
'image/webp': 'webp',
|
|
'application/pdf': 'pdf',
|
|
'application/vnd.openxmlformats-officedocument.wordprocessingml.document': 'docx',
|
|
'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet': 'xlsx',
|
|
};
|
|
return extensions[mimeType] || 'bin';
|
|
}
|
|
|
|
private resolveStoragePath(storageKey: string): string {
|
|
const safeKey = storageKey.replace(/[^a-zA-Z0-9/_.-]/g, '');
|
|
if (safeKey !== storageKey) throw new BadRequestException('附件路径无效');
|
|
const absolutePath = resolve(this.storageRoot, safeKey);
|
|
const relativePath = relative(this.storageRoot, absolutePath);
|
|
if (relativePath.startsWith('..') || isAbsolute(relativePath)) {
|
|
throw new BadRequestException('附件路径无效');
|
|
}
|
|
return absolutePath;
|
|
}
|
|
}
|