diff --git a/packages/api-client/src/posthog-client.ts b/packages/api-client/src/posthog-client.ts index 9e1fb0e3e9..b6330e579e 100644 --- a/packages/api-client/src/posthog-client.ts +++ b/packages/api-client/src/posthog-client.ts @@ -1370,6 +1370,9 @@ function previewTokenHeader( export class PostHogAPIClient { private api: ReturnType; 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; 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, + signal?: AbortSignal, + ): Promise { + 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/. diff --git a/packages/core/src/notebooks/identifiers.ts b/packages/core/src/notebooks/identifiers.ts index 4a88d6fd7d..9fe135a6e3 100644 --- a/packages/core/src/notebooks/identifiers.ts +++ b/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", +); diff --git a/packages/core/src/notebooks/notebookNodeAI.test.ts b/packages/core/src/notebooks/notebookNodeAI.test.ts new file mode 100644 index 0000000000..eea86f9efa --- /dev/null +++ b/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 { + 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); + }); +}); diff --git a/packages/core/src/notebooks/notebookNodeAIKnowledge.ts b/packages/core/src/notebooks/notebookNodeAIKnowledge.ts new file mode 100644 index 0000000000..3a16edfdeb --- /dev/null +++ b/packages/core/src/notebooks/notebookNodeAIKnowledge.ts @@ -0,0 +1,135 @@ +// Compact, prompt-ready descriptions of what each notebook component tag's +// props support. Fed to the model alongside the node's current props so it +// can summarize accurately and produce valid replacement props. Deliberately +// terse — the full vendored OpenAPI schema is megabytes; this is the working +// subset the notebook embeds actually render. + +const QUERY_KNOWLEDGE = `Props shape: { "query": , "title"?: string }. +The query node is a PostHog insight query. Supported kinds: + +- Wrappers: { "kind": "InsightVizNode", "source": } renders charts; + { "kind": "DataTableNode", "source": , "columns"?: string[] } renders tables. +- TrendsQuery: { "kind": "TrendsQuery", "series": [...], + "interval"?: "hour"|"day"|"week"|"month", + "dateRange"?: { "date_from": string, "date_to"?: string|null }, + "trendsFilter"?: { "display"?: "ActionsLineGraph"|"ActionsLineGraphCumulative"|"ActionsAreaGraph"|"ActionsBar"|"ActionsBarValue"|"ActionsPie"|"ActionsTable"|"BoldNumber"|"WorldMap" }, + "breakdownFilter"?: { "breakdown": string, "breakdown_type": "event"|"person"|"group" } or { "breakdowns": [{ "property": string, "type": "event"|"person" }] }, + "properties"?: [] }. +- EventsNode series entry: { "kind": "EventsNode", "event": string, "name"?: string, "custom_name"?: string, + "math"?: "total"|"dau"|"weekly_active"|"monthly_active"|"unique_session"|"sum"|"avg"|"min"|"max"|"median", + "math_property"?: string, "properties"?: [...] }. ActionsNode: { "kind": "ActionsNode", "id": number }. +- FunnelsQuery: { "kind": "FunnelsQuery", "series": [... in step order], "dateRange"?, "funnelsFilter"?: { "funnelVizType"?: "steps"|"time_to_convert"|"trends", "funnelWindowInterval"?: number, "funnelWindowIntervalUnit"?: "minute"|"hour"|"day"|"week"|"month" } }. +- RetentionQuery: { "kind": "RetentionQuery", "retentionFilter": { "period": "Day"|"Week"|"Month", "totalIntervals": number, "targetEntity": { "id": string, "name": string, "type": "events" }, "returningEntity": }, "dateRange"? }. +- PathsQuery: { "kind": "PathsQuery", "pathsFilter": { "includeEventTypes": ["$pageview"|"$screen"|"custom_event"...] }, "dateRange"? }. +- StickinessQuery / LifecycleQuery: { "kind": ..., "series": [...], "dateRange"? } (+ "stickinessFilter"/"lifecycleFilter"). +- HogQLQuery: { "kind": "HogQLQuery", "query": "" }. +- EventsQuery (raw events table, use inside DataTableNode): { "kind": "EventsQuery", "select": ["*","event","person","timestamp"...], "event"?: string, "after"?: "-24h", "limit"?: number, "properties"?: [...] }. +- SavedInsightNode: { "kind": "SavedInsightNode", "shortId": string } references an insight saved in PostHog. + +Date range strings: relative like "-24h", "-7d", "-30d", "-90d", "-1m", or "dStart" (today), "mStart" (this month), "yStart", "all", or absolute "YYYY-MM-DD". +Common events: "$pageview", "$pageleave", "$autocapture", "$identify". Common breakdown properties: "$browser", "$os", "$geoip_country_code", "$current_url", "$device_type". +Keep the wrapper kind (InsightVizNode/DataTableNode) unless the change requires switching it.`; + +const NOTEBOOK_NODE_AI_KNOWLEDGE: Record = { + Query: QUERY_KNOWLEDGE, + FeatureFlag: + 'Props shape: { "id": number | string, "title"?: string }. `id` is the feature flag\'s numeric id or its string key. The node renders the live flag (name, status, rollout) fetched from PostHog. Never invent an id or key.', + Experiment: + 'Props shape: { "id": number, "title"?: string }. `id` is the experiment\'s numeric id. The node renders the live experiment. Never invent an id.', + Survey: + 'Props shape: { "id": string, "title"?: string }. `id` is the survey\'s UUID. The node renders the live survey (questions, status). Never invent an id.', + EarlyAccessFeature: + 'Props shape: { "id": string, "title"?: string }. `id` is the early access feature\'s UUID. Never invent an id.', + Cohort: + 'Props shape: { "id": number, "title"?: string }. `id` is the cohort\'s numeric id. Never invent an id.', + Person: + 'Props shape: { "distinctId": string, "title"?: string }. `distinctId` identifies the person. Never invent a distinct id.', + Group: + 'Props shape: { "groupKey": string, "groupTypeIndex": number, "title"?: string }. Never invent a group key.', + Recording: + 'Props shape: { "sessionRecordingId": string, "title"?: string }. Never invent a recording id.', +}; + +const GENERIC_KNOWLEDGE = + 'Props are an arbitrary JSON object rendered by the notebook. A "title" string prop, when present, is shown as the block title.'; + +export function getNotebookNodeKnowledge(tagName: string): string { + return NOTEBOOK_NODE_AI_KNOWLEDGE[tagName] ?? GENERIC_KNOWLEDGE; +} + +export function buildNotebookNodeSummarySystemPrompt(): string { + return [ + "You summarize PostHog notebook blocks for a block-editing side panel.", + "Reply with ONLY the summary: one plain-text sentence (two at most), under 220 characters, no markdown, no quotes, no preamble.", + 'Describe what the block shows in product terms a PostHog user recognizes (e.g. "Trends line chart: pageviews vs. signups, last 30 days, broken down by browser").', + "Mention chart type, series/steps, date range, breakdowns, filters, or entity name/status when present. Do not mention JSON, props, or internal field names.", + ].join("\n"); +} + +export function buildNotebookNodeSummaryUserPrompt(input: { + tagName: string; + propsJson: string; + knowledge: string; + liveEntityJson?: string | null; +}): string { + const sections = [ + `Notebook block type: <${input.tagName} />`, + `What this block type supports:\n${input.knowledge}`, + `Current props JSON:\n${input.propsJson}`, + ]; + if (input.liveEntityJson) { + sections.push( + `Live object fetched from PostHog (for context):\n${input.liveEntityJson}`, + ); + } + sections.push("Summarize this block now."); + return sections.join("\n\n"); +} + +export function buildNotebookNodeChangeSystemPrompt(): string { + return [ + "You edit PostHog notebook blocks. Given a block's current props JSON and a user change request, produce the FULL replacement props and a fresh summary.", + 'Reply with ONLY one JSON object, no markdown fences, no commentary: {"props": { ...complete replacement props... }, "summary": "..."}.', + "Rules:", + "- `props` must be the complete new props object (not a diff). Preserve every existing prop the request does not ask to change.", + "- Only use fields the block type supports. Never invent entity ids, keys, or recording ids.", + "- `summary` is one plain-text sentence (two at most, under 220 characters) describing the block AFTER the change, same style as the current summary.", + "- If the request is impossible for this block type, return the original props unchanged and explain why in the summary, prefixed with 'Unchanged:'.", + ].join("\n"); +} + +export function buildNotebookNodeChangeUserPrompt(input: { + tagName: string; + propsJson: string; + knowledge: string; + request: string; + currentSummary?: string | null; + liveEntityJson?: string | null; + retryContext?: { previousResponse: string; problem: string } | null; +}): string { + const sections = [ + `Notebook block type: <${input.tagName} />`, + `What this block type supports:\n${input.knowledge}`, + `Current props JSON:\n${input.propsJson}`, + ]; + if (input.currentSummary) { + sections.push(`Current summary: ${input.currentSummary}`); + } + if (input.liveEntityJson) { + sections.push( + `Live object fetched from PostHog (for context):\n${input.liveEntityJson}`, + ); + } + sections.push(`User change request: ${input.request}`); + if (input.retryContext) { + sections.push( + [ + "Your previous reply could not be applied.", + `Previous reply:\n${input.retryContext.previousResponse}`, + `Problem: ${input.retryContext.problem}`, + 'Reply again with ONLY the corrected {"props": ..., "summary": ...} JSON object.', + ].join("\n"), + ); + } + return sections.join("\n\n"); +} diff --git a/packages/core/src/notebooks/notebookNodeAIModel.ts b/packages/core/src/notebooks/notebookNodeAIModel.ts new file mode 100644 index 0000000000..9b2ce9597a --- /dev/null +++ b/packages/core/src/notebooks/notebookNodeAIModel.ts @@ -0,0 +1,219 @@ +import type { PostHogAPIClient } from "@posthog/api-client/posthog-client"; +import { injectable } from "inversify"; + +// Transport seam for the notebook node AI editor. The service depends on this +// interface so tests can fake the model; the default implementation prefers +// the direct LLM-gateway completion (single non-agentic round trip) and falls +// back to a Max conversation turn when the gateway is unreachable for this +// login (e.g. self-hosted without a gateway, or a token the gateway rejects). + +export interface NotebookNodeAIModelRequest { + system: string; + user: string; + signal?: AbortSignal; + /** Called with the full accumulated text on every streamed delta. */ + onText?: (accumulatedText: string) => void; +} + +export interface NotebookNodeAIModel { + /** One-shot completion. Resolves with the full response text. */ + complete( + client: PostHogAPIClient, + request: NotebookNodeAIModelRequest, + ): Promise; +} + +export const NOTEBOOK_NODE_AI_MODEL = Symbol.for( + "posthog.notebooks.nodeAIModel", +); + +/** Fast + cheap; the same model the harness uses for one-shot summarization. */ +const GATEWAY_MODEL = "claude-haiku-4-5"; + +/** + * Minimal SSE parser shared by both transports: utf-8 decode, accumulate + * `event:`/`data:` field lines (multi-line data joined with "\n"), dispatch on + * a blank line, ignore `:` comments. Mirrors the parser proven against the + * conversations endpoint in the notebooks inline-AI controller. + */ +export async function readNotebookAISSEStream( + body: ReadableStream, + onEvent: (eventName: string, data: string) => void, +): Promise { + const reader = body.getReader(); + const decoder = new TextDecoder(); + let buffer = ""; + let eventName = ""; + let dataLines: string[] = []; + + const dispatch = () => { + if (dataLines.length) { + onEvent(eventName || "message", dataLines.join("\n")); + } + eventName = ""; + dataLines = []; + }; + + const handleLine = (rawLine: string) => { + const line = rawLine.endsWith("\r") ? rawLine.slice(0, -1) : rawLine; + if (!line) { + dispatch(); + return; + } + if (line.startsWith(":")) return; + const colonIndex = line.indexOf(":"); + const field = colonIndex === -1 ? line : line.slice(0, colonIndex); + let value = colonIndex === -1 ? "" : line.slice(colonIndex + 1); + if (value.startsWith(" ")) value = value.slice(1); + if (field === "event") { + eventName = value; + } else if (field === "data") { + dataLines.push(value); + } + }; + + try { + for (;;) { + const { done, value } = await reader.read(); + if (done) break; + buffer += decoder.decode(value, { stream: true }); + let newlineIndex = buffer.indexOf("\n"); + while (newlineIndex !== -1) { + handleLine(buffer.slice(0, newlineIndex)); + buffer = buffer.slice(newlineIndex + 1); + newlineIndex = buffer.indexOf("\n"); + } + } + buffer += decoder.decode(); + if (buffer) handleLine(buffer); + dispatch(); + } finally { + reader.releaseLock(); + } +} + +@injectable() +export class PostHogNotebookNodeAIModel implements NotebookNodeAIModel { + /** Sticky: once the gateway rejects this login, stop paying its round trip. */ + private gatewayUnavailable = false; + + async complete( + client: PostHogAPIClient, + request: NotebookNodeAIModelRequest, + ): Promise { + if (!this.gatewayUnavailable) { + try { + return await this.completeViaGateway(client, request); + } catch (error) { + if (request.signal?.aborted) throw error; + this.gatewayUnavailable = true; + } + } + return await this.completeViaMaxConversation(client, request); + } + + private async completeViaGateway( + client: PostHogAPIClient, + request: NotebookNodeAIModelRequest, + ): Promise { + const response = await client.llmGatewayChatCompletion( + { + model: GATEWAY_MODEL, + stream: true, + messages: [ + { role: "system", content: request.system }, + { role: "user", content: request.user }, + ], + }, + request.signal, + ); + if (!response.body) { + throw new Error("LLM gateway response had no body"); + } + let text = ""; + await readNotebookAISSEStream(response.body, (_eventName, data) => { + if (data === "[DONE]") return; + let chunk: { + choices?: { delta?: { content?: unknown } }[]; + error?: { message?: unknown }; + }; + try { + chunk = JSON.parse(data) as typeof chunk; + } catch { + return; + } + if (chunk.error) { + throw new Error( + typeof chunk.error.message === "string" + ? chunk.error.message + : "LLM gateway stream error", + ); + } + const delta = chunk.choices?.[0]?.delta?.content; + if (typeof delta === "string" && delta) { + text += delta; + request.onText?.(text); + } + }); + if (!text.trim()) { + throw new Error("LLM gateway returned an empty response"); + } + return text; + } + + /** + * Fallback: one Max conversation turn. Max is a server-side agent graph, so + * this is slower and chattier than the gateway, but it works with any + * PostHog login. Assistant messages re-send the full accumulated content + * each tick (same contract the notebook inline AI streams against). + */ + private async completeViaMaxConversation( + client: PostHogAPIClient, + request: NotebookNodeAIModelRequest, + ): Promise { + const abort = new AbortController(); + const onOuterAbort = () => abort.abort(); + request.signal?.addEventListener("abort", onOuterAbort, { once: true }); + try { + const response = await client.startConversationStream( + { + content: `${request.system}\n\n${request.user}`, + conversation: globalThis.crypto.randomUUID(), + contextual_tools: {}, + ui_context: {}, + trace_id: globalThis.crypto.randomUUID(), + }, + abort.signal, + ); + if (!response.body) { + throw new Error("Max conversation response had no body"); + } + let text = ""; + await readNotebookAISSEStream(response.body, (eventName, data) => { + if (eventName !== "message") return; + let message: { type?: unknown; content?: unknown }; + try { + message = JSON.parse(data) as typeof message; + } catch { + return; + } + if (message.type === "ai" && typeof message.content === "string") { + text = message.content; + request.onText?.(text); + } else if (message.type === "ai/failure") { + throw new Error( + typeof message.content === "string" && message.content + ? message.content + : "Max failed to generate an answer", + ); + } + }); + if (!text.trim()) { + throw new Error("Max returned an empty response"); + } + return text; + } finally { + request.signal?.removeEventListener("abort", onOuterAbort); + } + } +} diff --git a/packages/core/src/notebooks/notebookNodeAIService.ts b/packages/core/src/notebooks/notebookNodeAIService.ts new file mode 100644 index 0000000000..208e85868b --- /dev/null +++ b/packages/core/src/notebooks/notebookNodeAIService.ts @@ -0,0 +1,443 @@ +import type { PostHogAPIClient } from "@posthog/api-client/posthog-client"; +import { inject, injectable } from "inversify"; +import { + buildNotebookNodeChangeSystemPrompt, + buildNotebookNodeChangeUserPrompt, + buildNotebookNodeSummarySystemPrompt, + buildNotebookNodeSummaryUserPrompt, + getNotebookNodeKnowledge, +} from "./notebookNodeAIKnowledge"; +import { + NOTEBOOK_NODE_AI_MODEL, + type NotebookNodeAIModel, +} from "./notebookNodeAIModel"; +import type { + NotebookNodeJsonObject, + NotebookNodeJsonValue, +} from "./notebookNodeSummary"; +import { notebookNodeAIChangeResponseSchema } from "./schemas"; + +export interface NotebookNodeAIInput { + tagName: string; + props: NotebookNodeJsonObject; +} + +export interface NotebookNodeAIChange { + /** Complete replacement props (shell-managed props reattached). */ + props: NotebookNodeJsonObject; + summary: string; +} + +// Props the component shell manages (panel visibility toggles). They are +// noise to the model: stripped from prompts and cache keys, preserved +// verbatim across AI edits. +const SHELL_MANAGED_PROPS = new Set(["hideFilters", "hideResults"]); + +/** Longest live-entity JSON we are willing to spend prompt tokens on. */ +const MAX_ENTITY_CONTEXT_CHARS = 1200; + +/** + * AI summarizer + change-request engine behind the notebook node edit panel. + * Owns prompt construction, the model call, response parsing/validation with + * one retry, and an in-memory summary cache keyed by (tagName + props) so + * reopening a node is instant. Host-neutral: the authenticated api client is + * passed per call (the renderer container has no AuthService binding), the + * model transport is injected. + */ +@injectable() +export class NotebookNodeAIService { + private readonly summaryCache = new Map(); + /** Coalesces concurrent summarize calls for the same node state. */ + private readonly pendingSummaries = new Map>(); + + constructor( + @inject(NOTEBOOK_NODE_AI_MODEL) + private readonly model: NotebookNodeAIModel, + ) {} + + /** Synchronous cache probe so the UI can render a known summary instantly. */ + getCachedSummary(input: NotebookNodeAIInput): string | null { + return this.summaryCache.get(summaryCacheKey(input)) ?? null; + } + + /** Pre-seed the cache (e.g. with the summary returned by a change request). */ + primeSummary(input: NotebookNodeAIInput, summary: string): void { + this.summaryCache.set(summaryCacheKey(input), summary); + } + + async summarizeNode( + client: PostHogAPIClient, + input: NotebookNodeAIInput, + options?: { + signal?: AbortSignal; + /** Full accumulated summary text on every streamed delta. */ + onPartial?: (text: string) => void; + }, + ): Promise { + const key = summaryCacheKey(input); + const cached = this.summaryCache.get(key); + if (cached) return cached; + const pending = this.pendingSummaries.get(key); + if (pending) return pending; + + const run = (async () => { + const liveEntityJson = await this.fetchEntityContext(client, input); + const raw = await this.model.complete(client, { + system: buildNotebookNodeSummarySystemPrompt(), + user: buildNotebookNodeSummaryUserPrompt({ + tagName: input.tagName, + propsJson: promptPropsJson(input.props), + knowledge: getNotebookNodeKnowledge(input.tagName), + liveEntityJson, + }), + ...(options?.signal ? { signal: options.signal } : {}), + ...(options?.onPartial + ? { + onText: (text: string) => options.onPartial?.(cleanSummary(text)), + } + : {}), + }); + const summary = cleanSummary(raw); + if (!summary) throw new Error("AI returned an empty summary"); + this.summaryCache.set(key, summary); + return summary; + })(); + this.pendingSummaries.set(key, run); + try { + return await run; + } finally { + this.pendingSummaries.delete(key); + } + } + + /** + * Apply a natural-language change request to a node. One model call returns + * both the replacement props and the new summary; invalid replies get one + * retry with the parse problem quoted back. + */ + async requestNodeChange( + client: PostHogAPIClient, + input: NotebookNodeAIInput, + request: string, + options?: { signal?: AbortSignal }, + ): Promise { + const liveEntityJson = await this.fetchEntityContext(client, input); + const basePrompt = { + tagName: input.tagName, + propsJson: promptPropsJson(input.props), + knowledge: getNotebookNodeKnowledge(input.tagName), + request, + currentSummary: this.getCachedSummary(input), + liveEntityJson, + }; + + let retryContext: { previousResponse: string; problem: string } | null = + null; + let lastProblem = "unknown"; + for (let attempt = 0; attempt < 2; attempt++) { + const raw = await this.model.complete(client, { + system: buildNotebookNodeChangeSystemPrompt(), + user: buildNotebookNodeChangeUserPrompt({ + ...basePrompt, + retryContext, + }), + ...(options?.signal ? { signal: options.signal } : {}), + }); + const outcome = parseNodeChangeResponse(input.tagName, raw); + if (outcome.ok) { + const props = reattachShellProps(input.props, outcome.props); + const summary = cleanSummary(outcome.summary); + this.primeSummary({ tagName: input.tagName, props }, summary); + return { props, summary }; + } + lastProblem = outcome.problem; + retryContext = { + previousResponse: raw.slice(0, 2000), + problem: outcome.problem, + }; + } + throw new Error(`AI response could not be applied: ${lastProblem}`); + } + + /** + * Best-effort live-object fetch for entity nodes so summaries can name the + * flag/experiment/survey instead of parroting an id. Failures (bad id, no + * access, offline) degrade to props-only prompts. + */ + private async fetchEntityContext( + client: PostHogAPIClient, + input: NotebookNodeAIInput, + ): Promise { + try { + const picked = await pickEntityContext(client, input); + if (!picked) return null; + const json = JSON.stringify(picked); + return json.length > MAX_ENTITY_CONTEXT_CHARS + ? `${json.slice(0, MAX_ENTITY_CONTEXT_CHARS)}…` + : json; + } catch { + return null; + } + } +} + +async function pickEntityContext( + client: PostHogAPIClient, + input: NotebookNodeAIInput, +): Promise | null> { + const props = input.props; + const id = props.id; + switch (input.tagName) { + case "FeatureFlag": { + if (typeof id !== "string" && typeof id !== "number") return null; + const flag = await client.getFeatureFlag(String(id)); + return { key: flag.key, name: flag.name, active: flag.active }; + } + case "Experiment": { + if (typeof id !== "number") return null; + const experiment = await client.getExperiment(id); + return { + name: experiment.name, + description: experiment.description ?? undefined, + start_date: experiment.start_date ?? undefined, + end_date: experiment.end_date ?? undefined, + feature_flag_key: experiment.feature_flag_key ?? undefined, + }; + } + case "Survey": { + if (typeof id !== "string") return null; + const survey = await client.getSurvey(id); + return { + name: survey.name, + type: survey.type ?? undefined, + start_date: survey.start_date ?? undefined, + end_date: survey.end_date ?? undefined, + questions: (survey.questions ?? []) + .slice(0, 5) + .map((question) => + question && typeof question === "object" && "question" in question + ? (question as { question?: unknown }).question + : undefined, + ), + }; + } + case "EarlyAccessFeature": { + if (typeof id !== "string") return null; + const feature = await client.getEarlyAccessFeature(id); + return { + name: feature.name, + stage: feature.stage ?? undefined, + description: feature.description?.slice(0, 200), + }; + } + case "Cohort": { + if (typeof id !== "number") return null; + const cohort = await client.getCohort(id); + return { + name: cohort.name, + count: cohort.count ?? undefined, + is_static: cohort.is_static, + }; + } + case "Person": { + const distinctId = props.distinctId; + if (typeof distinctId !== "string") return null; + const person = await client.getPerson({ distinctId }); + return { + name: person.name ?? undefined, + email: person.properties?.email ?? person.properties?.$email, + }; + } + default: + return null; + } +} + +// --------------------------------------------------------------------------- +// Pure helpers (exported for tests and for the UI's apply step) +// --------------------------------------------------------------------------- + +/** + * Turn a full replacement props object into the partial-update shape the + * notebook shell's `updateProps` expects: keys missing from `next` are set to + * `undefined` so they get removed, shell-managed props are never deleted. + */ +export function buildNotebookNodePropsUpdate( + current: NotebookNodeJsonObject, + next: NotebookNodeJsonObject, +): Record { + const update: Record = { + ...next, + }; + for (const key of Object.keys(current)) { + if (!(key in next) && !SHELL_MANAGED_PROPS.has(key)) { + update[key] = undefined; + } + } + for (const key of SHELL_MANAGED_PROPS) { + delete update[key]; + } + return update; +} + +function reattachShellProps( + current: NotebookNodeJsonObject, + next: NotebookNodeJsonObject, +): NotebookNodeJsonObject { + const merged: NotebookNodeJsonObject = { ...next }; + for (const key of SHELL_MANAGED_PROPS) { + delete merged[key]; + if (key in current) { + merged[key] = current[key] as NotebookNodeJsonValue; + } + } + return merged; +} + +function stripShellProps( + props: NotebookNodeJsonObject, +): NotebookNodeJsonObject { + const stripped: NotebookNodeJsonObject = {}; + for (const [key, value] of Object.entries(props)) { + if (!SHELL_MANAGED_PROPS.has(key)) stripped[key] = value; + } + return stripped; +} + +function promptPropsJson(props: NotebookNodeJsonObject): string { + return JSON.stringify(stripShellProps(props), null, 1); +} + +/** FNV-1a over a stable (key-sorted) serialization — cache key material. */ +export function notebookNodeAICacheKey(input: NotebookNodeAIInput): string { + return summaryCacheKey(input); +} + +function summaryCacheKey(input: NotebookNodeAIInput): string { + const canonical = stableStringify(stripShellProps(input.props)); + let hash = 0x811c9dc5; + const text = `${input.tagName}${canonical}`; + for (let index = 0; index < text.length; index++) { + hash ^= text.charCodeAt(index); + hash = Math.imul(hash, 0x01000193); + } + return `${input.tagName}:${(hash >>> 0).toString(16)}`; +} + +function stableStringify(value: NotebookNodeJsonValue): string { + if (Array.isArray(value)) { + return `[${value.map(stableStringify).join(",")}]`; + } + if (value && typeof value === "object") { + const entries = Object.entries(value) + .sort(([a], [b]) => (a < b ? -1 : a > b ? 1 : 0)) + .map( + ([key, entry]) => `${JSON.stringify(key)}:${stableStringify(entry)}`, + ); + return `{${entries.join(",")}}`; + } + return JSON.stringify(value); +} + +/** Strip fences/quotes/whitespace the model may wrap a plain-text reply in. */ +function cleanSummary(text: string): string { + let cleaned = text.trim(); + const fence = /^```[a-z]*\n([\s\S]*?)\n?```$/.exec(cleaned); + if (fence?.[1] !== undefined) cleaned = fence[1].trim(); + if (cleaned.startsWith('"') && cleaned.endsWith('"') && cleaned.length > 1) { + cleaned = cleaned.slice(1, -1).trim(); + } + return cleaned.replace(/\s+/g, " "); +} + +type ParseOutcome = + | { ok: true; props: NotebookNodeJsonObject; summary: string } + | { ok: false; problem: string }; + +export function parseNodeChangeResponse( + tagName: string, + raw: string, +): ParseOutcome { + const jsonText = extractFirstJsonObject(raw); + if (!jsonText) { + return { ok: false, problem: "no JSON object found in the reply" }; + } + let parsed: unknown; + try { + parsed = JSON.parse(jsonText); + } catch (error) { + return { + ok: false, + problem: `invalid JSON: ${error instanceof Error ? error.message : String(error)}`, + }; + } + const validated = notebookNodeAIChangeResponseSchema.safeParse(parsed); + if (!validated.success) { + return { + ok: false, + problem: `reply must be {"props": object, "summary": string} — ${validated.error.issues + .map((issue) => issue.message) + .join("; ")}`, + }; + } + const props = validated.data.props as NotebookNodeJsonObject; + const shapeProblem = validateChangedProps(tagName, props); + if (shapeProblem) { + return { ok: false, problem: shapeProblem }; + } + return { ok: true, props, summary: validated.data.summary }; +} + +/** Per-tag structural sanity beyond "is a JSON object". */ +function validateChangedProps( + tagName: string, + props: NotebookNodeJsonObject, +): string | null { + if (tagName === "Query") { + const query = props.query; + if (!query || typeof query !== "object" || Array.isArray(query)) { + return 'props.query must be an object with a "kind" field'; + } + if (typeof (query as NotebookNodeJsonObject).kind !== "string") { + return 'props.query.kind must be a string (e.g. "InsightVizNode")'; + } + } + return null; +} + +/** + * Extract the first balanced top-level JSON object from free-form model + * output (handles fences and prose around the object, and braces inside + * strings). + */ +export function extractFirstJsonObject(text: string): string | null { + const start = text.indexOf("{"); + if (start === -1) return null; + let depth = 0; + let inString = false; + let escaped = false; + for (let index = start; index < text.length; index++) { + const char = text[index]; + if (inString) { + if (escaped) { + escaped = false; + } else if (char === "\\") { + escaped = true; + } else if (char === '"') { + inString = false; + } + continue; + } + if (char === '"') { + inString = true; + } else if (char === "{") { + depth++; + } else if (char === "}") { + depth--; + if (depth === 0) { + return text.slice(start, index + 1); + } + } + } + return null; +} diff --git a/packages/core/src/notebooks/notebookNodeSummary.test.ts b/packages/core/src/notebooks/notebookNodeSummary.test.ts new file mode 100644 index 0000000000..001f79e604 --- /dev/null +++ b/packages/core/src/notebooks/notebookNodeSummary.test.ts @@ -0,0 +1,138 @@ +import { describe, expect, it } from "vitest"; +import { + humanizeDateRange, + type NotebookNodeJsonObject, + summarizeNotebookNodeLocally, +} from "./notebookNodeSummary"; + +describe("humanizeDateRange", () => { + it.each([ + [{ date_from: "-30d" }, "last 30 days"], + [{ date_from: "-24h" }, "last 24 hours"], + [{ date_from: "-1w" }, "last 1 week"], + [{ date_from: "-90d" }, "last 90 days"], + [{ date_from: "all" }, "all time"], + [{ date_from: "dStart" }, "today"], + [{ date_from: "mStart" }, "this month"], + [ + { date_from: "2026-01-01", date_to: "2026-02-01" }, + "2026-01-01 to 2026-02-01", + ], + [{ date_from: "2026-01-01" }, "since 2026-01-01"], + [{}, null], + ])("maps %j to %j", (dateRange, expected) => { + expect(humanizeDateRange(dateRange as NotebookNodeJsonObject)).toBe( + expected, + ); + }); + + it("returns null for a null range", () => { + expect(humanizeDateRange(null)).toBeNull(); + }); +}); + +describe("summarizeNotebookNodeLocally", () => { + it("summarizes a trends query with breakdown, range and display", () => { + const props: NotebookNodeJsonObject = { + query: { + kind: "InsightVizNode", + source: { + kind: "TrendsQuery", + series: [ + { kind: "EventsNode", event: "$pageview" }, + { kind: "EventsNode", event: "signup", math: "dau" }, + ], + dateRange: { date_from: "-30d" }, + trendsFilter: { display: "ActionsBar" }, + breakdownFilter: { breakdown: "$browser", breakdown_type: "event" }, + }, + }, + }; + expect(summarizeNotebookNodeLocally("Query", props)).toBe( + "Trends bar chart, $pageview vs. signup (dau), last 30 days, broken down by $browser", + ); + }); + + it.each([ + [ + "funnel steps in order", + { + query: { + kind: "InsightVizNode", + source: { + kind: "FunnelsQuery", + series: [ + { kind: "EventsNode", event: "$pageview" }, + { kind: "EventsNode", event: "purchase" }, + ], + }, + }, + }, + "Funnel: $pageview → purchase", + ], + [ + "SQL query text", + { + query: { + kind: "HogQLQuery", + query: "SELECT event,\n count() FROM events", + }, + }, + "SQL: SELECT event, count() FROM events", + ], + [ + "events table", + { + query: { + kind: "DataTableNode", + source: { kind: "EventsQuery", after: "-24h", limit: 100 }, + }, + }, + "Events table, last 24 hours, limit 100", + ], + [ + "saved insight", + { query: { kind: "SavedInsightNode", shortId: "AbC123" } }, + "Saved insight AbC123", + ], + [ + "retention", + { + query: { + kind: "InsightVizNode", + source: { + kind: "RetentionQuery", + retentionFilter: { + period: "Week", + targetEntity: { name: "$pageview" }, + returningEntity: { name: "$pageview" }, + }, + }, + }, + }, + "Retention: $pageview then $pageview, by week", + ], + ["unconfigured query", {}, "Query (not configured)"], + ])("summarizes %s", (_label, props, expected) => { + expect( + summarizeNotebookNodeLocally("Query", props as NotebookNodeJsonObject), + ).toBe(expected); + }); + + it.each([ + ["FeatureFlag", { id: "my-flag" }, "Feature flag my-flag"], + ["FeatureFlag", {}, "Feature flag (no id)"], + ["Experiment", { id: 42 }, "Experiment 42"], + ["Survey", { id: "s-1" }, "Survey s-1"], + ["Cohort", { id: 7 }, "Cohort 7"], + ["Person", { distinctId: "user@x.com" }, "Person user@x.com"], + ["Group", { groupKey: "acme" }, "Group acme"], + ["Recording", { sessionRecordingId: "rec-9" }, "Session recording rec-9"], + ["EarlyAccessFeature", { id: "eaf" }, "Early access feature eaf"], + ["SomethingElse", { foo: 1 }, "Something Else (foo)"], + ])("summarizes a %s node", (tagName, props, expected) => { + expect( + summarizeNotebookNodeLocally(tagName, props as NotebookNodeJsonObject), + ).toBe(expected); + }); +}); diff --git a/packages/core/src/notebooks/notebookNodeSummary.ts b/packages/core/src/notebooks/notebookNodeSummary.ts new file mode 100644 index 0000000000..7c657146df --- /dev/null +++ b/packages/core/src/notebooks/notebookNodeSummary.ts @@ -0,0 +1,257 @@ +// Deterministic, instant summaries of notebook component nodes. These render +// immediately when a node enters edit mode; the AI summary (see +// NotebookNodeAIService) streams in afterwards to replace/enrich them. Pure +// functions only — no I/O, no React. + +export type NotebookNodeJsonValue = + | string + | number + | boolean + | null + | NotebookNodeJsonValue[] + | { [key: string]: NotebookNodeJsonValue }; + +export type NotebookNodeJsonObject = Record; + +function asObject(value: unknown): NotebookNodeJsonObject | null { + return value && typeof value === "object" && !Array.isArray(value) + ? (value as NotebookNodeJsonObject) + : null; +} + +function asString(value: unknown): string | null { + return typeof value === "string" && value.trim() ? value.trim() : null; +} + +function asIdentifier(value: unknown): string | null { + if (typeof value === "number") return String(value); + return asString(value); +} + +/** "TrendsQuery" → "Trends", "WebOverviewQuery" → "Web overview". */ +function cleanQueryKind(kind: string): string { + const spaced = kind.replace(/Query$/, "").replace(/([a-z])([A-Z])/g, "$1 $2"); + return spaced.charAt(0) + spaced.slice(1).toLowerCase(); +} + +const DISPLAY_LABELS: Record = { + ActionsLineGraph: "line chart", + ActionsLineGraphCumulative: "cumulative line chart", + ActionsAreaGraph: "area chart", + ActionsBar: "bar chart", + ActionsBarValue: "value bar chart", + ActionsPie: "pie chart", + ActionsTable: "table", + BoldNumber: "big number", + WorldMap: "world map", +}; + +/** "-30d" → "last 30 days", "-24h" → "last 24 hours", "mStart" → "this month". */ +export function humanizeDateRange( + dateRange: NotebookNodeJsonObject | null, +): string | null { + if (!dateRange) return null; + const from = asString(dateRange.date_from); + const to = asString(dateRange.date_to); + if (!from) return null; + if (from === "all") return "all time"; + if (from === "dStart") return "today"; + if (from === "-1dStart") return "yesterday"; + if (from === "mStart") return "this month"; + if (from === "yStart") return "this year"; + const relative = /^-(\d+)([hdwmy])$/.exec(from); + if (relative) { + const count = Number(relative[1]); + const unit = { h: "hour", d: "day", w: "week", m: "month", y: "year" }[ + relative[2] as "h" | "d" | "w" | "m" | "y" + ]; + return `last ${count} ${unit}${count === 1 ? "" : "s"}`; + } + if (to) return `${from} to ${to}`; + return `since ${from}`; +} + +function describeSeries(source: NotebookNodeJsonObject): string[] { + const series = Array.isArray(source.series) ? source.series : []; + const labels: string[] = []; + for (const entry of series) { + const item = asObject(entry); + if (!item) continue; + const label = + asString(item.custom_name) ?? + asString(item.name) ?? + asString(item.event) ?? + (item.kind === "ActionsNode" && item.id != null + ? `action ${String(item.id)}` + : null); + if (!label) continue; + const math = asString(item.math); + labels.push(math && math !== "total" ? `${label} (${math})` : label); + } + return labels; +} + +function describeBreakdown(source: NotebookNodeJsonObject): string | null { + const filter = asObject(source.breakdownFilter); + if (!filter) return null; + const multi = Array.isArray(filter.breakdowns) + ? filter.breakdowns + .map((entry) => asString(asObject(entry)?.property)) + .filter((property): property is string => property !== null) + : []; + const single = asString(filter.breakdown); + const properties = multi.length ? multi : single ? [single] : []; + return properties.length ? `broken down by ${properties.join(", ")}` : null; +} + +function summarizeQueryNode(query: NotebookNodeJsonObject): string { + const kind = asString(query.kind); + if (!kind) return "Query (no kind set)"; + if (kind === "SavedInsightNode") { + return `Saved insight ${asString(query.shortId) ?? ""}`.trim(); + } + + const wrapped = + (kind === "InsightVizNode" || kind === "DataTableNode") && + asObject(query.source) + ? asObject(query.source) + : null; + const source = wrapped ?? query; + const sourceKind = asString(source.kind); + if (!sourceKind) return "Query (no kind set)"; + + const dateRange = humanizeDateRange(asObject(source.dateRange)); + const parts: string[] = []; + + switch (sourceKind) { + case "TrendsQuery": { + const display = asString(asObject(source.trendsFilter)?.display); + const displayLabel = display ? DISPLAY_LABELS[display] : null; + const series = describeSeries(source); + parts.push(`Trends ${displayLabel ?? "line chart"}`); + if (series.length) parts.push(series.join(" vs. ")); + if (dateRange) parts.push(dateRange); + const breakdown = describeBreakdown(source); + if (breakdown) parts.push(breakdown); + const interval = asString(source.interval); + if (interval && interval !== "day") parts.push(`by ${interval}`); + break; + } + case "FunnelsQuery": { + const series = describeSeries(source); + parts.push( + series.length ? `Funnel: ${series.join(" → ")}` : "Funnel (no steps)", + ); + if (dateRange) parts.push(dateRange); + break; + } + case "RetentionQuery": { + const filter = asObject(source.retentionFilter); + const target = asString(asObject(filter?.targetEntity)?.name); + const returning = asString(asObject(filter?.returningEntity)?.name); + const period = asString(filter?.period); + parts.push( + target && returning + ? `Retention: ${target} then ${returning}` + : "Retention", + ); + if (period) parts.push(`by ${period.toLowerCase()}`); + if (dateRange) parts.push(dateRange); + break; + } + case "PathsQuery": + parts.push("User paths"); + if (dateRange) parts.push(dateRange); + break; + case "StickinessQuery": + case "LifecycleQuery": { + const series = describeSeries(source); + parts.push(sourceKind === "StickinessQuery" ? "Stickiness" : "Lifecycle"); + if (series.length) parts.push(series.join(", ")); + if (dateRange) parts.push(dateRange); + break; + } + case "HogQLQuery": { + const sql = asString(source.query)?.replace(/\s+/g, " ") ?? ""; + parts.push(sql ? `SQL: ${truncate(sql, 100)}` : "SQL query"); + break; + } + case "EventsQuery": { + parts.push("Events table"); + const after = asString(source.after); + if (after) { + parts.push(humanizeDateRange({ date_from: after }) ?? `after ${after}`); + } + if (typeof source.limit === "number") { + parts.push(`limit ${source.limit}`); + } + break; + } + case "ActorsQuery": + parts.push("People table"); + break; + default: + parts.push(cleanQueryKind(sourceKind)); + if (dateRange) parts.push(dateRange); + } + return parts.join(", "); +} + +function truncate(value: string, max: number): string { + return value.length > max ? `${value.slice(0, max - 1)}…` : value; +} + +/** + * Instant human-readable summary of a node's props. Never throws; unknown + * shapes fall back to a generic label so the edit panel always has something + * to show before (or without) the AI summary. + */ +export function summarizeNotebookNodeLocally( + tagName: string, + props: NotebookNodeJsonObject, +): string { + switch (tagName) { + case "Query": { + const query = asObject(props.query); + return query ? summarizeQueryNode(query) : "Query (not configured)"; + } + case "FeatureFlag": { + const id = asIdentifier(props.id); + return id ? `Feature flag ${id}` : "Feature flag (no id)"; + } + case "Experiment": { + const id = asIdentifier(props.id); + return id ? `Experiment ${id}` : "Experiment (no id)"; + } + case "Survey": { + const id = asIdentifier(props.id); + return id ? `Survey ${id}` : "Survey (no id)"; + } + case "EarlyAccessFeature": { + const id = asIdentifier(props.id); + return id ? `Early access feature ${id}` : "Early access feature (no id)"; + } + case "Cohort": { + const id = asIdentifier(props.id); + return id ? `Cohort ${id}` : "Cohort (no id)"; + } + case "Person": { + const id = asIdentifier(props.distinctId) ?? asIdentifier(props.id); + return id ? `Person ${id}` : "Person (no id)"; + } + case "Group": { + const key = asIdentifier(props.groupKey) ?? asIdentifier(props.id); + return key ? `Group ${key}` : "Group (no key)"; + } + case "Recording": { + const id = + asIdentifier(props.sessionRecordingId) ?? asIdentifier(props.id); + return id ? `Session recording ${id}` : "Session recording (no id)"; + } + default: { + const label = tagName.replace(/([a-z])([A-Z])/g, "$1 $2"); + const keys = Object.keys(props); + return keys.length ? `${label} (${keys.join(", ")})` : label; + } + } +} diff --git a/packages/core/src/notebooks/notebooks.module.ts b/packages/core/src/notebooks/notebooks.module.ts index a7a738d10e..ef88af1b9c 100644 --- a/packages/core/src/notebooks/notebooks.module.ts +++ b/packages/core/src/notebooks/notebooks.module.ts @@ -1,5 +1,10 @@ import { ContainerModule } from "inversify"; -import { NOTEBOOKS_SERVICE } from "./identifiers"; +import { NOTEBOOK_NODE_AI_SERVICE, NOTEBOOKS_SERVICE } from "./identifiers"; +import { + NOTEBOOK_NODE_AI_MODEL, + PostHogNotebookNodeAIModel, +} from "./notebookNodeAIModel"; +import { NotebookNodeAIService } from "./notebookNodeAIService"; import { NotebooksService } from "./notebooksService"; // Host-agnostic notebooks service (PostHog cloud Notebooks REST API). It only @@ -8,4 +13,9 @@ import { NotebooksService } from "./notebooksService"; export const notebooksCoreModule = new ContainerModule(({ bind }) => { bind(NotebooksService).toSelf().inSingletonScope(); bind(NOTEBOOKS_SERVICE).toService(NotebooksService); + bind(NOTEBOOK_NODE_AI_MODEL) + .to(PostHogNotebookNodeAIModel) + .inSingletonScope(); + bind(NotebookNodeAIService).toSelf().inSingletonScope(); + bind(NOTEBOOK_NODE_AI_SERVICE).toService(NotebookNodeAIService); }); diff --git a/packages/core/src/notebooks/schemas.ts b/packages/core/src/notebooks/schemas.ts index c024aca506..c7d8d35142 100644 --- a/packages/core/src/notebooks/schemas.ts +++ b/packages/core/src/notebooks/schemas.ts @@ -29,3 +29,14 @@ export const notebookListItemSchema = z last_modified_at: z.string(), }) .array(); + +// The single-call contract of the AI node editor: one model reply carries +// BOTH the full replacement props and the fresh human-readable summary. +export const notebookNodeAIChangeResponseSchema = z.object({ + props: z.record(z.string(), z.unknown()), + summary: z.string().trim().min(1), +}); + +export type NotebookNodeAIChangeResponse = z.infer< + typeof notebookNodeAIChangeResponseSchema +>; diff --git a/packages/ui/src/features/notebooks/ai/NotebookNodeAIEditPanel.tsx b/packages/ui/src/features/notebooks/ai/NotebookNodeAIEditPanel.tsx new file mode 100644 index 0000000000..026adbd7a4 --- /dev/null +++ b/packages/ui/src/features/notebooks/ai/NotebookNodeAIEditPanel.tsx @@ -0,0 +1,171 @@ +import { buildNotebookNodePropsUpdate } from "@posthog/core/notebooks/notebookNodeAIService"; +import type { NotebookNodeJsonObject } from "@posthog/core/notebooks/notebookNodeSummary"; +import { Button, Input, Textarea } from "@posthog/quill"; +import { ChevronRight, Sparkles } from "lucide-react"; +import { type JSX, type KeyboardEvent, useState } from "react"; +import type { NotebookComponentRenderProps } from "../markdown-notebook/types"; +import { + useNotebookNodeAIChange, + useNotebookNodeAISummary, +} from "./useNotebookNodeAI"; + +/** + * Universal AI edit panel for PostHog entity nodes: a human-readable summary + * of the node's current props (local heuristic instantly, AI streaming in + * over it), a prompt box that applies natural-language change requests via a + * single model call returning {props, summary}, and a collapsible raw-JSON + * escape hatch that commits through the same shell `updateProps` plumbing as + * the generic edit panel it replaces. + */ +export function NotebookNodeAIEditPanel({ + node, + updateProps, +}: NotebookComponentRenderProps): JSX.Element { + // NotebookComponentProps and NotebookNodeJsonObject are the same JSON shape; + // core just owns its own name for it to stay UI-independent. + const props = node.props as NotebookNodeJsonObject; + const summaryState = useNotebookNodeAISummary(node.tagName, props); + const change = useNotebookNodeAIChange(node.tagName); + const [prompt, setPrompt] = useState(""); + const [rawOpen, setRawOpen] = useState(false); + const [rawDraft, setRawDraft] = useState(null); + const [rawError, setRawError] = useState(null); + + const submitPrompt = (): void => { + const request = prompt.trim(); + if (!request || change.isPending) return; + change.mutate( + { props, request }, + { + onSuccess: (result) => { + updateProps(buildNotebookNodePropsUpdate(props, result.props)); + setPrompt(""); + setRawDraft(null); + setRawError(null); + }, + }, + ); + }; + + const handlePromptKeyDown = (event: KeyboardEvent) => { + if (event.key === "Enter") { + event.preventDefault(); + submitPrompt(); + } + }; + + const rawValue = rawDraft ?? JSON.stringify(props, null, 2); + const applyRawJson = (): void => { + try { + const parsed: unknown = JSON.parse(rawValue); + if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) { + setRawError("Props must be a JSON object."); + return; + } + updateProps( + buildNotebookNodePropsUpdate(props, parsed as NotebookNodeJsonObject), + ); + setRawDraft(null); + setRawError(null); + } catch (error) { + setRawError(error instanceof Error ? error.message : "Invalid JSON"); + } + }; + + const changeErrorMessage = + change.error instanceof Error ? change.error.message : null; + + return ( +
+
+ +

+ {summaryState.summary} + {summaryState.isStreaming ? ( + + ▍ + + ) : null} +

+
+ {summaryState.hasClient ? ( +
+ setPrompt(event.target.value)} + onKeyDown={handlePromptKeyDown} + placeholder='Describe a change, e.g. "split by country, last 90 days"' + aria-label="AI change request" + disabled={change.isPending} + className="flex-1" + /> + +
+ ) : ( +

+ Connect to PostHog to edit this block with AI. +

+ )} + {changeErrorMessage ? ( +

{changeErrorMessage}

+ ) : null} +
+ + {rawOpen ? ( +
+