yan-istart_IStart-Note-AI-Plugin/src/ai/QuestionClassifier.ts

import { DeepSeekSettings, QuestionClassification } from "../types";
import { LLMClient, parseJsonSafe } from "../core/llm";

const CLASSIFY_PROMPT = `你是一个知识图谱助手，负责对用户的问题进行分类和关联。

问题类型定义：
- new：全新问题，不依赖已有问题，引入新领域
- refinement：对某个已有问题的深入追问
- expansion：对某个已有问题的横向扩展

当前问题：
{{question}}

历史问题列表（最近20条）：
{{history}}

请严格按以下 JSON 格式返回，不要有任何其他内容：
{
  "category": "new | refinement | expansion",
  "parent": "最相关的历史问题标题，没有则为 null",
  "related": ["相关问题1", "相关问题2"],
  "confidence": 0.9,
  "refinements": ["推荐深化问题1", "推荐深化问题2"],
  "expansions": ["推荐扩展问题1", "推荐扩展问题2"]
}`;

export class QuestionClassifier {
  private llm: LLMClient;

  constructor(settings: DeepSeekSettings) {
    this.llm = new LLMClient(settings);
  }

  async classify(question: string, history: string[]): Promise<QuestionClassification> {
    try {
      this.llm.ensureApiKey();
    } catch {
      return this.defaultClassification();
    }

    const prompt = CLASSIFY_PROMPT
      .replace("{{question}}", question)
      .replace(
        "{{history}}",
        history.length > 0
          ? history.map((q, i) => `${i + 1}. ${q}`).join("\n")
          : "（无历史问题）"
      );

    try {
      const content = await this.llm.chat({ userPrompt: prompt, temperature: 0.3 });
      return this.parse(content);
    } catch {
      return this.defaultClassification();
    }
  }

  private parse(content: string): QuestionClassification {
    const p = parseJsonSafe<Record<string, unknown> | null>(content, null);
    if (!p) return this.defaultClassification();

    const category = ["new", "refinement", "expansion"].includes(p.category as string)
      ? (p.category as QuestionClassification["category"])
      : "new";

    return {
      category,
      parent: typeof p.parent === "string" ? p.parent : null,
      related: Array.isArray(p.related) ? (p.related as string[]) : [],
      confidence: typeof p.confidence === "number" ? p.confidence : 0.5,
      refinements: Array.isArray(p.refinements) ? (p.refinements as string[]) : [],
      expansions: Array.isArray(p.expansions) ? (p.expansions as string[]) : [],
    };
  }

  private defaultClassification(): QuestionClassification {
    return { category: "new", parent: null, related: [], confidence: 0, refinements: [], expansions: [] };
  }
}