-- 安装扩展（容器启动脚本或迁移文件里执行一次即可）
CREATE EXTENSION IF NOT EXISTS vector;

-- 文档表：记录每个上传文档的元数据和处理状态
CREATE TABLE documents (
  id          UUID PRIMARY KEY DEFAULT gen_random_uuid(),
  tenant_id   UUID NOT NULL REFERENCES tenants(id) ON DELETE CASCADE,
  name        VARCHAR(255) NOT NULL,
  -- 文档处理状态机：uploaded → parsing → chunking → embedding → indexed | failed
  status      VARCHAR(50) NOT NULL DEFAULT 'uploaded',
  file_path   TEXT,
  file_type   VARCHAR(50),          -- pdf / docx / md / txt
  version     INTEGER DEFAULT 1,
  chunk_count INTEGER DEFAULT 0,
  error_msg   TEXT,                 -- 失败时记录原因
  created_at  TIMESTAMPTZ DEFAULT NOW(),
  updated_at  TIMESTAMPTZ DEFAULT NOW()
);
 
-- 分块表：文档切块后的实际内容和向量
CREATE TABLE knowledge_chunks (
  id          UUID PRIMARY KEY DEFAULT gen_random_uuid(),
  tenant_id   UUID REFERENCES tenants(id) ON DELETE CASCADE,
  -- tenant_id IS NULL 表示平台内置知识库，所有租户可访问
  document_id UUID NOT NULL REFERENCES documents(id) ON DELETE CASCADE,
  content     TEXT NOT NULL,
  -- 全书统一 1024 维，对应默认 embedding 模型 voyage-3。
  -- 切换到 OpenAI text-embedding-3-small（1536 维）时需要把这列改成 vector(1536) 并重建索引。
  embedding   vector(1024),
  metadata    JSONB DEFAULT '{}',   -- 存储页码、标题等附加信息
  chunk_index INTEGER NOT NULL,     -- 在文档中的顺序位置
  is_deleted  BOOLEAN DEFAULT FALSE, -- 软删除标记，文档更新时使用
  created_at  TIMESTAMPTZ DEFAULT NOW()
);

-- IVFFlat 索引：lists 值建议约等于 sqrt(总行数)
-- 适合数据量中等、内存受限的场景
CREATE INDEX idx_chunks_ivfflat ON knowledge_chunks
  USING ivfflat (embedding vector_cosine_ops)
  WITH (lists = 100);

-- HNSW 索引：m 是每个节点的最大连接数，ef_construction 影响构建质量
-- 生产推荐配置
CREATE INDEX idx_chunks_hnsw ON knowledge_chunks
  USING hnsw (embedding vector_cosine_ops)
  WITH (m = 16, ef_construction = 64);

import { drizzle } from 'drizzle-orm/node-postgres';
import { sql } from 'drizzle-orm';
import pgvector from 'pgvector';
 
// 向量相似度检索，cosine distance（<=>) 越小表示越相似
// 1 - cosine_distance = cosine_similarity
async function vectorSearch(
  db: ReturnType<typeof drizzle>,
  tenantId: string,
  queryEmbedding: number[],
  limit: number = 5
) {
  const embeddingStr = pgvector.toSql(queryEmbedding);
 
  const results = await db.execute(sql`
    SELECT
      id,
      content,
      metadata,
      document_id,
      1 - (embedding <=> ${embeddingStr}::vector) AS similarity
    FROM knowledge_chunks
    WHERE
      (tenant_id = ${tenantId} OR tenant_id IS NULL)
      AND is_deleted = FALSE
    ORDER BY embedding <=> ${embeddingStr}::vector
    LIMIT ${limit}
  `);
 
  return results.rows;
}

// 入库流水线的 Job 数据结构
interface IngestionJobData {
  documentId: string;
  tenantId: string;
  filePath: string;
  fileType: 'pdf' | 'docx' | 'md' | 'txt';
}
 
// 在 Worker 里处理入库流水线
const ingestionWorker = new Worker<IngestionJobData>(
  'knowledge:ingestion',
  async (job) => {
    const { documentId, tenantId, filePath, fileType } = job.data;
 
    // 每个阶段更新状态，方便前端轮询进度
    await updateDocumentStatus(documentId, 'parsing');
    const text = await parseDocument(filePath, fileType);
 
    await updateDocumentStatus(documentId, 'chunking');
    const chunks = chunkText(text, { strategy: 'semantic', chunkSize: 512, overlap: 1 });
    // overlap: 1 表示保留上一个 chunk 的最后 1 句作为新 chunk 的开头（语义策略下单位为句子数）
 
    await updateDocumentStatus(documentId, 'embedding');
    // batchEmbed 的完整签名：batchEmbed(texts, provider, options)
    // provider 是 VoyageEmbeddingProvider 或 OpenAIEmbeddingProvider 实例
    const embeddings = await batchEmbed(chunks.map(c => c.content), provider);
 
    // 批量写入，更新 chunk_count，最后切换状态
    await saveChunks(documentId, tenantId, chunks, embeddings);
    await updateDocumentStatus(documentId, 'indexed');
  },
  { connection: redisConnection, concurrency: 3 }
);

