From a3b5eb72c4ef663c65bbdcfefdbf768296fae162 Mon Sep 17 00:00:00 2001 From: JeremyFunk Date: Wed, 2 Sep 2026 05:53:42 +0200 Subject: [PATCH] perf(agent-sessions): project the span read down to the keys the mapper reads MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The session spans read selected `SpanAttributes` and `ResourceAttributes` whole, then `mapAiSpan` read a fixed list of keys off the first and nothing off the second. Measured on the largest production sessions, the resource map was ~60% of the raw bytes and one unrelated key (`db.query.text`) was half of what remained — none of it reached the wire. `spanProjection` now filters the span map to `aiSpanAttributeKeys` — every source key of every integration plus what the refine hooks read, declared next to the hook as `refineKeys` — and the prompt-variable prefix, and drops the resource map. The byte cap on the read now measures what actually ships. `mapFilterKeys` is the builder primitive: `mapFilter((k, v) -> …)` with the key predicate written in the DSL's own conditions. --- .../routes/internal/ai-sessions.http.test.ts | 4 +-- .../src/ch/core-dsl.test.ts | 8 +++++ .../src/ch/functions/index.ts | 2 +- .../src/ch/functions/map.ts | 17 ++++++++++ lib/clickhouse-builder/src/ch/index.ts | 1 + .../src/__sql_baseline__/integrations.sql | 6 ++-- .../src/ai/ai-integrations.test.ts | 22 ------------- .../src/ai/ai-integrations.ts | 26 +++++++++++++++ .../src/ai/ai-sessions.test.ts | 33 ++++++++++++++----- .../src/ai/ai-sessions.ts | 23 ++++++++----- .../src/ai/ai-vendors.test.ts | 1 - .../src/ai/ai-vendors.ts | 3 ++ .../query-engine-integrations/src/ai/index.ts | 1 + 13 files changed, 100 insertions(+), 47 deletions(-) diff --git a/apps/api/src/routes/internal/ai-sessions.http.test.ts b/apps/api/src/routes/internal/ai-sessions.http.test.ts index 362df25a9..8d2a3cbf0 100644 --- a/apps/api/src/routes/internal/ai-sessions.http.test.ts +++ b/apps/api/src/routes/internal/ai-sessions.http.test.ts @@ -66,7 +66,6 @@ const spanRow = (index: number) => ({ statusMessage: "", timestamp: "2026-08-19 10:00:00.000000000", spanAttributes: { "gen_ai.operation.name": "chat", "maple_ai.session.id": SESSION_ID }, - resourceAttributes: {}, }) const makeHarness = (overrides: Partial) => { @@ -230,7 +229,8 @@ describe("POST /internal/ai-sessions/spans", () => { expect(windowSql).toContain(`TraceId = '${TRACE_ID}'`) expect(windowSql).not.toContain("maple_ai.session.id") expect(spansSql).toContain(`TraceId = '${TRACE_ID}'`) - expect(spansSql).not.toContain("maple_ai.session.id") + // The projection names the key; the predicate is what must be absent. + expect(spansSql).not.toContain("SpanAttributes['maple_ai.session.id']") // The bounds the window read handed back still prune the span read. expect(spansSql).toContain(`Timestamp >= '${resolved.startTime}'`) expect(spansSql).not.toContain("__PARAM_") diff --git a/lib/clickhouse-builder/src/ch/core-dsl.test.ts b/lib/clickhouse-builder/src/ch/core-dsl.test.ts index 66bcd17f2..e52180e07 100644 --- a/lib/clickhouse-builder/src/ch/core-dsl.test.ts +++ b/lib/clickhouse-builder/src/ch/core-dsl.test.ts @@ -88,6 +88,14 @@ describe("expression functions", () => { expect(sql).toContain("map('key1', Name, 'key2', 'val') AS m") }) + it("compiles mapFilterKeys with the DSL's own conditions on the key", () => { + const q = CH.from(TestTable).select(($) => ({ + m: CH.mapFilterKeys($.Attrs, (k) => k.in_("a", "b").or(k.like("x.%"))), + })) + const { sql } = compileCHUnsafe(q, {}) + expect(sql).toContain("mapFilter((k, v) -> (k IN ('a', 'b') OR k LIKE 'x.%'), Attrs) AS m") + }) + it("compiles empty mapLiteral", () => { const q = CH.from(TestTable).select(() => ({ m: CH.mapLiteral() })) const { sql } = compileCHUnsafe(q, {}) diff --git a/lib/clickhouse-builder/src/ch/functions/index.ts b/lib/clickhouse-builder/src/ch/functions/index.ts index 700d25d56..c527089c6 100644 --- a/lib/clickhouse-builder/src/ch/functions/index.ts +++ b/lib/clickhouse-builder/src/ch/functions/index.ts @@ -90,7 +90,7 @@ export { has, } from "./array" -export { mapContains, mapGet, mapKeys, mapValues, mapLiteral } from "./map" +export { mapContains, mapFilterKeys, mapGet, mapKeys, mapValues, mapLiteral } from "./map" export { toJSONString } from "./json" diff --git a/lib/clickhouse-builder/src/ch/functions/map.ts b/lib/clickhouse-builder/src/ch/functions/map.ts index a052671c1..e490869da 100644 --- a/lib/clickhouse-builder/src/ch/functions/map.ts +++ b/lib/clickhouse-builder/src/ch/functions/map.ts @@ -24,6 +24,23 @@ export function mapValues(mapExpr: Expr>): Expr , map)` — the entries whose KEY passes. + * + * The predicate is built from the lambda's key parameter, so it can use every + * condition the DSL has (`in_`, `like`, `or`, …). Values are not inspected. + */ +export function mapFilterKeys( + mapExpr: Expr>, + predicate: (key: Expr) => Condition, +): Expr> { + const key = makeExpr(raw("k"), T.string.schema) + return makeExpr( + raw(`mapFilter((k, v) -> ${compile(predicate(key).toFragment())}, ${compile(mapExpr.toFragment())})`), + STRING_MAP, + ) +} + export function mapLiteral(...pairs: Array<[string, Expr]>): Expr> { if (pairs.length === 0) return makeExpr(raw("map()"), STRING_MAP) const args = pairs.map(([k, v]) => `${compile(str(k))}, ${compile(v.toFragment())}`).join(", ") diff --git a/lib/clickhouse-builder/src/ch/index.ts b/lib/clickhouse-builder/src/ch/index.ts index 815a71264..c4d67fd80 100644 --- a/lib/clickhouse-builder/src/ch/index.ts +++ b/lib/clickhouse-builder/src/ch/index.ts @@ -203,6 +203,7 @@ export { mapGet, mapKeys, mapValues, + mapFilterKeys, mapLiteral, // JSON toJSONString, diff --git a/packages/query-engine-integrations/src/__sql_baseline__/integrations.sql b/packages/query-engine-integrations/src/__sql_baseline__/integrations.sql index 585d0f261..b7e91c567 100644 --- a/packages/query-engine-integrations/src/__sql_baseline__/integrations.sql +++ b/packages/query-engine-integrations/src/__sql_baseline__/integrations.sql @@ -129,8 +129,7 @@ SELECT StatusCode AS statusCode, StatusMessage AS statusMessage, toString(Timestamp) AS timestamp, - SpanAttributes AS spanAttributes, - ResourceAttributes AS resourceAttributes + mapFilter((k, v) -> (k IN ('maple_ai.session.id', 'maple_ai.vendor.id', 'maple_ai.vendor.version', 'gen_ai.operation.name', 'gen_ai.provider.name', 'gen_ai.system', 'gen_ai.request.model', 'gen_ai.request.max_tokens', 'gen_ai.request.choice.count', 'gen_ai.request.temperature', 'gen_ai.request.top_p', 'gen_ai.request.top_k', 'gen_ai.request.stop_sequences', 'gen_ai.request.frequency_penalty', 'gen_ai.request.presence_penalty', 'gen_ai.request.encoding_formats', 'gen_ai.request.seed', 'gen_ai.openai.request.seed', 'gen_ai.request.stream', 'gen_ai.request.reasoning.level', 'gen_ai.request.previous_response.id', 'gen_ai.request.stream_cursor', 'gen_ai.response.id', 'gen_ai.response.model', 'gen_ai.response.finish_reasons', 'gen_ai.response.finish_reason', 'gen_ai.response.status', 'gen_ai.response.time_to_first_chunk', 'gen_ai.output.type', 'gen_ai.usage.input_tokens', 'gen_ai.usage.prompt_tokens', 'gen_ai.usage.cache_read.input_tokens', 'gen_ai.usage.input_tokens.cached', 'gen_ai.usage.cache_creation.input_tokens', 'gen_ai.usage.cache_write.input_tokens', 'gen_ai.usage.output_tokens', 'gen_ai.usage.completion_tokens', 'gen_ai.usage.reasoning.output_tokens', 'gen_ai.usage.output_tokens.reasoning', 'gen_ai.usage.cost', 'gen_ai.usage.total_cost', 'gen_ai.conversation.id', 'gen_ai.conversation.compacted', 'gen_ai.agent.id', 'gen_ai.agent.name', 'gen_ai.agent.description', 'gen_ai.agent.version', 'gen_ai.tool.name', 'gen_ai.tool.call.id', 'gen_ai.tool.description', 'gen_ai.tool.type', 'gen_ai.tool.call.arguments', 'gen_ai.tool.call.result', 'gen_ai.tool.definitions', 'gen_ai.system_instructions', 'gen_ai.input.messages', 'gen_ai.prompt', 'gen_ai.output.messages', 'gen_ai.completion', 'gen_ai.data_source.id', 'gen_ai.retrieval.query.text', 'gen_ai.retrieval.top_k', 'gen_ai.retrieval.documents', 'gen_ai.memory.store.id', 'gen_ai.memory.record.id', 'gen_ai.memory.record.count', 'gen_ai.memory.query.text', 'gen_ai.memory.records', 'gen_ai.embeddings.dimension.count', 'gen_ai.evaluation.name', 'gen_ai.evaluation.score.value', 'gen_ai.evaluation.score.label', 'gen_ai.evaluation.explanation', 'gen_ai.prompt.name', 'gen_ai.prompt.version', 'gen_ai.workflow.name', 'error.type', 'server.address', 'server.port', 'ai.model.provider', 'ai.model.id', 'ai.response.id', 'ai.response.model', 'ai.response.finishReason', 'gen_ai.client.operation.time_to_first_chunk', 'ai.usage.inputTokens', 'ai.usage.promptTokens', 'ai.usage.cachedInputTokens', 'ai.usage.inputTokenDetails.cacheReadTokens', 'ai.usage.inputTokenDetails.cacheWriteTokens', 'ai.usage.outputTokens', 'ai.usage.completionTokens', 'ai.usage.reasoningTokens', 'ai.usage.outputTokenDetails.reasoningTokens', 'ai.telemetry.functionId', 'ai.toolCall.name', 'ai.toolCall.id', 'ai.toolCall.args', 'ai.toolCall.result', 'ai.prompt.tools', 'ai.prompt.messages', 'ai.prompt', 'llm.provider', 'llm.system', 'llm.model_name', 'llm.token_count.prompt', 'llm.token_count.prompt_details.cache_read', 'llm.token_count.completion', 'llm.token_count.completion_details.reasoning', 'llm.cost.total', 'tool.name', 'tool.description', 'llm.tools', 'llm.input_messages', 'input.value', 'llm.output_messages', 'output.value', 'openinference.span.kind', 'eve.turn.id', 'maple_ai.turn.id') OR k LIKE 'gen_ai.prompt.variable.%'), SpanAttributes) AS spanAttributes FROM trace_detail_spans WHERE OrgId = 'org_sql_catalog' AND Timestamp >= '2026-01-01 10:30:00' @@ -170,8 +169,7 @@ SELECT StatusCode AS statusCode, StatusMessage AS statusMessage, toString(Timestamp) AS timestamp, - SpanAttributes AS spanAttributes, - ResourceAttributes AS resourceAttributes + mapFilter((k, v) -> (k IN ('maple_ai.session.id', 'maple_ai.vendor.id', 'maple_ai.vendor.version', 'gen_ai.operation.name', 'gen_ai.provider.name', 'gen_ai.system', 'gen_ai.request.model', 'gen_ai.request.max_tokens', 'gen_ai.request.choice.count', 'gen_ai.request.temperature', 'gen_ai.request.top_p', 'gen_ai.request.top_k', 'gen_ai.request.stop_sequences', 'gen_ai.request.frequency_penalty', 'gen_ai.request.presence_penalty', 'gen_ai.request.encoding_formats', 'gen_ai.request.seed', 'gen_ai.openai.request.seed', 'gen_ai.request.stream', 'gen_ai.request.reasoning.level', 'gen_ai.request.previous_response.id', 'gen_ai.request.stream_cursor', 'gen_ai.response.id', 'gen_ai.response.model', 'gen_ai.response.finish_reasons', 'gen_ai.response.finish_reason', 'gen_ai.response.status', 'gen_ai.response.time_to_first_chunk', 'gen_ai.output.type', 'gen_ai.usage.input_tokens', 'gen_ai.usage.prompt_tokens', 'gen_ai.usage.cache_read.input_tokens', 'gen_ai.usage.input_tokens.cached', 'gen_ai.usage.cache_creation.input_tokens', 'gen_ai.usage.cache_write.input_tokens', 'gen_ai.usage.output_tokens', 'gen_ai.usage.completion_tokens', 'gen_ai.usage.reasoning.output_tokens', 'gen_ai.usage.output_tokens.reasoning', 'gen_ai.usage.cost', 'gen_ai.usage.total_cost', 'gen_ai.conversation.id', 'gen_ai.conversation.compacted', 'gen_ai.agent.id', 'gen_ai.agent.name', 'gen_ai.agent.description', 'gen_ai.agent.version', 'gen_ai.tool.name', 'gen_ai.tool.call.id', 'gen_ai.tool.description', 'gen_ai.tool.type', 'gen_ai.tool.call.arguments', 'gen_ai.tool.call.result', 'gen_ai.tool.definitions', 'gen_ai.system_instructions', 'gen_ai.input.messages', 'gen_ai.prompt', 'gen_ai.output.messages', 'gen_ai.completion', 'gen_ai.data_source.id', 'gen_ai.retrieval.query.text', 'gen_ai.retrieval.top_k', 'gen_ai.retrieval.documents', 'gen_ai.memory.store.id', 'gen_ai.memory.record.id', 'gen_ai.memory.record.count', 'gen_ai.memory.query.text', 'gen_ai.memory.records', 'gen_ai.embeddings.dimension.count', 'gen_ai.evaluation.name', 'gen_ai.evaluation.score.value', 'gen_ai.evaluation.score.label', 'gen_ai.evaluation.explanation', 'gen_ai.prompt.name', 'gen_ai.prompt.version', 'gen_ai.workflow.name', 'error.type', 'server.address', 'server.port', 'ai.model.provider', 'ai.model.id', 'ai.response.id', 'ai.response.model', 'ai.response.finishReason', 'gen_ai.client.operation.time_to_first_chunk', 'ai.usage.inputTokens', 'ai.usage.promptTokens', 'ai.usage.cachedInputTokens', 'ai.usage.inputTokenDetails.cacheReadTokens', 'ai.usage.inputTokenDetails.cacheWriteTokens', 'ai.usage.outputTokens', 'ai.usage.completionTokens', 'ai.usage.reasoningTokens', 'ai.usage.outputTokenDetails.reasoningTokens', 'ai.telemetry.functionId', 'ai.toolCall.name', 'ai.toolCall.id', 'ai.toolCall.args', 'ai.toolCall.result', 'ai.prompt.tools', 'ai.prompt.messages', 'ai.prompt', 'llm.provider', 'llm.system', 'llm.model_name', 'llm.token_count.prompt', 'llm.token_count.prompt_details.cache_read', 'llm.token_count.completion', 'llm.token_count.completion_details.reasoning', 'llm.cost.total', 'tool.name', 'tool.description', 'llm.tools', 'llm.input_messages', 'input.value', 'llm.output_messages', 'output.value', 'openinference.span.kind', 'eve.turn.id', 'maple_ai.turn.id') OR k LIKE 'gen_ai.prompt.variable.%'), SpanAttributes) AS spanAttributes FROM trace_detail_spans WHERE OrgId = 'org_sql_catalog' AND Timestamp >= '2026-01-01 10:30:00' diff --git a/packages/query-engine-integrations/src/ai/ai-integrations.test.ts b/packages/query-engine-integrations/src/ai/ai-integrations.test.ts index 7180c69f3..c04872012 100644 --- a/packages/query-engine-integrations/src/ai/ai-integrations.test.ts +++ b/packages/query-engine-integrations/src/ai/ai-integrations.test.ts @@ -18,7 +18,6 @@ const row = ( statusMessage: "", timestamp: "2026-08-12 15:18:42.207000000", spanAttributes, - resourceAttributes: {}, ...overrides, }) @@ -351,18 +350,6 @@ describe("span envelope", () => { }) }) - it("reads gen_ai keys from span attributes alone", () => { - // A resource-level `gen_ai.*` key describes the process, not the - // operation: honouring it would stamp every span of that service — - // Postgres, HTTP, everything — as an AI span. - const mapped = mapAiSpan( - row({}, { resourceAttributes: { "gen_ai.request.model": "resource-level" } }), - ) - - expect(mapped.genAi.requestModel).toBeUndefined() - expect(mapped.isAiSpan).toBe(false) - }) - it("maps a whole trace's worth of spans in order", () => { const mapped = mapAiSpans([ row(INVOKE_AGENT_ATTRS), @@ -397,15 +384,6 @@ describe("resolveAiIntegration", () => { }) describe("untrusted attribute keys", () => { - it("ignores a vendor stamp that arrives via a resource attribute", () => { - // The envelope is read from span attributes alone, so a resource-level - // stamp neither selects an integration nor marks the span. - const mapped = mapAiSpan(row({}, { resourceAttributes: { "maple_ai.vendor.id": "eve" } })) - - expect(mapped.vendorId).toBeUndefined() - expect(mapped.isAiSpan).toBe(false) - }) - it("keeps a prompt variable literally named __proto__ as AI signal", () => { expect(mapAiSpan(row({ "gen_ai.prompt.variable.__proto__": "kept" })).isAiSpan).toBe(true) }) diff --git a/packages/query-engine-integrations/src/ai/ai-integrations.ts b/packages/query-engine-integrations/src/ai/ai-integrations.ts index c49c1f68d..cb75d6dda 100644 --- a/packages/query-engine-integrations/src/ai/ai-integrations.ts +++ b/packages/query-engine-integrations/src/ai/ai-integrations.ts @@ -50,6 +50,12 @@ export interface AiIntegration { * field from something other than a single attribute. */ readonly refine?: (values: MutableAiGenAiValues, ctx: AiRefineContext) => void + /** + * Attribute keys `refine` reads that no source list names. The span read + * projects only the keys the mapper is known to read (`aiSpanAttributeKeys`), + * so a key missing here is a key `refine` never sees. + */ + readonly refineKeys?: readonly string[] } /** An integration carrying a source list for every catalog field. */ @@ -239,6 +245,26 @@ const resolvedIntegrations = new Map( ]), ) +/** + * Every attribute key the mapper can read off a span, across every integration: + * the envelope, each field's source keys, and what the refine hooks read. The + * span read projects the attribute map down to these (plus the + * `AI_PROMPT_VARIABLE_PREFIX` family, which has no fixed key), so a key not in + * this list never reaches `mapAiSpan` — in production the map's bulk is + * `db.query.text` and friends, which the mapper never looked at. + */ +export const aiSpanAttributeKeys: readonly string[] = [ + ...new Set([ + MAPLE_AI_SESSION_ID_ATTR, + MAPLE_AI_VENDOR_ID_ATTR, + MAPLE_AI_VENDOR_VERSION_ATTR, + ...[genAiIntegration, ...resolvedIntegrations.values()].flatMap((integration) => + Object.values(integration.sources).flat(), + ), + ...Object.values(AI_VENDOR_INTEGRATIONS).flatMap((vendor) => vendor.refineKeys ?? []), + ]), +] + /** The integration for a vendor stamp, or the default for a stamp with no entry. */ export const resolveAiIntegration = (vendorId: string | undefined): ResolvedAiIntegration => { const resolved = vendorId === undefined ? undefined : resolvedIntegrations.get(vendorId) diff --git a/packages/query-engine-integrations/src/ai/ai-sessions.test.ts b/packages/query-engine-integrations/src/ai/ai-sessions.test.ts index e4f3b078e..9d614fd9a 100644 --- a/packages/query-engine-integrations/src/ai/ai-sessions.test.ts +++ b/packages/query-engine-integrations/src/ai/ai-sessions.test.ts @@ -304,12 +304,31 @@ describe("aiSessionSpansQuery", () => { expect(sql).toContain("TraceId IN (SELECT") expect(sql).toContain("FROM traces") expect(sql).toContain("Duration / 1000000 AS durationMs") - expect(sql).toContain("SpanAttributes AS spanAttributes") - expect(sql).toContain("ResourceAttributes AS resourceAttributes") + expect(sql).toContain("mapFilter((k, v) -> (k IN ('maple_ai.session.id', ") + expect(sql).toContain("OR k LIKE 'gen_ai.prompt.variable.%'), SpanAttributes) AS spanAttributes") + expect(sql).not.toContain("ResourceAttributes") expect(sql).toContain("ORDER BY timestamp ASC") expect(sql).toContain("LIMIT 2000") }) + it("projects every key the mapper reads, across vendors", () => { + const { sql } = compileUnsafe(aiSessionSpansQuery(), spanParams) + + for (const key of [ + "maple_ai.vendor.id", + "gen_ai.input.messages", + "gen_ai.usage.prompt_tokens", // legacy alias + "ai.usage.inputTokens", // vercel_ai_sdk + "llm.token_count.prompt", // openinference + "openinference.span.kind", // read by a refine hook, not a source list + "eve.turn.id", + "maple_ai.turn.id", + "error.type", + ]) { + expect(sql, key).toContain(`'${key}'`) + } + }) + it("repeats the org predicate on every level that reads a table", () => { const { sql } = compileUnsafe(aiSessionSpansQuery(), spanParams) @@ -366,7 +385,6 @@ describe("aiSessionSpansQuery", () => { "maple_ai.vendor.id": "eve", "maple_ai.session.id": "wrun_01M0CSAEW96BH2W9185XZPRPKH", }, - resourceAttributes: { "service.name": "maple-slack-agent" }, }, ]) @@ -375,7 +393,6 @@ describe("aiSessionSpansQuery", () => { "maple_ai.vendor.id": "eve", "maple_ai.session.id": "wrun_01M0CSAEW96BH2W9185XZPRPKH", }) - expect(row?.resourceAttributes).toEqual({ "service.name": "maple-slack-agent" }) }) }) @@ -516,7 +533,8 @@ describe("aiTraceSpansQuery", () => { expect(sql).toContain(`TraceId = '${TRACE_ID}'`) expect(sql).not.toContain("TraceId IN (SELECT") expect(sql).not.toContain("FROM traces") - expect(sql).not.toContain("maple_ai.session.id") + // The projection still names the key; only the predicate is gone. + expect(sql).not.toContain("SpanAttributes['maple_ai.session.id']") }) it("keeps the projection and the order of the session form", () => { @@ -524,8 +542,8 @@ describe("aiTraceSpansQuery", () => { // One shape whichever kind of session the detail page opened. expect(sql).toContain("Duration / 1000000 AS durationMs") - expect(sql).toContain("SpanAttributes AS spanAttributes") - expect(sql).toContain("ResourceAttributes AS resourceAttributes") + expect(sql).toContain("SpanAttributes) AS spanAttributes") + expect(sql).not.toContain("ResourceAttributes") expect(sql).toContain("ORDER BY timestamp ASC, spanId ASC") expect(sql).toContain("LIMIT 2000") expect(compileUnsafe(aiTraceSpansQuery({ limit: 100 }), traceParams).sql).toContain("LIMIT 100") @@ -572,7 +590,6 @@ describe("aiTraceSpansQuery", () => { timestamp: "2026-08-19 10:33:25.825000000", // A sessionless vendor: the stamp is there, the session key is not. spanAttributes: { "maple_ai.vendor.id": "llamaindex" }, - resourceAttributes: { "service.name": "rag-service" }, }, ]) diff --git a/packages/query-engine-integrations/src/ai/ai-sessions.ts b/packages/query-engine-integrations/src/ai/ai-sessions.ts index 5c4accb0d..01fc28d4d 100644 --- a/packages/query-engine-integrations/src/ai/ai-sessions.ts +++ b/packages/query-engine-integrations/src/ai/ai-sessions.ts @@ -102,11 +102,13 @@ import { AiTraceIndex, TraceDetailSpans, Traces } from "@maple/query-engine/ch/t import { CHNumber } from "@maple/query-engine/ch/schema" import { AI_SESSION_SPANS_MAX_SPANS } from "@maple/domain/http" import { + AI_PROMPT_VARIABLE_PREFIX, MAPLE_AI_SESSION_ID_ATTR, MAPLE_AI_TRACE_SESSION_PREFIX, MAPLE_AI_VENDOR_ID_ATTR, MAPLE_AI_VENDOR_VERSION_ATTR, } from "@maple/domain/gen-ai" +import { aiSpanAttributeKeys } from "./ai-integrations" const SESSION_ID_ATTR = MAPLE_AI_SESSION_ID_ATTR const VENDOR_ID_ATTR = MAPLE_AI_VENDOR_ID_ATTR @@ -524,7 +526,6 @@ export interface AiSessionSpansOutput { readonly statusMessage: string readonly timestamp: string readonly spanAttributes: Record - readonly resourceAttributes: Record } export const aiSessionSpansRowSchema: CompiledQueryRowSchema = Schema.Struct({ @@ -542,7 +543,6 @@ export const aiSessionSpansRowSchema: CompiledQueryRowSchema) => ( statusCode: $.StatusCode, statusMessage: $.StatusMessage, timestamp: CH.toString_($.Timestamp), - spanAttributes: $.SpanAttributes, - resourceAttributes: $.ResourceAttributes, + // The map cut down to what `mapAiSpan` reads. Measured on production's + // largest sessions, the whole map is dominated by keys the mapper never + // touches (`db.query.text` alone was half of one session's bytes), and + // `ResourceAttributes` — which the mapper deliberately ignores, see + // `mapAiSpan` — was another 60% on top. Neither is read any more. + spanAttributes: CH.mapFilterKeys($.SpanAttributes, (key) => + key.in_(...aiSpanAttributeKeys).or(key.like(`${AI_PROMPT_VARIABLE_PREFIX}%`)), + ), }) /** @@ -568,11 +574,10 @@ const spanProjection = ($: ColumnAccessor) => ( * `sessionId` is a compile param rather than an opts field, so one compiled SQL * string serves every session. * - * Both attribute Maps come back whole: the integration layer that normalizes - * these into gen_ai form needs keys this query cannot know in advance. Projecting - * only the keys it wants is a later optimisation, and a real one — one production - * trace already carries 250 spans with up to ~17KB of attributes each, so callers - * should expect megabyte-scale payloads at the default limit. + * The attribute map is projected down to the keys the integration layer reads + * (`aiSpanAttributeKeys`); everything else on the span stays in the warehouse. + * Even so, a content-heavy vendor puts whole prompts in `gen_ai.input.messages`, + * so callers should still expect megabyte-scale payloads at the default limit. * * No scope columns: `trace_detail_spans` does not carry `ScopeName`/`ScopeVersion`, * and the read path does not need them. The ingest gateway already did the diff --git a/packages/query-engine-integrations/src/ai/ai-vendors.test.ts b/packages/query-engine-integrations/src/ai/ai-vendors.test.ts index 83be4f4fd..c7f3df5a5 100644 --- a/packages/query-engine-integrations/src/ai/ai-vendors.test.ts +++ b/packages/query-engine-integrations/src/ai/ai-vendors.test.ts @@ -15,7 +15,6 @@ const row = (vendorId: string, spanAttributes: Record): AiSessio statusMessage: "", timestamp: "2026-08-12 15:19:41.626000000", spanAttributes: { ...spanAttributes, "maple_ai.vendor.id": vendorId, "maple_ai.vendor.version": "0" }, - resourceAttributes: {}, }) describe("vercel_ai_sdk", () => { diff --git a/packages/query-engine-integrations/src/ai/ai-vendors.ts b/packages/query-engine-integrations/src/ai/ai-vendors.ts index 96f084a9a..1539410d0 100644 --- a/packages/query-engine-integrations/src/ai/ai-vendors.ts +++ b/packages/query-engine-integrations/src/ai/ai-vendors.ts @@ -115,6 +115,7 @@ const openInferenceIntegration: AiIntegration = { ) if (operation !== undefined) values.operationName = operation }, + refineKeys: ["openinference.span.kind"], } /** @@ -130,6 +131,7 @@ const eveIntegration: AiIntegration = { const turnId = ctx.attributes["eve.turn.id"] if (turnId !== undefined && turnId !== "") values.conversationId = turnId }, + refineKeys: ["eve.turn.id"], } /** @@ -147,6 +149,7 @@ const mapleIntegration: AiIntegration = { const turnId = ctx.attributes[MAPLE_NATIVE_TURN_ID_ATTR] if (turnId !== undefined && turnId !== "") values.conversationId = turnId }, + refineKeys: [MAPLE_NATIVE_TURN_ID_ATTR], } /** diff --git a/packages/query-engine-integrations/src/ai/index.ts b/packages/query-engine-integrations/src/ai/index.ts index 9db68d007..07936b635 100644 --- a/packages/query-engine-integrations/src/ai/index.ts +++ b/packages/query-engine-integrations/src/ai/index.ts @@ -24,6 +24,7 @@ export { } from "./ai-sessions" export { + aiSpanAttributeKeys, genAiIntegration, mapAiSpan, mapAiSpans,