feat: add telemetry to the chat API route

.vscode/settings.json

@@ -38,6 +38,7 @@
     "EHRC",
     "estree",
     "estruyf",
+    "eyfs",
     "fkey",
     "fontsource",
     "Geist",
@@ -82,6 +83,8 @@
     "Onboarded",
     "openai",
     "openapi",
+    "opentelemetry",
+    "otlp",
     "paragraphise",
     "pgvector",
     "Pinecone",
@@ -98,6 +101,7 @@
     "PSHE",
     "psql",
     "ratelimit",
+    "Regen",
     "remeda",
     "Rerank",
     "Retriable",

apps/nextjs/package.json

@@ -39,6 +39,13 @@
     "@oakai/prettier-config": "*",
     "@oaknational/oak-components": "^1.0.0",
     "@oaknational/oak-consent-client": "^2.1.0",
+    "@opentelemetry/api": "^1.9.0",
+    "@opentelemetry/auto-instrumentations-node": "^0.49.1",
+    "@opentelemetry/core": "^1.25.1",
+    "@opentelemetry/exporter-trace-otlp-http": "^0.52.1",
+    "@opentelemetry/sdk-node": "^0.52.1",
+    "@opentelemetry/sdk-trace-base": "^1.26.0",
+    "@opentelemetry/sdk-trace-web": "^1.26.0",
     "@prisma/client": "^5.13.0",
     "@prisma/extension-accelerate": "^1.0.0",
     "@radix-ui/react-accordion": "^1.1.2",
@@ -95,6 +102,7 @@
     "superjson": "^1.9.1",
     "tailwind-merge": "^2.2.1",
     "tailwindcss-animate": "^1.0.7",
+    "tiny-invariant": "^1.3.1",
     "trpc-openapi": "^1.2.0",
     "ts-node": "^10.9.2",
     "tsx": "^4.16.0",

apps/nextjs/src/app/api/chat/chatHandler.ts

