@goscribe/server 1.3.4 → 1.5.0

package/src/services/knowledge/knowledge-base.service.ts

@@ -0,0 +1,705 @@
+import type { PrismaClient } from '@prisma/client';
+import { TRPCError } from '@trpc/server';
+import { BaseService } from '../base.service.js';
+import { workspaceAccessWhere } from '../../repositories/workspace.repository.js';
+import { supabaseClient } from '../../lib/storage.js';
+import {
+  DEFAULT_EMBEDDING_DIM,
+  DEFAULT_EMBEDDING_MODEL,
+  embedQuery,
+  embedTexts,
+  toVectorLiteral,
+} from '../../lib/embeddings.js';
+import { extractPdfText } from '../../lib/pdf.js';
+import { chunkText } from '../../lib/chunking.js';
+import { ai } from '../../lib/ai/index.js';
+import { isWorkspaceManagedKb } from '../../lib/workspace-kb.js';
+import { buildKbMetaFromName } from '../../lib/kb-meta.js';
+
+const KB_BUCKET = process.env.SUPABASE_KB_BUCKET ?? 'media';
+const MAX_PDF_BYTES = 50 * 1024 * 1024;
+
+interface AdminCreateInput {
+  name: string;
+  description?: string;
+  isCurated?: boolean;
+}
+
+interface AdminUpdateInput {
+  id: string;
+  name?: string;
+  description?: string | null;
+  isCurated?: boolean;
+}
+
+interface RequestUploadUrlInput {
+  knowledgeBaseId: string;
+  filename: string;
+  contentType: string;
+  size: number;
+}
+
+interface ProcessDocumentInput {
+  knowledgeBaseId: string;
+  documentId: string;
+}
+
+interface QueryInput {
+  knowledgeBaseId: string;
+  query: string;
+  topK?: number;
+  generateAnswer?: boolean;
+}
+
+export class KnowledgeBaseService extends BaseService {
+  constructor(db: PrismaClient) {
+    super(db);
+  }
+
+  // ───────────────────────────────────────── access helpers
+
+  private async assertWorkspaceAccess(userId: string, workspaceId: string) {
+    const ws = await this.db.workspace.findFirst({
+      where: { id: workspaceId, ...workspaceAccessWhere(userId) },
+      select: { id: true },
+    });
+    if (!ws) {
+      throw new TRPCError({
+        code: 'NOT_FOUND',
+        message: 'Workspace not found or access denied',
+      });
+    }
+  }
+
+  /**
+   * Loads a knowledge base only if it's reachable to the calling user, i.e.
+   * either curated (visible in catalog) or attached to one of their accessible
+   * workspaces.
+   */
+  private async loadKbForUser(userId: string, kbId: string) {
+    const kb = await this.db.knowledgeBase.findFirst({
+      where: {
+        id: kbId,
+        OR: [
+          { isCurated: true },
+          { workspaces: { some: { workspace: workspaceAccessWhere(userId) } } },
+        ],
+      },
+    });
+    if (!kb) {
+      throw new TRPCError({
+        code: 'NOT_FOUND',
+        message: 'Knowledge base not found',
+      });
+    }
+    return kb;
+  }
+
+  /** Admin-only: load by id without membership/curation checks. */
+  private async loadKbAsAdmin(kbId: string) {
+    const kb = await this.db.knowledgeBase.findUnique({ where: { id: kbId } });
+    if (!kb) {
+      throw new TRPCError({
+        code: 'NOT_FOUND',
+        message: 'Knowledge base not found',
+      });
+    }
+    return kb;
+  }
+
+  // ───────────────────────────────────────── user-facing reads
+
+  /**
+   * KBs attached to a given workspace (the user's "library" for that
+   * workspace).
+   */
+  async listForWorkspace(userId: string, workspaceId: string) {
+    await this.assertWorkspaceAccess(userId, workspaceId);
+    const rows = await this.db.workspaceKnowledgeBase.findMany({
+      where: { workspaceId },
+      orderBy: { createdAt: 'desc' },
+      include: {
+        knowledgeBase: {
+          include: { _count: { select: { documents: true } } },
+        },
+      },
+    });
+    return rows
+      .map((r) => ({
+        ...r.knowledgeBase,
+        attachmentId: r.id,
+        attachedAt: r.createdAt,
+        attachedById: r.addedById,
+      }))
+      .filter((kb) => !isWorkspaceManagedKb(kb));
+  }
+
+  /**
+   * Search the curated catalog. Optionally pass a workspaceId to also tag
+   * each result with whether the workspace already has it attached.
+   */
+  async searchCatalog(
+    userId: string,
+    input: { query?: string; workspaceId?: string; limit?: number },
+  ) {
+    const limit = Math.min(Math.max(input.limit ?? 50, 1), 100);
+    const q = (input.query ?? '').trim();
+
+    if (input.workspaceId) {
+      await this.assertWorkspaceAccess(userId, input.workspaceId);
+    }
+
+    const kbs = await this.db.knowledgeBase.findMany({
+      where: {
+        isCurated: true,
+        ...(q
+          ? {
+              OR: [
+                { name: { contains: q, mode: 'insensitive' as const } },
+                { description: { contains: q, mode: 'insensitive' as const } },
+              ],
+            }
+          : {}),
+      },
+      orderBy: { updatedAt: 'desc' },
+      take: limit,
+      include: { _count: { select: { documents: true } } },
+    });
+
+    if (!input.workspaceId) {
+      return kbs.map((kb) => ({ ...kb, attached: false }));
+    }
+
+    const attachments = await this.db.workspaceKnowledgeBase.findMany({
+      where: {
+        workspaceId: input.workspaceId,
+        knowledgeBaseId: { in: kbs.map((k) => k.id) },
+      },
+      select: { knowledgeBaseId: true },
+    });
+    const attachedSet = new Set(attachments.map((a) => a.knowledgeBaseId));
+    return kbs.map((kb) => ({ ...kb, attached: attachedSet.has(kb.id) }));
+  }
+
+  async get(userId: string, id: string) {
+    const kb = await this.loadKbForUser(userId, id);
+    const documents = await this.db.knowledgeBaseDocument.findMany({
+      where: { knowledgeBaseId: id },
+      orderBy: { createdAt: 'desc' },
+    });
+    return { ...kb, documents };
+  }
+
+  async listDocuments(userId: string, knowledgeBaseId: string) {
+    await this.loadKbForUser(userId, knowledgeBaseId);
+    return this.db.knowledgeBaseDocument.findMany({
+      where: { knowledgeBaseId },
+      orderBy: { createdAt: 'desc' },
+    });
+  }
+
+  // ───────────────────────────────────────── user-facing writes
+
+  /** Attach a curated KB to one of the user's workspaces. */
+  async addToWorkspace(
+    userId: string,
+    input: { workspaceId: string; knowledgeBaseId: string },
+  ) {
+    await this.assertWorkspaceAccess(userId, input.workspaceId);
+
+    const kb = await this.db.knowledgeBase.findUnique({
+      where: { id: input.knowledgeBaseId },
+      select: { id: true, isCurated: true },
+    });
+    if (!kb) throw new TRPCError({ code: 'NOT_FOUND' });
+    if (!kb.isCurated) {
+      throw new TRPCError({
+        code: 'FORBIDDEN',
+        message:
+          'This knowledge base is private. Only an administrator can attach it to a workspace.',
+      });
+    }
+
+    const existing = await this.db.workspaceKnowledgeBase.findUnique({
+      where: {
+        workspaceId_knowledgeBaseId: {
+          workspaceId: input.workspaceId,
+          knowledgeBaseId: kb.id,
+        },
+      },
+    });
+    if (existing) return existing;
+
+    return this.db.workspaceKnowledgeBase.create({
+      data: {
+        workspaceId: input.workspaceId,
+        knowledgeBaseId: kb.id,
+        addedById: userId,
+      },
+    });
+  }
+
+  /** Detach a KB from one of the user's workspaces. */
+  async removeFromWorkspace(
+    userId: string,
+    input: { workspaceId: string; knowledgeBaseId: string },
+  ) {
+    await this.assertWorkspaceAccess(userId, input.workspaceId);
+
+    const kb = await this.db.knowledgeBase.findUnique({
+      where: { id: input.knowledgeBaseId },
+      select: { meta: true },
+    });
+    if (kb && isWorkspaceManagedKb(kb)) {
+      throw new TRPCError({
+        code: 'BAD_REQUEST',
+        message: 'The workspace knowledge base is managed automatically and cannot be removed.',
+      });
+    }
+
+    await this.db.workspaceKnowledgeBase.deleteMany({
+      where: {
+        workspaceId: input.workspaceId,
+        knowledgeBaseId: input.knowledgeBaseId,
+      },
+    });
+    return true;
+  }
+
+  // ───────────────────────────────────────── retrieval
+
+  async query(userId: string, input: QueryInput) {
+    const kb = await this.loadKbForUser(userId, input.knowledgeBaseId);
+    return this.queryForKb(kb, input);
+  }
+
+  private async queryForKb(
+    kb: { id: string; embeddingModel: string },
+    input: QueryInput,
+  ) {
+    const topK = Math.min(Math.max(input.topK ?? 5, 1), 20);
+
+    const queryVec = await embedQuery(input.query, { model: kb.embeddingModel });
+    const vecLiteral = toVectorLiteral(queryVec);
+
+    const rows = await this.db.$queryRawUnsafe<
+      Array<{
+        id: string;
+        documentId: string;
+        chunkIndex: number;
+        content: string;
+        distance: number;
+        documentName: string;
+      }>
+    >(
+      `SELECT c."id", c."documentId", c."chunkIndex", c."content",
+              (c."embedding" <=> $1::vector) AS "distance",
+              d."name" AS "documentName"
+         FROM "KnowledgeBaseChunk" c
+         JOIN "KnowledgeBaseDocument" d ON d."id" = c."documentId"
+        WHERE c."knowledgeBaseId" = $2
+          AND c."embedding" IS NOT NULL
+        ORDER BY c."embedding" <=> $1::vector
+        LIMIT ${topK}`,
+      vecLiteral,
+      kb.id,
+    );
+
+    const matches = rows.map((r) => ({
+      chunkId: r.id,
+      documentId: r.documentId,
+      documentName: r.documentName,
+      chunkIndex: r.chunkIndex,
+      content: r.content,
+      score: 1 - Number(r.distance),
+    }));
+
+    let answer: string | null = null;
+    if (input.generateAnswer && matches.length > 0) {
+      try {
+        const context = matches
+          .map(
+            (m, i) =>
+              `[${i + 1}] (${m.documentName}, chunk ${m.chunkIndex})\n${m.content}`,
+          )
+          .join('\n\n---\n\n');
+        const completion = await ai.llm.complete([
+          {
+            role: 'system',
+            content:
+              'You are a helpful assistant. Answer the user\'s question using ONLY the provided context. Cite sources as [1], [2], etc. If the answer is not in the context, say so.',
+          },
+          {
+            role: 'user',
+            content: `Context:\n${context}\n\nQuestion: ${input.query}`,
+          },
+        ]);
+        answer = completion.choices?.[0]?.message?.content ?? null;
+      } catch (err) {
+        this.logger.error('KB answer generation failed', err);
+      }
+    }
+
+    return { matches, answer };
+  }
+
+  // ───────────────────────────────────────── document processing core
+
+  private async processDocumentForKb(
+    kb: { id: string; embeddingModel: string },
+    documentId: string,
+  ) {
+    const doc = await this.db.knowledgeBaseDocument.findFirst({
+      where: { id: documentId, knowledgeBaseId: kb.id },
+    });
+    if (!doc) throw new TRPCError({ code: 'NOT_FOUND' });
+    if (!doc.bucket || !doc.objectKey) {
+      throw new TRPCError({
+        code: 'BAD_REQUEST',
+        message: 'Document is missing storage location',
+      });
+    }
+
+    if (doc.status === 'PROCESSING' || doc.status === 'READY') {
+      return doc;
+    }
+
+    await this.db.knowledgeBaseDocument.update({
+      where: { id: doc.id },
+      data: { status: 'PROCESSING', errorMessage: null },
+    });
+    await this.db.knowledgeBase.update({
+      where: { id: kb.id },
+      data: { status: 'INDEXING' },
+    });
+
+    try {
+      const { data, error } = await supabaseClient.storage
+        .from(doc.bucket)
+        .download(doc.objectKey);
+
+      if (error || !data) {
+        throw new Error(error?.message ?? 'Failed to download document');
+      }
+
+      const buf = Buffer.from(await data.arrayBuffer());
+      let text: string;
+      let numPages: number | null = null;
+
+      if (doc.mimeType === 'application/pdf' || doc.name.toLowerCase().endsWith('.pdf')) {
+        const parsed = await extractPdfText(buf);
+        text = parsed.text;
+        numPages = parsed.numPages;
+      } else {
+        text = buf.toString('utf8');
+      }
+
+
+      if (!text.trim()) {
+        throw new Error('No text could be extracted from the document');
+      }
+
+      const chunks = chunkText(text, { targetTokens: 500, overlapTokens: 60 });
+
+      if (chunks.length === 0) {
+        throw new Error('Document produced no chunks');
+      }
+        console.log('chunks', chunks);
+      const embeddings = await embedTexts(
+        chunks.map((c) => c.content),
+        { model: kb.embeddingModel },
+      );
+
+      console.log('embeddings', embeddings);
+
+      await this.db.knowledgeBaseChunk.deleteMany({
+        where: { documentId: doc.id },
+      });
+
+      for (let i = 0; i < chunks.length; i += 1) {
+        const chunk = chunks[i];
+        const vec = toVectorLiteral(embeddings[i] ?? []);
+        await this.db.$executeRawUnsafe(
+          `INSERT INTO "KnowledgeBaseChunk"
+            ("id", "documentId", "knowledgeBaseId", "chunkIndex",
+             "pageNumber", "content", "tokenCount", "embedding", "createdAt")
+           VALUES (
+             gen_random_uuid()::text,
+             $1, $2, $3, $4, $5, $6, $7::vector, NOW()
+           )`,
+          doc.id,
+          kb.id,
+          chunk.index,
+          null,
+          chunk.content,
+          chunk.approxTokens,
+          vec,
+        );
+      }
+
+      await this.db.knowledgeBaseDocument.update({
+        where: { id: doc.id },
+        data: {
+          status: 'READY',
+          numChunks: chunks.length,
+          numPages: numPages ?? undefined,
+        },
+      });
+      
+
+      const stillProcessing = await this.db.knowledgeBaseDocument.count({
+        where: { knowledgeBaseId: kb.id, status: 'PROCESSING' },
+      });
+      if (stillProcessing === 0) {
+        await this.db.knowledgeBase.update({
+          where: { id: kb.id },
+          data: { status: 'READY' },
+        });
+      }
+
+      return this.db.knowledgeBaseDocument.findUnique({
+        where: { id: doc.id },
+      });
+    } catch (err: any) {
+      this.logger.error(`KB processing failed for doc ${doc.id}`, err);
+      await this.db.knowledgeBaseDocument.update({
+        where: { id: doc.id },
+        data: {
+          status: 'FAILED',
+          errorMessage: err?.message?.slice(0, 500) ?? 'Unknown error',
+        },
+      });
+      await this.db.knowledgeBase.update({
+        where: { id: kb.id },
+        data: { status: 'ERROR' },
+      });
+      throw new TRPCError({
+        code: 'INTERNAL_SERVER_ERROR',
+        message: err?.message ?? 'Failed to process document',
+      });
+    }
+  }
+
+  // ───────────────────────────────────────── admin entrypoints
+
+  async adminListAll() {
+    return this.db.knowledgeBase.findMany({
+      orderBy: { updatedAt: 'desc' },
+      include: {
+        _count: {
+          select: { documents: true, workspaces: true },
+        },
+      },
+    });
+  }
+
+  async adminGet(id: string) {
+    const kb = await this.loadKbAsAdmin(id);
+    const [documents, workspaces] = await Promise.all([
+      this.db.knowledgeBaseDocument.findMany({
+        where: { knowledgeBaseId: id },
+        orderBy: { createdAt: 'desc' },
+      }),
+      this.db.workspaceKnowledgeBase.findMany({
+        where: { knowledgeBaseId: id },
+        orderBy: { createdAt: 'desc' },
+        include: {
+          workspace: { select: { id: true, title: true, icon: true } },
+        },
+      }),
+    ]);
+    return { ...kb, documents, workspaces };
+  }
+
+  async adminListDocuments(knowledgeBaseId: string) {
+    await this.loadKbAsAdmin(knowledgeBaseId);
+    return this.db.knowledgeBaseDocument.findMany({
+      where: { knowledgeBaseId },
+      orderBy: { createdAt: 'desc' },
+    });
+  }
+
+  async adminCreate(actorUserId: string, input: AdminCreateInput) {
+    const meta = buildKbMetaFromName(input.name);
+    return this.db.knowledgeBase.create({
+      data: {
+        createdById: actorUserId,
+        name: input.name,
+        description: input.description,
+        isCurated: input.isCurated ?? false,
+        embeddingModel: DEFAULT_EMBEDDING_MODEL,
+        embeddingDim: DEFAULT_EMBEDDING_DIM,
+        meta,
+      },
+    });
+  }
+
+  async adminUpdate(input: AdminUpdateInput) {
+    const existing = await this.loadKbAsAdmin(input.id);
+    const nextName = input.name ?? existing.name;
+    const meta =
+      input.name !== undefined
+        ? buildKbMetaFromName(nextName, (existing.meta as Record<string, unknown> | null) ?? null)
+        : undefined;
+
+    return this.db.knowledgeBase.update({
+      where: { id: input.id },
+      data: {
+        name: input.name ?? undefined,
+        description:
+          input.description === undefined ? undefined : input.description,
+        isCurated: input.isCurated ?? undefined,
+        ...(meta !== undefined ? { meta } : {}),
+      },
+    });
+  }
+
+  async adminDelete(id: string) {
+    const kb = await this.loadKbAsAdmin(id);
+    const docs = await this.db.knowledgeBaseDocument.findMany({
+      where: { knowledgeBaseId: kb.id },
+      select: { bucket: true, objectKey: true },
+    });
+    for (const d of docs) {
+      if (d.bucket && d.objectKey) {
+        supabaseClient.storage
+          .from(d.bucket)
+          .remove([d.objectKey])
+          .catch((err: unknown) => {
+            this.logger.error(
+              `Failed to remove KB doc ${d.objectKey} from ${d.bucket}`,
+              err,
+            );
+          });
+      }
+    }
+    await this.db.knowledgeBase.delete({ where: { id: kb.id } });
+    return true;
+  }
+
+  /** Admin-only: attach a KB to a workspace (curated or not). */
+  async adminAttachToWorkspace(
+    actorUserId: string,
+    input: { workspaceId: string; knowledgeBaseId: string },
+  ) {
+    const [ws, kb] = await Promise.all([
+      this.db.workspace.findUnique({
+        where: { id: input.workspaceId },
+        select: { id: true },
+      }),
+      this.db.knowledgeBase.findUnique({
+        where: { id: input.knowledgeBaseId },
+        select: { id: true },
+      }),
+    ]);
+    if (!ws) throw new TRPCError({ code: 'NOT_FOUND', message: 'Workspace not found' });
+    if (!kb) throw new TRPCError({ code: 'NOT_FOUND', message: 'Knowledge base not found' });
+
+    const existing = await this.db.workspaceKnowledgeBase.findUnique({
+      where: {
+        workspaceId_knowledgeBaseId: {
+          workspaceId: ws.id,
+          knowledgeBaseId: kb.id,
+        },
+      },
+    });
+    if (existing) return existing;
+
+    return this.db.workspaceKnowledgeBase.create({
+      data: {
+        workspaceId: ws.id,
+        knowledgeBaseId: kb.id,
+        addedById: actorUserId,
+      },
+    });
+  }
+
+  async adminDetachFromWorkspace(input: {
+    workspaceId: string;
+    knowledgeBaseId: string;
+  }) {
+    await this.db.workspaceKnowledgeBase.deleteMany({
+      where: {
+        workspaceId: input.workspaceId,
+        knowledgeBaseId: input.knowledgeBaseId,
+      },
+    });
+    return true;
+  }
+
+  async adminRequestUploadUrl(input: RequestUploadUrlInput) {
+    if (input.size > MAX_PDF_BYTES) {
+      throw new TRPCError({
+        code: 'PAYLOAD_TOO_LARGE',
+        message: `Document exceeds the ${MAX_PDF_BYTES / (1024 * 1024)}MB limit`,
+      });
+    }
+    const kb = await this.loadKbAsAdmin(input.knowledgeBaseId);
+
+    const objectKey = `kb/${kb.id}/${Date.now()}-${input.filename.replace(
+      /[^a-zA-Z0-9._-]+/g,
+      '_',
+    )}`;
+
+    const doc = await this.db.knowledgeBaseDocument.create({
+      data: {
+        knowledgeBaseId: kb.id,
+        name: input.filename,
+        mimeType: input.contentType,
+        size: input.size,
+        bucket: KB_BUCKET,
+        objectKey,
+        status: 'PENDING',
+      },
+    });
+
+    const { data: signed, error } = await supabaseClient.storage
+      .from(KB_BUCKET)
+      .createSignedUploadUrl(objectKey);
+
+    if (error || !signed) {
+      this.logger.error('KB upload URL error', error);
+      await this.db.knowledgeBaseDocument.delete({ where: { id: doc.id } });
+      throw new TRPCError({
+        code: 'INTERNAL_SERVER_ERROR',
+        message: 'Failed to create upload URL',
+      });
+    }
+
+    return {
+      documentId: doc.id,
+      uploadUrl: signed.signedUrl,
+      bucket: KB_BUCKET,
+      objectKey,
+    };
+  }
+
+  async adminProcessDocument(input: ProcessDocumentInput) {
+    const kb = await this.loadKbAsAdmin(input.knowledgeBaseId);
+    return this.processDocumentForKb(kb, input.documentId);
+  }
+
+  async adminDeleteDocument(documentId: string) {
+    const doc = await this.db.knowledgeBaseDocument.findUnique({
+      where: { id: documentId },
+    });
+    if (!doc) throw new TRPCError({ code: 'NOT_FOUND' });
+
+    if (doc.bucket && doc.objectKey) {
+      supabaseClient.storage
+        .from(doc.bucket)
+        .remove([doc.objectKey])
+        .catch((err: unknown) => {
+          this.logger.error('Failed to remove KB doc object', err);
+        });
+    }
+
+    await this.db.knowledgeBaseDocument.delete({ where: { id: documentId } });
+    return true;
+  }
+
+  async adminQuery(input: QueryInput) {
+    const kb = await this.loadKbAsAdmin(input.knowledgeBaseId);
+    return this.queryForKb(kb, input);
+  }
+}