teatea-pension/modules/ai/server/knowledge.ts

import "server-only";

import { and, desc, eq, isNull, or } from "drizzle-orm";

import type {
  AiKnowledgeScope,
  KnowledgeEntry,
  KnowledgeEntryInput,
  KnowledgeRetrievalResult,
} from "@/modules/ai/types";
import { getDatabase } from "@/modules/core/server/db";
import { aiKnowledgeChunks, aiKnowledgeEntries } from "@/modules/core/server/schema";
import type { AuthContext } from "@/modules/core/types";

type KnowledgeRow = typeof aiKnowledgeEntries.$inferSelect;

type ServiceResult<T> = { success: true; data: T } | { success: false; reason: string; status: number };

const CHUNK_TARGET_SIZE = 900;
const CHUNK_OVERLAP_SIZE = 160;
const DEFAULT_RETRIEVAL_LIMIT = 6;

function toIsoString(value: Date): string {
  return value.toISOString();
}

function toKnowledgeEntry(row: KnowledgeRow): KnowledgeEntry {
  return {
    id: row.id,
    scope: row.scope,
    organizationId: row.organizationId ?? undefined,
    title: row.title,
    category: row.category,
    tags: row.tags,
    body: row.body,
    status: row.status,
    createdAt: toIsoString(row.createdAt),
    updatedAt: toIsoString(row.updatedAt),
  };
}

function normalizeSearchText(value: string): string[] {
  const tokens = new Set<string>();
  const normalizedTokens = value.toLowerCase().match(/[\p{L}\p{N}]+/gu) ?? [];
  for (const token of normalizedTokens) {
    if (token.length >= 2) {
      tokens.add(token);
    }
    const cjkSequences = token.match(/[\u3400-\u9fff]+/gu) ?? [];
    for (const sequence of cjkSequences) {
      for (let index = 0; index < sequence.length - 1; index += 1) {
        tokens.add(sequence.slice(index, index + 2));
      }
    }
  }
  return Array.from(tokens);
}

function chunkKnowledgeText(input: KnowledgeEntryInput): string[] {
  const prefix = [input.title, input.category, input.tags].filter(Boolean).join("\n");
  const text = `${prefix}\n\n${input.body}`.trim();
  if (text.length <= CHUNK_TARGET_SIZE) {
    return [text];
  }

  const chunks: string[] = [];
  let start = 0;
  while (start < text.length) {
    const end = Math.min(start + CHUNK_TARGET_SIZE, text.length);
    chunks.push(text.slice(start, end).trim());
    if (end >= text.length) {
      break;
    }
    start = Math.max(end - CHUNK_OVERLAP_SIZE, start + 1);
  }
  return chunks.filter(Boolean);
}

function ensureWritableScope(context: AuthContext, input: KnowledgeEntryInput): ServiceResult<{ organizationId: string | null }> {
  if (input.scope === "platform") {
    if (!context.permissions.includes("platform:manage")) {
      return { success: false, reason: "只有平台管理员可以维护平台知识", status: 403 };
    }
    return { success: true, data: { organizationId: null } };
  }

  const organizationId = context.organization?.id;
  if (!organizationId) {
    return { success: false, reason: "请选择机构后维护知识库", status: 400 };
  }
  if (input.organizationId && input.organizationId !== organizationId) {
    return { success: false, reason: "不能维护其他机构的知识库", status: 403 };
  }
  return { success: true, data: { organizationId } };
}

function getVisibleKnowledgeWhere(organizationId?: string) {
  const platformFilter = and(eq(aiKnowledgeEntries.scope, "platform"), isNull(aiKnowledgeEntries.organizationId));
  if (!organizationId) {
    return platformFilter;
  }
  return or(
    platformFilter,
    and(eq(aiKnowledgeEntries.scope, "organization"), eq(aiKnowledgeEntries.organizationId, organizationId)),
  );
}

