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

import { beforeEach, describe, expect, it, vi } from "vitest";

import type { AiRuntimeConfigResult } from "@/modules/ai/server/config";
import { getAiRuntimeConfig } from "@/modules/ai/server/config";
import {
  createKnowledgeEntry,
  listKnowledgeEntries,
  retrieveKnowledge,
  updateKnowledgeEntry,
} from "@/modules/ai/server/knowledge";
import type { KnowledgeEntryInput } from "@/modules/ai/types";
import type { AppDatabase } from "@/modules/core/server/db";
import { getDatabase } from "@/modules/core/server/db";
import type { AuthContext, Permission } from "@/modules/core/types";

const embeddingMocks = vi.hoisted(() => ({
  embedDocuments: vi.fn(),
}));

const drizzleDsl = vi.hoisted(() => {
  type TableName = "entries" | "chunks";
  type ColumnRef = { table: TableName; key: string };
  const entries = {
    id: { table: "entries", key: "id" } as ColumnRef,
    scope: { table: "entries", key: "scope" } as ColumnRef,
    organizationId: { table: "entries", key: "organizationId" } as ColumnRef,
    status: { table: "entries", key: "status" } as ColumnRef,
    updatedAt: { table: "entries", key: "updatedAt" } as ColumnRef,
  };
  const chunks = {
    id: { table: "chunks", key: "id" } as ColumnRef,
    entryId: { table: "chunks", key: "entryId" } as ColumnRef,
    scope: { table: "chunks", key: "scope" } as ColumnRef,
    organizationId: { table: "chunks", key: "organizationId" } as ColumnRef,
    embedding: { table: "chunks", key: "embedding" } as ColumnRef,
    sourceTitle: { table: "chunks", key: "sourceTitle" } as ColumnRef,
    sourceCategory: { table: "chunks", key: "sourceCategory" } as ColumnRef,
    content: { table: "chunks", key: "content" } as ColumnRef,
  };
  return { entries, chunks };
});

vi.mock("server-only", () => ({}));

vi.mock("@langchain/openai", () => ({
  OpenAIEmbeddings: vi.fn(function MockOpenAIEmbeddings() {
    return {
      embedDocuments: embeddingMocks.embedDocuments,
    };
  }),
}));

vi.mock("@/modules/ai/server/config", () => ({
  getAiRuntimeConfig: vi.fn(),
}));

vi.mock("@/modules/core/server/db", () => ({
  getDatabase: vi.fn(),
}));

vi.mock("@/modules/core/server/schema", () => ({
  aiKnowledgeEntries: drizzleDsl.entries,
  aiKnowledgeChunks: drizzleDsl.chunks,
}));

vi.mock("drizzle-orm", () => {
  type JoinedRow = {
    entries?: Record<string, unknown>;
    chunks?: Record<string, unknown>;
  };
  type ColumnRef = { table: "entries" | "chunks"; key: string };
  type Predicate = (row: JoinedRow) => boolean;

  function readColumn(row: JoinedRow, column: ColumnRef): unknown {
    return row[column.table]?.[column.key];
  }

  return {
    eq: (column: ColumnRef, expected: unknown): Predicate => (row) => readColumn(row, column) === expected,
    isNull: (column: ColumnRef): Predicate => (row) => readColumn(row, column) === null || readColumn(row, column) === undefined,
    and: (...predicates: Predicate[]): Predicate => (row) => predicates.every((predicate) => predicate(row)),
    or: (...predicates: Predicate[]): Predicate => (row) => predicates.some((predicate) => predicate(row)),
    desc: (column: ColumnRef) => ({ type: "desc", column }),
    sql: () => ({ type: "distance" }),
  };
});

type KnowledgeScope = "platform" | "organization";
type KnowledgeStatus = "enabled" | "disabled";

type KnowledgeRow = {
  id: string;
  scope: KnowledgeScope;
  organizationId: string | null;
  title: string;
  category: string;
  tags: string;
  body: string;
  status: KnowledgeStatus;
  createdByAccountId: string | null;
  updatedByAccountId: string | null;
  createdAt: Date;
  updatedAt: Date;
};

