experiment(notebooks): AI summarizer + prompter edit panel

packages/api-client/src/posthog-client.ts

@@ -1370,6 +1370,9 @@ function previewTokenHeader(
 export class PostHogAPIClient {
   private api: ReturnType<typeof createApiClient>;
   private _teamId: number | null = null;
+  // Kept for callers that hit non-API origins with the same bearer (the LLM
+  // gateway); the fetcher only covers requests against the API base URL.
+  private readonly getAccessToken: () => Promise<string>;
 
   constructor(
     apiHost: string,
@@ -1378,6 +1381,7 @@ export class PostHogAPIClient {
     teamId?: number,
   ) {
     const baseUrl = apiHost.endsWith("/") ? apiHost.slice(0, -1) : apiHost;
+    this.getAccessToken = getAccessToken;
     this.api = createApiClient(
       buildApiFetcher({
         getAccessToken,
@@ -2042,6 +2046,56 @@ export class PostHogAPIClient {
     });
   }
 
+  /**
+   * PostHog LLM gateway base for this cloud region, scoped to the
+   * `posthog_code` product (mirrors `@posthog/agent`'s `getLlmGatewayUrl`,
+   * inlined here so the api-client stays dependency-free). The gateway
+   * accepts the same OAuth bearer this client already holds.
+   */
+  get llmGatewayUrl(): string {
+    const url = new URL(this.api.baseUrl);
+    const hostname = url.hostname;
+    let base: string;
+    if (hostname === "localhost" || hostname === "127.0.0.1") {
+      base = `${url.protocol}//localhost:3308`;
+    } else if (hostname === "app.dev.posthog.dev") {
+      base = "https://gateway.dev.posthog.dev";
+    } else {
+      const region = hostname.match(/^(us|eu)\.posthog\.com$/)?.[1] ?? "us";
+      base = `https://gateway.${region}.posthog.com`;
+    }
+    return `${base}/posthog_code`;
+  }
+
+  // One-shot OpenAI-compatible chat completion against the PostHog LLM
+  // gateway (non-agentic — much lower latency than a Max conversation turn).
+  // Returns the raw Response; with `stream: true` in the body it streams
+  // `chat.completion.chunk` SSE events. The gateway is a different origin
+  // from the API, so this goes through plain fetch with the bearer attached
+  // rather than the shared fetcher.
+  async llmGatewayChatCompletion(
+    body: Record<string, unknown>,
+    signal?: AbortSignal,
+  ): Promise<Response> {
+    const token = await this.getAccessToken();
+    const response = await fetch(`${this.llmGatewayUrl}/v1/chat/completions`, {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        Authorization: `Bearer ${token}`,
+      },
+      body: JSON.stringify(body),
+      ...(signal ? { signal } : {}),
+    });
+    if (!response.ok) {
+      const detail = await response.text().catch(() => "");
+      throw new Error(
+        `LLM gateway request failed: [${response.status}] ${detail.slice(0, 300)}`,
+      );
+    }
+    return response;
+  }
+
   // Cancel an in-progress Max AI conversation turn (best-effort; the caller
   // usually also aborts its stream fetch). Same verb/path as the PostHog web
   // app: PATCH .../conversations/{id}/cancel/.

packages/core/src/notebooks/identifiers.ts

@@ -2,3 +2,9 @@
 // @posthog/core so host routers and containers can reference it without
 // importing the concrete class's module.
 export const NOTEBOOKS_SERVICE = Symbol.for("posthog.notebooks.service");
+
+// AI summarizer + change-request engine for notebook component nodes (the
+// model-transport token lives beside its interface in notebookNodeAIModel.ts).
+export const NOTEBOOK_NODE_AI_SERVICE = Symbol.for(
+  "posthog.notebooks.nodeAIService",
+);

packages/core/src/notebooks/notebookNodeAI.test.ts

@@ -0,0 +1,231 @@
+import type { PostHogAPIClient } from "@posthog/api-client/posthog-client";
+import { describe, expect, it, vi } from "vitest";
+import type {
+  NotebookNodeAIModel,
+  NotebookNodeAIModelRequest,
+} from "./notebookNodeAIModel";
+import {
+  buildNotebookNodePropsUpdate,
+  extractFirstJsonObject,
+  NotebookNodeAIService,
+  notebookNodeAICacheKey,
+  parseNodeChangeResponse,
+} from "./notebookNodeAIService";
+import type { NotebookNodeJsonObject } from "./notebookNodeSummary";
+
+function fakeClient(
+  overrides: Partial<PostHogAPIClient> = {},
+): PostHogAPIClient {
+  return overrides as PostHogAPIClient;
+}
+
+/** Scripted fake model: returns queued replies in order, records requests. */
+function fakeModel(replies: string[]): {
+  model: NotebookNodeAIModel;
+  requests: NotebookNodeAIModelRequest[];
+} {
+  const requests: NotebookNodeAIModelRequest[] = [];
+  let index = 0;
+  return {
+    requests,
+    model: {
+      complete: vi.fn(
+        async (
+          _client: PostHogAPIClient,
+          request: NotebookNodeAIModelRequest,
+        ) => {
+          requests.push(request);
+          const reply = replies[Math.min(index, replies.length - 1)];
+          index++;
+          if (reply === undefined) throw new Error("no scripted reply");
+          request.onText?.(reply);
+          return reply;
+        },
+      ),
+    },
+  };
+}
+
+const TRENDS_PROPS: NotebookNodeJsonObject = {
+  query: {
+    kind: "InsightVizNode",
+    source: {
+      kind: "TrendsQuery",
+      series: [{ kind: "EventsNode", event: "$pageview" }],
+    },
+  },
+  hideResults: true,
+};
+
+describe("NotebookNodeAIService.summarizeNode", () => {
+  it("returns the cleaned model text and streams partials", async () => {
+    const { model } = fakeModel(['  "Trends of pageviews, last 7 days."  ']);
+    const service = new NotebookNodeAIService(model);
+    const partials: string[] = [];
+    const summary = await service.summarizeNode(
+      fakeClient(),
+      { tagName: "Query", props: TRENDS_PROPS },
+      { onPartial: (text) => partials.push(text) },
+    );
+    expect(summary).toBe("Trends of pageviews, last 7 days.");
+    expect(partials).toEqual(["Trends of pageviews, last 7 days."]);
+  });
+
+  it("caches by (tagName + props) so the second call skips the model", async () => {
+    const { model } = fakeModel(["A summary."]);
+    const service = new NotebookNodeAIService(model);
+    const client = fakeClient();
+    const input = { tagName: "Query", props: TRENDS_PROPS };
+    await service.summarizeNode(client, input);
+    const again = await service.summarizeNode(client, input);
+    expect(again).toBe("A summary.");
+    expect(model.complete).toHaveBeenCalledTimes(1);
+    expect(service.getCachedSummary(input)).toBe("A summary.");
+  });
+
+  it("ignores shell-managed props in the cache key", () => {
+    const withPanels = notebookNodeAICacheKey({
+      tagName: "Query",
+      props: { ...TRENDS_PROPS, hideFilters: true },
+    });
+    const withoutPanels = notebookNodeAICacheKey({
+      tagName: "Query",
+      props: { query: TRENDS_PROPS.query },
+    });
+    expect(withPanels).toBe(withoutPanels);
+  });
+
+  it("enriches entity nodes with the live object and survives fetch failure", async () => {
+    const { model, requests } = fakeModel(["Flag summary."]);
+    const service = new NotebookNodeAIService(model);
+    const getFeatureFlag = vi.fn().mockResolvedValue({
+      id: 1,
+      key: "my-flag",
+      name: "My flag",
+      active: true,
+      filters: {},
+    });
+    await service.summarizeNode(fakeClient({ getFeatureFlag }), {
+      tagName: "FeatureFlag",
+      props: { id: "my-flag" },
+    });
+    expect(getFeatureFlag).toHaveBeenCalledWith("my-flag");
+    expect(requests[0]?.user).toContain('"key":"my-flag"');
+
+    const failing = vi.fn().mockRejectedValue(new Error("403"));
+    const summary = await service.summarizeNode(
+      fakeClient({ getFeatureFlag: failing }),
+      { tagName: "FeatureFlag", props: { id: "other-flag" } },
+    );
+    expect(summary).toBe("Flag summary.");
+  });
+});
+
+describe("NotebookNodeAIService.requestNodeChange", () => {
+  const change = {
+    props: {
+      query: {
+        kind: "InsightVizNode",
+        source: {
+          kind: "TrendsQuery",
+          series: [{ kind: "EventsNode", event: "$pageview" }],
+          trendsFilter: { display: "ActionsBar" },
+        },
+      },
+    },
+    summary: "Trends bar chart of pageviews.",
+  };
+
+  it("applies a change from one model call and preserves shell props", async () => {
+    const { model, requests } = fakeModel([
+      `Here you go:\n\`\`\`json\n${JSON.stringify(change)}\n\`\`\``,
+    ]);
+    const service = new NotebookNodeAIService(model);
+    const result = await service.requestNodeChange(
+      fakeClient(),
+      { tagName: "Query", props: TRENDS_PROPS },
+      "make it a bar chart",
+    );
+    expect(model.complete).toHaveBeenCalledTimes(1);
+    expect(result.summary).toBe("Trends bar chart of pageviews.");
+    expect(result.props.hideResults).toBe(true);
+    expect(result.props.query).toEqual(change.props.query);
+    // The new summary is primed so reopening the node is instant.
+    expect(
+      service.getCachedSummary({ tagName: "Query", props: result.props }),
+    ).toBe("Trends bar chart of pageviews.");
+    // Shell props stay out of the prompt.
+    expect(requests[0]?.user).not.toContain("hideResults");
+  });
+
+  it.each([
+    ["prose with no JSON", "Sorry, I cannot help with that."],
+    ["unparsable JSON", '{"props": {"query": }}'],
+    ["missing summary", '{"props": {"query": {"kind": "TrendsQuery"}}}'],
+    ["query without kind", '{"props": {"query": {}}, "summary": "x"}'],
+  ])("retries once after %s, then succeeds", async (_label, badReply) => {
+    const { model, requests } = fakeModel([badReply, JSON.stringify(change)]);
+    const service = new NotebookNodeAIService(model);
+    const result = await service.requestNodeChange(
+      fakeClient(),
+      { tagName: "Query", props: TRENDS_PROPS },
+      "make it a bar chart",
+    );
+    expect(model.complete).toHaveBeenCalledTimes(2);
+    expect(requests[1]?.user).toContain("could not be applied");
+    expect(result.summary).toBe("Trends bar chart of pageviews.");
+  });
+
+  it("throws after two unusable replies", async () => {
+    const { model } = fakeModel(["nope", "still nope"]);
+    const service = new NotebookNodeAIService(model);
+    await expect(
+      service.requestNodeChange(
+        fakeClient(),
+        { tagName: "Query", props: TRENDS_PROPS },
+        "do a thing",
+      ),
+    ).rejects.toThrow(/no JSON object/);
+    expect(model.complete).toHaveBeenCalledTimes(2);
+  });
+});
+
+describe("buildNotebookNodePropsUpdate", () => {
+  it("replaces, deletes dropped keys, and never touches shell props", () => {
+    const current: NotebookNodeJsonObject = {
+      query: { kind: "HogQLQuery", query: "SELECT 1" },
+      title: "Old",
+      stale: "remove-me",
+      hideFilters: true,
+    };
+    const next: NotebookNodeJsonObject = {
+      query: { kind: "HogQLQuery", query: "SELECT 2" },
+      title: "New",
+    };
+    expect(buildNotebookNodePropsUpdate(current, next)).toEqual({
+      query: { kind: "HogQLQuery", query: "SELECT 2" },
+      title: "New",
+      stale: undefined,
+    });
+  });
+});
+
+describe("parse helpers", () => {
+  it.each([
+    ['{"a": 1}', '{"a": 1}'],
+    ['pre {"a": {"b": "}"}} post', '{"a": {"b": "}"}}'],
+    ['```json\n{"a": 1}\n```', '{"a": 1}'],
+    ["no braces here", null],
+    ['{"unbalanced": ', null],
+  ])("extractFirstJsonObject(%j) -> %j", (input, expected) => {
+    expect(extractFirstJsonObject(input)).toBe(expected);
+  });
+
+  it("rejects a non-object props payload", () => {
+    const outcome = parseNodeChangeResponse(
+      "Query",
+      '{"props": [1, 2], "summary": "x"}',
+    );
+    expect(outcome.ok).toBe(false);
+  });
+});