24.09-知识库管理

-- 知识库表
CREATE TABLE knowledge_bases (
  id TEXT PRIMARY KEY,
  tenant_id TEXT NOT NULL,
  workspace_id TEXT,
  name TEXT NOT NULL,
  description TEXT,
  document_count INTEGER DEFAULT 0,
  total_tokens INTEGER DEFAULT 0,
  status TEXT DEFAULT 'active',       -- active/archived
  embedding_model TEXT DEFAULT 'text-embedding-3-small',
  created_at INTEGER NOT NULL,
  updated_at INTEGER NOT NULL,
  FOREIGN KEY (tenant_id) REFERENCES tenants(id),
  FOREIGN KEY (workspace_id) REFERENCES workspaces(id)
)
 
-- 文档表
CREATE TABLE kb_documents (
  id TEXT PRIMARY KEY,
  knowledge_base_id TEXT NOT NULL,
  file_name TEXT NOT NULL,
  file_type TEXT NOT NULL,            -- pdf/docx/md/txt
  file_size INTEGER NOT NULL,
  r2_key TEXT NOT NULL,               -- R2 存储路径
  content TEXT,                       -- 解析后的文本
  chunk_count INTEGER DEFAULT 0,
  token_count INTEGER DEFAULT 0,
  status TEXT DEFAULT 'pending',      -- pending/processing/ready/failed
  error_message TEXT,
  created_at INTEGER NOT NULL,
  updated_at INTEGER NOT NULL,
  FOREIGN KEY (knowledge_base_id) REFERENCES knowledge_bases(id)
)
 
-- 文档 chunks 表
CREATE TABLE kb_chunks (
  id TEXT PRIMARY KEY,
  document_id TEXT NOT NULL,
  chunk_index INTEGER NOT NULL,
  content TEXT NOT NULL,
  token_count INTEGER NOT NULL,
  metadata TEXT,                      -- JSON: 页码、标题等
  created_at INTEGER NOT NULL,
  FOREIGN KEY (document_id) REFERENCES kb_documents(id)
)
 
-- 向量存储（如果使用外部 Vector DB，则不需要此表）
-- 这里以 D1 存储为例，实际生产中建议使用 Pinecone/Weaviate
CREATE TABLE kb_embeddings (
  id TEXT PRIMARY KEY,
  chunk_id TEXT NOT NULL,
  embedding TEXT NOT NULL,            -- JSON 数组
  metadata TEXT,
  FOREIGN KEY (chunk_id) REFERENCES kb_chunks(id)
)
 
-- 索引
CREATE INDEX idx_chunks_document ON kb_chunks(document_id, chunk_index)
CREATE INDEX idx_embeddings_chunk ON kb_embeddings(chunk_id)

// src/types/knowledge-base.ts
export interface KnowledgeBase {
  id: string
  tenantId: string
  workspaceId?: string
  name: string
  description?: string
  documentCount: number
  totalTokens: number
  status: 'active' | 'archived'
  embeddingModel: string
  createdAt: number
  updatedAt: number
}
 
export interface KBDocument {
  id: string
  knowledgeBaseId: string
  fileName: string
  fileType: 'pdf' | 'docx' | 'md' | 'txt'
  fileSize: number
  r2Key: string
  content?: string
  chunkCount: number
  tokenCount: number
  status: 'pending' | 'processing' | 'ready' | 'failed'
  errorMessage?: string
  createdAt: number
  updatedAt: number
}
 
export interface KBChunk {
  id: string
  documentId: string
  chunkIndex: number
  content: string
  tokenCount: number
  metadata?: {
    page?: number
    section?: string
    [key: string]: any
  }
}

// src/routes/knowledge-bases.ts
import { Hono } from 'hono'
import { zValidator } from '@hono/zod-validator'
import { z } from 'zod'
import { generateId } from '../lib/utils'
import { requireRole } from '../middleware/permission'
 
const app = new Hono()
 
app.post('/', requireRole('admin', 'member'), zValidator('json', z.object({
  name: z.string().min(1).max(100),
  description: z.string().optional(),
  workspaceId: z.string().optional(),
  embeddingModel: z.string().optional(),
})), async (c) => {
  const tenantId = c.get('tenantId')
  const userId = c.get('userId')
  const { name, description, workspaceId, embeddingModel = 'text-embedding-3-small' } = c.req.valid('json')
 
  // 检查知识库数量限制
  const plan = c.get('plan')
  const kbCount = await c.env.DB.prepare(`
    SELECT COUNT(*) as count FROM knowledge_bases WHERE tenant_id = ?
  `).bind(tenantId).first()
 
  if (plan.maxKnowledgeBases > 0 && kbCount.count >= plan.maxKnowledgeBases) {
    return c.json({ error: '知识库数量已达上限' }, 403)
  }
 
  const id = generateId()
  const now = Date.now()
 
  await c.env.DB.prepare(`
    INSERT INTO knowledge_bases (id, tenant_id, workspace_id, name, description, embedding_model, created_at, updated_at)
    VALUES (?, ?, ?, ?, ?, ?, ?, ?)
  `).bind(
    id,
    tenantId,
    workspaceId || null,
    name,
    description || null,
    embeddingModel,
    now,
    now
  ).run()
 
  return c.json({
    id,
    name,
    description,
    embeddingModel,
  }, 201)
})