type KnowledgeChunkRow = {
  id: string;
  entryId: string;
  organizationId: string | null;
  scope: KnowledgeScope;
  chunkIndex: number;
  content: string;
  embedding: number[];
  sourceTitle: string;
  sourceCategory: string;
  createdAt: Date;
};

type JoinedRow = {
  entries?: KnowledgeRow;
  chunks?: KnowledgeChunkRow;
};

type Predicate = (row: JoinedRow) => boolean;
type SelectRow = KnowledgeRow | RetrievedChunkRow;

type RetrievedChunkRow = {
  entryId: string;
  chunkId: string;
  title: string;
  category: string;
  content: string;
  distance: number;
};

type KnowledgeDatabaseFake = {
  entries: KnowledgeRow[];
  chunks: KnowledgeChunkRow[];
  transaction: ReturnlessFunction;
  select: ReturnlessFunction;
};

type ReturnlessFunction = (...args: unknown[]) => unknown;

const runtimeConfig: AiRuntimeConfigResult = {
  success: true,
  config: {
    baseUrl: "",
    apiKey: "test-api-key",
    chatModel: "gpt-test",
    embeddingModel: "embedding-test",
    embeddingDimensions: 1536,
  },
};

const organization: NonNullable<AuthContext["organization"]> = {
  id: "org-1",
  name: "TeaCare Home",
  slug: "teacare",
  status: "active",
  registrationEnabled: true,
  oidcEnabled: false,
  oidcIssuerUrl: "",
  oidcClientId: "",
  oidcHasClientSecret: false,
  oidcScopes: "",
  oidcRedirectUri: "",
  oidcAvatarClaim: "",
  oidcAutoProvision: false,
  createdAt: "2026-07-02T00:00:00.000Z",
  updatedAt: "2026-07-02T00:00:00.000Z",
};

const account: AuthContext["account"] = {
  id: "account-1",
  name: "Knowledge Admin",
  email: "knowledge@example.com",
  avatarUrl: "",
  role: "org_admin",
  status: "active",
  createdAt: "2026-07-02T00:00:00.000Z",
  updatedAt: "2026-07-02T00:00:00.000Z",
};

const baseInput: KnowledgeEntryInput = {
  scope: "organization",
  organizationId: "org-1",
  title: "Fall prevention",
  category: "Safety",
  tags: "fall,night",
  body: "Keep walkways clear and increase night rounds after wandering events.",
  status: "enabled",
};

const organizationContext: NonNullable<AuthContext["organization"]> = organization;

function createAuthContext(input: { permissions?: Permission[]; activeOrganization?: boolean } = {}): AuthContext {
  const activeOrganization = input.activeOrganization ?? true;
  const permissions = input.permissions ?? ["knowledge:read", "knowledge:manage"];
  return {
    account,
    organization: activeOrganization ? organizationContext : undefined,
    organizations: [
      {
        id: organizationContext.id,
        name: organizationContext.name,
        slug: organizationContext.slug,
        status: organizationContext.status,
        registrationEnabled: organizationContext.registrationEnabled,
        oidcEnabled: organizationContext.oidcEnabled,
        roleLabel: "Admin",
        isActive: activeOrganization,
      },
    ],
    membership: activeOrganization
      ? {
          id: "membership-1",
          accountId: account.id,
          organizationId: organizationContext.id,
          roleId: "role-1",
          roleKey: "org_admin",
          roleLabel: "Admin",
          status: "active",
          createdAt: "2026-07-02T00:00:00.000Z",
          updatedAt: "2026-07-02T00:00:00.000Z",
        }
      : undefined,
    permissions,
    session: {
      id: "session-1",
      accountId: account.id,
      activeOrganizationId: activeOrganization ? organizationContext.id : undefined,
      createdAt: "2026-07-02T00:00:00.000Z",
      expiresAt: "2026-07-09T00:00:00.000Z",
    },
  };
}