import * as fs from 'fs/promises';
 
async function parseMarkdown(filePath: string): Promise<string> {
  // Markdown 直接读取，保留原始格式供分块器使用
  return fs.readFile(filePath, 'utf-8');
}
 
async function parsePlainText(filePath: string): Promise<string> {
  return fs.readFile(filePath, 'utf-8');
}

import pdfParse from 'pdf-parse';
import * as fs from 'fs/promises';
 
async function parsePDF(filePath: string): Promise<string> {
  const buffer = await fs.readFile(filePath);
  const data = await pdfParse(buffer);
  // data.text 是所有页面的文本拼接，data.numpages 是总页数
  return data.text;
}

import mammoth from 'mammoth';
 
async function parseDocx(filePath: string): Promise<string> {
  // 转为 Markdown 格式，便于后续按标题分块
  const result = await mammoth.convertToMarkdown({ path: filePath });
 
  if (result.messages.length > 0) {
    // messages 里是转换警告，比如"不支持的样式 xxx"
    // 不是错误，记录日志即可
    console.warn('DOCX 转换警告:', result.messages);
  }
 
  return result.value;
}

interface ChunkOptions {
  chunkSize: number;   // 块的目标大小（字符数）
  overlap: number;     // 相邻块的重叠字符数，保留上下文连贯性
}
 
function fixedSizeChunk(text: string, options: ChunkOptions): string[] {
  const { chunkSize, overlap } = options;
  const chunks: string[] = [];
  let start = 0;
 
  while (start < text.length) {
    const end = Math.min(start + chunkSize, text.length);
    chunks.push(text.slice(start, end));
    // 下一个 chunk 从 (start + chunkSize - overlap) 开始
    // overlap 确保块边界处的语义不会完全丢失
    start += chunkSize - overlap;
  }
 
  return chunks;
}

interface SemanticChunk {
  content: string;
  index: number;
}
 
function semanticChunk(text: string, maxChunkSize: number, overlap: number = 1): SemanticChunk[] {
  // 按中文句末标点分割句子
  // 同时处理换行段落
  const sentences = text
    .split(/(?<=[。！？\n])\s*/)
    .map(s => s.trim())
    .filter(s => s.length > 0);
 
  const chunks: SemanticChunk[] = [];
  let currentChunk = '';
  let lastSentence = '';  // 用于 overlap：保留上一个 chunk 的最后一句
 
  for (const sentence of sentences) {
    // 如果加上当前句子会超过 maxChunkSize，先保存当前 chunk
    if (currentChunk.length + sentence.length > maxChunkSize && currentChunk.length > 0) {
      chunks.push({ content: currentChunk.trim(), index: chunks.length });
      // 新 chunk 以上一个 chunk 的最后一句开头（overlap）
      currentChunk = overlap > 0 ? lastSentence + '\n' : '';
    }
    lastSentence = sentence;
    currentChunk += sentence + '\n';
  }
 
  // 处理最后一个 chunk
  if (currentChunk.trim().length > 0) {
    chunks.push({ content: currentChunk.trim(), index: chunks.length });
  }
 
  return chunks;
}

interface MarkdownChunk {
  content: string;
  heading: string;   // 对应的标题，写入 metadata 方便调试
  level: number;     // 标题层级（1-6）
  index: number;
}
 