@@ -0,0 +1,255 @@
+import {
+  Aila,
+  AilaAuthenticationError,
+  AilaThreatDetectionError,
+} from "@oakai/aila";
+import type { AilaOptions, AilaPublicChatOptions, Message } from "@oakai/aila";
+import { LooseLessonPlan } from "@oakai/aila/src/protocol/schema";
+import { handleHeliconeError } from "@oakai/aila/src/utils/moderation/moderationErrorHandling";
+import { tracer } from "@oakai/core/src/tracing/serverTracing";
+import { PrismaClientWithAccelerate, prisma as globalPrisma } from "@oakai/db";
+import {
+  SpanContext,
+  Span,
+  SpanStatusCode,
+  TraceFlags,
+} from "@opentelemetry/api";
+import { TraceState } from "@opentelemetry/core";
+import { StreamingTextResponse } from "ai";
+import { NextRequest } from "next/server";
+import invariant from "tiny-invariant";
+
+import { Config } from "./config";
+import { streamingJSON } from "./protocol";
+
+export const maxDuration = 300;
+
+const prisma: PrismaClientWithAccelerate = globalPrisma;
+
+export async function GET() {
+  return new Response("Chat API is working", { status: 200 });
+}
+
+async function setupChatHandler(req: NextRequest) {
+  const json = await req.json();
+  const {
+    id: chatId,
+    messages,
+    lessonPlan = {},
+    options: chatOptions = {},
+    traceContext: serializedTraceContext,
+  }: {
+    id: string;
+    messages: Message[];
+    lessonPlan?: LooseLessonPlan;
+    options?: AilaPublicChatOptions;
+    traceContext?: {
+      traceId: string;
+      spanId: string;
+      traceFlags: number;
+      traceState?: string;
+    };
+  } = json;
+
+  let traceContext: SpanContext | undefined;
+  if (serializedTraceContext) {
+    traceContext = {
+      traceId: serializedTraceContext.traceId,
+      spanId: serializedTraceContext.spanId,
+      traceFlags: serializedTraceContext.traceFlags as TraceFlags,
+      traceState: serializedTraceContext.traceState
+        ? new TraceState(serializedTraceContext.traceState)
+        : undefined,
+    };
+  }
+
+  const options: AilaOptions = {
+    useRag: chatOptions.useRag ?? true,
+    temperature: chatOptions.temperature ?? 0.7,
+    numberOfLessonPlansInRag: chatOptions.numberOfLessonPlansInRag ?? 5,
+    usePersistence: true,
+    useModeration: true,
+  };
+
+  return { chatId, messages, lessonPlan, options, traceContext };
+}
+
+function reportErrorTelemetry(
+  span: Span,
+  error: Error,
+  errorType: string,
+  statusMessage: string,
+  additionalAttributes: Record<string, unknown> = {},
+) {
+  span.recordException(error);
+  span.setStatus({ code: SpanStatusCode.ERROR, message: statusMessage });
+  span.setAttributes({
+    errorType,
+    errorStack: error.stack,
+    ...additionalAttributes,
+  });
+}
+
+function setTelemetryMetadata(
+  span: Span,
+  id: string,
+  messages: Message[],
+  lessonPlan: LooseLessonPlan,
+  options: AilaOptions,
+) {
+  span.setAttributes({
+    chatId: id,
+    messageCount: messages.length,
+    hasLessonPlan: Object.keys(lessonPlan).length > 0,
+    useRag: options.useRag,
+    temperature: options.temperature,
+    numberOfLessonPlansInRag: options.numberOfLessonPlansInRag,
+    usePersistence: options.usePersistence,
+    useModeration: options.useModeration,
+  });
+}
+
+function handleConnectionAborted(req: NextRequest) {
+  const abortController = new AbortController();
+
+  req.signal.addEventListener("abort", () => {
+    console.log("Client has disconnected");
+    abortController.abort();
+  });
+  return abortController;
+}
+
+async function handleThreatDetectionError(
+  span: Span,
+  e: AilaThreatDetectionError,
+  userId: string,
+  id: string,
+  prisma: PrismaClientWithAccelerate,
+) {
+  const heliconeErrorMessage = await handleHeliconeError(userId, id, e, prisma);
+  reportErrorTelemetry(span, e, "AilaThreatDetectionError", "Threat detected");
+  return streamingJSON(heliconeErrorMessage);
+}
+
+async function handleAilaAuthenticationError(
+  span: Span,
+  e: AilaAuthenticationError,
+) {
+  reportErrorTelemetry(span, e, "AilaAuthenticationError", "Unauthorized");
+  return new Response("Unauthorized", { status: 401 });
+}
+
+async function handleGenericError(span: Span, e: Error) {
+  reportErrorTelemetry(span, e, e.name, e.message);
+  return streamingJSON({
+    type: "error",
+    message: e.message,
+    value: `Sorry, an error occurred: ${e.message}`,
+  });
+}
+
+async function getUserId(
+  config: Config,
+  span: Span,
+  id: string,
+): Promise<string> {
+  if (config.shouldPerformUserLookup) {
+    const userLookup = await config.handleUserLookup(span, id);
+
+    if (!userLookup) {
+      throw new Error("User lookup failed");
+    }
+
+    if ("failureResponse" in userLookup) {
+      if (userLookup.failureResponse) {
+        throw new Error("User lookup failed: failureResponse received");
+      }
+    }
+
+    if ("userId" in userLookup) {
+      return userLookup.userId;
+    }
+
+    throw new Error("User lookup failed: userId not found");
+  }
+  invariant(config.mockUserId, "User ID is required");
+
+  return config.mockUserId;
+}
+
+async function generateChatStream(
+  aila: Aila,
+  abortController: AbortController,
+) {
+  return tracer.startActiveSpan("chat-aila-generate", async (generateSpan) => {
+    invariant(aila, "Aila instance is required");
+    const result = await aila.generate({ abortController });
+    generateSpan.end();
+    return result;
+  });
+}
+
+async function handleChatException(
+  span: Span,
+  e: unknown,
+  userId: string | undefined,
+  chatId: string,
+  prisma: PrismaClientWithAccelerate,
+): Promise<Response> {
+  if (e instanceof AilaAuthenticationError) {
+    return handleAilaAuthenticationError(span, e);
+  }
+
+  if (e instanceof AilaThreatDetectionError && userId) {
+    return handleThreatDetectionError(span, e, userId, chatId, prisma);
+  }
+
+  if (e instanceof Error) {
+    return handleGenericError(span, e);
+  }
+
+  throw e;
+}
+
+export async function handleChatPostRequest(
+  req: NextRequest,
+  config: Config,
+): Promise<Response> {
+  const { chatId, messages, lessonPlan, options, traceContext } =
+    await setupChatHandler(req);
+
+  const span = tracer.startSpan("chat-api", {
+    links: traceContext ? [{ context: traceContext }] : [],
+  });
+
+  let userId: string | undefined;
+  let aila: Aila | undefined;
+
+  try {
+    setTelemetryMetadata(span, chatId, messages, lessonPlan, options);
+
+    userId = await getUserId(config, span, chatId);
+
+    aila = await config.createAila({
+      options,
+      chat: {
+        id: chatId,
+        userId,
+        messages,
+      },
+      lessonPlan,
+    });
+
+    const abortController = handleConnectionAborted(req);
+    const stream = await generateChatStream(aila, abortController);
+    console.log("Completed handleChatPostRequest. Returning response");
+    return new StreamingTextResponse(stream);
+  } catch (e) {
+    return handleChatException(span, e, userId, chatId, prisma);
+  } finally {
+    if (aila) {
+      await aila.ensureShutdown();
+    }
+    span.end();
+  }
+}

apps/nextjs/src/app/api/chat/config.ts

@@ -0,0 +1,43 @@
+import { Aila, AilaInitializationOptions } from "@oakai/aila";
+import {
+  prisma as globalPrisma,
+  type PrismaClientWithAccelerate,
+} from "@oakai/db";
+import { Span } from "@opentelemetry/api";
+import { nanoid } from "ai";
+
+import { handleUserLookup as defaultHandleUserLookup } from "./user";
+import { createWebActionsPlugin } from "./webActionsPlugin";
+
+export interface Config {
+  shouldPerformUserLookup: boolean;
+  mockUserId?: string;
+  handleUserLookup: (
+    span: Span,
+    id: string,
+  ) => Promise<
+    | {
+        userId: string;
+      }
+    | {
+        failureResponse: Response;
+      }
+  >;
+  prisma: PrismaClientWithAccelerate;
+  createAila: (options: Partial<AilaInitializationOptions>) => Promise<Aila>;
+}
+
+export const defaultConfig: Config = {
+  shouldPerformUserLookup: true,
+  handleUserLookup: defaultHandleUserLookup,
+  prisma: globalPrisma,
+  createAila: async (options) => {
+    const webActionsPlugin = createWebActionsPlugin(globalPrisma);
+    return new Aila({
+      ...options,
+      plugins: [...(options.plugins || []), webActionsPlugin],
+      prisma: options.prisma ?? globalPrisma,
+      chat: options.chat || { id: nanoid(), userId: undefined },
+    });
+  },
+};