@goscribe/server 1.3.4 → 1.6.0

package/dist/services/knowledge/knowledge-base.service.js

@@ -0,0 +1,544 @@
+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;
+export class KnowledgeBaseService extends BaseService {
+    constructor(db) {
+        super(db);
+    }
+    // ───────────────────────────────────────── access helpers
+    async assertWorkspaceAccess(userId, workspaceId) {
+        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.
+     */
+    async loadKbForUser(userId, kbId) {
+        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. */
+    async loadKbAsAdmin(kbId) {
+        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, workspaceId) {
+        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, input) {
+        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' } },
+                            { description: { contains: q, mode: 'insensitive' } },
+                        ],
+                    }
+                    : {}),
+            },
+            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, id) {
+        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, knowledgeBaseId) {
+        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, input) {
+        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, input) {
+        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, input) {
+        const kb = await this.loadKbForUser(userId, input.knowledgeBaseId);
+        return this.queryForKb(kb, input);
+    }
+    async queryForKb(kb, input) {
+        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(`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 = 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
+    async processDocumentForKb(kb, documentId) {
+        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;
+            let numPages = 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) {
+            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) {
+        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) {
+        await this.loadKbAsAdmin(knowledgeBaseId);
+        return this.db.knowledgeBaseDocument.findMany({
+            where: { knowledgeBaseId },
+            orderBy: { createdAt: 'desc' },
+        });
+    }
+    async adminCreate(actorUserId, input) {
+        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) {
+        const existing = await this.loadKbAsAdmin(input.id);
+        const nextName = input.name ?? existing.name;
+        const meta = input.name !== undefined
+            ? buildKbMetaFromName(nextName, existing.meta ?? 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) {
+        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) => {
+                    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, input) {
+        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) {
+        await this.db.workspaceKnowledgeBase.deleteMany({
+            where: {
+                workspaceId: input.workspaceId,
+                knowledgeBaseId: input.knowledgeBaseId,
+            },
+        });
+        return true;
+    }
+    async adminRequestUploadUrl(input) {
+        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) {
+        const kb = await this.loadKbAsAdmin(input.knowledgeBaseId);
+        return this.processDocumentForKb(kb, input.documentId);
+    }
+    async adminDeleteDocument(documentId) {
+        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) => {
+                this.logger.error('Failed to remove KB doc object', err);
+            });
+        }
+        await this.db.knowledgeBaseDocument.delete({ where: { id: documentId } });
+        return true;
+    }
+    async adminQuery(input) {
+        const kb = await this.loadKbAsAdmin(input.knowledgeBaseId);
+        return this.queryForKb(kb, input);
+    }
+}