function markdownChunk(text: string, maxChunkSize: number = 1500): MarkdownChunk[] {
  const lines = text.split('\n');
  const chunks: MarkdownChunk[] = [];
  let currentContent = '';
  let currentHeading = '（无标题）';
  let currentLevel = 0;
 
  for (const line of lines) {
    const headingMatch = line.match(/^(#{1,6})\s+(.+)/);
 
    if (headingMatch) {
      const level = headingMatch[1].length;
      const heading = headingMatch[2];
 
      // 遇到新标题时，保存当前积累的内容
      if (currentContent.trim().length > 0) {
        // 如果当前块超过 maxChunkSize，递归用语义分块处理
        if (currentContent.length > maxChunkSize) {
          const subChunks = semanticChunk(currentContent, maxChunkSize);
          for (const sub of subChunks) {
            chunks.push({
              content: sub.content,
              heading: currentHeading,
              level: currentLevel,
              index: chunks.length,
            });
          }
        } else {
          chunks.push({
            content: currentContent.trim(),
            heading: currentHeading,
            level: currentLevel,
            index: chunks.length,
          });
        }
      }
 
      currentHeading = heading;
      currentLevel = level;
      currentContent = line + '\n';  // 标题本身也包含在 chunk 里
    } else {
      currentContent += line + '\n';
    }
  }
 
  // 处理最后一个节
  if (currentContent.trim().length > 0) {
    chunks.push({
      content: currentContent.trim(),
      heading: currentHeading,
      level: currentLevel,
      index: chunks.length,
    });
  }
 
  return chunks;
}

import { HierarchicalNodeParser, getNodesFromDocument, Document } from 'llamaindex';
 
const parser = new HierarchicalNodeParser({
  chunkSizes: [2048, 512, 128],  // 从粗到细三个粒度
});
 
const document = new Document({ text: rawText });
const nodes = getNodesFromDocument(document, parser);
// nodes 里同时包含三个粒度的 chunk，并记录了父子关系

import Anthropic from '@anthropic-ai/sdk';
 
// 注意：voyage embedding 通过单独的 voyageai 包调用
// 这里用 OpenAI 格式展示逻辑，实际可替换为任意 embedding 提供商
async function batchEmbed(
  texts: string[],
  batchSize: number = 20,
  delayMs: number = 100
): Promise<number[][]> {
  const embeddings: number[][] = [];
 
  for (let i = 0; i < texts.length; i += batchSize) {
    const batch = texts.slice(i, i + batchSize);
 
    // 调用 embedding API
    const batchEmbeddings = await embedBatch(batch);
    embeddings.push(...batchEmbeddings);
 
    // 批次间等待，避免触发速率限制
    // 生产环境应该用指数退避重试，而不是固定等待
    if (i + batchSize < texts.length) {
      await new Promise(resolve => setTimeout(resolve, delayMs));
    }
 
    console.log(`向量化进度: ${Math.min(i + batchSize, texts.length)}/${texts.length}`);
  }
 
  return embeddings;
}

-- 给 knowledge_chunks 表加全文搜索列（生成列，自动维护）
ALTER TABLE knowledge_chunks
  ADD COLUMN tsv TSVECTOR
  GENERATED ALWAYS AS (to_tsvector('simple', content)) STORED;
 
-- GIN 索引加速全文搜索
CREATE INDEX idx_chunks_tsv ON knowledge_chunks USING gin(tsv);

-- 混合检索查询（使用 RRF 融合向量检索和全文检索结果）
WITH
vector_results AS (
  SELECT
    id,
    ROW_NUMBER() OVER (ORDER BY embedding <=> $1::vector) AS rank
  FROM knowledge_chunks
  WHERE
    (tenant_id = $2 OR tenant_id IS NULL)
    AND is_deleted = FALSE
  LIMIT 20
),
bm25_results AS (
  SELECT
    id,
    ROW_NUMBER() OVER (ORDER BY ts_rank(tsv, query) DESC) AS rank
  FROM knowledge_chunks,
       to_tsquery('simple', $3) AS query
  WHERE
    (tenant_id = $2 OR tenant_id IS NULL)
    AND is_deleted = FALSE
    AND tsv @@ query
  LIMIT 20
)
SELECT
  kc.id,
  kc.content,
  kc.metadata,
  kc.document_id,
  -- RRF 分数：k=60 是行业常用默认值
  (COALESCE(1.0 / (60.0 + vr.rank), 0.0) +
   COALESCE(1.0 / (60.0 + br.rank), 0.0)) AS rrf_score
FROM knowledge_chunks kc
LEFT JOIN vector_results vr ON kc.id = vr.id
LEFT JOIN bm25_results br ON kc.id = br.id
WHERE (vr.id IS NOT NULL OR br.id IS NOT NULL)
  AND kc.is_deleted = FALSE
ORDER BY rrf_score DESC
LIMIT $4;

function buildTsquery(text: string): string {
  // 简单实现：按空格和标点切分，取有意义的词（长度 > 1）
  // 生产环境应使用 jieba 或其他中文分词工具
  const tokens = text
    .split(/[\s，。！？、：；""''【】（）\(\)\[\]]+/)
    .map(t => t.trim())
    .filter(t => t.length > 1);
 
  if (tokens.length === 0) return '';
 
  // 用 | 连接表示 OR，& 表示 AND
  // 对于用户问题，OR 召回更多结果，配合 RRF 重排
  //
  // 注意：to_tsquery 对特殊字符（& | ! : 等）会报 SQL 语法错误。
  // 生产中应使用 plainto_tsquery($config, $text) 代替手工拼接，
  // 它会自动将自然语言转为安全的 tsquery 表达式。
  return tokens.join(' | ');
}

WHERE (tenant_id = $tenantId OR tenant_id IS NULL)

interface SearchResult {
  id: string;
  content: string;
  metadata: Record<string, unknown>;
  documentId: string;
  rrfScore: number;
  tenantId: string | null;   // null 表示平台内置
}
 
function boostTenantResults(
  results: SearchResult[],
  tenantId: string,
  boostFactor: number = 1.2
): SearchResult[] {
  return results
    .map(r => ({
      ...r,
      // 租户自己的内容乘以提升系数
      rrfScore: r.tenantId === tenantId ? r.rrfScore * boostFactor : r.rrfScore,
    }))
    .sort((a, b) => b.rrfScore - a.rrfScore);
}

interface KnowledgeSearchOptions {
  tenantId: string;
  query: string;
  queryEmbedding: number[];
  includeBuiltin: boolean;      // false 时只查租户自己的知识库
  tenantContentBoost: number;   // 租户内容排名提升系数，默认 1.2
  limit: number;
  strategy: 'vector' | 'bm25' | 'hybrid';  // 检索策略
}

async function updateDocument(
  db: ReturnType<typeof drizzle>,
  documentId: string,
  tenantId: string,
  newContent: string,
  fileType: DocumentFileType
): Promise<void> {
  // 1. 解析和分块
  const text = await parseDocument('/tmp/updated-file', fileType);
  const chunks = chunkText(text, { strategy: 'semantic', chunkSize: 512, overlap: 1 });
  const embeddings = await batchEmbed(chunks.map(c => c.content));
 
  // 2. 在事务里完成：写新 chunks + 软删除旧 chunks
  await db.transaction(async (tx) => {
    // 写入新 chunks
    const newChunkIds = await insertChunks(tx, documentId, tenantId, chunks, embeddings);
 
    // 软删除旧 chunks（排除刚写入的新 chunks）
    await tx.execute(sql`
      UPDATE knowledge_chunks
      SET is_deleted = TRUE
      WHERE document_id = ${documentId}
        AND id != ALL(${newChunkIds}::uuid[])
        AND is_deleted = FALSE
    `);
 
    // 更新文档元数据
    await tx.execute(sql`
      UPDATE documents
      SET chunk_count = ${chunks.length},
          version = version + 1,
          updated_at = NOW()
      WHERE id = ${documentId}
    `);
  });
}

-- 定时清理软删除的 chunks（保留 7 天以上的才清理）
DELETE FROM knowledge_chunks
WHERE is_deleted = TRUE
  AND created_at < NOW() - INTERVAL '7 days';

// Agent 处理用户消息时的 RAG 步骤
async function answerWithRAG(
  userMessage: string,
  tenantId: string
): Promise<string> {
  // 1. 将用户问题向量化
  const queryEmbedding = await embedSingleText(userMessage);
 
  // 2. 混合检索相关 chunks
  const chunks = await hybridSearch({
    tenantId,
    query: userMessage,
    queryEmbedding,
    includeBuiltin: true,
    tenantContentBoost: 1.2,
    limit: 5,
    strategy: 'hybrid',
  });
 
  if (chunks.length === 0) {
    // 没有检索到相关内容，让 LLM 明确表示不知道
    return await callLLM(userMessage, '');
  }
 
  // 3. 拼装 context，传入 LLM
  const context = chunks
    .map((c, i) => `[来源 ${i + 1}]\n${c.content}`)
    .join('\n\n---\n\n');
 
  return await callLLM(userMessage, context);
}
 
async function callLLM(userMessage: string, context: string): Promise<string> {
  const client = new Anthropic();
 
  const systemPrompt = context
    ? `你是一个企业智能客服助手。请基于以下知识库内容回答用户问题。
如果知识库内容不能回答问题，请明确说明，不要编造信息。
 
知识库内容：
${context}`
    : '你是一个企业智能客服助手。你没有找到与该问题相关的知识库内容，请如实告知用户。';
 
  const response = await client.messages.create({
    model: 'claude-sonnet-4-5',
    max_tokens: 1024,
    system: systemPrompt,
    messages: [{ role: 'user', content: userMessage }],
  });
 
  return response.content[0].type === 'text' ? response.content[0].text : '';
}