function getEnabledChunkWhere(organizationId: string) {
  return and(
    eq(aiKnowledgeEntries.status, "enabled"),
    or(
      and(eq(aiKnowledgeChunks.scope, "platform"), isNull(aiKnowledgeChunks.organizationId)),
      and(eq(aiKnowledgeChunks.scope, "organization"), eq(aiKnowledgeChunks.organizationId, organizationId)),
    ),
  );
}

function buildKnowledgeChunks(input: KnowledgeEntryInput): string[] {
  return chunkKnowledgeText(input);
}

export async function listKnowledgeEntries(context: AuthContext): Promise<KnowledgeEntry[]> {
  const database = getDatabase();
  const rows = await database
    .select()
    .from(aiKnowledgeEntries)
    .where(getVisibleKnowledgeWhere(context.organization?.id))
    .orderBy(desc(aiKnowledgeEntries.updatedAt));
  return rows.map(toKnowledgeEntry);
}

export async function createKnowledgeEntry(
  context: AuthContext,
  input: KnowledgeEntryInput,
): Promise<ServiceResult<{ entry: KnowledgeEntry }>> {
  const scope = ensureWritableScope(context, input);
  if (!scope.success) {
    return scope;
  }

  const chunks = buildKnowledgeChunks(input);
  const database = getDatabase();

  const row = await database.transaction(async (transaction) => {
    const rows = await transaction
      .insert(aiKnowledgeEntries)
      .values({
        scope: input.scope,
        organizationId: scope.data.organizationId,
        title: input.title,
        category: input.category,
        tags: input.tags,
        body: input.body,
        status: input.status,
        createdByAccountId: context.account.id,
        updatedByAccountId: context.account.id,
        updatedAt: new Date(),
      })
      .returning();
    const entry = rows[0];
    if (!entry) {
      return null;
    }
    if (chunks.length > 0) {
      await transaction.insert(aiKnowledgeChunks).values(
        chunks.map((content, index) => ({
          entryId: entry.id,
          organizationId: entry.organizationId,
          scope: entry.scope,
          chunkIndex: index,
          content,
          embedding: [],
          sourceTitle: entry.title,
          sourceCategory: entry.category,
        })),
      );
    }
    return entry;
  });

  if (!row) {
    return { success: false, reason: "知识库创建失败", status: 500 };
  }

  return { success: true, data: { entry: toKnowledgeEntry(row) } };
}

async function updateEntryWithChunks(
  id: string,
  context: AuthContext,
  input: KnowledgeEntryInput,
  organizationId: string | null,
): Promise<KnowledgeRow | null> {
  const chunks = buildKnowledgeChunks(input);

  const database = getDatabase();
  return database.transaction(async (transaction) => {
    const rows = await transaction
      .update(aiKnowledgeEntries)
      .set({
        scope: input.scope,
        organizationId,
        title: input.title,
        category: input.category,
        tags: input.tags,
        body: input.body,
        status: input.status,
        updatedByAccountId: context.account.id,
        updatedAt: new Date(),
      })
      .where(eq(aiKnowledgeEntries.id, id))
      .returning();
    const row = rows[0];
    if (!row) {
      return null;
    }

    await transaction.delete(aiKnowledgeChunks).where(eq(aiKnowledgeChunks.entryId, row.id));
    if (chunks.length > 0) {
      await transaction.insert(aiKnowledgeChunks).values(
        chunks.map((content, index) => ({
          entryId: row.id,
          organizationId: row.organizationId,
          scope: row.scope,
          chunkIndex: index,
          content,
          embedding: [],
          sourceTitle: row.title,
          sourceCategory: row.category,
        })),
      );
    }

    return row;
  });
}

