@goscribe/server 1.1.7 → 1.3.0

package/src/routers/copilot.ts

@@ -0,0 +1,719 @@
+import { TRPCError } from "@trpc/server";
+import { z } from "zod";
+import { router, verifiedProcedure } from "../trpc.js";
+import inference from "../lib/inference.js";
+import { sanitizeString } from "../lib/validation.js";
+import { workspaceAccessFilter } from "../lib/workspace-access.js";
+import PusherService from "../lib/pusher.js";
+import { ArtifactType } from "../lib/constants.js";
+import type { Context } from "../context.js";
+
+const copilotArtifactType = z.enum([
+  "study-guide",
+  "worksheet",
+  "flashcards",
+  "notes",
+]);
+
+const copilotContextSchema = z.object({
+  workspaceId: z.string().min(1),
+  artifactId: z.string().min(1),
+  artifactType: copilotArtifactType,
+  documentContent: z.string().min(1),
+  selectedText: z.string().optional(),
+  viewportText: z.string().optional(),
+  cursorPosition: z
+    .object({
+      start: z.number().int().min(0),
+      end: z.number().int().min(0),
+    })
+    .optional(),
+  metadata: z
+    .object({
+      flashcardId: z.string().optional(),
+      questionId: z.string().optional(),
+      documentId: z.string().optional(),
+    })
+    .optional(),
+});
+
+const highlightInstructionSchema = z.object({
+  start: z.number().int().min(0),
+  end: z.number().int().min(0),
+  label: z.string().optional(),
+});
+
+const flashcardSchema = z.object({
+  front: z.string().min(1),
+  back: z.string().min(1),
+});
+
+type CopilotContextInput = z.infer<typeof copilotContextSchema>;
+
+const RATE_LIMIT_WINDOW_MS = 60_000;
+const RATE_LIMIT_MAX_REQUESTS = 20;
+const requestBuckets = new Map<string, number[]>();
+
+function enforceRateLimit(userId: string, workspaceId: string) {
+  const now = Date.now();
+  const key = `${userId}:${workspaceId}`;
+  const existing = requestBuckets.get(key) ?? [];
+  const recent = existing.filter((timestamp) => now - timestamp < RATE_LIMIT_WINDOW_MS);
+
+  if (recent.length >= RATE_LIMIT_MAX_REQUESTS) {
+    throw new TRPCError({
+      code: "TOO_MANY_REQUESTS",
+      message: "Too many copilot requests. Please wait a moment.",
+    });
+  }
+
+  recent.push(now);
+  requestBuckets.set(key, recent);
+}
+
+function contextWindow(context: CopilotContextInput): string {
+  // Pass the full document so the copilot has full context and does not make mistakes
+  return sanitizeString(context.documentContent, 5_000_000);
+}
+
+function extractJson(content: string): Record<string, unknown> | null {
+  const fencedMatch = content.match(/```json\s*([\s\S]*?)```/i);
+  if (fencedMatch?.[1]) {
+    try {
+      return JSON.parse(fencedMatch[1]);
+    } catch {
+      // fall through
+    }
+  }
+
+  const objectMatch = content.match(/\{[\s\S]*\}/);
+  if (objectMatch?.[0]) {
+    try {
+      return JSON.parse(objectMatch[0]);
+    } catch {
+      return null;
+    }
+  }
+
+  return null;
+}
+
+function getAnswerFromParsedJSON(parsed: Record<string, any> | null): string | null {
+  if (!parsed) return null;
+
+  const potentialKeys = [
+    "answer",
+    "summary",
+    "explanation",
+    "content",
+    "explanationText",
+    "response",
+    "result",
+  ];
+
+  for (const key of potentialKeys) {
+    if (typeof parsed[key] === "string" && parsed[key].length > 0) {
+      return parsed[key];
+    }
+  }
+
+  return null;
+}
+
+async function assertWorkspaceAccess(
+  ctx: Context & { session: { user: { id: string } } },
+  workspaceId: string
+) {
+  const workspace = await ctx.db.workspace.findFirst({
+    where: {
+      id: workspaceId,
+      ...workspaceAccessFilter(ctx.session.user.id),
+    },
+    select: { id: true },
+  });
+
+  if (!workspace) {
+    throw new TRPCError({
+      code: "FORBIDDEN",
+      message: "You do not have access to this workspace",
+    });
+  }
+}
+
+async function assertConversationAccess(params: {
+  ctx: Context & { session: { user: { id: string } } };
+  workspaceId: string;
+  conversationId: string;
+}) {
+  const conversation = await params.ctx.db.copilotConversation.findFirst({
+    where: {
+      id: params.conversationId,
+      workspaceId: params.workspaceId,
+      userId: params.ctx.session.user.id,
+    },
+    select: {
+      id: true,
+    },
+  });
+
+  if (!conversation) {
+    throw new TRPCError({
+      code: "NOT_FOUND",
+      message: "Conversation not found",
+    });
+  }
+}
+
+async function persistConversationExchange(params: {
+  ctx: Context & { session: { user: { id: string } } };
+  workspaceId: string;
+  conversationId?: string;
+  userMessage: string;
+  assistantMessage: string;
+}) {
+  if (!params.conversationId) return;
+
+  await assertConversationAccess({
+    ctx: params.ctx,
+    workspaceId: params.workspaceId,
+    conversationId: params.conversationId,
+  });
+
+  const generatedTitle = sanitizeString(
+    params.userMessage.replace(/\s+/g, " ").trim(),
+    80
+  );
+
+  const conversation = await params.ctx.db.copilotConversation.findUnique({
+    where: { id: params.conversationId },
+    select: { title: true },
+  });
+
+  const shouldUpdateTitle =
+    !!conversation &&
+    conversation.title.trim().toLowerCase() === "new chat" &&
+    generatedTitle.length > 0;
+
+  await params.ctx.db.copilotMessage.createMany({
+    data: [
+      {
+        conversationId: params.conversationId,
+        role: "USER",
+        content: sanitizeString(params.userMessage, 8_000),
+      },
+      {
+        conversationId: params.conversationId,
+        role: "ASSISTANT",
+        content: sanitizeString(params.assistantMessage, 20_000),
+      },
+    ],
+  });
+
+  await params.ctx.db.copilotConversation.update({
+    where: { id: params.conversationId },
+    data: {
+      updatedAt: new Date(),
+      ...(shouldUpdateTitle ? { title: generatedTitle } : {}),
+    },
+  });
+}
+
+async function callCopilotModel(params: {
+  mode: "ask" | "explainSelection" | "suggestHighlights" | "generateFlashcards";
+  context: CopilotContextInput;
+  message: string;
+  history?: { role: "user" | "assistant"; content: string }[];
+}) {
+  const documentSnippet = contextWindow(params.context);
+  const selectedText = params.context.selectedText
+    ? sanitizeString(params.context.selectedText, 2_000)
+    : "";
+  const viewportText = params.context.viewportText
+    ? sanitizeString(params.context.viewportText, 4_000)
+    : "";
+  const message = sanitizeString(params.message, 4_000);
+
+  const systemPrompt = [
+    "You are an AI study assistant for the Scribe learning platform.",
+    "Be concise, correct, and safe. Do not claim edits were applied unless asked and confirmed by the user.",
+    "Return valid JSON only.",
+    "",
+    "OUTPUT FORMAT:",
+    "{",
+    '  "answer": "markdown response for the user",',
+    '  "highlights": [{ "start": 0, "end": 10, "label": "optional" }],',
+    '  "flashcards": [{ "front": "question", "back": "answer" }]',
+    "}",
+    "Include only relevant keys for the task.",
+    "",
+    `MODE: ${params.mode}`,
+    `ARTIFACT_TYPE: ${params.context.artifactType}`,
+  ].join("\n");
+
+  const contextPrompt = [
+    "DOCUMENT:",
+    documentSnippet,
+    "",
+    "SELECTION:",
+    selectedText || "None",
+    "",
+    "VIEWPORT_TEXT:",
+    viewportText || "None",
+  ].join("\n");
+
+  const messages: { role: "system" | "user" | "assistant"; content: string }[] = [
+    { role: "system", content: systemPrompt },
+    { role: "user", content: contextPrompt },
+  ];
+
+  if (params.history && params.history.length > 0) {
+    params.history.forEach((msg) => {
+      messages.push({
+        role: msg.role === "user" ? "user" : "assistant",
+        content: msg.content,
+      });
+    });
+  }
+
+  messages.push({ role: "user", content: message });
+
+  const response = await inference(messages);
+  const content = response.choices?.[0]?.message?.content ?? "";
+  const parsed = extractJson(content);
+
+  return {
+    rawContent: content,
+    parsed,
+  };
+}
+
+export const copilot = router({
+  listConversations: verifiedProcedure
+    .input(
+      z.object({
+        workspaceId: z.string().min(1),
+      })
+    )
+    .query(async ({ ctx, input }) => {
+      await assertWorkspaceAccess(ctx, input.workspaceId);
+
+      const conversations = await ctx.db.copilotConversation.findMany({
+        where: {
+          workspaceId: input.workspaceId,
+          userId: ctx.session.user.id,
+        },
+        include: {
+          messages: {
+            orderBy: { createdAt: "desc" },
+            take: 1,
+          },
+        },
+        orderBy: { updatedAt: "desc" },
+      });
+
+      return conversations.map((conversation) => ({
+        id: conversation.id,
+        title: conversation.title,
+        updatedAt: conversation.updatedAt,
+        preview: conversation.messages[0]?.content ?? "",
+      }));
+    }),
+
+  getConversation: verifiedProcedure
+    .input(
+      z.object({
+        workspaceId: z.string().min(1),
+        conversationId: z.string().min(1),
+      })
+    )
+    .query(async ({ ctx, input }) => {
+      await assertWorkspaceAccess(ctx, input.workspaceId);
+      await assertConversationAccess({
+        ctx,
+        workspaceId: input.workspaceId,
+        conversationId: input.conversationId,
+      });
+
+      const conversation = await ctx.db.copilotConversation.findUnique({
+        where: { id: input.conversationId },
+        include: {
+          messages: {
+            orderBy: { createdAt: "asc" },
+          },
+        },
+      });
+
+      if (!conversation) {
+        throw new TRPCError({
+          code: "NOT_FOUND",
+          message: "Conversation not found",
+        });
+      }
+
+      return {
+        id: conversation.id,
+        title: conversation.title,
+        messages: conversation.messages.map((message) => ({
+          id: message.id,
+          role: message.role === "USER" ? "user" : "assistant",
+          content: message.content,
+          createdAt: message.createdAt,
+        })),
+      };
+    }),
+
+  createConversation: verifiedProcedure
+    .input(
+      z.object({
+        workspaceId: z.string().min(1),
+        title: z.string().optional(),
+      })
+    )
+    .mutation(async ({ ctx, input }) => {
+      await assertWorkspaceAccess(ctx, input.workspaceId);
+
+      const conversation = await ctx.db.copilotConversation.create({
+        data: {
+          workspaceId: input.workspaceId,
+          userId: ctx.session.user.id,
+          title: sanitizeString(input.title || "New Chat", 120),
+        },
+      });
+
+      return {
+        id: conversation.id,
+        title: conversation.title,
+        updatedAt: conversation.updatedAt,
+      };
+    }),
+
+  deleteConversation: verifiedProcedure
+    .input(
+      z.object({
+        workspaceId: z.string().min(1),
+        conversationId: z.string().min(1),
+      })
+    )
+    .mutation(async ({ ctx, input }) => {
+      await assertWorkspaceAccess(ctx, input.workspaceId);
+      await assertConversationAccess({
+        ctx,
+        workspaceId: input.workspaceId,
+        conversationId: input.conversationId,
+      });
+
+      await ctx.db.copilotConversation.delete({
+        where: { id: input.conversationId },
+      });
+
+      return { success: true };
+    }),
+
+  ask: verifiedProcedure
+    .input(
+      z.object({
+        context: copilotContextSchema,
+        message: z.string().min(1),
+        conversationId: z.string().optional(),
+      })
+    )
+    .mutation(async ({ ctx, input }) => {
+      enforceRateLimit(ctx.session.user.id, input.context.workspaceId);
+      await assertWorkspaceAccess(ctx, input.context.workspaceId);
+
+      await PusherService.emitTaskComplete(input.context.workspaceId, "copilot:thinking", {
+        status: "started",
+      });
+
+      let history: { role: "user" | "assistant"; content: string }[] = [];
+      if (input.conversationId) {
+        const messages = await ctx.db.copilotMessage.findMany({
+          where: { conversationId: input.conversationId },
+          orderBy: { createdAt: "asc" },
+          take: 50,
+        });
+        history = messages.map((m) => ({
+          role: m.role.toLowerCase() as "user" | "assistant",
+          content: m.content,
+        }));
+      }
+
+      const model = await callCopilotModel({
+        mode: "ask",
+        context: input.context,
+        message: input.message,
+        history,
+      });
+
+      const answer =
+        getAnswerFromParsedJSON(model.parsed) ||
+        model.rawContent ||
+        "I could not generate a response.";
+
+      const highlights = z.array(highlightInstructionSchema).safeParse(model.parsed?.highlights);
+      const safeHighlights = highlights.success
+        ? highlights.data.filter(
+          (item) =>
+            item.start < item.end &&
+            item.end <= Math.max(input.context.documentContent.length, 0)
+        )
+        : [];
+
+      await PusherService.emitTaskComplete(input.context.workspaceId, "copilot:response", {
+        status: "completed",
+      });
+
+      await persistConversationExchange({
+        ctx,
+        workspaceId: input.context.workspaceId,
+        conversationId: input.conversationId,
+        userMessage: input.message,
+        assistantMessage: answer,
+      });
+
+      return {
+        answer,
+        highlights: safeHighlights,
+      };
+    }),
+
+  explainSelection: verifiedProcedure
+    .input(
+      z.object({
+        context: copilotContextSchema,
+        message: z.string().optional(),
+        conversationId: z.string().optional(),
+      })
+    )
+    .mutation(async ({ ctx, input }) => {
+      enforceRateLimit(ctx.session.user.id, input.context.workspaceId);
+      await assertWorkspaceAccess(ctx, input.context.workspaceId);
+
+      await PusherService.emitTaskComplete(input.context.workspaceId, "copilot:thinking", {
+        status: "started",
+      });
+
+      const question =
+        input.message && input.message.trim().length > 0
+          ? input.message
+          : "Explain the selected text in simple study-friendly terms and include one practical example.";
+
+      let history: { role: "user" | "assistant"; content: string }[] = [];
+      if (input.conversationId) {
+        const messages = await ctx.db.copilotMessage.findMany({
+          where: { conversationId: input.conversationId },
+          orderBy: { createdAt: "asc" },
+          take: 50,
+        });
+        history = messages.map((m) => ({
+          role: m.role.toLowerCase() as "user" | "assistant",
+          content: m.content,
+        }));
+      }
+
+      const model = await callCopilotModel({
+        mode: "explainSelection",
+        context: input.context,
+        message: question,
+        history,
+      });
+
+      const answer =
+        getAnswerFromParsedJSON(model.parsed) ||
+        model.rawContent ||
+        "I could not explain this selection.";
+
+      await PusherService.emitTaskComplete(input.context.workspaceId, "copilot:response", {
+        status: "completed",
+      });
+
+      await persistConversationExchange({
+        ctx,
+        workspaceId: input.context.workspaceId,
+        conversationId: input.conversationId,
+        userMessage: question,
+        assistantMessage: answer,
+      });
+
+      return { answer };
+    }),
+
+  suggestHighlights: verifiedProcedure
+    .input(
+      z.object({
+        context: copilotContextSchema,
+        message: z.string().optional(),
+        conversationId: z.string().optional(),
+      })
+    )
+    .mutation(async ({ ctx, input }) => {
+      enforceRateLimit(ctx.session.user.id, input.context.workspaceId);
+      await assertWorkspaceAccess(ctx, input.context.workspaceId);
+
+      await PusherService.emitTaskComplete(input.context.workspaceId, "copilot:thinking", {
+        status: "started",
+      });
+
+      const instruction =
+        input.message && input.message.trim().length > 0
+          ? input.message
+          : "Suggest key highlights for the provided section.";
+
+      let history: { role: "user" | "assistant"; content: string }[] = [];
+      if (input.conversationId) {
+        const messages = await ctx.db.copilotMessage.findMany({
+          where: { conversationId: input.conversationId },
+          orderBy: { createdAt: "asc" },
+          take: 50,
+        });
+        history = messages.map((m) => ({
+          role: m.role.toLowerCase() as "user" | "assistant",
+          content: m.content,
+        }));
+      }
+
+      const model = await callCopilotModel({
+        mode: "suggestHighlights",
+        context: input.context,
+        message: instruction,
+        history,
+      });
+
+      const parsedHighlights = z.array(highlightInstructionSchema).safeParse(
+        model.parsed?.highlights
+      );
+
+      const safeHighlights = parsedHighlights.success
+        ? parsedHighlights.data.filter(
+          (item) =>
+            item.start < item.end &&
+            item.end <= Math.max(input.context.documentContent.length, 0)
+        )
+        : [];
+
+      await PusherService.emitTaskComplete(
+        input.context.workspaceId,
+        "copilot:highlight_suggestions",
+        {
+          highlights: safeHighlights,
+        }
+      );
+
+      const answer =
+        getAnswerFromParsedJSON(model.parsed) ||
+        "Here are suggested highlights.";
+
+      await persistConversationExchange({
+        ctx,
+        workspaceId: input.context.workspaceId,
+        conversationId: input.conversationId,
+        userMessage: instruction,
+        assistantMessage: answer,
+      });
+
+      return {
+        answer,
+        highlights: safeHighlights,
+      };
+    }),
+
+  generateFlashcards: verifiedProcedure
+    .input(
+      z.object({
+        context: copilotContextSchema,
+        message: z.string().optional(),
+        numCards: z.number().int().min(1).max(30).default(10),
+        conversationId: z.string().optional(),
+      })
+    )
+    .mutation(async ({ ctx, input }) => {
+      enforceRateLimit(ctx.session.user.id, input.context.workspaceId);
+      await assertWorkspaceAccess(ctx, input.context.workspaceId);
+
+      await PusherService.emitTaskComplete(input.context.workspaceId, "copilot:thinking", {
+        status: "started",
+      });
+
+      const instruction =
+        input.message && input.message.trim().length > 0
+          ? input.message
+          : `Generate ${input.numCards} concise study flashcards from this context.`;
+
+      let history: { role: "user" | "assistant"; content: string }[] = [];
+      if (input.conversationId) {
+        const messages = await ctx.db.copilotMessage.findMany({
+          where: { conversationId: input.conversationId },
+          orderBy: { createdAt: "asc" },
+          take: 50,
+        });
+        history = messages.map((m) => ({
+          role: m.role.toLowerCase() as "user" | "assistant",
+          content: m.content,
+        }));
+      }
+
+      const model = await callCopilotModel({
+        mode: "generateFlashcards",
+        context: input.context,
+        message: instruction,
+        history,
+      });
+
+      const parsedFlashcards = z.array(flashcardSchema).safeParse(model.parsed?.flashcards);
+      const cards = (parsedFlashcards.success ? parsedFlashcards.data : []).slice(
+        0,
+        input.numCards
+      );
+
+      if (cards.length === 0) {
+        throw new TRPCError({
+          code: "BAD_REQUEST",
+          message: "Copilot did not return valid flashcards.",
+        });
+      }
+
+      const artifact = await ctx.db.artifact.create({
+        data: {
+          workspaceId: input.context.workspaceId,
+          type: ArtifactType.FLASHCARD_SET,
+          title: `Copilot Flashcards - ${new Date().toLocaleString()}`,
+          createdById: ctx.session.user.id,
+          generating: false,
+          flashcards: {
+            create: cards.map((card, index) => ({
+              front: sanitizeString(card.front, 200),
+              back: sanitizeString(card.back, 500),
+              order: index,
+              tags: ["copilot-generated"],
+            })),
+          },
+        },
+        include: {
+          flashcards: true,
+        },
+      });
+
+      await PusherService.emitTaskComplete(input.context.workspaceId, "copilot:response", {
+        status: "completed",
+        flashcardSetId: artifact.id,
+      });
+
+      const answer =
+        getAnswerFromParsedJSON(model.parsed) ||
+        `Generated ${artifact.flashcards.length} flashcards.`;
+
+      await persistConversationExchange({
+        ctx,
+        workspaceId: input.context.workspaceId,
+        conversationId: input.conversationId,
+        userMessage: instruction,
+        assistantMessage: answer,
+      });
+
+      return {
+        answer,
+        artifactId: artifact.id,
+        flashcards: artifact.flashcards,
+      };
+    }),
+});