app.get('/', async (c) => {
  const tenantId = c.get('tenantId')
  const workspaceId = c.req.query('workspaceId')
 
  let query = `SELECT * FROM knowledge_bases WHERE tenant_id = ?`
  const params: any[] = [tenantId]
 
  if (workspaceId) {
    query += ` AND workspace_id = ?`
    params.push(workspaceId)
  }
 
  query += ` ORDER BY updated_at DESC`
 
  const knowledgeBases = await c.env.DB.prepare(query).bind(...params).all()
 
  return c.json({
    knowledgeBases: knowledgeBases.results.map(kb => ({
      ...kb,
      workspaceId: kb.workspace_id,
      documentCount: kb.document_count,
      totalTokens: kb.total_tokens,
      embeddingModel: kb.embedding_model,
    })),
  })
})
 
app.get('/:id', async (c) => {
  const kbId = c.req.param('id')
  const tenantId = c.get('tenantId')
 
  const kb = await c.env.DB.prepare(`
    SELECT * FROM knowledge_bases WHERE id = ? AND tenant_id = ?
  `).bind(kbId, tenantId).first()
 
  if (!kb) {
    return c.json({ error: '知识库不存在' }, 404)
  }
 
  // 获取文档列表
  const documents = await c.env.DB.prepare(`
    SELECT id, file_name, file_type, file_size, chunk_count, token_count, status, created_at
    FROM kb_documents
    WHERE knowledge_base_id = ?
    ORDER BY created_at DESC
  `).bind(kbId).all()
 
  return c.json({
    knowledgeBase: {
      ...kb,
      workspaceId: kb.workspace_id,
      documentCount: kb.document_count,
      totalTokens: kb.total_tokens,
      embeddingModel: kb.embedding_model,
    },
    documents: documents.results.map(doc => ({
      ...doc,
      knowledgeBaseId: doc.knowledge_base_id,
      fileName: doc.file_name,
      fileType: doc.file_type,
      fileSize: doc.file_size,
      chunkCount: doc.chunk_count,
      tokenCount: doc.token_count,
      errorMessage: doc.error_message,
    })),
  })
})

// 上传文档
app.post('/:id/documents', async (c) => {
  const kbId = c.req.param('id')
  const tenantId = c.get('tenantId')
  const body = await c.req.formData()
  const file = body.get('file') as File
 
  if (!file) {
    return c.json({ error: '请上传文件' }, 400)
  }
 
  // 验证文件类型
  const allowedTypes = ['application/pdf', 'application/vnd.openxmlformats-officedocument.wordprocessingml.document', 'text/markdown', 'text/plain']
  if (!allowedTypes.includes(file.type)) {
    return c.json({ error: '不支持的文件类型' }, 400)
  }
 
  // 验证文件大小（最大 50MB）
  if (file.size > 50 * 1024 * 1024) {
    return c.json({ error: '文件大小不能超过 50MB' }, 400)
  }
 
  // 确定文件类型
  const fileType = getFileType(file.type, file.name)
 
  // 上传到 R2
  const docId = generateId()
  const r2Key = `kb/${tenantId}/${kbId}/${docId}/${file.name}`
 
  await c.env.R2.put(r2Key, file.stream(), {
    httpMetadata: {
      contentType: file.type,
    },
  })
 
  // 创建文档记录
  const now = Date.now()
  await c.env.DB.prepare(`
    INSERT INTO kb_documents (id, knowledge_base_id, file_name, file_type, file_size, r2_key, status, created_at, updated_at)
    VALUES (?, ?, ?, ?, ?, ?, 'pending', ?, ?)
  `).bind(docId, kbId, file.name, fileType, file.size, r2Key, now, now).run()
 
  // 异步处理文档
  c.executionCtx.waitUntil(processDocument(docId, c.env))
 
  return c.json({
    id: docId,
    fileName: file.name,
    fileType,
    fileSize: file.size,
    status: 'pending',
  }, 201)
})
 
function getFileType(mimeType: string, fileName: string): string {
  if (mimeType === 'application/pdf') return 'pdf'
  if (mimeType.includes('wordprocessingml')) return 'docx'
  if (mimeType === 'text/markdown' || fileName.endsWith('.md')) return 'md'
  return 'txt'
}