function createEntry(overrides: Partial<KnowledgeRow> = {}): KnowledgeRow {
  return {
    id: "entry-org-1",
    scope: "organization",
    organizationId: "org-1",
    title: "Org fall prevention",
    category: "Safety",
    tags: "fall",
    body: "Organization-specific fall prevention guidance.",
    status: "enabled",
    createdByAccountId: "account-1",
    updatedByAccountId: "account-1",
    createdAt: new Date("2026-07-01T00:00:00.000Z"),
    updatedAt: new Date("2026-07-02T00:00:00.000Z"),
    ...overrides,
  };
}

function createChunk(entry: KnowledgeRow, overrides: Partial<KnowledgeChunkRow> = {}): KnowledgeChunkRow {
  return {
    id: `${entry.id}-chunk-1`,
    entryId: entry.id,
    organizationId: entry.organizationId,
    scope: entry.scope,
    chunkIndex: 0,
    content: entry.body,
    embedding: [0.1, 0.2, 0.3],
    sourceTitle: entry.title,
    sourceCategory: entry.category,
    createdAt: new Date("2026-07-02T00:00:00.000Z"),
    ...overrides,
  };
}

class SelectBuilder implements PromiseLike<SelectRow[]> {
  private table: "entries" | "chunks" = "entries";
  private predicate: Predicate = () => true;

  constructor(
    private readonly entries: KnowledgeRow[],
    private readonly chunks: KnowledgeChunkRow[],
  ) {}

  from(table: unknown): SelectBuilder {
    this.table = table === drizzleDsl.chunks ? "chunks" : "entries";
    return this;
  }

  innerJoin(): SelectBuilder {
    return this;
  }

  where(predicate: Predicate): SelectBuilder {
    this.predicate = predicate;
    return this;
  }

  orderBy(): SelectBuilder {
    return this;
  }

  limit(count: number): Promise<SelectRow[]> {
    return Promise.resolve(this.execute().slice(0, count));
  }

  then<TResult1 = SelectRow[], TResult2 = never>(
    onfulfilled?: ((value: SelectRow[]) => TResult1 | PromiseLike<TResult1>) | null,
    onrejected?: ((reason: unknown) => TResult2 | PromiseLike<TResult2>) | null,
  ): Promise<TResult1 | TResult2> {
    return Promise.resolve(this.execute()).then(onfulfilled, onrejected);
  }

  private execute(): SelectRow[] {
    if (this.table === "entries") {
      return this.entries
        .filter((entry) => this.predicate({ entries: entry }))
        .sort((left, right) => right.updatedAt.getTime() - left.updatedAt.getTime());
    }

    const rows: RetrievedChunkRow[] = [];
    for (const chunk of this.chunks) {
      const entry = this.entries.find((candidate) => candidate.id === chunk.entryId);
      if (entry && this.predicate({ entries: entry, chunks: chunk })) {
        rows.push({
          entryId: chunk.entryId,
          chunkId: chunk.id,
          title: chunk.sourceTitle,
          category: chunk.sourceCategory,
          content: chunk.content,
          distance: 0.1 + rows.length * 0.1,
        });
      }
    }
    return rows;
  }
}

function createDatabaseFake(input: { entries?: KnowledgeRow[]; chunks?: KnowledgeChunkRow[] } = {}): KnowledgeDatabaseFake {
  const entries = input.entries ?? [];
  const chunks = input.chunks ?? [];
  return {
    entries,
    chunks,
    transaction: vi.fn(),
    select: vi.fn(() => new SelectBuilder(entries, chunks)),
  };
}

function useDatabase(fake: KnowledgeDatabaseFake): void {
  vi.mocked(getDatabase).mockReturnValue(fake as unknown as AppDatabase);
}

beforeEach(() => {
  vi.clearAllMocks();
  vi.mocked(getAiRuntimeConfig).mockReturnValue(runtimeConfig);
  embeddingMocks.embedDocuments.mockResolvedValue([[0.4, 0.5, 0.6]]);
  useDatabase(createDatabaseFake());
});