export async function updateKnowledgeEntry(
  context: AuthContext,
  id: string,
  input: KnowledgeEntryInput,
): Promise<ServiceResult<{ entry: KnowledgeEntry }>> {
  const scope = ensureWritableScope(context, input);
  if (!scope.success) {
    return scope;
  }

  const database = getDatabase();
  const existingRows = await database
    .select()
    .from(aiKnowledgeEntries)
    .where(and(eq(aiKnowledgeEntries.id, id), getVisibleKnowledgeWhere(context.organization?.id)))
    .limit(1);
  const existing = existingRows[0];
  if (!existing) {
    return { success: false, reason: "知识库条目不存在", status: 404 };
  }

  if (existing.scope === "platform" && !context.permissions.includes("platform:manage")) {
    return { success: false, reason: "只有平台管理员可以维护平台知识", status: 403 };
  }

  let row: KnowledgeRow | null;
  try {
    row = await updateEntryWithChunks(id, context, input, scope.data.organizationId);
  } catch {
    return { success: false, reason: "知识库更新失败", status: 500 };
  }

  if (!row) {
    return { success: false, reason: "知识库更新失败", status: 500 };
  }
  return { success: true, data: { entry: toKnowledgeEntry(row) } };
}

export async function deleteKnowledgeEntry(context: AuthContext, id: string): Promise<ServiceResult<{ entry: KnowledgeEntry }>> {
  const database = getDatabase();
  const existingRows = await database
    .select()
    .from(aiKnowledgeEntries)
    .where(and(eq(aiKnowledgeEntries.id, id), getVisibleKnowledgeWhere(context.organization?.id)))
    .limit(1);
  const existing = existingRows[0];
  if (!existing) {
    return { success: false, reason: "知识库条目不存在", status: 404 };
  }
  if (existing.scope === "platform" && !context.permissions.includes("platform:manage")) {
    return { success: false, reason: "只有平台管理员可以维护平台知识", status: 403 };
  }

  const rows = await database.delete(aiKnowledgeEntries).where(eq(aiKnowledgeEntries.id, id)).returning();
  const row = rows[0];
  if (!row) {
    return { success: false, reason: "知识库删除失败", status: 500 };
  }
  return { success: true, data: { entry: toKnowledgeEntry(row) } };
}

function scoreKnowledgeChunk(
  chunk: Pick<KnowledgeRetrievalResult, "title" | "category" | "content">,
  queryTokens: readonly string[],
): number {
  if (queryTokens.length === 0) {
    return 0;
  }

  const chunkTokens = new Set(normalizeSearchText(`${chunk.title}\n${chunk.category}\n${chunk.content}`));
  let matchedTokens = 0;
  for (const token of queryTokens) {
    if (chunkTokens.has(token)) {
      matchedTokens += 1;
    }
  }
  return matchedTokens / queryTokens.length;
}

export async function retrieveKnowledge(
  organizationId: string,
  query: string,
  options?: { limit?: number },
): Promise<ServiceResult<{ results: KnowledgeRetrievalResult[] }>> {
  const queryTokens = normalizeSearchText(query);

  const database = getDatabase();
  const rows = await database
    .select({
      entryId: aiKnowledgeChunks.entryId,
      chunkId: aiKnowledgeChunks.id,
      title: aiKnowledgeChunks.sourceTitle,
      category: aiKnowledgeChunks.sourceCategory,
      content: aiKnowledgeChunks.content,
    })
    .from(aiKnowledgeChunks)
    .innerJoin(aiKnowledgeEntries, eq(aiKnowledgeEntries.id, aiKnowledgeChunks.entryId))
    .where(getEnabledChunkWhere(organizationId));

  return {
    success: true,
    data: {
      results: rows
        .map((row) => ({
          entryId: row.entryId,
          chunkId: row.chunkId,
          title: row.title,
          category: row.category,
          content: row.content,
          score: scoreKnowledgeChunk(row, queryTokens),
        }))
        .filter((row) => row.score > 0)
        .sort((left, right) => right.score - left.score)
        .slice(0, options?.limit ?? DEFAULT_RETRIEVAL_LIMIT),
    },
  };
}

export function normalizeKnowledgeScope(scope: AiKnowledgeScope, organizationId?: string): string {
  return scope === "platform" ? "platform" : `organization:${organizationId ?? ""}`;
}