// src/lib/document-processor.ts
export async function processDocument(documentId: string, env: Env) {
  // 1. 获取文档信息
  const doc = await env.DB.prepare(`
    SELECT * FROM kb_documents WHERE id = ?
  `).bind(documentId).first()
 
  if (!doc) return
 
  // 2. 更新状态为处理中
  await env.DB.prepare(`
    UPDATE kb_documents SET status = 'processing', updated_at = ?
    WHERE id = ?
  `).bind(Date.now(), documentId).run()
 
  try {
    // 3. 从 R2 下载文件
    const object = await env.R2.get(doc.r2_key)
    if (!object) throw new Error('File not found in R2')
 
    const buffer = await object.arrayBuffer()
 
    // 4. 解析文档内容
    const content = await parseDocument(buffer, doc.file_type)
 
    // 5. 切分成 chunks
    const chunks = splitIntoChunks(content, {
      maxTokens: 500,
      overlap: 50,
    })
 
    // 6. 生成 embedding
    const embeddings = await generateEmbeddings(
      chunks.map(c => c.content),
      env
    )
 
    // 7. 保存到数据库
    const now = Date.now()
    let totalTokens = 0
 
    for (let i = 0; i < chunks.length; i++) {
      const chunk = chunks[i]
      const embedding = embeddings[i]
      const chunkId = generateId()
 
      await env.DB.prepare(`
        INSERT INTO kb_chunks (id, document_id, chunk_index, content, token_count, metadata, created_at)
        VALUES (?, ?, ?, ?, ?, ?, ?)
      `).bind(
        chunkId,
        documentId,
        i,
        chunk.content,
        chunk.tokenCount,
        JSON.stringify(chunk.metadata),
        now
      ).run()
 
      await env.DB.prepare(`
        INSERT INTO kb_embeddings (id, chunk_id, embedding)
        VALUES (?, ?, ?)
      `).bind(generateId(), chunkId, JSON.stringify(embedding)).run()
 
      totalTokens += chunk.tokenCount
    }
 
    // 8. 更新文档状态
    await env.DB.prepare(`
      UPDATE kb_documents
      SET content = ?, chunk_count = ?, token_count = ?, status = 'ready', updated_at = ?
      WHERE id = ?
    `).bind(content, chunks.length, totalTokens, now, documentId).run()
 
    // 9. 更新知识库统计
    await env.DB.prepare(`
      UPDATE knowledge_bases
      SET document_count = document_count + 1, total_tokens = total_tokens + ?, updated_at = ?
      WHERE id = ?
    `).bind(totalTokens, now, doc.knowledge_base_id).run()
 
  } catch (error) {
    console.error('Document processing failed:', error)
 
    await env.DB.prepare(`
      UPDATE kb_documents
      SET status = 'failed', error_message = ?, updated_at = ?
      WHERE id = ?
    `).bind(error.message, Date.now(), documentId).run()
  }
}
 
// 解析文档
async function parseDocument(buffer: ArrayBuffer, fileType: string): Promise<string> {
  switch (fileType) {
    case 'pdf':
      return parsePdf(buffer)
    case 'docx':
      return parseDocx(buffer)
    case 'md':
    case 'txt':
      return new TextDecoder().decode(buffer)
    default:
      throw new Error('Unsupported file type')
  }
}
 
// 切分成 chunks
function splitIntoChunks(content: string, options: { maxTokens: number; overlap: number }) {
  const sentences = content.split(/(?<=[.!?])\s+/)
  const chunks: Array<{ content: string; tokenCount: number; metadata?: any }> = []
  let currentChunk = ''
  let currentTokens = 0
 
  for (const sentence of sentences) {
    const sentenceTokens = Math.ceil(sentence.length / 4)
 
    if (currentTokens + sentenceTokens > options.maxTokens && currentChunk) {
      chunks.push({
        content: currentChunk.trim(),
        tokenCount: currentTokens,
      })
 
      // 保留重叠
      const overlapWords = currentChunk.split(' ').slice(-options.overlap / 4)
      currentChunk = overlapWords.join(' ') + ' ' + sentence
      currentTokens = Math.ceil(currentChunk.length / 4)
    } else {
      currentChunk += (currentChunk ? ' ' : '') + sentence
      currentTokens += sentenceTokens
    }
  }
 
  if (currentChunk) {
    chunks.push({
      content: currentChunk.trim(),
      tokenCount: currentTokens,
    })
  }
 
  return chunks
}
 