describe("AI knowledge service", () => {
  it("lists platform and active-organization entries while excluding other organizations", async () => {
    const platformEntry = createEntry({
      id: "entry-platform",
      scope: "platform",
      organizationId: null,
      title: "Platform safety standard",
      updatedAt: new Date("2026-07-03T00:00:00.000Z"),
    });
    const orgEntry = createEntry({ id: "entry-org-1", title: "Org safety note" });
    const otherOrgEntry = createEntry({ id: "entry-org-2", organizationId: "org-2", title: "Other org note" });
    useDatabase(createDatabaseFake({ entries: [orgEntry, otherOrgEntry, platformEntry] }));

    const entries = await listKnowledgeEntries(createAuthContext());

    expect(entries.map((entry) => entry.id)).toEqual(["entry-platform", "entry-org-1"]);
    expect(entries).toEqual([
      expect.objectContaining({ id: "entry-platform", scope: "platform", organizationId: undefined }),
      expect.objectContaining({ id: "entry-org-1", scope: "organization", organizationId: "org-1" }),
    ]);
  });

  it("retrieves only enabled platform and active-organization chunks", async () => {
    const platformEntry = createEntry({ id: "entry-platform", scope: "platform", organizationId: null, title: "Platform standard" });
    const orgEntry = createEntry({ id: "entry-org-1", title: "Org guidance" });
    const disabledEntry = createEntry({ id: "entry-disabled", status: "disabled", title: "Disabled guidance" });
    const otherOrgEntry = createEntry({ id: "entry-org-2", organizationId: "org-2", title: "Other org guidance" });
    const database = createDatabaseFake({
      entries: [platformEntry, orgEntry, disabledEntry, otherOrgEntry],
      chunks: [
        createChunk(platformEntry),
        createChunk(orgEntry),
        createChunk(disabledEntry),
        createChunk(otherOrgEntry, { organizationId: "org-2", scope: "organization" }),
      ],
    });
    useDatabase(database);

    const result = await retrieveKnowledge("org-1", "night fall risk", { limit: 10 });

    expect(result).toEqual({
      success: true,
      data: {
        results: [
          expect.objectContaining({ entryId: "entry-platform", title: "Platform standard" }),
          expect.objectContaining({ entryId: "entry-org-1", title: "Org guidance" }),
        ],
      },
    });
    expect(embeddingMocks.embedDocuments).toHaveBeenCalledWith(["night fall risk"]);
  });

  it("requires platform management permission before mutating platform knowledge", async () => {
    const existingPlatformEntry = createEntry({ id: "entry-platform", scope: "platform", organizationId: null });
    useDatabase(createDatabaseFake({ entries: [existingPlatformEntry] }));

    const result = await updateKnowledgeEntry(
      createAuthContext({ permissions: ["knowledge:manage"] }),
      "entry-platform",
      baseInput,
    );

    expect(result).toEqual({ success: false, reason: "只有平台管理员可以维护平台知识", status: 403 });
    expect(embeddingMocks.embedDocuments).not.toHaveBeenCalled();
  });

  it("requires an active organization before mutating organization knowledge", async () => {
    const result = await createKnowledgeEntry(
      createAuthContext({ activeOrganization: false, permissions: ["knowledge:manage"] }),
      baseInput,
    );

    expect(result).toEqual({ success: false, reason: "请选择机构后维护知识库", status: 400 });
    expect(embeddingMocks.embedDocuments).not.toHaveBeenCalled();
  });

  it("rejects attempts to mutate another organization's knowledge", async () => {
    const result = await createKnowledgeEntry(createAuthContext(), {
      ...baseInput,
      organizationId: "org-2",
    });

    expect(result).toEqual({ success: false, reason: "不能维护其他机构的知识库", status: 403 });
    expect(embeddingMocks.embedDocuments).not.toHaveBeenCalled();
  });

  it("surfaces embedding failures without opening a write transaction", async () => {
    const database = createDatabaseFake();
    useDatabase(database);
    embeddingMocks.embedDocuments.mockRejectedValue(new Error("embedding provider unavailable"));

    const result = await createKnowledgeEntry(createAuthContext(), baseInput);

    expect(result).toEqual({ success: false, reason: "知识库向量生成失败", status: 500 });
    expect(database.transaction).not.toHaveBeenCalled();
  });
});