dukang/server/dukang-api/src/integrations/llm/knowledge-retrieval.service.ts

import { Injectable } from '@nestjs/common';
import { PrismaService } from '../../common/prisma/prisma.module';

const MAX_CONTEXT_CHARS = 6000;

@Injectable()
export class KnowledgeRetrievalService {
  constructor(private readonly prisma: PrismaService) {}

  /** 简易关键词命中：按文档正文包含查询词打分，拼进上下文 */
  async buildContext(knowledgeBaseId: bigint, query: string): Promise<string> {
    const kb = await this.prisma.knowledgeBase.findUnique({
      where: { id: knowledgeBaseId },
      select: { id: true, enabled: true, name: true },
    });
    if (!kb?.enabled) return '';

    const docs = await this.prisma.knowledgeDocument.findMany({
      where: {
        knowledgeBaseId,
        status: 'READY',
        contentText: { not: null },
      },
      select: { title: true, contentText: true },
      take: 50,
    });
    if (!docs.length) return '';

    const tokens = tokenize(query);
    const scored = docs
      .map((d) => {
        const body = d.contentText || '';
        let score = 0;
        for (const t of tokens) {
          if (body.includes(t) || d.title.includes(t)) score += 1;
        }
        if (!tokens.length) score = 1;
        return { title: d.title, body, score };
      })
      .filter((x) => x.score > 0)
      .sort((a, b) => b.score - a.score);

    const picked = scored.length ? scored.slice(0, 5) : docs.slice(0, 3).map((d) => ({
      title: d.title,
      body: d.contentText || '',
      score: 0,
    }));

    let out = `【知识库：${kb.name}】\n`;
    for (const p of picked) {
      const chunk = `### ${p.title}\n${p.body}\n\n`;
      if (out.length + chunk.length > MAX_CONTEXT_CHARS) {
        out += chunk.slice(0, Math.max(0, MAX_CONTEXT_CHARS - out.length));
        break;
      }
      out += chunk;
    }
    return out.trim();
  }
}

function tokenize(q: string): string[] {
  return q
    .split(/[\s,，。；;、！？!?\-_/\\]+/)
    .map((s) => s.trim().toLowerCase())
    .filter((s) => s.length >= 2)
    .slice(0, 12);
}