// 生成 embedding
async function generateEmbeddings(texts: string[], env: Env): Promise<number[][]> {
  const response = await fetch('https://api.openai.com/v1/embeddings', {
    method: 'POST',
    headers: {
      'Content-Type': 'application/json',
      'Authorization': `Bearer ${env.OPENAI_API_KEY}`,
    },
    body: JSON.stringify({
      input: texts,
      model: 'text-embedding-3-small',
    }),
  })
 
  const data = await response.json()
  return data.data.map((item: any) => item.embedding)
}

// src/routes/knowledge-bases.ts（补充）
app.post('/:id/search', zValidator('json', z.object({
  query: z.string(),
  topK: z.number().min(1).max(20).optional(),
})), async (c) => {
  const kbId = c.req.param('id')
  const tenantId = c.get('tenantId')
  const { query, topK = 5 } = c.req.valid('json')
 
  // 验证知识库属于该租户
  const kb = await c.env.DB.prepare(`
    SELECT * FROM knowledge_bases WHERE id = ? AND tenant_id = ?
  `).bind(kbId, tenantId).first()
 
  if (!kb) {
    return c.json({ error: '知识库不存在' }, 404)
  }
 
  // 生成查询向量
  const queryEmbedding = await generateEmbeddings([query], c.env)
 
  // 检索相似 chunks（简化版，实际应使用 Vector DB）
  const results = await searchSimilarChunks(
    queryEmbedding[0],
    kbId,
    topK,
    c.env
  )
 
  return c.json({
    query,
    results: results.map(r => ({
      content: r.content,
      score: r.score,
      document: {
        id: r.documentId,
        fileName: r.fileName,
      },
      metadata: r.metadata,
    })),
  })
})
 
// 相似度搜索（使用余弦相似度）
async function searchSimilarChunks(
  queryEmbedding: number[],
  knowledgeBaseId: string,
  topK: number,
  env: Env
) {
  // 获取知识库下所有 chunks 的 embeddings
  const embeddings = await env.DB.prepare(`
    SELECT e.id, e.chunk_id, e.embedding, c.content, c.token_count, c.metadata,
           d.id as document_id, d.file_name
    FROM kb_embeddings e
    JOIN kb_chunks c ON e.chunk_id = c.id
    JOIN kb_documents d ON c.document_id = d.id
    WHERE d.knowledge_base_id = ?
  `).bind(knowledgeBaseId).all()
 
  // 计算余弦相似度
  const results = embeddings.results.map(e => {
    const embedding = JSON.parse(e.embedding)
    const score = cosineSimilarity(queryEmbedding, embedding)
 
    return {
      content: e.content,
      score,
      documentId: e.document_id,
      fileName: e.file_name,
      metadata: JSON.parse(e.metadata || '{}'),
    }
  })
 
  // 按相似度排序
  results.sort((a, b) => b.score - a.score)
 
  return results.slice(0, topK)
}
 
function cosineSimilarity(a: number[], b: number[]): number {
  let dotProduct = 0
  let normA = 0
  let normB = 0
 
  for (let i = 0; i < a.length; i++) {
    dotProduct += a[i] * b[i]
    normA += a[i] * a[i]
    normB += b[i] * b[i]
  }
 
  return dotProduct / (Math.sqrt(normA) * Math.sqrt(normB))
}

app.delete('/:id/documents/:docId', async (c) => {
  const kbId = c.req.param('id')
  const docId = c.req.param('docId')
  const tenantId = c.get('tenantId')
 
  // 验证知识库
  const kb = await c.env.DB.prepare(`
    SELECT * FROM knowledge_bases WHERE id = ? AND tenant_id = ?
  `).bind(kbId, tenantId).first()
 
  if (!kb) {
    return c.json({ error: '知识库不存在' }, 404)
  }
 
  // 获取文档
  const doc = await c.env.DB.prepare(`
    SELECT * FROM kb_documents WHERE id = ? AND knowledge_base_id = ?
  `).bind(docId, kbId).first()
 
  if (!doc) {
    return c.json({ error: '文档不存在' }, 404)
  }
 
  // 删除 R2 文件
  await c.env.R2.delete(doc.r2_key)
 
  // 删除数据库记录
  await c.env.DB.batch([
    c.env.DB.prepare(`DELETE FROM kb_embeddings WHERE chunk_id IN (SELECT id FROM kb_chunks WHERE document_id = ?)`).bind(docId),
    c.env.DB.prepare(`DELETE FROM kb_chunks WHERE document_id = ?`).bind(docId),
    c.env.DB.prepare(`DELETE FROM kb_documents WHERE id = ?`).bind(docId),
  ])
 
  // 更新知识库统计
  await c.env.DB.prepare(`
    UPDATE knowledge_bases
    SET document_count = document_count - 1, total_tokens = total_tokens - ?, updated_at = ?
    WHERE id = ?
  `).bind(doc.token_count, Date.now(), kbId).run()
 
  return c.json({ success: true })
})