diff --git a/docs/sdk-conformance.md b/docs/sdk-conformance.md index 853d4ec..4eb6eeb 100644 --- a/docs/sdk-conformance.md +++ b/docs/sdk-conformance.md @@ -41,9 +41,10 @@ repository, not this SDK workspace. Paths below refer to that product repository | error | span status ERROR + `error.type` + `exception` event (`exception.type/message/stacktrace`, `log.severity_number: 17`) | | | `log()` | OTLP log record to `/v1/logs` (always; never span events) | severity debug=5/info=9/warn=13/error=17; trace-correlated via active context | | `withPrompt(prompt, fn)` / `with with_prompt(prompt):` → `prompt.name` | `td.prompt.name` | on model-call spans that start in scope: operation `chat`, `text_completion`, or `generate_content` | +| outside a scope: a message object that `compile()` of a chat prompt returned, in the request messages of a provider integration | `td.prompt.name`, `td.prompt.version`, `td.prompt.label` | on the `chat` spans of the provider integrations listed in [Prompts](#prompts); the scope wins | | `prompt.version` | `td.prompt.version` (int) | omitted for a fallback prompt (`version` null / None) | | `prompt.label` | `td.prompt.label` | only for a fetch by label; a fetch without label or version is the `production` label | -| variables of the last `prompt.compile()` before `withPrompt` / `with_prompt` starts, or in its scope | `td.prompt.variables` (JSON object) | mask → JSON stringify → truncate; gated by the effective captureInput of the span | +| variables of the last `prompt.compile()` before `withPrompt` / `with_prompt` starts, or in its scope; outside a scope, of the newest `compile()` that made a message of the request | `td.prompt.variables` (JSON object) | mask → JSON stringify → truncate; gated by the effective captureInput of the span | | `runExperiment` / `run_experiment` → each span and log record of an item run | `td.experiment.id` | the id that `POST /v1/experiments` returns; set like a propagated attribute, so the provider integration spans get it | | the item of the run | `td.dataset.item_id` | the `id` of the dataset item: its key, the same across dataset versions | | the item run | `gen_ai.conversation.id = "::"` | one session and one trace for each run; replaces an ambient or process session, and the run drops the active span | @@ -187,6 +188,71 @@ cannot tell a concurrent task from an awaited helper, so in TypeScript all spans its last `compile()`; concurrent tasks that compile their own variables call `withPrompt` in each task. +Outside a scope, the SDKs find the prompt from the message objects of the request, not from +their text: + +- `compile()` of a chat prompt tags each message object that it makes with the prompt, a copy of + the variables of that call and of each list in them (not of the objects in a list), and a serial + number that grows with each chat compile in the process. With input capture off in the client + at that time, the tag holds no variables. It does not tag the messages of a placeholder + variable: the app passes those, and they keep the tag of their own compile. TypeScript: a + non-enumerable symbol property; the keys and the JSON of a message do not change. Python: the + messages are instances of a private `dict` subclass that holds the tag in a slot. + `isinstance(message, dict)` is true and the keys and the JSON do not change, but a check of the + exact type (`type(message) is dict`), `yaml.safe_dump`, and `marshal` reject it. Its + constructor is the one of `dict`. +- A copy has no tag: a spread, `structuredClone`, or a JSON round trip (TypeScript); `dict()`, + `.copy()`, `copy.copy`, `copy.deepcopy`, pickle, a JSON round trip, `type(message)(message)`, or + `dataclasses.asdict` (Python). A new message with the same text has no tag either, thus a + call never records the prompt of a message from another call, task, or prompt version. +- `promptAttributesFor(messages)` / `prompt_attributes_for(messages)` take a list of messages or + one message (Python: a `list` or `tuple`, not a subclass, so the read runs no code of the app). + They look at the first 1000 and the last 1000 items and give the name, version, label, and + variables of the tag with the highest serial number: the newest compile. Thus a message that a + placeholder brings from an older compile does not win over the compile that put it in the list. + They do not look inside a message. +- In a `withPrompt` / `with_prompt` scope they give nothing. A model-call span whose parent context + is in a scope gets the attributes of the scope and drops each `td.prompt.*` attribute that the + app passes, at start and in each update (`update`, `end`, `updateActiveSpan` / + `update_current_span`), also when the app calls these functions outside the scope and passes an + explicit parent. They never throw: an error goes to `onError` / `on_error`. +- They give `td.prompt.variables` only when input capture is on in the client, at the compile and + at the call. Then they serialize the variables, which calls the `mask` of the client. A span + whose effective captureInput is off drops `td.prompt.variables` from its attributes, also when + the app passes the result of these functions to its own span. +- A tag keeps the variables of its compile for 5 minutes. A call of these functions keeps them + for 5 more minutes when the newest compile of the call is that compile. The SDKs measure this + time in elapsed time, so a change of the system clock does not move it. After that time, + they give the attributes of the tag without `td.prompt.variables`, and the next compile or call + removes the variables from the tag, at most one minute late, also when the app keeps the + message. Other compiles do not change this time. Thus a chat history keeps the variables of the + compiles of the last 5 minutes, for a history of any length or window. A guard prompt, a queue, + or a compile in other code does not remove the variables of a request before a call in these 5 + minutes. A message that more than one request sends keeps its variables while a call names its + compile at least one time in 5 minutes. A request with only messages of earlier compiles gives + the variables of the newest of them, when they are less than 5 minutes old. +- The SDKs keep a weak reference to each tag with variables, by the minute of its time, thus the + reference does not keep a tag in memory. In a runtime without `WeakRef` (Cloudflare Workers with a + compatibility date before 2025-05-05, or with `disable_weak_ref`), the TypeScript tag keeps no + variables. The SDKs keep no other reference to the variables, except those of the last compile + of each prompt object and of each running scope, for `withPrompt` / `with_prompt`. +- `compile()` of a text prompt returns a string, which cannot carry a tag. Outside a scope, the + SDKs do not record a text prompt: use `withPrompt` / `with_prompt`. +- The provider integrations call them with the request messages before the span starts, also + when input capture is off: OpenAI chat completions (`messages`) and responses (`input`), + Anthropic messages (`messages`; the `system` string has no tag), Amazon Bedrock Converse + (`messages`), Google GenAI `generateContent` (`contents`), LiteLLM completion (`messages`) and + responses (`input`) (Python), and OpenRouter chat (`messages`) and responses (`input`). Bedrock + Converse and Google GenAI take another message shape: a compiled message that the app changes + in place keeps its tag, and a new message does not. Bedrock InvokeModel takes a serialized body + and is not recorded outside a scope. The other spans of the integrations (embeddings, images, + audio, video, and OpenAI Realtime) do not look at the request. +- Python Anthropic `messages.stream()` starts its span at `__enter__`. It decides there, with the + messages that `.stream()` received, thus a scope around `__enter__` wins. +- A provider integration reads the function from the core module when it runs, thus it also + operates with an older core that does not have it. With such a core, the prompt is recorded + only in a scope. + ## Experiments and scores `runExperiment(options)` / `run_experiment(...)` (and `run_experiment_async`) and `score(options)` / diff --git a/packages/anthropic/src/index.ts b/packages/anthropic/src/index.ts index d9cb1f6..1082844 100644 --- a/packages/anthropic/src/index.ts +++ b/packages/anthropic/src/index.ts @@ -5,6 +5,7 @@ import { type SpanHandle, type StartSpanOptions, } from "@telemetry-dev/sdk"; +import * as telemetrySdk from "@telemetry-dev/sdk"; import { Stream } from "@anthropic-ai/sdk/core/streaming"; import { Messages as BetaMessages } from "@anthropic-ai/sdk/resources/beta/messages"; import { Messages } from "@anthropic-ai/sdk/resources/messages"; @@ -147,6 +148,16 @@ function toolDefinitions(tools: T): Record | undefined { : undefined; } +// Cores older than promptAttributesFor record a prompt only inside withPrompt. A namespace read +// keeps them loadable without the export. +function promptAttributes(messages: unknown): SpanFields["attributes"] { + const read = ( + telemetrySdk as { promptAttributesFor?: (messages: unknown) => SpanFields["attributes"] } + ).promptAttributesFor; + + return typeof read === "function" ? read(messages) : undefined; +} + function messagesRequest(body: ValueRecord): RequestMapping { const model = readString(body.model); @@ -161,7 +172,10 @@ function messagesRequest(body: ValueRecord): RequestMapping { type: "generation", model, input, - attributes: toolDefinitions(body.tools), + attributes: { + ...toolDefinitions(body.tools), + ...promptAttributes(body.messages), + }, systemInstructions: body.system, temperature: readNumber(body.temperature), topP: readNumber(body.top_p), diff --git a/packages/anthropic/tests/anthropic.test.ts b/packages/anthropic/tests/anthropic.test.ts index e53d040..6516205 100644 --- a/packages/anthropic/tests/anthropic.test.ts +++ b/packages/anthropic/tests/anthropic.test.ts @@ -4,7 +4,14 @@ import { type PushMetricExporter, } from "@opentelemetry/sdk-metrics"; import { InMemorySpanExporter, type ReadableSpan } from "@opentelemetry/sdk-trace-base"; -import { flush, init, shutdown, type TelemetryOptions } from "@telemetry-dev/sdk"; +import { + type ChatMessage, + flush, + getPrompt, + init, + shutdown, + type TelemetryOptions, +} from "@telemetry-dev/sdk"; import * as telemetrySdk from "@telemetry-dev/sdk"; import Anthropic from "@anthropic-ai/sdk"; import { createRequire } from "node:module"; @@ -274,6 +281,72 @@ afterEach(async () => { vi.restoreAllMocks(); }); +test.each([ + { messages: "compiled", olderCore: false }, + { messages: "copied", olderCore: false }, + { messages: "compiled", olderCore: true }, +])( + "messages.create outside withPrompt records the prompt of $messages messages, older core $olderCore", + async ({ messages: kind, olderCore }) => { + if (olderCore) { + vi.spyOn( + telemetrySdk as { promptAttributesFor: unknown }, + "promptAttributesFor", + "get", + ).mockReturnValue(undefined); + } + + const prompt = await getPrompt("support/reply", { + apiKey: "td_live_a", + baseUrl: "https://prompts.anthropic.test", + fetch: async () => + Response.json({ + name: "support/reply", + type: "chat", + label: "production", + version: 12, + messages: [ + { role: "system", content: "You answer for the {{team}} desk." }, + { role: "user", content: "Where is order {{order}}?" }, + ], + text: null, + model: { provider: "anthropic", name: "claude-sonnet-4-6" }, + parameters: {}, + tools: [], + responseFormat: { type: "text" }, + rollout: null, + }), + }); + const [system, user] = prompt.compile({ team: "billing", order: "#4" }) as ChatMessage[]; + const message = user as { role: "user"; content: string }; + const spans = setupSpans(); + const client = clientWith(createFakeFetch(jsonResponse(messagePayload())).fetch); + + await client.messages.create({ + model: "claude-sonnet-4-6", + max_tokens: 64, + system: system!.content, + messages: [kind === "compiled" ? message : { ...message }], + }); + + const { attributes } = await exportedSpan(spans); + expect(attributes["gen_ai.input.messages"]).toBeDefined(); + + if (olderCore || kind === "copied") { + expect(attributes["td.prompt.name"]).toBeUndefined(); + + return; + } + + expect(attributes).toMatchObject({ + "td.prompt.name": "support/reply", + "td.prompt.version": 12, + "td.prompt.label": "production", + "td.prompt.variables": '{"team":"billing","order":"#4"}', + }); + }, +); + test("messages.create maps request, response, usage, finish reason, provider, and sampling attributes", async () => { const spans = setupSpans(); const messages = [{ role: "user" as const, content: "Say hello" }]; diff --git a/packages/bedrock/src/converse.ts b/packages/bedrock/src/converse.ts index 074d1b1..fd9c8e1 100644 --- a/packages/bedrock/src/converse.ts +++ b/packages/bedrock/src/converse.ts @@ -9,6 +9,7 @@ import { numberValue, omitUndefined, parseJson, + promptAttributes, PROVIDER, stringValue, StreamRetention, @@ -34,6 +35,7 @@ export function converseRequestFields(input: JsonRecord): SpanFields { model: stringValue(input.modelId), input: normalizeMessages(input.messages), systemInstructions: normalizeSystem(input.system), + attributes: promptAttributes(input.messages), temperature: numberValue(inference.temperature), topP: numberValue(inference.topP), maxTokens: numberValue(inference.maxTokens), diff --git a/packages/bedrock/src/internal.ts b/packages/bedrock/src/internal.ts index 104f595..2a89e01 100644 --- a/packages/bedrock/src/internal.ts +++ b/packages/bedrock/src/internal.ts @@ -192,6 +192,16 @@ export function captureMaskedBeforeStart(): boolean { } } +// Cores older than promptAttributesFor record a prompt only inside withPrompt. A namespace read +// keeps them loadable without the export. +export function promptAttributes(messages: unknown): SpanFields["attributes"] { + const read = ( + telemetrySdk as { promptAttributesFor?: (messages: unknown) => SpanFields["attributes"] } + ).promptAttributesFor; + + return typeof read === "function" ? read(messages) : undefined; +} + // The decision is made before the span starts so samplers see the same attributes as before. function startRequestSpan(name: string, options: StartSpanOptions): SpanHandle { const withheld = options.input !== undefined && isCut(options) && captureMaskedBeforeStart(); diff --git a/packages/bedrock/tests/bedrock.test.ts b/packages/bedrock/tests/bedrock.test.ts index 8e94b8f..0eb0e7f 100644 --- a/packages/bedrock/tests/bedrock.test.ts +++ b/packages/bedrock/tests/bedrock.test.ts @@ -9,7 +9,8 @@ import { } from "@aws-sdk/client-bedrock-runtime"; import { afterEach, expect, test, vi } from "vitest"; -import type { MaskFn } from "@telemetry-dev/sdk"; +import { type ChatMessage, getPrompt, type MaskFn } from "@telemetry-dev/sdk"; +import * as telemetrySdk from "@telemetry-dev/sdk"; import { instrumentBedrock, uninstrumentBedrock, wrapBedrock } from "../src/index.ts"; import { @@ -423,6 +424,90 @@ test("ConverseStream records stream errors with partial output", async () => { expect(span.attributes["telemetry.dev.capture.truncated"]).toBe(true); }); +test.each([ + { messages: "changed in place", olderCore: false }, + { messages: "new", olderCore: false }, + { messages: "changed in place", olderCore: true }, +])( + "Converse outside withPrompt records the prompt of $messages messages, older core $olderCore", + async ({ messages: kind, olderCore }) => { + const spans = setup(); + + if (olderCore) { + vi.spyOn( + telemetrySdk as { promptAttributesFor: unknown }, + "promptAttributesFor", + "get", + ).mockReturnValue(undefined); + } + + const prompt = await getPrompt("support/reply", { + apiKey: "td_live_a", + baseUrl: "https://prompts.bedrock.test", + fetch: async () => + Response.json({ + name: "support/reply", + type: "chat", + label: "production", + version: 12, + messages: [ + { role: "system", content: "You answer for the {{team}} desk." }, + { role: "user", content: "Where is order {{order}}?" }, + ], + text: null, + model: { provider: "amazon-bedrock", name: "anthropic.claude-3-haiku" }, + parameters: {}, + tools: [], + responseFormat: { type: "text" }, + rollout: null, + }), + }); + const [system, user] = prompt.compile({ team: "billing", order: "#4" }) as ChatMessage[]; + const blocks = [{ text: user!.content }]; + // Converse takes content blocks: a message changed in place keeps its tag, a new one has none. + const message = + kind === "new" + ? { role: "user", content: blocks } + : Object.assign(user!, { content: blocks }); + + const client = wrapBedrock( + new BedrockRuntimeClient({ + region: "us-east-1", + credentials: { accessKeyId: "test", secretAccessKey: "test" }, + maxAttempts: 1, + requestHandler: jsonRequestHandler("prompt-1", { + output: { message: { role: "assistant", content: [{ text: "Soon." }] } }, + stopReason: "end_turn", + usage: { inputTokens: 1, outputTokens: 1, totalTokens: 2 }, + }), + }), + ); + + await client.send( + new ConverseCommand({ + modelId: "anthropic.claude-3-haiku", + system: [{ text: system!.content }], + messages: [message] as ConverseCommandInput["messages"], + }), + ); + + const [span] = spans.getFinishedSpans(); + expect(span!.attributes["gen_ai.response.id"]).toBe("prompt-1"); + expect( + Object.entries(span!.attributes).filter(([key]) => key.startsWith("td.prompt.")), + ).toEqual( + olderCore || kind === "new" + ? [] + : [ + ["td.prompt.name", "support/reply"], + ["td.prompt.version", 12], + ["td.prompt.label", "production"], + ["td.prompt.variables", '{"team":"billing","order":"#4"}'], + ], + ); + }, +); + test("InvokeModel captures provider-native bodies and embedding spans", async () => { const spans = setup(); diff --git a/packages/google-genai/src/index.ts b/packages/google-genai/src/index.ts index 96ffef5..341d5e8 100644 --- a/packages/google-genai/src/index.ts +++ b/packages/google-genai/src/index.ts @@ -567,7 +567,11 @@ function responseAttributes(response: UnknownRecord, metadataContent: boolean) { function generateRequestFields(params: UnknownRecord): RequestMapping { const model = readString(params.model); const config = asRecord(params.config); - const attributes = requestAttributesFromConfig(config); + + const attributes = { + ...requestAttributesFromConfig(config), + ...promptAttributes(params.contents), + }; const input = captureEnabled("input") ? normalizeContentInput(params.contents) : undefined; const systemInstructions = captureEnabled("input") @@ -940,6 +944,16 @@ function captureMaskedBeforeStart(): boolean { } } +// Cores older than promptAttributesFor record a prompt only inside withPrompt. A namespace read +// keeps them loadable without the export. +function promptAttributes(messages: unknown): SpanFields["attributes"] { + const read = ( + telemetrySdk as { promptAttributesFor?: (messages: unknown) => SpanFields["attributes"] } + ).promptAttributesFor; + + return typeof read === "function" ? read(messages) : undefined; +} + // The decision is made before the span starts so samplers see the same attributes as before. function startRequestSpan(request: RequestMapping, provider: string): SpanHandle { const fields = { ...request.fields, provider }; diff --git a/packages/google-genai/tests/google-genai.test.ts b/packages/google-genai/tests/google-genai.test.ts index d1d291a..19d7fbf 100644 --- a/packages/google-genai/tests/google-genai.test.ts +++ b/packages/google-genai/tests/google-genai.test.ts @@ -10,7 +10,8 @@ import { type PushMetricExporter, type ResourceMetrics, } from "@opentelemetry/sdk-metrics"; -import { flush, init, shutdown } from "@telemetry-dev/sdk"; +import { type ChatMessage, flush, getPrompt, init, shutdown } from "@telemetry-dev/sdk"; +import * as telemetrySdk from "@telemetry-dev/sdk"; import { FunctionCallingConfigMode, GoogleGenAI, @@ -234,6 +235,86 @@ afterEach(async () => { vi.restoreAllMocks(); }); +test.each([ + { contents: "changed in place", captureInput: true, olderCore: false }, + { contents: "changed in place", captureInput: false, olderCore: false }, + { contents: "new", captureInput: true, olderCore: false }, + { contents: "changed in place", captureInput: true, olderCore: true }, +])( + "generateContent outside withPrompt records the prompt of $contents contents, input capture $captureInput, older core $olderCore", + async ({ contents: kind, captureInput, olderCore }) => { + if (olderCore) { + vi.spyOn( + telemetrySdk as { promptAttributesFor: unknown }, + "promptAttributesFor", + "get", + ).mockReturnValue(undefined); + } + + const prompt = await getPrompt("support/reply", { + apiKey: "td_live_a", + baseUrl: "https://prompts.google-genai.test", + fetch: async () => + Response.json({ + name: "support/reply", + type: "chat", + label: "production", + version: 12, + messages: [ + { role: "system", content: "You answer for the {{team}} desk." }, + { role: "user", content: "Where is order {{order}}?" }, + ], + text: null, + model: { provider: "gcp.gemini", name: "gemini-2.5-flash" }, + parameters: {}, + tools: [], + responseFormat: { type: "text" }, + rollout: null, + }), + }); + const [system, user] = prompt.compile({ team: "billing", order: "#4" }) as ChatMessage[]; + const parts = [{ text: user!.content }]; + // Gemini takes parts: a message changed in place keeps its tag, a new one has none. + const content = kind === "new" ? { role: "user", parts } : Object.assign(user!, { parts }); + const spans = setupSpans({ captureInput }); + const fake = createFakeFetch( + jsonResponse({ + candidates: [ + { content: { role: "model", parts: [{ text: "Soon." }] }, finishReason: "STOP" }, + ], + }), + ); + vi.stubGlobal("fetch", fake.fetch); + + await clientWith().models.generateContent({ + model: "gemini-2.5-flash", + contents: [content], + config: { temperature: 0.2, systemInstruction: system!.content }, + }); + + const { attributes } = await exportedSpan(spans); + expect(attributes["gen_ai.request.temperature"]).toBe(0.2); + expect(fake.requests[0]!.body).toMatchObject({ + contents: [{ role: "user", parts: [{ text: "Where is order #4?" }] }], + }); + + if (olderCore || kind === "new") { + expect(attributes["td.prompt.name"]).toBeUndefined(); + + return; + } + + expect(attributes).toMatchObject({ + "td.prompt.name": "support/reply", + "td.prompt.version": 12, + "td.prompt.label": "production", + }); + expect(attributes["td.prompt.variables"]).toBe( + captureInput ? '{"team":"billing","order":"#4"}' : undefined, + ); + }, +); + test("generateContent maps request, response, usage, finish reason, provider, and sampling attributes", async () => { const spans = setupSpans(); diff --git a/packages/openai/src/index.ts b/packages/openai/src/index.ts index 3512476..9b86e73 100644 --- a/packages/openai/src/index.ts +++ b/packages/openai/src/index.ts @@ -156,6 +156,7 @@ function stopSequences(value: T): string[] | undefined { function chatRequest(body: JsonRecord): RequestMapping { const model = readString(body.model); const capture = captureEnabled("input") ? sanitizeResponsesCapture(body.messages) : undefined; + const prompt = promptAttributes(body.messages); return { name: `chat ${model ?? "unknown"}`, @@ -170,7 +171,9 @@ function chatRequest(body: JsonRecord): RequestMapping { seed: readNumber(body.seed), frequencyPenalty: readNumber(body.frequency_penalty), presencePenalty: readNumber(body.presence_penalty), - attributes: capture?.truncated ? { "telemetry.dev.capture.truncated": true } : undefined, + attributes: capture?.truncated + ? { ...prompt, "telemetry.dev.capture.truncated": true } + : prompt, }, }; } @@ -243,6 +246,7 @@ function chatResponse(response: T): SpanFields { function responsesRequest(body: JsonRecord): RequestMapping { const model = readString(body.model); const capture = captureEnabled("input") ? sanitizeResponsesCapture(body.input) : undefined; + const prompt = promptAttributes(body.input); return { name: `chat ${model ?? "unknown"}`, @@ -254,7 +258,9 @@ function responsesRequest(body: JsonRecord): RequestMapping { temperature: readNumber(body.temperature), topP: readNumber(body.top_p), maxTokens: readNumber(body.max_output_tokens), - attributes: capture?.truncated ? { "telemetry.dev.capture.truncated": true } : undefined, + attributes: capture?.truncated + ? { ...prompt, "telemetry.dev.capture.truncated": true } + : prompt, }, }; } @@ -1626,6 +1632,16 @@ function captureMaskedBeforeStart(): boolean { } } +// Cores older than promptAttributesFor record a prompt only inside withPrompt. A namespace read +// keeps them loadable without the export. +function promptAttributes(messages: unknown): SpanFields["attributes"] { + const read = ( + telemetrySdk as { promptAttributesFor?: (messages: unknown) => SpanFields["attributes"] } + ).promptAttributesFor; + + return typeof read === "function" ? read(messages) : undefined; +} + // The decision is made before the span starts so samplers see the same attributes as before. function startRequestSpan(name: string, fields: StartSpanOptions): SpanHandle { const cut = diff --git a/packages/openai/tests/openai.test.ts b/packages/openai/tests/openai.test.ts index 4467aa3..099ec37 100644 --- a/packages/openai/tests/openai.test.ts +++ b/packages/openai/tests/openai.test.ts @@ -8,7 +8,14 @@ import { InMemoryMetricExporter, type PushMetricExporter, } from "@opentelemetry/sdk-metrics"; -import { flush, init, shutdown, type TelemetryOptions } from "@telemetry-dev/sdk"; +import { + type ChatMessage, + flush, + getPrompt, + init, + shutdown, + type TelemetryOptions, +} from "@telemetry-dev/sdk"; import * as telemetrySdk from "@telemetry-dev/sdk"; import { EventEmitter } from "node:events"; import { createRequire } from "node:module"; @@ -189,7 +196,7 @@ function setupSpans( metricExporter?: PushMetricExporter, options: Pick< TelemetryOptions, - "captureOutput" | "mask" | "maxAttributeLength" | "onError" | "sampler" + "captureInput" | "captureOutput" | "mask" | "maxAttributeLength" | "onError" | "sampler" > = {}, ): InMemorySpanExporter { const spanExporter = new InMemorySpanExporter(); @@ -2887,6 +2894,102 @@ test("chat streams preserve caller-requested usage chunks", async () => { expect(span.attributes["gen_ai.usage.total_tokens"]).toBe(7); }); +test.each([ + { api: "chat completions", messages: "compiled", captureInput: false, olderCore: false }, + { api: "responses", messages: "compiled", captureInput: true, olderCore: false }, + { api: "chat completions", messages: "copied", captureInput: true, olderCore: false }, + { api: "chat completions", messages: "compiled", captureInput: true, olderCore: true }, +])( + "$api outside withPrompt records the prompt of $messages messages, input capture $captureInput, older core $olderCore", + async ({ api, messages: kind, captureInput, olderCore }) => { + if (olderCore) { + vi.spyOn( + telemetrySdk as { promptAttributesFor: unknown }, + "promptAttributesFor", + "get", + ).mockReturnValue(undefined); + } + + const prompt = await getPrompt("support/reply", { + apiKey: "td_live_a", + baseUrl: "https://prompts.openai.test", + fetch: async () => + Response.json({ + name: "support/reply", + type: "chat", + label: "production", + version: 12, + messages: [ + { role: "system", content: "You answer for the {{team}} desk." }, + { role: "user", content: "Where is order {{order}}?" }, + ], + text: null, + model: { provider: "openai", name: "gpt-4o" }, + parameters: {}, + tools: [], + responseFormat: { type: "text" }, + rollout: null, + }), + }); + const compiled = prompt.compile({ team: "billing", order: "#4" }) as ChatMessage[]; + const messages = kind === "compiled" ? compiled : compiled.map((message) => ({ ...message })); + const spans = setupSpans(undefined, { captureInput }); + const client = clientWith( + createFakeFetch( + jsonResponse( + api === "responses" + ? { + id: "resp_1", + object: "response", + status: "completed", + model: "gpt-4.1", + output: [], + usage: { input_tokens: 1, output_tokens: 1, total_tokens: 2 }, + } + : { + id: "chatcmpl_1", + object: "chat.completion", + created: 1, + model: "gpt-4o", + choices: [ + { + index: 0, + message: { role: "assistant", content: "Soon." }, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }, + ), + ).fetch, + ); + + if (api === "responses") { + await client.responses.create({ model: "gpt-4.1", input: messages }); + } else { + await client.chat.completions.create({ model: "gpt-4o", messages }); + } + + const { attributes } = await exportedSpan(spans); + expect(attributes["gen_ai.input.messages"] !== undefined).toBe(captureInput); + + if (olderCore || kind === "copied") { + expect(attributes["td.prompt.name"]).toBeUndefined(); + + return; + } + + expect(attributes).toMatchObject({ + "td.prompt.name": "support/reply", + "td.prompt.version": 12, + "td.prompt.label": "production", + }); + expect(attributes["td.prompt.variables"]).toBe( + captureInput ? '{"team":"billing","order":"#4"}' : undefined, + ); + }, +); + test("responses create maps instructions to system instructions", async () => { const spans = setupSpans(); diff --git a/packages/openrouter/src/index.ts b/packages/openrouter/src/index.ts index bfb0076..74e2c76 100644 --- a/packages/openrouter/src/index.ts +++ b/packages/openrouter/src/index.ts @@ -4,6 +4,7 @@ import { type SpanHandle, type StartSpanOptions, } from "@telemetry-dev/sdk"; +import * as telemetrySdk from "@telemetry-dev/sdk"; import { activeContext } from "@telemetry-dev/otel"; import { Chat } from "@openrouter/sdk/sdk/chat.js"; import { Embeddings } from "@openrouter/sdk/sdk/embeddings.js"; @@ -384,6 +385,16 @@ function usageCost(usage: T): number | undefined { return cost ?? upstreamCost; } +// Cores older than promptAttributesFor record a prompt only inside withPrompt. A namespace read +// keeps them loadable without the export. +function promptAttributes(messages: unknown): SpanFields["attributes"] { + const read = ( + telemetrySdk as { promptAttributesFor?: (messages: unknown) => SpanFields["attributes"] } + ).promptAttributesFor; + + return typeof read === "function" ? read(messages) : undefined; +} + function chatRequest(request: JsonRecord): RequestMapping { const body = asRecord(request.chatRequest) ?? {}; const model = readString(body.model); @@ -403,6 +414,7 @@ function chatRequest(request: JsonRecord): RequestMapping { seed: readNumber(body.seed), frequencyPenalty: readNumber(body.frequencyPenalty), presencePenalty: readNumber(body.presencePenalty), + attributes: promptAttributes(body.messages), }, }; } @@ -472,6 +484,7 @@ function responsesRequest(request: JsonRecord): RequestMapping { maxTokens: readNumber(body.maxOutputTokens), frequencyPenalty: readNumber(body.frequencyPenalty), presencePenalty: readNumber(body.presencePenalty), + attributes: promptAttributes(body.input), }, }; } diff --git a/packages/openrouter/tests/openrouter.test.ts b/packages/openrouter/tests/openrouter.test.ts index 2dc9baf..0d50a1b 100644 --- a/packages/openrouter/tests/openrouter.test.ts +++ b/packages/openrouter/tests/openrouter.test.ts @@ -5,13 +5,16 @@ import { Embeddings } from "@openrouter/sdk/sdk/embeddings.js"; import { EventStream } from "@openrouter/sdk/lib/event-streams.js"; import { Responses } from "@openrouter/sdk/sdk/responses.js"; import { + type ChatMessage, flush, + getPrompt, init, type MaskFn, shutdown, startActiveSpan, type TelemetryOptions, } from "@telemetry-dev/sdk"; +import * as telemetrySdk from "@telemetry-dev/sdk"; import { afterEach, expect, test, vi } from "vitest"; import { instrumentOpenRouter, uninstrumentOpenRouter, wrapOpenRouter } from "../src/index.ts"; @@ -308,6 +311,82 @@ afterEach(async () => { vi.restoreAllMocks(); }); +test.each([ + { api: "chat", messages: "compiled", olderCore: false }, + { api: "streamed responses", messages: "compiled", olderCore: false }, + { api: "chat", messages: "copied", olderCore: false }, + { api: "chat", messages: "compiled", olderCore: true }, +])( + "$api outside withPrompt records the prompt of $messages messages, older core $olderCore", + async ({ api, messages: kind, olderCore }) => { + if (olderCore) { + vi.spyOn( + telemetrySdk as { promptAttributesFor: unknown }, + "promptAttributesFor", + "get", + ).mockReturnValue(undefined); + } + + const prompt = await getPrompt("support/reply", { + apiKey: "td_live_a", + baseUrl: "https://prompts.openrouter.test", + fetch: async () => + Response.json({ + name: "support/reply", + type: "chat", + label: "production", + version: 12, + messages: [ + { role: "system", content: "You answer for the {{team}} desk." }, + { role: "user", content: "Where is order {{order}}?" }, + ], + text: null, + model: { provider: "openrouter", name: "openai/gpt-4o" }, + parameters: {}, + tools: [], + responseFormat: { type: "text" }, + rollout: null, + }), + }); + const compiled = prompt.compile({ team: "billing", order: "#4" }) as ChatMessage[]; + const messages = kind === "compiled" ? compiled : compiled.map((message) => ({ ...message })); + const spans = setupSpans(); + + if (api === "chat") { + const client = clientWith(createFakeFetcher(jsonResponse(chatBody("chat_1"))).fetcher); + await client.chat.send({ chatRequest: { model: "openai/gpt-4o", messages } }); + } else { + const client = clientWith( + createFakeFetcher( + sseResponse([{ type: "response.completed", response: responsesBody("resp_1") }]), + ).fetcher, + ); + const stream = await client.responses.send({ + responsesRequest: { model: "openai/gpt-4o", input: messages, stream: true }, + }); + + if (!(stream instanceof ReadableStream)) throw new Error("expected a readable stream"); + await collectStream(stream); + } + + const { attributes } = await exportedSpan(spans); + expect(attributes["gen_ai.request.model"]).toBe("openai/gpt-4o"); + + if (olderCore || kind === "copied") { + expect(attributes["td.prompt.name"]).toBeUndefined(); + + return; + } + + expect(attributes).toMatchObject({ + "td.prompt.name": "support/reply", + "td.prompt.version": 12, + "td.prompt.label": "production", + "td.prompt.variables": '{"team":"billing","order":"#4"}', + }); + }, +); + test("chat maps request, response, usage, primary cost, provider, and sampling fields", async () => { const spans = setupSpans(); diff --git a/packages/sdk/README.md b/packages/sdk/README.md index c3b9a5e..ed35293 100644 --- a/packages/sdk/README.md +++ b/packages/sdk/README.md @@ -179,8 +179,43 @@ compiles its own. Model calls are the spans of type `generation` (operations `ch `text_completion`, and `generate_content`), including the spans of the provider integrations that use this SDK: OpenAI, Anthropic, Amazon Bedrock, Google GenAI, and OpenRouter. Spans from the AI SDK and TanStack AI integrations, and from other OpenTelemetry tracers, do not get the attributes. -A model call outside `withPrompt` does not get the attributes, even when it sends messages from -`compile()`; with input capture on, the call still records those messages as its input. + +Outside `withPrompt`, those provider integrations record the prompt of the message objects that +`compile()` returned. `compile()` of a chat prompt tags each message that it makes with its prompt +and a copy of its variables and of each list in them, in a non-enumerable symbol property: the keys +and the JSON of the message do not change. The objects in a list are not copied, so a change to one +of them after `compile()` changes the variables that a call records. A chat call that sends tagged +messages gets `td.prompt.name`, `td.prompt.version`, and `td.prompt.label` of the newest `compile()` +among them, also when input capture is off, and `td.prompt.variables` when it is on. The SDK looks +at the first 1000 and the last 1000 messages of the call. A message from a placeholder variable +keeps the tag of its own, older `compile()`, so it does not win over the `compile()` that put it in +the list. A copy of a message (a spread, `structuredClone`, or JSON) and a new message with the same +text have no tag, so a call that sends only such messages gets no `td.prompt.*` attributes; with +input capture on, the call still records them as its input. Bedrock Converse and Google GenAI take +another message shape: change the compiled messages in place, or use `withPrompt`. A text prompt +compiles to a string, which cannot carry a tag: use `withPrompt`. `withPrompt` wins over the tag. The +integrations do it for OpenAI chat completions and responses, Anthropic messages, Bedrock Converse, +Google GenAI `generateContent`, and OpenRouter chat and responses, but not for Bedrock InvokeModel +(its body is bytes), embeddings, images, audio, video, or OpenAI Realtime. + +With input capture on in the client when `compile()` runs, a tagged message keeps the variables of +its `compile()` for 5 minutes. A call keeps them for 5 more minutes when the newest `compile()` of +the call is that one. After that time, a call records the prompt without `td.prompt.variables`, +and the SDK removes the variables from the tag, also when you keep the message. Other compiles do +not change this time. Thus a chat history keeps the variables of the compiles of its last 5 +minutes, and a guard prompt or a queue before the call does not remove them. A message that you +compile one time and send in many calls keeps its variables while a call names its `compile()` at +least one time in 5 minutes. A call that sends only messages of earlier turns records the prompt of +the newest `compile()` among them, and its variables when they are less than 5 minutes old. With +input capture off, or in a runtime without `WeakRef` (Cloudflare Workers with a compatibility date +before 2025-05-05), the tag keeps no variables. + +For your own spans, `promptAttributesFor(messages)` gives the attributes to pass as `attributes`. +It gives nothing inside `withPrompt`, and a span with `captureInput: false` drops +`td.prompt.variables`. A model call whose parent is in `withPrompt` gets the prompt of the +`withPrompt` and none of the `td.prompt.*` attributes that you pass, at start or in `update`, +`end`, or `updateActiveSpan`, also when its `parent` option leads into a `withPrompt` that the call +is not in. ## Experiments diff --git a/packages/sdk/src/attrs.ts b/packages/sdk/src/attrs.ts index 75afcb4..6dbf34f 100644 --- a/packages/sdk/src/attrs.ts +++ b/packages/sdk/src/attrs.ts @@ -179,6 +179,9 @@ export function fieldsToAttributes( if (fields.attributes) { Object.assign(attrs, omitUndefined(fields.attributes)); + + // The variables are input, also when the app passes them from promptAttributesFor(). + if (!cfg.captureInput) delete attrs["td.prompt.variables"]; } return attrs; diff --git a/packages/sdk/src/index.ts b/packages/sdk/src/index.ts index 2c81aee..d4055c7 100644 --- a/packages/sdk/src/index.ts +++ b/packages/sdk/src/index.ts @@ -68,6 +68,7 @@ export { type GetPromptOptions, type GetTextPromptOptions, type Prompt, + promptAttributesFor, PromptError, type PromptMessage, type PromptModel, diff --git a/packages/sdk/src/prompts.ts b/packages/sdk/src/prompts.ts index 6cb830d..93b50d9 100644 --- a/packages/sdk/src/prompts.ts +++ b/packages/sdk/src/prompts.ts @@ -1,4 +1,9 @@ -import { type Attributes, type Context, createContextKey } from "@opentelemetry/api"; +import { + type Attributes, + type AttributeValue, + type Context, + createContextKey, +} from "@opentelemetry/api"; import { activeContext, jsonAttr, @@ -91,6 +96,12 @@ export interface ChatPrompt extends PromptBase { /** * Fills in the variables and replaces each placeholder message with the messages in its * variable. `withPrompt` records the variables of the last call before it starts or inside it. + * It tags each message object that it makes with this prompt and a copy of these variables and + * of each list in them, so that a model call outside `withPrompt` that sends the message records + * them; see `promptAttributesFor`. The objects in a list are not copied. The tag keeps the + * variables for 5 minutes after this compile and after each call that names it. Other compiles + * do not change this. The tag keeps no variables when input capture is off in the client, or + * when the runtime has no `WeakRef`. */ compile(variables?: PromptVariables): (ChatMessage | M)[]; } @@ -101,7 +112,8 @@ export interface TextPrompt extends PromptBase { readonly text: string; /** * Fills in the variables. `withPrompt` records the variables of the last call before it starts - * or inside it. + * or inside it. A string cannot carry a tag, so a model call outside `withPrompt` does not + * record a text prompt. */ compile(variables?: PromptVariables): string; } @@ -195,6 +207,20 @@ interface PromptScope { variables: object | undefined; } +/** The `compile()` call that made a chat message. */ +interface CompileTag { + readonly prompt: Prompt; + /** Undefined when input capture was off, or after `expires`; see expire. */ + variables: PromptVariables | undefined; + /** Counts the compiles of the process, so a higher one is newer. */ + readonly serial: number; + /** + * When the tag removes its variables, in `performance.now()` time. A call that names the tag + * moves it. + */ + expires: number; +} + const FETCH_TIMEOUT_MS = 10_000; /** The longest Retry-After that getPrompt honors, so that a bad header cannot stop fetches. */ const MAX_RETRY_AFTER_MS = 3_600_000; @@ -204,11 +230,25 @@ const MAX_CACHE_ENTRIES = 1_000; const VARIABLE = /\{\{\s*([A-Za-z_][A-Za-z0-9_]*)\s*\}\}/g; const MODEL_CALL_OPERATIONS = new Set(["chat", "text_completion", "generate_content"]); const PROMPT_KEY = createContextKey("telemetry.dev prompt"); +// Not enumerable, so JSON.stringify, spreads, and structuredClone leave the tag out of a copy. +const COMPILED_BY = Symbol("telemetry.dev compiled by"); +/** How many messages promptAttributesFor looks at, at each end of the list. */ +const TAG_LOOKUP = 1_000; +/** How long a tag keeps its variables after its compile, or after a call that names it. */ +const VARIABLES_TTL_MS = 300_000; +const MINUTE_MS = 60_000; const cache = new Map(); const compiledVariables = new WeakMap(); /** The Retry-After of the response that failed with this error, in milliseconds. */ const retryAfter = new WeakMap(); +let compiles = 0; +/** + * The tags that keep variables, by the minute of their deadline. A weak reference, so that the + * variables of a message that the app dropped do not stay until the deadline. The minutes are not + * sorted, since a tag that a call named goes back in at its new minute. + */ +const deadlines = new Map[]>(); /** * Fetches a prompt version from telemetry.dev and keeps it. When the copy is older than @@ -329,19 +369,82 @@ export function promptAttributes( | undefined; if (!scope) return undefined; - const { prompt, variables } = scope; + return attributesOf(scope.prompt, scope.variables, capture); +} + +/** + * The td.prompt.* attributes for a model call outside `withPrompt` that sends `messages` (a list + * of messages, or one message): those of the newest `compile()` of a chat prompt that made one of + * the messages. It looks at the first 1000 and the last 1000 messages. The tag is on the message + * object, so a copy of it gives nothing. Gives undefined inside `withPrompt`, whose prompt the span + * gets. `td.prompt.variables` needs input capture on in the client, at the compile and now, and + * the span drops it when its own input capture is off. A compile gives no variables 5 minutes + * after it or after the last call that named it; see `compile()`. It never throws. + */ +export function promptAttributesFor(messages: unknown): Record | undefined { + const config = currentClient().core?.config; + + try { + if (activeContext().getValue(PROMPT_KEY) !== undefined) return undefined; + + const now = performance.now(); + expire(now); + const tag = newestTag(messages); + + if (!tag) return undefined; + + if (tag.expires > now) tag.expires = now + VARIABLES_TTL_MS; + else tag.variables = undefined; + + return attributesOf(tag.prompt, tag.variables, config) as Record; + } catch (error) { + reportError(config?.onError, error); + + return undefined; + } +} + +function attributesOf( + prompt: Prompt, + variables: object | undefined, + capture: FieldCaptureConfig | undefined, +): Attributes { return omitUndefined({ "td.prompt.name": prompt.name, "td.prompt.version": prompt.version ?? undefined, "td.prompt.label": prompt.label ?? undefined, "td.prompt.variables": - variables !== undefined && capture.captureInput + variables !== undefined && capture?.captureInput ? prepareCaptureValue("td.prompt.variables", variables, capture) : undefined, }); } +// The messages of a placeholder come from an older compile, so the newest compile, not the last +// tagged message, names the prompt of the request. Its own messages are at the start or the end of +// the list, unless placeholders put more than 1000 messages on each side of them. +function newestTag(messages: unknown): CompileTag | undefined { + if (!Array.isArray(messages)) return tagOf(messages); + const { length } = messages; + let newest: CompileTag | undefined; + + for (let index = 0; index < length; index += 1) { + if (index === TAG_LOOKUP && length > 2 * TAG_LOOKUP) index = length - TAG_LOOKUP; + const tag = tagOf(messages[index]); + + if (tag && (newest === undefined || tag.serial > newest.serial)) newest = tag; + } + + return newest; +} + +function tagOf(message: unknown): CompileTag | undefined { + return typeof message === "object" && message !== null + ? (message as { [COMPILED_BY]?: CompileTag })[COMPILED_BY] + : undefined; +} + function checkOptions(name: string, options: GetPromptOptions, type: string): void { if (typeof name !== "string" || name.length === 0) { throw new TypeError("name must be a non-empty string"); @@ -528,27 +631,90 @@ function createPrompt(input: { responseFormat: content.responseFormat, isFallback: input.isFallback, compile(variables: PromptVariables = {}) { - compiledVariables.set(prompt, variables); + const values = copyVariables(variables); + compiledVariables.set(prompt, values); const scope = activeContext().getValue(PROMPT_KEY) as PromptScope | undefined; - if (scope?.prompt === prompt) scope.variables = variables; - - return type === "chat" - ? compileMessages(content.messages, variables) - : fill(content.text ?? "", variables); + if (scope?.prompt === prompt) scope.variables = values; + + if (type === "text") return fill(content.text ?? "", values); + compiles += 1; + const now = performance.now(); + expire(now); + const tag: CompileTag = { + prompt, + variables: undefined, + serial: compiles, + expires: now + VARIABLES_TTL_MS, + }; + + // promptAttributesFor cannot record the variables when input capture is off in the client. + // Without WeakRef (Cloudflare Workers before 2025-05-05), expire could not remove them. + if ((currentClient().core?.config.captureInput ?? true) && typeof WeakRef === "function") { + tag.variables = values; + addDeadline(new WeakRef(tag), tag.expires); + } + + return compileMessages(content.messages, values, tag); }, }) as Prompt; return prompt; } +/** + * Removes the variables of each expired tag, also when the app keeps its messages. Thus a chat + * history keeps the variables of the last 5 minutes only, and other compiles do not change the time + * of a request. A tag that a call named gets its new deadline. + */ +function expire(now: number): void { + const minute = Math.floor(now / MINUTE_MS); + + for (const [due, refs] of deadlines) { + if (due >= minute) continue; + deadlines.delete(due); + + for (const ref of refs) { + const tag = ref.deref(); + + if (!tag) continue; + + if (tag.expires > now) addDeadline(ref, tag.expires); + else tag.variables = undefined; + } + } +} + +function addDeadline(ref: WeakRef, at: number): void { + const minute = Math.floor(at / MINUTE_MS); + const refs = deadlines.get(minute); + + if (refs) refs.push(ref); + else deadlines.set(minute, [ref]); +} + +// A copy of each list too, so that a change of the app's history after compile() does not change +// what a span records. +function copyVariables(variables: PromptVariables): PromptVariables { + return Object.fromEntries( + Object.entries(variables).map(([name, value]) => [ + name, + Array.isArray(value) ? [...value] : value, + ]), + ); +} + function compileMessages( messages: readonly PromptMessage[], variables: PromptVariables, + tag: CompileTag, ): unknown[] { return messages.flatMap((message) => { if (message.role !== "placeholder") { - return [{ role: message.role, content: fill(message.content, variables) }]; + const compiled = { role: message.role, content: fill(message.content, variables) }; + Object.defineProperty(compiled, COMPILED_BY, { value: tag }); + + return [compiled]; } const value = Object.hasOwn(variables, message.name) ? variables[message.name] : undefined; diff --git a/packages/sdk/src/span.ts b/packages/sdk/src/span.ts index 4c4c38a..3ed9cd9 100644 --- a/packages/sdk/src/span.ts +++ b/packages/sdk/src/span.ts @@ -1,4 +1,5 @@ import { + type Attributes, type Context, defaultTextMapGetter, defaultTextMapSetter, @@ -98,6 +99,8 @@ interface SpanMeta { type: SpanType; capture: FieldCaptureConfig; onError?: (cause: Error) => void; + /** The span got its prompt from withPrompt. */ + scopedPrompt?: boolean; } // Remembers each span's type + capture flags so update()/end()/updateActiveSpan map @@ -190,11 +193,22 @@ function applyFields(span: Span, meta: SpanMeta, fields: SpanFields): void { if (fields.name) span.updateName(fields.name); const attrs = fieldsToAttributes(fields, meta.type, meta.capture); + if (meta.scopedPrompt) dropPromptAttributes(attrs); + if (Object.keys(attrs).length > 0) span.setAttributes(attrs); if (fields.error !== undefined) applyError(span, fields.error, meta.capture); } +// The prompt of withPrompt replaces each td.prompt.* attribute that the app passes, at start and in +// each update, so that a span does not mix two prompts when the app calls promptAttributesFor +// outside withPrompt. +function dropPromptAttributes(attrs: Attributes): void { + for (const key of Object.keys(attrs)) { + if (key.startsWith("td.prompt.")) delete attrs[key]; + } +} + function metaFor(core: ClientCore, options?: StartSpanOptions): SpanMeta { const cfg = core.config; @@ -231,6 +245,10 @@ export function createSpanHandle( const parentCtx = withSessionParent(sessionContext, sessionId, core.config.apiKey); const attrs = fieldsToAttributes(options, meta.type, meta.capture); const operation = attrs["gen_ai.operation.name"] ?? SPAN_TYPE_OPERATIONS[meta.type]; + const prompt = promptAttributes(parentCtx, operation, meta.capture); + meta.scopedPrompt = prompt !== undefined; + + if (meta.scopedPrompt) dropPromptAttributes(attrs); const span = core.tracer.startSpan( name, @@ -241,7 +259,7 @@ export function createSpanHandle( attributes: { "gen_ai.operation.name": SPAN_TYPE_OPERATIONS[meta.type], ...propagatedFromContext(parentCtx), - ...promptAttributes(parentCtx, operation, meta.capture), + ...prompt, ...attrs, }, }, diff --git a/packages/sdk/tests/prompts.test.ts b/packages/sdk/tests/prompts.test.ts index 6048d16..d75d116 100644 --- a/packages/sdk/tests/prompts.test.ts +++ b/packages/sdk/tests/prompts.test.ts @@ -1,16 +1,22 @@ import { createServer, type RequestListener } from "node:http"; import type { AddressInfo } from "node:net"; +import { setFlagsFromString } from "node:v8"; +import { runInNewContext } from "node:vm"; import { afterEach, beforeEach, expect, test, vi } from "vitest"; import { + type ChatMessage, flush, getPrompt, + type Prompt, + promptAttributesFor, propagateAttributes, PromptError, shutdown, startActiveSpan, startSpan, + updateActiveSpan, withPrompt, } from "../src/index.ts"; import { setup } from "./helpers.ts"; @@ -107,6 +113,7 @@ afterEach(async () => { await shutdown(); vi.useRealTimers(); vi.unstubAllEnvs(); + vi.unstubAllGlobals(); vi.restoreAllMocks(); }); @@ -729,6 +736,736 @@ test("variables are omitted when input capture is off; a version fetch has no la expect(off.spans.getFinishedSpans()[0]!.attributes["td.prompt.variables"]).toBeUndefined(); }); +async function promptOf(overrides: Record): Promise { + const { fetch } = fakeFetch(reply(overrides)); + + return getPrompt("support/reply", { + apiKey: "td_live_a", + baseUrl: newBaseUrl(), + fetch, + type: overrides.type as "chat", + }); +} + +const DESK_PROMPT = { + messages: [ + { role: "system", content: "Answer as the {{team}} desk." }, + { role: "user", content: "Order {{order}}: where is it?" }, + ], +}; + +test("promptAttributesFor names the compile of a message object; a copy or a string names nothing", async () => { + setup(); + const prompt = await promptOf(DESK_PROMPT); + const compiled = prompt.compile({ team: "billing", order: "#4" }) as ChatMessage[]; + const attributes = { + "td.prompt.name": "support/reply", + "td.prompt.version": 12, + "td.prompt.label": "production", + "td.prompt.variables": '{"team":"billing","order":"#4"}', + }; + + expect(promptAttributesFor(compiled)).toEqual(attributes); + expect(promptAttributesFor(compiled[0])).toEqual(attributes); + expect(JSON.stringify(compiled[0])).toBe( + '{"role":"system","content":"Answer as the billing desk."}', + ); + + for (const copy of [ + compiled.map((message) => ({ ...message })), + JSON.parse(JSON.stringify(compiled)) as unknown, + structuredClone(compiled), + compiled.map((message) => message.content), + ]) { + expect(promptAttributesFor(copy)).toBeUndefined(); + } + + const text = await promptOf({ type: "text", messages: [], text: "Summarize {{doc}}." }); + expect(promptAttributesFor([text.compile({ doc: "the policy" })])).toBeUndefined(); +}); + +test("an unrelated message with the text of a compiled message is not tagged", async () => { + setup(); + const prompt = await promptOf({ + messages: [{ role: "system", content: "You are a helpful assistant." }], + }); + prompt.compile(); + + expect( + promptAttributesFor([ + { role: "system", content: "You are a helpful assistant." }, + { role: "user", content: "What is 2 + 2?" }, + ]), + ).toBeUndefined(); +}); + +test("two tasks with one text and different other variables each record their own", async () => { + const { spans } = setup(); + const prompt = await promptOf(DESK_PROMPT); + + // Both compiles make the system text "Answer as the billing desk.", and each task sends only it. + const [systemA] = prompt.compile({ team: "billing", order: "#A" }) as ChatMessage[]; + const [systemB] = prompt.compile({ team: "billing", order: "#B" }) as ChatMessage[]; + startSpan("a", { type: "generation", attributes: promptAttributesFor([systemA]) }).end(); + startSpan("b", { type: "generation", attributes: promptAttributesFor([systemB]) }).end(); + await flush(); + + expect( + Object.fromEntries( + spans.getFinishedSpans().map((span) => [span.name, span.attributes["td.prompt.variables"]]), + ), + ).toEqual({ + a: '{"team":"billing","order":"#A"}', + b: '{"team":"billing","order":"#B"}', + }); +}); + +test("a rollout with two versions of the same messages records the version of each call", async () => { + setup(); + const { fetch } = fakeFetch( + reply({ + ...DESK_PROMPT, + rollout: { + percent: 16, + version: 13, + ...DESK_PROMPT, + text: null, + model: { provider: "openai", name: "gpt-4o-mini" }, + parameters: {}, + tools: [], + responseFormat: { type: "text" }, + }, + }), + ); + const options = { apiKey: "td_live_a", baseUrl: newBaseUrl(), fetch }; + vi.spyOn(Math, "random").mockReturnValueOnce(0.159).mockReturnValueOnce(0.16); + + const rollout = await getPrompt("support/reply", options); + const production = await getPrompt("support/reply", options); + const fromRollout = rollout.compile({ team: "billing", order: "#4" }); + const fromProduction = production.compile({ team: "billing", order: "#4" }); + + expect(promptAttributesFor(fromRollout)?.["td.prompt.version"]).toBe(13); + expect(promptAttributesFor(fromProduction)?.["td.prompt.version"]).toBe(12); +}); + +test("the newest compiled message wins, also after 170 messages of history", async () => { + setup(); + const prompt = await promptOf({ + messages: [ + { role: "system", content: "Answer as the {{team}} desk." }, + { role: "placeholder", name: "history" }, + { role: "user", content: "{{question}}" }, + ], + }); + const history = Array.from({ length: 170 }, (_, index) => ({ + role: index % 2 === 0 ? ("user" as const) : ("assistant" as const), + content: `message ${index}`, + })); + + const firstTurn = prompt.compile({ + team: "billing", + question: "Where is #4?", + history, + }) as ChatMessage[]; + const answer = { role: "assistant", content: "On its way." }; + const secondTurn = prompt.compile({ + team: "returns", + question: "Can I return it?", + }) as ChatMessage[]; + const chat = [...firstTurn, answer, ...secondTurn.slice(1)]; + + expect(promptAttributesFor(firstTurn)?.["td.prompt.variables"]).toContain('"team":"billing"'); + expect(promptAttributesFor([firstTurn[0], ...history])?.["td.prompt.variables"]).toContain( + '"team":"billing"', + ); + expect(promptAttributesFor(chat)?.["td.prompt.variables"]).toBe( + '{"team":"returns","question":"Can I return it?"}', + ); +}); + +test("a message that a placeholder brings from an older compile does not win", async () => { + setup(); + const router = await promptOf({ + name: "router", + version: 1, + messages: [{ role: "user", content: "{{question}}" }], + }); + const answer = await promptOf({ + name: "answer", + version: 7, + messages: [ + { role: "system", content: "Answer for {{team}}." }, + { role: "placeholder", name: "history" }, + ], + }); + const chat = await promptOf({ + name: "chat", + messages: [{ role: "user", content: "{{question}}" }], + }); + const summary = await promptOf({ + name: "summary", + messages: [ + { role: "system", content: "Summarize the chat." }, + { role: "placeholder", name: "history" }, + ], + }); + + const question = router.compile({ question: "Where is #4?" }) as ChatMessage[]; + const routed = answer.compile({ team: "billing", history: question }); + // The last 1000 messages are all from chat; the summary message is the first. + const conversation = Array.from( + { length: 1500 }, + (_, index) => chat.compile({ question: `question ${index}` }) as ChatMessage[], + ).flat(); + const summarized = summary.compile({ history: conversation }); + + expect(promptAttributesFor(routed)).toEqual({ + "td.prompt.name": "answer", + "td.prompt.version": 7, + "td.prompt.label": "production", + "td.prompt.variables": + '{"team":"billing","history":[{"role":"user","content":"Where is #4?"}]}', + }); + expect(promptAttributesFor(summarized)?.["td.prompt.name"]).toBe("summary"); +}); + +test("promptAttributesFor looks at the first and the last 1000 messages", async () => { + setup(); + const prompt = await promptOf(DESK_PROMPT); + const [system] = prompt.compile({ team: "billing", order: "#4" }) as ChatMessage[]; + const filler = (count: number) => Array.from({ length: count }, () => ({ role: "user" })); + + expect(promptAttributesFor([...filler(1000), system, ...filler(999)])).toBeDefined(); + expect(promptAttributesFor([...filler(999), system, ...filler(1001)])).toBeDefined(); + expect(promptAttributesFor([...filler(1001), system, ...filler(999)])).toBeDefined(); + expect(promptAttributesFor([...filler(1000), system, ...filler(1000)])).toBeUndefined(); +}); + +test("compile records its variables, not later changes of the object or of its lists", async () => { + const { spans } = setup(); + const prompt = await promptOf({ + messages: [ + { role: "placeholder", name: "history" }, + { role: "user", content: "Order {{order}}" }, + ], + }); + const history: ChatMessage[] = [{ role: "user", content: "Hi" }]; + const variables = { order: "#A", history }; + + const first = prompt.compile(variables); + variables.order = "#B"; + history.push({ role: "assistant", content: "Hello" }); + const second = prompt.compile(variables); + variables.order = "#C"; + history.push({ role: "user", content: "Added after compile" }); + withPrompt(prompt, () => startSpan("scope", { type: "generation" }).end()); + await flush(); + + const secondVariables = + '{"order":"#B","history":[{"role":"user","content":"Hi"},{"role":"assistant","content":"Hello"}]}'; + expect(promptAttributesFor(first)?.["td.prompt.variables"]).toBe( + '{"order":"#A","history":[{"role":"user","content":"Hi"}]}', + ); + expect(promptAttributesFor(second)?.["td.prompt.variables"]).toBe(secondVariables); + expect(spans.getFinishedSpans()[0]!.attributes["td.prompt.variables"]).toBe(secondVariables); +}); + +// A WeakRef keeps its target until the current job ends, so collect in a new one. +async function collect(): Promise { + setFlagsFromString("--expose-gc"); + const gc = runInNewContext("gc") as () => void; + await new Promise((resolve) => setTimeout(resolve, 0)); + gc(); +} + +const RAG_PROMPT = { + messages: [ + { role: "system", content: "Answer from {{documents}}." }, + { role: "user", content: "{{question}}" }, + ], +}; + +test("the variables of a compile are freed with its messages", async () => { + setup(); + const prompt = await promptOf(RAG_PROMPT); + + let document: ChatMessage | undefined = { role: "user", content: "The refund policy" }; + const freed = new WeakRef(document); + let messages = prompt.compile({ documents: [document], question: "q" }) as + | ChatMessage[] + | undefined; + document = undefined; + // withPrompt keeps the variables of the last compile of a prompt object. + prompt.compile({ documents: [], question: "other" }); + await collect(); + expect(freed.deref()).toBeDefined(); + + expect(messages).toHaveLength(2); + messages = undefined; + await collect(); + expect(freed.deref()).toBeUndefined(); +}); + +// A row removes the global WeakRef, so the test keeps its own. +const TrackedRef = WeakRef; + +test.each<[string, () => void]>([ + ["input capture is off", () => setup({ captureInput: false })], + [ + "the runtime has no WeakRef (Cloudflare Workers before 2025-05-05)", + () => { + setup(); + vi.stubGlobal("WeakRef", undefined); + }, + ], +])("when %s, compiled messages keep no variables", async (_case, start) => { + start(); + const prompt = await promptOf(RAG_PROMPT); + const documents: WeakRef[] = []; + const questions: ChatMessage[] = []; + + for (let turn = 0; turn < 200; turn += 1) { + const document: ChatMessage = { role: "user", content: `document ${turn}` }; + documents.push(new TrackedRef(document)); + const compiled = prompt.compile({ documents: [document], question: `question ${turn}` }); + questions.push(compiled.at(-1) as ChatMessage); + } + await collect(); + + // withPrompt keeps the variables of the last compile of a prompt object. + expect(documents.filter((document) => document.deref()).length).toBe(1); + + await shutdown(); + vi.unstubAllGlobals(); + setup(); + expect(promptAttributesFor(questions)).toEqual({ + "td.prompt.name": "support/reply", + "td.prompt.version": 12, + "td.prompt.label": "production", + }); +}); + +const CHAT_PROMPT = { + messages: [ + { role: "system", content: "Answer as the {{team}} desk." }, + { role: "placeholder", name: "history" }, + { role: "user", content: "{{question}}" }, + ], +}; + +interface ChatPrompts { + chat: Prompt; + question: Prompt; + rag: Prompt; + reply: Prompt; +} + +async function chatPrompts(): Promise { + return { + chat: await promptOf(CHAT_PROMPT), + question: await promptOf({ + messages: [{ role: "user", content: "Context: {{documents}}\n\n{{question}}" }], + }), + rag: await promptOf({ + messages: [ + { role: "system", content: "Answer as the {{team}} desk." }, + { role: "placeholder", name: "history" }, + { role: "placeholder", name: "current" }, + ], + }), + reply: await promptOf({ messages: [{ role: "assistant", content: "{{answer}}" }] }), + }; +} + +function answer(turn: number): ChatMessage { + return { role: "assistant", content: `answer ${turn}` }; +} + +// A row runs one turn of a chat. It puts `marker` in the variables of the compile that made the +// message that the chat keeps. It gives the request, the history for the next turn, and that message. +type ChatTurn = ( + prompts: ChatPrompts, + history: ChatMessage[], + turn: number, + marker: ChatMessage, +) => [request: ChatMessage[], history: ChatMessage[], kept: ChatMessage]; + +test.each<[string, ChatTurn]>([ + [ + "a new list each turn", + ({ chat }, history, turn, marker) => { + const request = chat.compile({ + team: "billing", + history, + question: `question ${turn}`, + marker: [marker], + }) as ChatMessage[]; + const kept = request.at(-1)!; + + return [request, [...history, kept, answer(turn)], kept]; + }, + ], + [ + "a new list of the last 20 messages", + ({ chat }, history, turn, marker) => { + const request = chat.compile({ + team: "billing", + history, + question: `question ${turn}`, + marker: [marker], + }) as ChatMessage[]; + const kept = request.at(-1)!; + + return [request, [...history, kept, answer(turn)].slice(-20), kept]; + }, + ], + [ + "a full list that compile gets as copies", + ({ chat }, history, turn, marker) => { + const request = chat.compile({ + team: "billing", + history: history.map((message) => ({ ...message })), + question: `question ${turn}`, + marker: [marker], + }) as ChatMessage[]; + const kept = request.at(-1)!; + + return [request, [...history, kept, answer(turn)], kept]; + }, + ], + [ + "a question prompt for each turn, in a placeholder of the chat prompt", + ({ question, rag }, history, turn, marker) => { + const [kept] = question.compile({ + documents: `document ${turn}`, + question: `question ${turn}`, + marker: [marker], + }) as ChatMessage[]; + const request = rag.compile({ team: "billing", history, current: [kept!] }) as ChatMessage[]; + + return [request, [...history, kept!, answer(turn)], kept!]; + }, + ], + [ + "raw questions and a reply prompt for each answer", + ({ chat, reply }, history, turn, marker) => { + const request = chat.compile({ + team: "billing", + history, + question: `question ${turn}`, + }) as ChatMessage[]; + const [kept] = reply.compile({ answer: `answer ${turn}`, marker: [marker] }) as ChatMessage[]; + + return [request, [...history, { role: "user", content: `question ${turn}` }, kept!], kept!]; + }, + ], +])("a chat history in %s keeps the variables of the last 5 minutes", async (_shape, run) => { + vi.useFakeTimers({ toFake: ["Date", "performance"] }); + setup(); + const prompts = await chatPrompts(); + const markers: WeakRef[] = []; + let history: ChatMessage[] = []; + let request: ChatMessage[] = []; + let older: ChatMessage | undefined; + + // One turn in 2 minutes, for 300 turns. An unused variable marks the variables of each turn. + for (let turn = 0; turn < 300; turn += 1) { + if (turn > 0) vi.advanceTimersByTime(120_000); + const marker: ChatMessage = { role: "user", content: `turn ${turn}` }; + markers.push(new WeakRef(marker)); + let kept: ChatMessage; + [request, history, kept] = run(prompts, history, turn, marker); + // The model call of the turn. + promptAttributesFor(request); + + if (turn === 296) older = kept; + } + await collect(); + + // Turns 297, 298, and 299: 4, 2, and 0 minutes old. + expect(markers.filter((marker) => marker.deref()).length).toBe(3); + expect(promptAttributesFor(request)?.["td.prompt.variables"]).toContain("question 299"); + expect(promptAttributesFor(older)).toEqual({ + "td.prompt.name": "support/reply", + "td.prompt.version": 12, + "td.prompt.label": "production", + }); +}); + +function variablesOf(messages: unknown): unknown { + return promptAttributesFor(messages)?.["td.prompt.variables"]; +} + +const REFUND_VARIABLES = '{"documents":"The refund policy","question":"Where is #4?"}'; + +interface LivePrompts { + answer: Prompt; + guard: Prompt; + judge: Prompt; + other: Prompt; + chat: Prompt; + system: Prompt; + withSystem: Prompt; +} + +// A row compiles one or more requests and runs other compiles before their calls. It gives the +// variables that each call records and the variables of the compile of its request. +type LiveRequests = (prompts: LivePrompts) => Promise<[recorded: unknown, compiled: string][]>; + +test.each<[string, LiveRequests]>([ + [ + "150 compiles of other requests, then a guard", + async ({ answer, guard, other }) => { + const request = answer.compile({ documents: "The refund policy", question: "Where is #4?" }); + for (let index = 0; index < 150; index += 1) other.compile({ q: `${index}` }); + guard.compile({ conversation: request }); + + return [[variablesOf(request), REFUND_VARIABLES]]; + }, + ], + [ + "a guard before each of 150 agent steps", + async ({ answer, guard }) => { + const conversation = answer.compile({ + documents: "The refund policy", + question: "Where is #4?", + }) as ChatMessage[]; + + for (let step = 0; step < 150; step += 1) { + guard.compile({ conversation }); + conversation.push({ role: "assistant", content: `step ${step}` }); + } + + return [[variablesOf(conversation), REFUND_VARIABLES]]; + }, + ], + [ + "a guard, then a judge of the messages of the guard", + async ({ answer, guard, judge }) => { + const request = answer.compile({ documents: "The refund policy", question: "Where is #4?" }); + judge.compile({ conversation: guard.compile({ conversation: request }) }); + + return [[variablesOf(request), REFUND_VARIABLES]]; + }, + ], + [ + "three turns of a session that add their question to the history before their calls", + async ({ chat }) => { + const history: ChatMessage[] = []; + let open!: () => void; + const queue = new Promise((resolve) => { + open = resolve; + }); + const turn = async (question: string) => { + const request = chat.compile({ team: "billing", history, question }) as ChatMessage[]; + history.push(request.at(-1)!); + await queue; + + return variablesOf(request); + }; + const turns = ["q1", "q2", "q3"].map(turn); + open(); + const [first, second, third] = await Promise.all(turns); + + return [ + [first, '{"team":"billing","history":[],"question":"q1"}'], + [second, '{"team":"billing","history":[{"role":"user","content":"q1"}],"question":"q2"}'], + [ + third, + '{"team":"billing","history":[{"role":"user","content":"q1"},{"role":"user","content":"q2"}],"question":"q3"}', + ], + ]; + }, + ], + [ + "a system message compiled at startup, after 3 turns of a chat that gets it", + async ({ system, withSystem }) => { + const startup = system.compile({ product: "Acme" }) as ChatMessage[]; + let history: ChatMessage[] = []; + + for (let turn = 0; turn < 3; turn += 1) { + const request = withSystem.compile({ + system: startup, + history, + question: `question ${turn}`, + }) as ChatMessage[]; + promptAttributesFor(request); + history = [...history, request.at(-1)!, answer(turn)]; + } + + return [ + [ + variablesOf([...startup, { role: "user", content: "Where is #4?" }]), + '{"product":"Acme"}', + ], + ]; + }, + ], +])("a request keeps its variables until its call: %s", async (_case, run) => { + setup(); + const calls = await run({ + answer: await promptOf(RAG_PROMPT), + guard: await promptOf({ + name: "guard", + messages: [ + { role: "system", content: "Is this safe?" }, + { role: "placeholder", name: "conversation" }, + ], + }), + judge: await promptOf({ + name: "judge", + messages: [ + { role: "system", content: "Check the answer of the guard." }, + { role: "placeholder", name: "conversation" }, + ], + }), + other: await promptOf({ name: "other", messages: [{ role: "user", content: "{{q}}" }] }), + chat: await promptOf(CHAT_PROMPT), + system: await promptOf({ + name: "system", + messages: [{ role: "system", content: "You support {{product}}." }], + }), + withSystem: await promptOf({ + messages: [ + { role: "placeholder", name: "system" }, + { role: "placeholder", name: "history" }, + { role: "user", content: "{{question}}" }, + ], + }), + }); + + for (const [recorded, compiled] of calls) expect(recorded).toBe(compiled); +}); + +test("a compile keeps its variables for 5 minutes after it or after the last call that names it", async () => { + vi.useFakeTimers({ toFake: ["Date", "performance"] }); + setup(); + const chat = await promptOf(CHAT_PROMPT); + let history: ChatMessage[] = []; + + for (let turn = 0; turn < 5; turn += 1) { + const request = chat.compile({ team: "billing", history, question: `q${turn}` }); + promptAttributesFor(request); + history = [...history, request.at(-1) as ChatMessage, answer(turn)]; + } + // These calls send only messages of the chat, so they name the newest compile among them. + const summary = [...history, { role: "user", content: "Summarize the chat." }]; + const earlier = history.slice(0, -2); + + expect(variablesOf(summary)).toContain('"question":"q4"'); + expect(variablesOf(earlier)).toContain('"question":"q3"'); + + vi.advanceTimersByTime(299_999); + expect(variablesOf(summary)).toContain('"question":"q4"'); + vi.advanceTimersByTime(1); + expect(variablesOf(earlier)).toBeUndefined(); + // The call 1 ms before moved the deadline of q4. + expect(variablesOf(summary)).toContain('"question":"q4"'); + // The system clock moves back an hour; 5 minutes still pass. + vi.setSystemTime(Date.now() - 3_600_000); + vi.advanceTimersByTime(300_000); + expect(promptAttributesFor(summary)).toEqual({ + "td.prompt.name": "support/reply", + "td.prompt.version": 12, + "td.prompt.label": "production", + }); +}); + +test("withPrompt wins over a compiled message of another prompt, through a parent and in updates", async () => { + const { spans } = setup(); + const scoped = await promptOf({ name: "scoped/prompt" }); + // Fetched by version, so it has no label that the label of the scope could hide. + const { fetch } = fakeFetch( + reply({ + name: "other/prompt", + label: null, + version: 3, + messages: [{ role: "user", content: "Cancel order {{order}}." }], + }), + ); + const other = await getPrompt("other/prompt", { + apiKey: "td_live_a", + baseUrl: newBaseUrl(), + version: 3, + fetch, + }); + const compiled = other.compile({ order: "#9" }); + + const parent = withPrompt(scoped, () => { + startSpan("in scope", { type: "generation", attributes: promptAttributesFor(compiled) }).end(); + + return startSpan("parent"); + }); + const attributes = promptAttributesFor(compiled); + startSpan("with parent", { type: "generation", parent, attributes }).end(); + parent.end(); + withPrompt(scoped, () => { + startSpan("updated", { type: "generation" }).update({ attributes }).end(); + startSpan("ended", { type: "generation" }).end({ attributes }); + startActiveSpan("active", { type: "generation" }, () => updateActiveSpan({ attributes })); + }); + await flush(); + + expect(attributes?.["td.prompt.name"]).toBe("other/prompt"); + for (const name of ["in scope", "with parent", "updated", "ended", "active"]) { + const span = spans.getFinishedSpans().find((finished) => finished.name === name)!; + expect( + Object.fromEntries( + Object.entries(span.attributes).filter(([key]) => key.startsWith("td.prompt.")), + ), + ).toEqual({ + "td.prompt.name": "scoped/prompt", + "td.prompt.version": 12, + "td.prompt.label": "production", + }); + } +}); + +test("td.prompt.variables follows the input capture of the client and of the span", async () => { + const prompt = await promptOf({ + messages: [{ role: "user", content: "My card is {{card}}." }], + }); + const compiled = prompt.compile({ card: "4111 1111" }); + + setup({ captureInput: false }); + expect(promptAttributesFor(compiled)).toEqual({ + "td.prompt.name": "support/reply", + "td.prompt.version": 12, + "td.prompt.label": "production", + }); + + await shutdown(); + const { spans } = setup({ + mask: (value, { key }) => (key === "td.prompt.variables" ? { card: "[masked]" } : value), + }); + const attributes = promptAttributesFor(compiled); + expect(attributes?.["td.prompt.variables"]).toBe('{"card":"[masked]"}'); + + startSpan("private", { type: "generation", captureInput: false, attributes }).end(); + startSpan("updated", { type: "generation", captureInput: false }).update({ attributes }).end(); + await flush(); + + for (const span of spans.getFinishedSpans()) { + expect(span.attributes["td.prompt.name"]).toBe("support/reply"); + expect(span.attributes["td.prompt.variables"]).toBeUndefined(); + } +}); + +test("promptAttributesFor reports an error and gives undefined", async () => { + const errors: Error[] = []; + setup({ onError: (error) => errors.push(error) }); + + const trap = new Proxy([], { + get() { + throw new Error("trap"); + }, + }); + + expect(promptAttributesFor(trap)).toBeUndefined(); + expect(errors.map((error) => error.message)).toEqual(["trap"]); +}); + test("a fallback prompt records its name and label but no version", async () => { vi.spyOn(console, "error").mockImplementation(() => {}); const { spans } = setup(); diff --git a/sdks/python-anthropic/src/telemetry_dev_anthropic/__init__.py b/sdks/python-anthropic/src/telemetry_dev_anthropic/__init__.py index 0f257aa..a0dd315 100644 --- a/sdks/python-anthropic/src/telemetry_dev_anthropic/__init__.py +++ b/sdks/python-anthropic/src/telemetry_dev_anthropic/__init__.py @@ -244,6 +244,18 @@ def _offered_tools(tools: Any) -> dict[str, str] | None: return {"gen_ai.tool.definitions": json_dumps(definitions)} if definitions else None +def _add_prompt(fields: Mapping[str, Any], messages: object) -> dict[str, Any]: + """``fields`` plus the prompt tag of ``messages``, the list that the app passed: the copy that + ``_normalize_request_params`` makes has no tag. Call it where the span starts, so that a + ``with_prompt`` scope there wins.""" + # Older cores have no prompt_attributes_for and record a prompt only inside with_prompt. + read = getattr(telemetry_dev, "prompt_attributes_for", None) + prompt = cast("dict[str, Any]", read(messages)) if callable(read) else {} + if not prompt: + return dict(fields) + return {**fields, "attributes": {**(fields.get("attributes") or {}), **prompt}} + + def _messages_request(params: Mapping[str, Any]) -> tuple[str, dict[str, Any]]: params = _sent_params(params) model = _string(params.get("model")) @@ -1452,16 +1464,21 @@ def __getattr__(self, name: str) -> Any: class _InstrumentedMessageStreamManager: - def __init__(self, inner: Any, name: str, fields: Mapping[str, Any], provider: str) -> None: + def __init__( + self, inner: Any, name: str, fields: Mapping[str, Any], messages: object, provider: str + ) -> None: self._inner = inner self._name = name self._fields = fields + self._messages = messages self._provider = provider self._proxy: _InstrumentedStream | None = None def __enter__(self) -> Any: handle = telemetry_dev.start_span( - self._name, provider=self._provider, **_clean_fields(self._fields) + self._name, + provider=self._provider, + **_clean_fields(_add_prompt(self._fields, self._messages)), ) end = _end_once(handle) started_at = time.perf_counter() @@ -1491,16 +1508,21 @@ def __getattr__(self, name: str) -> Any: class _InstrumentedAsyncMessageStreamManager: - def __init__(self, inner: Any, name: str, fields: Mapping[str, Any], provider: str) -> None: + def __init__( + self, inner: Any, name: str, fields: Mapping[str, Any], messages: object, provider: str + ) -> None: self._inner = inner self._name = name self._fields = fields + self._messages = messages self._provider = provider self._proxy: _InstrumentedAsyncStream | None = None async def __aenter__(self) -> Any: handle = telemetry_dev.start_span( - self._name, provider=self._provider, **_clean_fields(self._fields) + self._name, + provider=self._provider, + **_clean_fields(_add_prompt(self._fields, self._messages)), ) end = _end_once(handle) started_at = time.perf_counter() @@ -1533,9 +1555,10 @@ def __getattr__(self, name: str) -> Any: def _start_span( - params: Mapping[str, Any], provider: str + params: Mapping[str, Any], messages: object, provider: str ) -> tuple[telemetry_dev.SpanHandle, Callable[..., None], float]: name, fields = _messages_request(params) + fields = _add_prompt(fields, messages) handle = telemetry_dev.start_span(name, provider=provider, **_clean_fields(fields)) return handle, _end_once(handle), time.perf_counter() @@ -1548,7 +1571,9 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: if _raw_response_requested(call_kwargs): return original(*args, **call_kwargs) streaming = call_kwargs.get("stream") is True - handle, end, started_at = _start_span(call_kwargs, provider(resource)) + handle, end, started_at = _start_span( + call_kwargs, _sent_params(kwargs).get("messages"), provider(resource) + ) try: result = original(*args, **call_kwargs) except BaseException as exc: @@ -1572,7 +1597,9 @@ async def wrapper(*args: Any, **kwargs: Any) -> Any: if _raw_response_requested(call_kwargs): return await original(*args, **call_kwargs) streaming = call_kwargs.get("stream") is True - handle, end, started_at = _start_span(call_kwargs, provider(resource)) + handle, end, started_at = _start_span( + call_kwargs, _sent_params(kwargs).get("messages"), provider(resource) + ) try: result = await original(*args, **call_kwargs) except BaseException as exc: @@ -1596,8 +1623,11 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: resource = args[0] if args else None call_kwargs = _normalize_request_params(kwargs) name, fields = _messages_request(call_kwargs) + messages = _sent_params(kwargs).get("messages") manager = original(*args, **call_kwargs) - return _InstrumentedMessageStreamManager(manager, name, fields, provider(resource)) + return _InstrumentedMessageStreamManager( + manager, name, fields, messages, provider(resource) + ) setattr(wrapper, _WRAPPED_ATTR, True) setattr(wrapper, _ORIGINAL_ATTR, original) @@ -1612,8 +1642,11 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: resource = args[0] if args else None call_kwargs = _normalize_request_params(kwargs) name, fields = _messages_request(call_kwargs) + messages = _sent_params(kwargs).get("messages") manager = original(*args, **call_kwargs) - return _InstrumentedAsyncMessageStreamManager(manager, name, fields, provider(resource)) + return _InstrumentedAsyncMessageStreamManager( + manager, name, fields, messages, provider(resource) + ) setattr(wrapper, _WRAPPED_ATTR, True) setattr(wrapper, _ORIGINAL_ATTR, original) diff --git a/sdks/python-anthropic/tests/test_anthropic.py b/sdks/python-anthropic/tests/test_anthropic.py index 3911ffb..c22c555 100644 --- a/sdks/python-anthropic/tests/test_anthropic.py +++ b/sdks/python-anthropic/tests/test_anthropic.py @@ -421,6 +421,126 @@ def handler(request: httpx.Request) -> httpx.Response: assert a["gen_ai.usage.cache_read.input_tokens"] == 1000 +SUPPORT_PROMPT = telemetry_dev.ChatPrompt( + name="support/reply", + version=12, + label="production", + messages=[ + {"role": "system", "content": "You answer for the {{team}} desk."}, + {"role": "user", "content": "Where is order {{order}}?"}, + ], + model=None, + parameters={}, + tools=[], + response_format={"type": "text"}, + is_fallback=False, +) + + +SCOPED_PROMPT = telemetry_dev.ChatPrompt( + name="scoped/prompt", + version=3, + label="production", + messages=[{"role": "user", "content": "Hi"}], + model=None, + parameters={}, + tools=[], + response_format={"type": "text"}, + is_fallback=False, +) + + +@pytest.mark.parametrize( + ("messages", "capture_input", "older_core"), + [("compiled", False, False), ("copied", True, False), ("compiled", True, True)], +) +def test_create_outside_with_prompt_records_the_prompt_of_compiled_messages( + make: Any, + monkeypatch: pytest.MonkeyPatch, + messages: str, + capture_input: bool, + older_core: bool, +) -> None: + memory = make(capture_input=capture_input) + if older_core: + monkeypatch.delattr(telemetry_dev, "prompt_attributes_for") + system, user = SUPPORT_PROMPT.compile(team="billing", order="#4") + client = wrapped_sync_client(lambda _request: json_response(message_payload())) + + client.messages.create( + model="claude-sonnet-4-6", + max_tokens=64, + system=system["content"], + messages=[user if messages == "compiled" else dict(user)], + ) + + a = attrs(only_span(memory)) + assert a["gen_ai.request.model"] == "claude-sonnet-4-6" + assert {key: value for key, value in a.items() if key.startswith("td.prompt.")} == ( + {} + if older_core or messages == "copied" + else { + "td.prompt.name": "support/reply", + "td.prompt.version": 12, + "td.prompt.label": "production", + **({"td.prompt.variables": '{"team":"billing","order":"#4"}'} if capture_input else {}), + } + ) + + +@pytest.mark.parametrize("scope_around", ["stream()", "__enter__"]) +def test_messages_stream_takes_the_prompt_from_the_context_where_its_span_starts( + memory: SimpleNamespace, scope_around: str +) -> None: + _system, user = SUPPORT_PROMPT.compile(team="billing", order="#4") + client = wrapped_sync_client(lambda _request: named_sse_response(stream_events())) + request: dict[str, Any] = {"model": "claude-sonnet-4-6", "max_tokens": 64, "messages": [user]} + + if scope_around == "stream()": + with telemetry_dev.with_prompt(SCOPED_PROMPT): + manager = client.messages.stream(**request) + with manager as stream: + "".join(stream.text_stream) + else: + manager = client.messages.stream(**request) + with telemetry_dev.with_prompt(SCOPED_PROMPT), manager as stream: + "".join(stream.text_stream) + + # The span starts at __enter__, so the scope there wins, and without one the message names it. + assert attrs(only_span(memory))["td.prompt.name"] == ( + "support/reply" if scope_around == "stream()" else "scoped/prompt" + ) + + +@pytest.mark.parametrize("scope_around", ["stream()", "__aenter__"]) +async def test_async_messages_stream_takes_the_prompt_from_the_context_where_its_span_starts( + memory: SimpleNamespace, scope_around: str +) -> None: + _system, user = SUPPORT_PROMPT.compile(team="billing", order="#4") + + async def handler(_request: httpx.Request) -> httpx.Response: + return named_sse_response(stream_events()) + + client = wrapped_async_client(handler) + request: dict[str, Any] = {"model": "claude-sonnet-4-6", "max_tokens": 64, "messages": [user]} + + if scope_around == "stream()": + with telemetry_dev.with_prompt(SCOPED_PROMPT): + manager = client.messages.stream(**request) + async with manager as stream: + [text async for text in stream.text_stream] + else: + manager = client.messages.stream(**request) + with telemetry_dev.with_prompt(SCOPED_PROMPT): + async with manager as stream: + [text async for text in stream.text_stream] + await client.close() + + assert attrs(only_span(memory))["td.prompt.name"] == ( + "support/reply" if scope_around == "stream()" else "scoped/prompt" + ) + + def test_create_maps_native_messages_system_params_usage_finish_and_provider( memory: SimpleNamespace, ) -> None: diff --git a/sdks/python-bedrock/src/telemetry_dev_bedrock/_fields.py b/sdks/python-bedrock/src/telemetry_dev_bedrock/_fields.py index 1017467..f6fc07d 100644 --- a/sdks/python-bedrock/src/telemetry_dev_bedrock/_fields.py +++ b/sdks/python-bedrock/src/telemetry_dev_bedrock/_fields.py @@ -1,7 +1,8 @@ from __future__ import annotations -from typing import Any +from typing import Any, cast +import telemetry_dev from telemetry_dev._serialize import json_dumps PROVIDER = "amazon-bedrock" @@ -11,6 +12,12 @@ def clean(fields: dict[str, Any]) -> dict[str, Any]: return {key: value for key, value in fields.items() if value is not None} +def prompt_attributes(messages: object) -> dict[str, Any]: + # Older cores have no prompt_attributes_for and record a prompt only inside with_prompt. + read = getattr(telemetry_dev, "prompt_attributes_for", None) + return cast("dict[str, Any]", read(messages)) if callable(read) else {} + + def metadata_fields(response: dict[str, Any] | BaseException | None) -> dict[str, Any]: metadata: dict[str, Any] | None = None if isinstance(response, dict) and isinstance(response.get("ResponseMetadata"), dict): diff --git a/sdks/python-bedrock/src/telemetry_dev_bedrock/_registry.py b/sdks/python-bedrock/src/telemetry_dev_bedrock/_registry.py index 6d09848..d27304d 100644 --- a/sdks/python-bedrock/src/telemetry_dev_bedrock/_registry.py +++ b/sdks/python-bedrock/src/telemetry_dev_bedrock/_registry.py @@ -6,7 +6,14 @@ import telemetry_dev -from ._fields import PROVIDER, clean, merge_fields, metadata_fields, usage_from_converse +from ._fields import ( + PROVIDER, + clean, + merge_fields, + metadata_fields, + prompt_attributes, + usage_from_converse, +) from ._invoke_model import is_embedding_model from ._invoke_model import request_fields as invoke_request_fields from ._invoke_model import response_fields as invoke_response_fields @@ -50,6 +57,7 @@ def _converse_request(params: dict[str, Any]) -> dict[str, Any]: "model": _model(params), "input": normalize_messages(params.get("messages")), "system_instructions": normalize_content_list(params.get("system")) or None, + "attributes": prompt_attributes(params.get("messages")) or None, "temperature": inference.get("temperature"), "top_p": inference.get("topP"), "max_tokens": inference.get("maxTokens"), diff --git a/sdks/python-bedrock/tests/test_bedrock.py b/sdks/python-bedrock/tests/test_bedrock.py index 3b6acac..505d4d6 100644 --- a/sdks/python-bedrock/tests/test_bedrock.py +++ b/sdks/python-bedrock/tests/test_bedrock.py @@ -14,6 +14,7 @@ from botocore.client import BaseClient from botocore.exceptions import ClientError, EventStreamError from botocore.response import StreamingBody +from botocore.stub import Stubber import telemetry_dev_bedrock._streams as bedrock_streams from telemetry_dev_bedrock import instrument_bedrock, uninstrument_bedrock, wrap_bedrock @@ -114,6 +115,74 @@ def _json_attr(span: Any, key: str) -> Any: return json.loads(value) +SUPPORT_PROMPT = telemetry_dev.ChatPrompt( + name="support/reply", + version=12, + label="production", + messages=[ + {"role": "system", "content": "You answer for the {{team}} desk."}, + {"role": "user", "content": "Where is order {{order}}?"}, + ], + model=None, + parameters={}, + tools=[], + response_format={"type": "text"}, + is_fallback=False, +) + + +@pytest.mark.parametrize( + ("messages", "older_core"), + [("changed in place", False), ("new", False), ("changed in place", True)], +) +def test_converse_outside_with_prompt_records_the_prompt_of_compiled_messages( + memory: SimpleNamespace, monkeypatch: pytest.MonkeyPatch, messages: str, older_core: bool +) -> None: + if older_core: + monkeypatch.delattr(telemetry_dev, "prompt_attributes_for") + system, user = SUPPORT_PROMPT.compile(team="billing", order="#4") + blocks = [{"text": user["content"]}] + # Converse takes content blocks: a message changed in place keeps its tag, a new one has none. + if messages == "new": + message = {"role": "user", "content": blocks} + else: + user["content"] = blocks + message = user + client = _client("bedrock-runtime") + # The stubber runs botocore's validation of the request. + stubber = Stubber(client) + stubber.add_response( + "converse", + { + "output": {"message": {"role": "assistant", "content": [{"text": "Soon."}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2}, + "metrics": {"latencyMs": 1}, + }, + ) + + with stubber: + cast(Any, wrap_bedrock(client)).converse( + modelId="anthropic.claude-3-haiku", + system=[{"text": system["content"]}], + messages=[message], + ) + + [span] = memory.span_exporter.get_finished_spans() + a = dict(span.attributes or {}) + assert a["gen_ai.request.model"] == "anthropic.claude-3-haiku" + assert {key: value for key, value in a.items() if key.startswith("td.prompt.")} == ( + {} + if older_core or messages == "new" + else { + "td.prompt.name": "support/reply", + "td.prompt.version": 12, + "td.prompt.label": "production", + "td.prompt.variables": '{"team":"billing","order":"#4"}', + } + ) + + def test_converse_happy_path(memory: SimpleNamespace, monkeypatch: pytest.MonkeyPatch) -> None: client = _client("bedrock-runtime") request = { diff --git a/sdks/python-google-genai/src/telemetry_dev_google_genai/__init__.py b/sdks/python-google-genai/src/telemetry_dev_google_genai/__init__.py index 1d8e0ae..e094f83 100644 --- a/sdks/python-google-genai/src/telemetry_dev_google_genai/__init__.py +++ b/sdks/python-google-genai/src/telemetry_dev_google_genai/__init__.py @@ -353,6 +353,12 @@ def _request_attributes(config: Any) -> dict[str, Any]: return attrs +def _prompt_attributes(messages: object) -> dict[str, Any]: + # Older cores have no prompt_attributes_for and record a prompt only inside with_prompt. + read = getattr(telemetry_dev, "prompt_attributes_for", None) + return cast("dict[str, Any]", read(messages)) if callable(read) else {} + + def _generate_request_fields(model: Any, contents: Any, config: Any) -> tuple[str, dict[str, Any]]: model_name = _string(model) or "unknown" fields: dict[str, Any] = { @@ -383,7 +389,7 @@ def _generate_request_fields(model: Any, contents: Any, config: Any) -> tuple[st output_type = _output_type(config) if output_type is not None: fields["output_type"] = output_type - attrs = _request_attributes(config) + attrs = {**_request_attributes(config), **_prompt_attributes(contents)} if attrs: fields["attributes"] = attrs return f"chat {model_name}", fields diff --git a/sdks/python-google-genai/tests/test_google_genai.py b/sdks/python-google-genai/tests/test_google_genai.py index f045e68..e397763 100644 --- a/sdks/python-google-genai/tests/test_google_genai.py +++ b/sdks/python-google-genai/tests/test_google_genai.py @@ -656,6 +656,75 @@ async def run() -> Any: assert attrs(span)["google_genai.response.operation_done"] is True +SUPPORT_PROMPT = telemetry_dev.ChatPrompt( + name="support/reply", + version=12, + label="production", + messages=[ + {"role": "system", "content": "You answer for the {{team}} desk."}, + {"role": "user", "content": "Where is order {{order}}?"}, + ], + model=None, + parameters={}, + tools=[], + response_format={"type": "text"}, + is_fallback=False, +) + + +@pytest.mark.parametrize( + ("contents", "capture_input", "older_core"), + [ + ("changed in place", True, False), + ("changed in place", False, False), + ("new", True, False), + ("changed in place", True, True), + ], +) +def test_generate_content_outside_with_prompt_records_the_prompt_of_compiled_messages( + make: Any, + monkeypatch: pytest.MonkeyPatch, + contents: str, + capture_input: bool, + older_core: bool, +) -> None: + memory = make(capture_input=capture_input) + if older_core: + monkeypatch.delattr(telemetry_dev, "prompt_attributes_for") + system, user = SUPPORT_PROMPT.compile(team="billing", order="#4") + parts = [{"text": user["content"]}] + # Gemini takes parts: a message changed in place keeps its tag, a new one has none. + if contents == "new": + content: dict[str, Any] = {"role": "user", "parts": parts} + else: + del user["content"] + user["parts"] = parts + content = user + client, transport = client_with_transport([happy_payload()]) + + wrap_google_genai(client).models.generate_content( + model="gemini-2.5-flash", + contents=[content], + config=types.GenerateContentConfig(system_instruction=system["content"]), + ) + + assert cast(Any, transport.calls[0])["contents"] == [ + {"role": "user", "parts": [{"text": "Where is order #4?"}]} + ] + a = attrs(only_span(memory)) + assert a["gen_ai.request.model"] == "gemini-2.5-flash" + assert {key: value for key, value in a.items() if key.startswith("td.prompt.")} == ( + {} + if older_core or contents == "new" + else { + "td.prompt.name": "support/reply", + "td.prompt.version": 12, + "td.prompt.label": "production", + **({"td.prompt.variables": '{"team":"billing","order":"#4"}'} if capture_input else {}), + } + ) + + @pytest.mark.parametrize("async_mode", [False, True], ids=["sync", "async"]) def test_generate_content_maps_request_response_usage_and_sampling( memory: SimpleNamespace, async_mode: bool diff --git a/sdks/python-litellm/src/telemetry_dev_litellm/__init__.py b/sdks/python-litellm/src/telemetry_dev_litellm/__init__.py index b2f6680..49ec964 100644 --- a/sdks/python-litellm/src/telemetry_dev_litellm/__init__.py +++ b/sdks/python-litellm/src/telemetry_dev_litellm/__init__.py @@ -455,6 +455,12 @@ def _output_type(response_format: Any) -> str | None: return kind +def _prompt_attributes(messages: object) -> dict[str, Any]: + # Older cores have no prompt_attributes_for and record a prompt only inside with_prompt. + read = getattr(telemetry_dev, "prompt_attributes_for", None) + return cast("dict[str, Any]", read(messages)) if callable(read) else {} + + def _completion_request( args: tuple[Any, ...], kwargs: Mapping[str, Any] ) -> tuple[str, dict[str, Any]]: @@ -478,6 +484,7 @@ def _completion_request( "presence_penalty": _number(kwargs.get("presence_penalty")), "output_type": _output_type(kwargs.get("response_format")), "metadata": _request_metadata(kwargs), + "attributes": _prompt_attributes(_arg(args, kwargs, "messages", 1)) or None, }, ) @@ -545,7 +552,10 @@ def _responses_request( input_budget: telemetry_dev.CaptureBudget | None = None if capture_input: input_value, input_budget = _bounded_responses_native(_arg(args, kwargs, "input", 0)) - attributes: dict[str, Any] = {"gen_ai.operation.name": "chat"} + attributes: dict[str, Any] = { + "gen_ai.operation.name": "chat", + **_prompt_attributes(_arg(args, kwargs, "input", 0)), + } if input_budget is not None and input_budget.truncated: attributes["telemetry.dev.capture.truncated"] = True return ( diff --git a/sdks/python-litellm/tests/test_litellm.py b/sdks/python-litellm/tests/test_litellm.py index f7b08a1..a4c8dd1 100644 --- a/sdks/python-litellm/tests/test_litellm.py +++ b/sdks/python-litellm/tests/test_litellm.py @@ -57,6 +57,69 @@ def finished_spans(env: SimpleNamespace) -> list[ReadableSpan]: return list(env.span_exporter.get_finished_spans()) +SUPPORT_PROMPT = telemetry_dev.ChatPrompt( + name="support/reply", + version=12, + label="production", + messages=[ + {"role": "system", "content": "You answer for the {{team}} desk."}, + {"role": "user", "content": "Where is order {{order}}?"}, + ], + model=None, + parameters={}, + tools=[], + response_format={"type": "text"}, + is_fallback=False, +) + + +@pytest.mark.parametrize( + ("api", "messages", "capture_input", "older_core"), + [ + ("completion", "compiled", False, False), + ("responses", "compiled", True, False), + ("completion", "copied", True, False), + ("completion", "compiled", True, True), + ], +) +def test_calls_outside_with_prompt_record_the_prompt_of_compiled_messages( + make: Any, + monkeypatch: pytest.MonkeyPatch, + api: str, + messages: str, + capture_input: bool, + older_core: bool, +) -> None: + memory = make(capture_input=capture_input) + if older_core: + monkeypatch.delattr(telemetry_dev, "prompt_attributes_for") + compiled = SUPPORT_PROMPT.compile(team="billing", order="#4") + sent = compiled if messages == "compiled" else [dict(message) for message in compiled] + + if api == "completion": + telemetry_dev_litellm.completion(model="gpt-4o-mini", messages=sent, mock_response="Soon.") + else: + + def fake_responses(*_args: Any, **_kwargs: Any) -> SimpleNamespace: + return responses_result() + + monkeypatch.setattr(litellm, "responses", fake_responses) + telemetry_dev_litellm.responses(sent, "openai/gpt-4.1-mini") + + a = attrs(only_span(memory)) + assert a["gen_ai.request.model"] == ("gpt-4o-mini" if api == "completion" else "gpt-4.1-mini") + assert {key: value for key, value in a.items() if key.startswith("td.prompt.")} == ( + {} + if older_core or messages == "copied" + else { + "td.prompt.name": "support/reply", + "td.prompt.version": 12, + "td.prompt.label": "production", + **({"td.prompt.variables": '{"team":"billing","order":"#4"}'} if capture_input else {}), + } + ) + + def test_completion_maps_model_messages_usage_finish_provider_and_sampling( memory: SimpleNamespace, ) -> None: diff --git a/sdks/python-openai/src/telemetry_dev_openai/__init__.py b/sdks/python-openai/src/telemetry_dev_openai/__init__.py index af2a498..3cbe610 100644 --- a/sdks/python-openai/src/telemetry_dev_openai/__init__.py +++ b/sdks/python-openai/src/telemetry_dev_openai/__init__.py @@ -349,9 +349,18 @@ def _capture_text(value: str, budget: telemetry_dev.CaptureBudget) -> str: return value +def _prompt_attributes(messages: object) -> dict[str, Any]: + # Older cores have no prompt_attributes_for and record a prompt only inside with_prompt. + read = getattr(telemetry_dev, "prompt_attributes_for", None) + return cast("dict[str, Any]", read(messages)) if callable(read) else {} + + def _chat_request(params: Mapping[str, Any]) -> tuple[str, dict[str, Any]]: model = _string(params.get("model")) + attributes = _prompt_attributes(params.get("messages")) captured, truncated = _bounded_responses_capture(params.get("messages"), "input") + if truncated: + attributes["telemetry.dev.capture.truncated"] = True return ( f"chat {model or 'unknown'}", { @@ -366,7 +375,7 @@ def _chat_request(params: Mapping[str, Any]) -> tuple[str, dict[str, Any]]: "seed": _number(params.get("seed")), "frequency_penalty": _number(params.get("frequency_penalty")), "presence_penalty": _number(params.get("presence_penalty")), - "attributes": {"telemetry.dev.capture.truncated": True} if truncated else None, + "attributes": attributes or None, }, ) @@ -432,7 +441,10 @@ def _chat_response(response: Any) -> dict[str, Any]: def _responses_request(params: Mapping[str, Any]) -> tuple[str, dict[str, Any]]: model = _string(params.get("model")) + attributes = _prompt_attributes(params.get("input")) captured, truncated = _bounded_responses_capture(params.get("input"), "input") + if truncated: + attributes["telemetry.dev.capture.truncated"] = True return ( f"chat {model or 'unknown'}", { @@ -443,7 +455,7 @@ def _responses_request(params: Mapping[str, Any]) -> tuple[str, dict[str, Any]]: "temperature": _number(params.get("temperature")), "top_p": _number(params.get("top_p")), "max_tokens": _number(params.get("max_output_tokens")), - "attributes": {"telemetry.dev.capture.truncated": True} if truncated else None, + "attributes": attributes or None, }, ) diff --git a/sdks/python-openai/tests/test_openai.py b/sdks/python-openai/tests/test_openai.py index 238c7e6..55ab3c6 100644 --- a/sdks/python-openai/tests/test_openai.py +++ b/sdks/python-openai/tests/test_openai.py @@ -1206,6 +1206,70 @@ def text(self) -> str: assert a["gen_ai.usage.text.output_tokens"] == 2 +SUPPORT_PROMPT = telemetry_dev.ChatPrompt( + name="support/reply", + version=12, + label="production", + messages=[ + {"role": "system", "content": "You answer for the {{team}} desk."}, + {"role": "user", "content": "Where is order {{order}}?"}, + ], + model=None, + parameters={}, + tools=[], + response_format={"type": "text"}, + is_fallback=False, +) + + +@pytest.mark.parametrize( + ("api", "messages", "capture_input", "older_core"), + [ + ("chat", "compiled", False, False), + ("responses", "compiled", True, False), + ("chat", "copied", True, False), + ("chat", "compiled", True, True), + ], +) +def test_calls_outside_with_prompt_record_the_prompt_of_compiled_messages( + make: Any, + monkeypatch: pytest.MonkeyPatch, + api: str, + messages: str, + capture_input: bool, + older_core: bool, +) -> None: + memory = make(capture_input=capture_input) + if older_core: + monkeypatch.delattr(telemetry_dev, "prompt_attributes_for") + compiled = SUPPORT_PROMPT.compile(team="billing", order="#4") + sent = compiled if messages == "compiled" else [dict(message) for message in compiled] + + def handler(_request: httpx.Request) -> httpx.Response: + return json_response(chat_completion() if api == "chat" else response_payload()) + + client = wrap_openai(sync_client(handler)) + if api == "chat": + cast(Any, client.chat.completions.create)(model="gpt-4o", messages=sent) + else: + cast(Any, client.responses.create)(model="gpt-4.1", input=sent) + + span = only_span(memory) + a = attrs(span) + assert a["gen_ai.request.model"] == ("gpt-4o" if api == "chat" else "gpt-4.1") + assert {key: value for key, value in a.items() if key.startswith("td.prompt.")} == ( + {} + if older_core or messages == "copied" + else { + "td.prompt.name": "support/reply", + "td.prompt.version": 12, + "td.prompt.label": "production", + **({"td.prompt.variables": '{"team":"billing","order":"#4"}'} if capture_input else {}), + } + ) + assert ("gen_ai.input.messages" in a) is capture_input + + def test_chat_completion_maps_native_messages_usage_finish_provider_and_sampling( memory: SimpleNamespace, ) -> None: diff --git a/sdks/python-openrouter/src/telemetry_dev_openrouter/__init__.py b/sdks/python-openrouter/src/telemetry_dev_openrouter/__init__.py index 8b46b90..c15bfb3 100644 --- a/sdks/python-openrouter/src/telemetry_dev_openrouter/__init__.py +++ b/sdks/python-openrouter/src/telemetry_dev_openrouter/__init__.py @@ -174,6 +174,12 @@ def _cost_usd(raw_usage: Any) -> float | None: return cost if cost is not None else upstream_cost +def _prompt_attributes(messages: object) -> dict[str, Any]: + # Older cores have no prompt_attributes_for and record a prompt only inside with_prompt. + read = getattr(telemetry_dev, "prompt_attributes_for", None) + return cast("dict[str, Any]", read(messages)) if callable(read) else {} + + def _chat_request(params: Mapping[str, Any]) -> tuple[str, dict[str, Any]]: model = _string(params.get("model")) max_completion_tokens = _number(params.get("max_completion_tokens")) @@ -193,6 +199,7 @@ def _chat_request(params: Mapping[str, Any]) -> tuple[str, dict[str, Any]]: "seed": _number(params.get("seed")), "frequency_penalty": _number(params.get("frequency_penalty")), "presence_penalty": _number(params.get("presence_penalty")), + "attributes": _prompt_attributes(params.get("messages")) or None, }, ) @@ -260,6 +267,7 @@ def _responses_request(params: Mapping[str, Any]) -> tuple[str, dict[str, Any]]: "max_tokens": _number(params.get("max_output_tokens")), "frequency_penalty": _number(params.get("frequency_penalty")), "presence_penalty": _number(params.get("presence_penalty")), + "attributes": _prompt_attributes(params.get("input")) or None, }, ) diff --git a/sdks/python-openrouter/tests/test_openrouter.py b/sdks/python-openrouter/tests/test_openrouter.py index 1efe3a5..400c4f5 100644 --- a/sdks/python-openrouter/tests/test_openrouter.py +++ b/sdks/python-openrouter/tests/test_openrouter.py @@ -221,6 +221,68 @@ def embeddings_payload() -> dict[str, Any]: } +SUPPORT_PROMPT = telemetry_dev.ChatPrompt( + name="support/reply", + version=12, + label="production", + messages=[ + {"role": "system", "content": "You answer for the {{team}} desk."}, + {"role": "user", "content": "Where is order {{order}}?"}, + ], + model=None, + parameters={}, + tools=[], + response_format={"type": "text"}, + is_fallback=False, +) + + +@pytest.mark.parametrize( + ("api", "messages", "capture_input", "older_core"), + [ + ("chat", "compiled", False, False), + ("responses", "compiled", True, False), + ("chat", "copied", True, False), + ("chat", "compiled", True, True), + ], +) +def test_calls_outside_with_prompt_record_the_prompt_of_compiled_messages( + make: Any, + monkeypatch: pytest.MonkeyPatch, + api: str, + messages: str, + capture_input: bool, + older_core: bool, +) -> None: + memory = make(capture_input=capture_input) + if older_core: + monkeypatch.delattr(telemetry_dev, "prompt_attributes_for") + compiled = SUPPORT_PROMPT.compile(team="billing", order="#4") + sent = compiled if messages == "compiled" else [dict(message) for message in compiled] + + def handler(_request: httpx.Request) -> httpx.Response: + return json_response(chat_payload() if api == "chat" else responses_payload()) + + client = wrap_open_router(sync_client(handler)) + if api == "chat": + client.chat.send(model="openai/gpt-4o-mini", messages=cast(Any, sent)) + else: + client.responses.send(model="openai/gpt-4o-mini", input=cast(Any, sent)) + + a = attrs(only_span(memory)) + assert a["gen_ai.request.model"] == "openai/gpt-4o-mini" + assert {key: value for key, value in a.items() if key.startswith("td.prompt.")} == ( + {} + if older_core or messages == "copied" + else { + "td.prompt.name": "support/reply", + "td.prompt.version": 12, + "td.prompt.label": "production", + **({"td.prompt.variables": '{"team":"billing","order":"#4"}'} if capture_input else {}), + } + ) + + def test_chat_send_maps_messages_usage_cost_and_sampling(memory: SimpleNamespace) -> None: def handler(request: httpx.Request) -> httpx.Response: return json_response(chat_payload()) diff --git a/sdks/python/README.md b/sdks/python/README.md index 8fe3dd3..c7eb87c 100644 --- a/sdks/python/README.md +++ b/sdks/python/README.md @@ -260,9 +260,51 @@ ran in it, or before it started in the code that started it; with none, it recor for code that has not compiled in its own task. Model calls are the spans of type `generation` (operations `chat`, `text_completion`, and `generate_content`), including the spans of the provider integrations: OpenAI, Anthropic, Amazon Bedrock, Google GenAI, LiteLLM, and OpenRouter. Spans from -other OpenTelemetry tracers do not get the attributes. A model call outside `with_prompt` does not -get the attributes, even when it sends messages from `compile()`; with input capture on, the call -still records those messages as its input. +other OpenTelemetry tracers do not get the attributes. + +Outside `with_prompt`, those provider integrations record the prompt of the message dicts that +`compile()` returned. `compile()` of a chat prompt returns instances of a private `dict` subclass +that also hold its prompt and a copy of its variables and of each list in them. The objects in a +list are not copied, so a change to one of them after `compile()` changes the variables that a call +records. A chat call that sends such messages gets `td.prompt.name`, `td.prompt.version`, and +`td.prompt.label` of the newest `compile()` among them, also when input capture is off, and +`td.prompt.variables` when it is on. The SDK looks at the first 1000 and the last 1000 messages of +the call. A message from a placeholder variable keeps the prompt of its own, older `compile()`, so +it does not win over the `compile()` that put it in the list. A copy of a message (`dict()`, +`.copy()`, `copy.copy`, `copy.deepcopy`, pickle, JSON, `type(message)(message)`, or +`dataclasses.asdict`) and a new message with the same text do not hold the prompt, so a call that +sends only such messages gets no `td.prompt.*` attributes; with input capture on, the call still +records them as its input. Bedrock Converse and Google GenAI take another message shape: change the +compiled dicts in place, or use `with_prompt`. A text prompt compiles to a string, which cannot hold +its prompt: use `with_prompt`. `with_prompt` wins over the message. The integrations do it for OpenAI +chat completions and responses, Anthropic messages (for `messages.stream()`, at `__enter__`), +Bedrock Converse, Google GenAI `generate_content`, LiteLLM completion and responses, and OpenRouter +chat and responses, but not for Bedrock InvokeModel (its body is bytes), embeddings, images, audio, +video, or OpenAI Realtime. With an older `telemetry-dev` that has no `prompt_attributes_for`, a +provider integration records the prompt only in `with_prompt`. + +A compiled message is a `dict` for `isinstance`, and its keys and its JSON (`json`, `orjson`, +pydantic) do not change. A check of the exact type (`type(message) is dict`), `yaml.safe_dump`, +and `marshal` reject it: give them `dict(message)`. + +With input capture on in the client when `compile()` runs, a compiled message keeps the variables +of its `compile()` for 5 minutes. A call keeps them for 5 more minutes when the newest `compile()` +of the call is that one. After that time, a call records the prompt without +`td.prompt.variables`, and the SDK removes the variables from the message, also when you keep the +message. Other compiles do not change this time. Thus a chat history keeps the variables of the +compiles of its last 5 minutes, and a guard prompt, a queue, or a compile in another thread before +the call does not remove them. A message that you compile one time and send in many calls keeps its +variables while a call names its `compile()` at least one time in 5 minutes. A call that sends only +messages of earlier turns records the prompt of the newest `compile()` among them, and its +variables when they are less than 5 minutes old. With input capture off, the message keeps no +variables. + +For your own spans, `prompt_attributes_for(messages)` gives the attributes to pass as +`attributes`. It gives `{}` inside `with_prompt`, and a span with `capture_input=False` drops +`td.prompt.variables`. A model call whose parent is in `with_prompt` gets the prompt of the +`with_prompt` and none of the `td.prompt.*` attributes that you pass, at start or in `update`, +`end`, or `update_current_span`, also when its `parent` context comes from a `with_prompt` that the +call is not in. ## Experiments diff --git a/sdks/python/src/telemetry_dev/__init__.py b/sdks/python/src/telemetry_dev/__init__.py index bbcb230..f6ba26c 100644 --- a/sdks/python/src/telemetry_dev/__init__.py +++ b/sdks/python/src/telemetry_dev/__init__.py @@ -39,6 +39,7 @@ TextPrompt, get_prompt, get_prompt_async, + prompt_attributes_for, with_prompt, ) from ._scores import ScoreError, score, score_async @@ -89,6 +90,7 @@ "init", "log", "observe", + "prompt_attributes_for", "propagate_attributes", "run_experiment", "run_experiment_async", diff --git a/sdks/python/src/telemetry_dev/_prompts.py b/sdks/python/src/telemetry_dev/_prompts.py index 52ba5e7..c7e9e80 100644 --- a/sdks/python/src/telemetry_dev/_prompts.py +++ b/sdks/python/src/telemetry_dev/_prompts.py @@ -4,6 +4,7 @@ import contextlib import copy import hashlib +import itertools import json import os import random @@ -20,7 +21,7 @@ from dataclasses import dataclass, field from time import monotonic from typing import Any, Literal, TypeAlias, TypedDict, cast, overload -from weakref import WeakKeyDictionary +from weakref import WeakKeyDictionary, ref from opentelemetry import context as otel_context from opentelemetry.context import Context @@ -47,6 +48,11 @@ _VARIABLE = re.compile(r"\{\{\s*([A-Za-z_][A-Za-z0-9_]*)\s*\}\}") _MODEL_CALL_OPERATIONS = frozenset({"chat", "text_completion", "generate_content"}) _PROMPT_KEY = otel_context.create_key("telemetry_dev.prompt") +# How many messages prompt_attributes_for looks at, at each end of the list. +_TAG_LOOKUP = 1000 +# How long a tag keeps its variables after its compile, or after a call that names it. +_VARIABLES_TTL_S = 300.0 +_MINUTE_S = 60.0 class PromptModel(TypedDict): @@ -95,14 +101,31 @@ def compile( ) -> list[dict[str, Any]]: """Fill in the variables and replace each placeholder message with the list of messages in its variable. ``with_prompt`` records the variables of the last call before it starts - or inside it.""" + or inside it. Each message that it makes is a ``dict`` subclass that also holds this + prompt and a copy of these variables and of each list in them, so that a model call + outside ``with_prompt`` that sends the message records them; see + ``prompt_attributes_for``. The objects in a list are not copied. The message keeps the + variables for 5 minutes after this compile and after each call that names it. Other + compiles do not change this. The message holds no variables when input capture is off in + the client. A check of the exact type (``type(message) is dict``), ``yaml.safe_dump``, and + ``marshal`` reject it: give them ``dict(message)``.""" values = _record(self, variables, kwargs) + now = monotonic() + _expire(now) + tag = _CompileTag(self, None, next(_compiles), now + _VARIABLES_TTL_S) + client = get_client() + # prompt_attributes_for cannot record the variables when input capture is off in the client. + if client is None or client.capture_input: + tag.variables = values + _add_deadline(ref(tag), tag.expires) compiled: list[dict[str, Any]] = [] for message in self.messages: if message.get("role") != "placeholder": - compiled.append( - {"role": message["role"], "content": _fill(str(message["content"]), values)} + tagged = _CompiledMessage( + role=message["role"], content=_fill(str(message["content"]), values) ) + tagged.compiled_by = tag + compiled.append(tagged) continue name = str(message["name"]) value = values.get(name) @@ -138,7 +161,8 @@ def messages(self) -> list[dict[str, Any]]: def compile(self, variables: Mapping[str, Any] | None = None, /, **kwargs: Any) -> str: """Fill in the variables. ``with_prompt`` records the variables of the last call before it - starts or inside it.""" + starts or inside it. A string cannot carry a tag, so a model call outside ``with_prompt`` + does not record a text prompt.""" return _fill(self.text, _record(self, variables, kwargs)) @@ -279,16 +303,59 @@ def prompt_attributes( scope = otel_context.get_value(_PROMPT_KEY, context) or otel_context.get_value(_PROMPT_KEY) if not isinstance(scope, _PromptScope): return {} - prompt = scope.prompt + task = _task_variables.get() + variables = task[1] if task is not None and task[0] is scope else scope.variables + return _attributes_of(scope.prompt, variables, client if capture_input else None) + + +def prompt_attributes_for(messages: object) -> dict[str, AttributeValue]: + """The td.prompt.* attributes for a model call outside ``with_prompt`` that sends ``messages`` + (a list or tuple of messages, or one message): those of the newest ``compile()`` of a chat + prompt that made one of the messages. It looks at the first 1000 and the last 1000 messages. + The tag is on the dict that ``compile()`` made, so a copy of it gives ``{}``. Gives ``{}`` + inside ``with_prompt``, whose prompt the span gets. ``td.prompt.variables`` needs input + capture on in the client, at the compile and now, and the span drops it when its own input + capture is off. A compile gives no variables 5 minutes after it or after the last call that + named it; see ``ChatPrompt.compile``. + + It reads only a ``list`` or ``tuple``, not a subclass, so the read runs no code of the app. + With input capture on, it serializes the variables, which calls the ``mask`` of the client and + can run code of a value, such as ``__repr__``. It never raises: it reports an error to + ``on_error`` and gives ``{}``.""" + client = get_client() + try: + if isinstance(otel_context.get_value(_PROMPT_KEY), _PromptScope): + return {} + now = monotonic() + _expire(now) + tag = _newest_tag(messages) + if tag is None: + return {} + # One read: _expire in another thread can remove the variables of the tag. + variables = tag.variables + if tag.expires > now: + tag.expires = now + _VARIABLES_TTL_S + else: + variables = tag.variables = None + capture = client if client is not None and client.capture_input else None + return _attributes_of(tag.prompt, variables, capture) + except Exception as exc: + if client is not None: + client.report("prompt_attributes_for failed", exc) + return {} + + +def _attributes_of( + prompt: Prompt, variables: dict[str, Any] | None, capture: Client | None +) -> dict[str, AttributeValue]: + """The attributes of ``prompt``; ``capture`` is the client when input capture is on.""" attrs: dict[str, AttributeValue] = {ATTR_PROMPT_NAME: prompt.name} if prompt.version is not None: attrs[ATTR_PROMPT_VERSION] = prompt.version if prompt.label is not None: attrs[ATTR_PROMPT_LABEL] = prompt.label - task = _task_variables.get() - variables = task[1] if task is not None and task[0] is scope else scope.variables - if variables is not None and capture_input: - serialized = client.serialize(variables, ATTR_PROMPT_VARIABLES) + if variables is not None and capture is not None: + serialized = capture.serialize(variables, ATTR_PROMPT_VARIABLES) if serialized is not None: attrs[ATTR_PROMPT_VARIABLES] = serialized return attrs @@ -386,7 +453,12 @@ class _PromptScope: def _record( prompt: Prompt, variables: Mapping[str, Any] | None, kwargs: Mapping[str, Any] ) -> dict[str, Any]: - values = {**(variables or {}), **kwargs} + # A copy of each list too, so that a change of the app's history after compile() does not + # change what a span records. + values: dict[str, Any] = { + name: list(cast("list[Any]", value)) if isinstance(value, list) else value + for name, value in {**(variables or {}), **kwargs}.items() + } _compiled_variables[prompt] = values scope = otel_context.get_value(_PROMPT_KEY) if isinstance(scope, _PromptScope) and scope.prompt is prompt: @@ -395,6 +467,81 @@ def _record( return values +@dataclass(eq=False) +class _CompileTag: + """The ``compile()`` call that made a chat message.""" + + prompt: ChatPrompt + variables: dict[str, Any] | None + """None when input capture was off, or after ``expires``; see _expire.""" + serial: int + """Counts the compiles of the process, so a higher one is newer.""" + expires: float + """When the tag removes its variables (monotonic). A call that names the tag moves it.""" + + +_compiles = itertools.count() +# The tags that keep variables, by the minute of their deadline. A weak reference, so that the +# variables of a message that the app dropped do not stay until the deadline. +_deadlines: dict[int, list[ref[_CompileTag]]] = {} + + +class _CompiledMessage(dict[str, Any]): + """A message that ``compile()`` made: a dict with the tag in a slot. The constructor is the one + of ``dict``, so ``type(message)(message)`` and ``dataclasses.asdict`` rebuild it, without the + tag. A copy, pickle, or ``copy.deepcopy`` of it is a plain dict without the tag.""" + + __slots__ = ("compiled_by",) + compiled_by: _CompileTag + + def __reduce__(self) -> tuple[type[dict[str, Any]], tuple[dict[str, Any]]]: + return dict, (dict(self),) + + +def _expire(now: float) -> None: + """Remove the variables of each expired tag, also when the app keeps its messages. Thus a chat + history keeps the variables of the last 5 minutes only, and other compiles do not change the + time of a request. A tag that a call named gets its new deadline. Threads add only to minutes + that are not due, and ``list()`` and ``pop()`` of a dict are atomic, so the threads share the + minutes without a lock.""" + minute = int(now // _MINUTE_S) + for due in list(_deadlines): + if due >= minute: + continue + for tag_ref in _deadlines.pop(due, ()): + tag = tag_ref() + if tag is None: + continue + if tag.expires > now: + _add_deadline(tag_ref, tag.expires) + else: + tag.variables = None + + +def _add_deadline(tag_ref: ref[_CompileTag], at: float) -> None: + _deadlines.setdefault(int(at // _MINUTE_S), []).append(tag_ref) + + +def _newest_tag(messages: object) -> _CompileTag | None: + """The messages of a placeholder come from an older compile, so the newest compile, not the + last tagged message, names the prompt of the request. Its own messages are at the start or the + end of the list, unless placeholders put more than 1000 messages on each side of them.""" + if type(messages) is _CompiledMessage: + return getattr(messages, "compiled_by", None) + # Exact types: a subclass of list or tuple could run code of the app on each read. + if type(messages) is not list and type(messages) is not tuple: + return None + items = cast("list[object]", messages) + if len(items) > 2 * _TAG_LOOKUP: + items = [*items[:_TAG_LOOKUP], *items[-_TAG_LOOKUP:]] + newest: _CompileTag | None = None + for item in items: + tag = getattr(item, "compiled_by", None) if type(item) is _CompiledMessage else None + if tag is not None and (newest is None or tag.serial > newest.serial): + newest = tag + return newest + + def _fill(text: str, values: Mapping[str, Any]) -> str: def replace(match: re.Match[str]) -> str: value = values.get(match.group(1)) diff --git a/sdks/python/src/telemetry_dev/_spans.py b/sdks/python/src/telemetry_dev/_spans.py index 0ac1b26..6328fa3 100644 --- a/sdks/python/src/telemetry_dev/_spans.py +++ b/sdks/python/src/telemetry_dev/_spans.py @@ -38,6 +38,7 @@ ATTR_FINISH_REASONS, ATTR_OPERATION, ATTR_OUTPUT_TYPE, + ATTR_PROMPT_VARIABLES, ATTR_PROVIDER, ATTR_REQUEST_MODEL, ATTR_RESPONSE_ID, @@ -102,6 +103,7 @@ class _SpanState: operation: str capture_input: bool capture_output: bool + prompt_from_scope: bool = False output_chunk_last_ms: float | None = None output_chunk_count: int = 0 output_chunk_sum_s: float = 0 @@ -279,6 +281,9 @@ def _fields_to_attributes( raw = fields.get("attributes") if raw is not None: for key, value in raw.items(): + # The variables are input, also when the app passes them from prompt_attributes_for(). + if key == ATTR_PROMPT_VARIABLES and not state.capture_input: + continue attr = coerce_attr_value( value, max_len=sys.maxsize, @@ -296,6 +301,8 @@ def _apply_fields(client: Client, span: Span, state: _SpanState, fields: dict[st if name is not None: span.update_name(name) attrs = _fields_to_attributes(client, state, fields) + if state.prompt_from_scope: + attrs = _without_prompt(attrs) if attrs: span.set_attributes(attrs) @@ -304,6 +311,15 @@ def _apply_fields(client: Client, span: Span, state: _SpanState, fields: dict[st _record_error_on_span(client, span, state, error) +def _without_prompt( + attrs: dict[str, otel_types.AttributeValue], +) -> dict[str, otel_types.AttributeValue]: + """The prompt of with_prompt replaces each td.prompt.* attribute that the app passes, at start + and in each update, so that a span does not mix two prompts when the app calls + prompt_attributes_for outside with_prompt.""" + return {key: value for key, value in attrs.items() if not key.startswith("td.prompt.")} + + def _collect_fields( *, input: Any = NOT_GIVEN, @@ -744,11 +760,13 @@ def start_span( ctx, session_id if isinstance(session_id, str) else None, client.config.api_key ) initial.update(propagated) - initial.update( - prompt_attributes( - client, ctx, attrs.get(ATTR_OPERATION, operation), state.capture_input - ) + prompt = prompt_attributes( + client, ctx, attrs.get(ATTR_OPERATION, operation), state.capture_input ) + if prompt: + state.prompt_from_scope = True + attrs = _without_prompt(attrs) + initial.update(prompt) initial.update(attrs) span = client.tracer.start_span( name, diff --git a/sdks/python/tests/test_prompts.py b/sdks/python/tests/test_prompts.py index 764c4ef..74416e7 100644 --- a/sdks/python/tests/test_prompts.py +++ b/sdks/python/tests/test_prompts.py @@ -2,9 +2,14 @@ import asyncio import contextvars +import copy +import dataclasses +import gc import importlib import itertools +import json import os +import pickle import signal import socket import sys @@ -13,15 +18,17 @@ import traceback import urllib.error import urllib.request +import weakref from collections.abc import Callable, Generator from datetime import datetime, timedelta, timezone from email.utils import format_datetime from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from types import FrameType, SimpleNamespace -from typing import Any, NoReturn +from typing import Any, NoReturn, TypeAlias from urllib.parse import parse_qsl, urlsplit import pytest +from opentelemetry import context as otel_context from opentelemetry import trace import telemetry_dev @@ -29,8 +36,10 @@ from telemetry_dev import ( ChatPrompt, PromptError, + TextPrompt, get_prompt, get_prompt_async, + prompt_attributes_for, propagate_attributes, start_span, with_prompt, @@ -92,6 +101,34 @@ def empty_cache() -> None: prompts_module._cache.clear() # pyright: ignore[reportPrivateUsage] +def chat_prompt(*messages: dict[str, Any], name: str = "support/reply") -> ChatPrompt: + return ChatPrompt( + name=name, + version=12, + label="production", + messages=list(messages), + model=None, + parameters={}, + tools=[], + response_format={"type": "text"}, + is_fallback=False, + ) + + +def text_prompt(text: str) -> TextPrompt: + return TextPrompt( + name="support/reply", + version=12, + label="production", + text=text, + model=None, + parameters={}, + tools=[], + response_format={"type": "text"}, + is_fallback=False, + ) + + def query_of(path: str) -> dict[str, str]: return dict(parse_qsl(urlsplit(path).query)) @@ -1094,6 +1131,716 @@ def mask(value: Any, context: telemetry_dev.MaskContext) -> Any: assert "td.prompt.variables" not in (span.attributes or {}) +DESK_MESSAGES = [ + {"role": "system", "content": "Answer as the {{team}} desk."}, + {"role": "user", "content": "Order {{order}}: where is it?"}, +] + + +def test_prompt_attributes_for_names_the_compile_of_a_message_dict_and_a_copy_names_nothing( + make: MakeClient, +) -> None: + make() + compiled = chat_prompt(*DESK_MESSAGES).compile(team="billing", order="#4") + attributes = { + "td.prompt.name": "support/reply", + "td.prompt.version": 12, + "td.prompt.label": "production", + "td.prompt.variables": '{"team":"billing","order":"#4"}', + } + + assert prompt_attributes_for(compiled) == attributes + assert prompt_attributes_for(tuple(compiled)) == attributes + assert prompt_attributes_for(compiled[0]) == attributes + assert json.dumps(compiled[0]) == '{"role": "system", "content": "Answer as the billing desk."}' + + copies: list[object] = [ + [dict(message) for message in compiled], + [message.copy() for message in compiled], + [copy.copy(message) for message in compiled], + copy.deepcopy(compiled), + pickle.loads(pickle.dumps(compiled)), + json.loads(json.dumps(compiled)), + [message["content"] for message in compiled], + ] + for messages in copies: + assert prompt_attributes_for(messages) == {} + + summary = text_prompt("Summarize {{doc}}.").compile(doc="the policy") + assert prompt_attributes_for([summary]) == {} + + +def test_an_unrelated_message_with_the_text_of_a_compiled_message_is_not_tagged( + make: MakeClient, +) -> None: + make() + chat_prompt({"role": "system", "content": "You are a helpful assistant."}).compile() + + unrelated = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "What is 2 + 2?"}, + ] + assert prompt_attributes_for(unrelated) == {} + + +def test_two_tasks_with_one_text_and_different_other_variables_each_record_their_own( + make: MakeClient, +) -> None: + memory = make() + prompt = chat_prompt(*DESK_MESSAGES) + + # Both compiles make the system text "Answer as the billing desk.", and each task sends only it. + system_a = prompt.compile(team="billing", order="#A")[0] + system_b = prompt.compile(team="billing", order="#B")[0] + start_span("a", type="generation", attributes=prompt_attributes_for([system_a])).end() + start_span("b", type="generation", attributes=prompt_attributes_for([system_b])).end() + + assert { + span.name: (span.attributes or {})["td.prompt.variables"] + for span in memory.span_exporter.get_finished_spans() + } == { + "a": '{"team":"billing","order":"#A"}', + "b": '{"team":"billing","order":"#B"}', + } + + +def test_a_rollout_with_two_versions_of_the_same_messages_records_the_version_of_each_call( + server: CaptureServer, make: MakeClient, monkeypatch: pytest.MonkeyPatch +) -> None: + make() + server.replies = [ + ( + 200, + reply( + messages=DESK_MESSAGES, + rollout={ + "percent": 16, + "version": 13, + "messages": DESK_MESSAGES, + "text": None, + "model": {"provider": "openai", "name": "gpt-4o-mini"}, + "parameters": {}, + "tools": [], + "responseFormat": {"type": "text"}, + }, + ), + ) + ] + draws = iter([15, 16]) + + def draw(_stop: int) -> int: + return next(draws) + + monkeypatch.setattr(prompts_module.random, "randrange", draw) + rollout = get_prompt("support/reply", api_key="td_live_a", base_url=server.url) + production = get_prompt("support/reply", api_key="td_live_a", base_url=server.url) + from_rollout = rollout.compile(team="billing", order="#4") + from_production = production.compile(team="billing", order="#4") + + assert prompt_attributes_for(from_rollout)["td.prompt.version"] == 13 + assert prompt_attributes_for(from_production)["td.prompt.version"] == 12 + + +def test_the_newest_compiled_message_wins_also_after_170_messages_of_history( + make: MakeClient, +) -> None: + make() + prompt = chat_prompt( + {"role": "system", "content": "Answer as the {{team}} desk."}, + {"role": "placeholder", "name": "history"}, + {"role": "user", "content": "{{question}}"}, + ) + history = [ + {"role": "user" if index % 2 == 0 else "assistant", "content": f"message {index}"} + for index in range(170) + ] + + first_turn = prompt.compile(team="billing", question="Where is #4?", history=history) + second_turn = prompt.compile(team="returns", question="Can I return it?") + chat = [*first_turn, {"role": "assistant", "content": "On its way."}, *second_turn[1:]] + + assert '"team":"billing"' in str(prompt_attributes_for(first_turn)["td.prompt.variables"]) + assert '"team":"billing"' in str( + prompt_attributes_for([first_turn[0], *history])["td.prompt.variables"] + ) + assert ( + prompt_attributes_for(chat)["td.prompt.variables"] + == '{"team":"returns","question":"Can I return it?"}' + ) + + +def test_a_message_that_a_placeholder_brings_from_an_older_compile_does_not_win( + make: MakeClient, +) -> None: + make() + router = chat_prompt({"role": "user", "content": "{{question}}"}, name="router") + answer = chat_prompt( + {"role": "system", "content": "Answer for {{team}}."}, + {"role": "placeholder", "name": "history"}, + name="answer", + ) + chat = chat_prompt({"role": "user", "content": "{{question}}"}, name="chat") + summary = chat_prompt( + {"role": "system", "content": "Summarize the chat."}, + {"role": "placeholder", "name": "history"}, + name="summary", + ) + + question = router.compile(question="Where is #4?") + routed = answer.compile(team="billing", history=question) + # The last 1000 messages are all from chat; the summary message is the first. + conversation = [chat.compile(question=f"question {index}")[0] for index in range(1500)] + summarized = summary.compile(history=conversation) + + assert prompt_attributes_for(routed) == { + "td.prompt.name": "answer", + "td.prompt.version": 12, + "td.prompt.label": "production", + "td.prompt.variables": ( + '{"team":"billing","history":[{"role":"user","content":"Where is #4?"}]}' + ), + } + assert prompt_attributes_for(summarized)["td.prompt.name"] == "summary" + + +def test_prompt_attributes_for_looks_at_the_first_and_the_last_1000_messages( + make: MakeClient, +) -> None: + make() + system = chat_prompt(*DESK_MESSAGES).compile(team="billing", order="#4")[0] + + def around(before: int, after: int) -> list[dict[str, Any]]: + return [*[{"role": "user"}] * before, system, *[{"role": "user"}] * after] + + assert prompt_attributes_for(around(1000, 999)) != {} + assert prompt_attributes_for(around(999, 1001)) != {} + assert prompt_attributes_for(tuple(around(1001, 999))) != {} + assert prompt_attributes_for(around(1000, 1000)) == {} + + +def test_compile_records_its_variables_not_later_changes_of_the_mapping_or_of_its_lists( + make: MakeClient, +) -> None: + memory = make() + prompt = chat_prompt( + {"role": "placeholder", "name": "history"}, + {"role": "user", "content": "Order {{order}}"}, + ) + history = [{"role": "user", "content": "Hi"}] + variables: dict[str, Any] = {"order": "#A", "history": history} + + first = prompt.compile(variables) + variables["order"] = "#B" + history.append({"role": "assistant", "content": "Hello"}) + second = prompt.compile(variables) + variables["order"] = "#C" + history.append({"role": "user", "content": "Added after compile"}) + with with_prompt(prompt): + start_span("scope", type="generation").end() + + second_variables = ( + '{"order":"#B","history":[{"role":"user","content":"Hi"},' + '{"role":"assistant","content":"Hello"}]}' + ) + assert ( + prompt_attributes_for(first)["td.prompt.variables"] + == '{"order":"#A","history":[{"role":"user","content":"Hi"}]}' + ) + assert prompt_attributes_for(second)["td.prompt.variables"] == second_variables + [span] = memory.span_exporter.get_finished_spans() + assert (span.attributes or {})["td.prompt.variables"] == second_variables + + +def test_a_compiled_message_rebuilds_like_a_dict_and_the_rebuilt_one_has_no_tag( + make: MakeClient, +) -> None: + make() + compiled = chat_prompt(*DESK_MESSAGES).compile(team="billing", order="#4") + + @dataclasses.dataclass + class Conversation: + messages: list[dict[str, Any]] + + rebuilt = dataclasses.asdict(Conversation(messages=compiled))["messages"] + same_type = [type(message)(message) for message in compiled] + plain = [ + {"role": "system", "content": "Answer as the billing desk."}, + {"role": "user", "content": "Order #4: where is it?"}, + ] + + assert rebuilt == plain + assert same_type == plain + assert prompt_attributes_for(rebuilt) == {} + assert prompt_attributes_for(same_type) == {} + + +def test_with_prompt_wins_over_a_compiled_message_of_another_prompt_through_a_parent_and_in_updates( + make: MakeClient, +) -> None: + memory = make() + scoped = chat_prompt({"role": "user", "content": "Hi"}, name="scoped/prompt") + # Fetched by version, so it has no label that the label of the scope could hide. + other = dataclasses.replace( + chat_prompt({"role": "user", "content": "Cancel order {{order}}."}, name="other/prompt"), + version=3, + label=None, + ) + compiled = other.compile(order="#9") + + with with_prompt(scoped): + start_span("in scope", type="generation", attributes=prompt_attributes_for(compiled)).end() + parent = otel_context.get_current() + attributes = prompt_attributes_for(compiled) + start_span("with parent", type="generation", parent=parent, attributes=attributes).end() + with with_prompt(scoped): + start_span("updated", type="generation").update(attributes=attributes).end() + start_span("ended", type="generation").end(attributes=attributes) + with start_span("current", type="generation"): + telemetry_dev.update_current_span(attributes=attributes) + + assert attributes["td.prompt.name"] == "other/prompt" + spans = memory.span_exporter.get_finished_spans() + assert sorted(span.name for span in spans) == [ + "current", + "ended", + "in scope", + "updated", + "with parent", + ] + for span in spans: + assert { + key: value + for key, value in dict(span.attributes or {}).items() + if key.startswith("td.prompt.") + } == { + "td.prompt.name": "scoped/prompt", + "td.prompt.version": 12, + "td.prompt.label": "production", + } + + +def test_prompt_variables_follow_the_input_capture_of_the_client_and_of_the_span( + make: MakeClient, +) -> None: + compiled = chat_prompt({"role": "user", "content": "My card is {{card}}."}).compile( + card="4111 1111" + ) + + make(capture_input=False) + assert prompt_attributes_for(compiled) == { + "td.prompt.name": "support/reply", + "td.prompt.version": 12, + "td.prompt.label": "production", + } + + telemetry_dev.shutdown() + + def mask(value: Any, context: telemetry_dev.MaskContext) -> Any: + return {"card": "[masked]"} if context.key == "td.prompt.variables" else value + + memory = make(mask=mask) + attributes = prompt_attributes_for(compiled) + assert attributes["td.prompt.variables"] == '{"card":"[masked]"}' + + start_span("private", type="generation", capture_input=False, attributes=attributes).end() + with start_span("updated", type="generation", capture_input=False): + telemetry_dev.update_current_span(attributes=attributes) + + for span in memory.span_exporter.get_finished_spans(): + assert (span.attributes or {})["td.prompt.name"] == "support/reply" + assert "td.prompt.variables" not in (span.attributes or {}) + + +def test_prompt_attributes_for_runs_no_code_of_the_app(make: MakeClient) -> None: + make() + [message] = chat_prompt({"role": "user", "content": "Refund order {{order}}."}).compile( + order="#3" + ) + reads: list[str] = [] + + class Messages(list[object]): + def __getitem__(self, index: Any) -> Any: + reads.append("__getitem__") + return super().__getitem__(index) + + def __iter__(self) -> Any: + reads.append("__iter__") + return super().__iter__() + + def __reversed__(self) -> Any: + reads.append("__reversed__") + return super().__reversed__() + + assert prompt_attributes_for(Messages([message])) == {} + assert reads == [] + + +RAG_MESSAGES = [ + {"role": "system", "content": "Answer from {{documents}}."}, + {"role": "user", "content": "{{question}}"}, +] + + +class Document(dict[str, str]): + """A dict that a weak reference can point to.""" + + +def test_the_variables_of_a_compile_are_freed_with_its_messages(make: MakeClient) -> None: + make() + prompt = chat_prompt(*RAG_MESSAGES) + + document = Document(content="The refund policy") + freed = weakref.ref(document) + messages = prompt.compile(documents=[document], question="q") + del document + # with_prompt keeps the variables of the last compile of a prompt object. + prompt.compile(documents=[], question="other") + gc.collect() + assert freed() is not None + + assert len(messages) == 2 + del messages + gc.collect() + assert freed() is None + + +def test_with_input_capture_off_compiled_messages_keep_no_variables(make: MakeClient) -> None: + make(capture_input=False) + prompt = chat_prompt(*RAG_MESSAGES) + documents: list[weakref.ref[Document]] = [] + questions: list[dict[str, Any]] = [] + + for turn in range(200): + document = Document(content=f"document {turn}") + documents.append(weakref.ref(document)) + questions.append(prompt.compile(documents=[document], question=f"question {turn}")[-1]) + gc.collect() + + # with_prompt keeps the variables of the last compile of a prompt object. + assert sum(1 for document in documents if document() is not None) == 1 + + telemetry_dev.shutdown() + make() + assert prompt_attributes_for(questions) == { + "td.prompt.name": "support/reply", + "td.prompt.version": 12, + "td.prompt.label": "production", + } + + +CHAT_MESSAGES = ( + {"role": "system", "content": "Answer as the {{team}} desk."}, + {"role": "placeholder", "name": "history"}, + {"role": "user", "content": "{{question}}"}, +) + + +@dataclasses.dataclass +class ChatPrompts: + chat: ChatPrompt + question: ChatPrompt + rag: ChatPrompt + reply: ChatPrompt + + +def _answer(turn: int) -> dict[str, Any]: + return {"role": "assistant", "content": f"answer {turn}"} + + +# What a row gives for one turn of a chat: the request, the history for the next turn, and the kept +# message. The row puts `marker` in the variables of the compile that made the kept message. +ChatTurn: TypeAlias = tuple[list[dict[str, Any]], list[dict[str, Any]], dict[str, Any]] + + +def _new_list( + p: ChatPrompts, history: list[dict[str, Any]], turn: int, marker: Document +) -> ChatTurn: + request = p.chat.compile( + team="billing", history=history, question=f"question {turn}", marker=marker + ) + return request, [*history, request[-1], _answer(turn)], request[-1] + + +def _new_list_of_20( + p: ChatPrompts, history: list[dict[str, Any]], turn: int, marker: Document +) -> ChatTurn: + request = p.chat.compile( + team="billing", history=history, question=f"question {turn}", marker=marker + ) + return request, [*history, request[-1], _answer(turn)][-20:], request[-1] + + +def _full_list_as_copies( + p: ChatPrompts, history: list[dict[str, Any]], turn: int, marker: Document +) -> ChatTurn: + request = p.chat.compile( + team="billing", + history=[dict(message) for message in history], + question=f"question {turn}", + marker=marker, + ) + return request, [*history, request[-1], _answer(turn)], request[-1] + + +def _question_prompt_in_a_placeholder( + p: ChatPrompts, history: list[dict[str, Any]], turn: int, marker: Document +) -> ChatTurn: + [kept] = p.question.compile( + documents=f"document {turn}", question=f"question {turn}", marker=marker + ) + request = p.rag.compile(team="billing", history=history, current=[kept]) + return request, [*history, kept, _answer(turn)], kept + + +def _raw_questions_and_a_reply_prompt( + p: ChatPrompts, history: list[dict[str, Any]], turn: int, marker: Document +) -> ChatTurn: + request = p.chat.compile(team="billing", history=history, question=f"question {turn}") + [kept] = p.reply.compile(answer=f"answer {turn}", marker=marker) + return request, [*history, {"role": "user", "content": f"question {turn}"}, kept], kept + + +@pytest.mark.parametrize( + "run", + [ + _new_list, + _new_list_of_20, + _full_list_as_copies, + _question_prompt_in_a_placeholder, + _raw_questions_and_a_reply_prompt, + ], + ids=[ + "new-list", + "new-list-of-20", + "full-list-as-copies", + "question-prompt-in-a-placeholder", + "raw-questions-and-a-reply-prompt", + ], +) +def test_a_chat_history_keeps_the_variables_of_the_last_5_minutes( + make: MakeClient, + clock: list[float], + run: Callable[[ChatPrompts, list[dict[str, Any]], int, Document], ChatTurn], +) -> None: + make() + prompts = ChatPrompts( + chat=chat_prompt(*CHAT_MESSAGES), + question=chat_prompt({"role": "user", "content": "Context: {{documents}}\n\n{{question}}"}), + rag=chat_prompt( + {"role": "system", "content": "Answer as the {{team}} desk."}, + {"role": "placeholder", "name": "history"}, + {"role": "placeholder", "name": "current"}, + ), + reply=chat_prompt({"role": "assistant", "content": "{{answer}}"}), + ) + markers: list[weakref.ref[Document]] = [] + history: list[dict[str, Any]] = [] + request: list[dict[str, Any]] = [] + older: dict[str, Any] = {} + + # One turn in 2 minutes, for 300 turns. An unused variable marks the variables of each turn. + for turn in range(300): + if turn > 0: + clock[0] += 120 + marker = Document(content=f"turn {turn}") + markers.append(weakref.ref(marker)) + request, history, kept = run(prompts, history, turn, marker) + # The model call of the turn. + prompt_attributes_for(request) + if turn == 296: + older = kept + gc.collect() + + # Turns 297, 298, and 299: 4, 2, and 0 minutes old. + assert sum(1 for marker in markers if marker() is not None) == 3 + assert "question 299" in str(prompt_attributes_for(request)["td.prompt.variables"]) + assert prompt_attributes_for(older) == { + "td.prompt.name": "support/reply", + "td.prompt.version": 12, + "td.prompt.label": "production", + } + + +def _variables_of(messages: object) -> object: + return prompt_attributes_for(messages).get("td.prompt.variables") + + +REFUND_VARIABLES = '{"documents":"The refund policy","question":"Where is #4?"}' + + +@dataclasses.dataclass +class LivePrompts: + answer: ChatPrompt + guard: ChatPrompt + judge: ChatPrompt + other: ChatPrompt + chat: ChatPrompt + system: ChatPrompt + with_system: ChatPrompt + + +# A row compiles one or more requests and runs other compiles before their calls. It gives the +# variables that each call records and the variables of the compile of its request. +LiveCalls: TypeAlias = list[tuple[object, str]] + + +def _other_compiles_then_a_guard(p: LivePrompts) -> LiveCalls: + request = p.answer.compile(documents="The refund policy", question="Where is #4?") + for index in range(150): + p.other.compile(q=str(index)) + p.guard.compile(conversation=request) + return [(_variables_of(request), REFUND_VARIABLES)] + + +def _a_guard_before_each_agent_step(p: LivePrompts) -> LiveCalls: + conversation = p.answer.compile(documents="The refund policy", question="Where is #4?") + for step in range(150): + p.guard.compile(conversation=conversation) + conversation.append({"role": "assistant", "content": f"step {step}"}) + return [(_variables_of(conversation), REFUND_VARIABLES)] + + +def _a_guard_then_a_judge_of_the_guard(p: LivePrompts) -> LiveCalls: + request = p.answer.compile(documents="The refund policy", question="Where is #4?") + p.judge.compile(conversation=p.guard.compile(conversation=request)) + return [(_variables_of(request), REFUND_VARIABLES)] + + +def _turns_in_threads_that_share_a_history(p: LivePrompts) -> LiveCalls: + history: list[dict[str, Any]] = [] + turns = [threading.Event() for _ in range(4)] + compiled = threading.Barrier(3, timeout=5) + recorded: list[object] = [None] * 3 + + def turn(index: int) -> None: + turns[index].wait(5) + request = p.chat.compile(team="billing", history=history, question=f"q{index + 1}") + history.append(request[-1]) + turns[index + 1].set() + # A guard call or a queue: the call comes after the other turns compiled. + compiled.wait() + recorded[index] = _variables_of(request) + + threads = [threading.Thread(target=turn, args=(index,)) for index in range(3)] + for thread in threads: + thread.start() + turns[0].set() + for thread in threads: + thread.join(5) + return list( + zip( + recorded, + [ + '{"team":"billing","history":[],"question":"q1"}', + '{"team":"billing","history":[{"role":"user","content":"q1"}],"question":"q2"}', + '{"team":"billing","history":[{"role":"user","content":"q1"},' + '{"role":"user","content":"q2"}],"question":"q3"}', + ], + strict=True, + ) + ) + + +def _a_system_message_from_startup_after_a_chat_that_gets_it(p: LivePrompts) -> LiveCalls: + startup = p.system.compile(product="Acme") + history: list[dict[str, Any]] = [] + for turn in range(3): + request = p.with_system.compile( + system=startup, history=history, question=f"question {turn}" + ) + prompt_attributes_for(request) + history = [*history, request[-1], _answer(turn)] + return [ + ( + _variables_of([*startup, {"role": "user", "content": "Where is #4?"}]), + '{"product":"Acme"}', + ) + ] + + +@pytest.mark.parametrize( + "run", + [ + _other_compiles_then_a_guard, + _a_guard_before_each_agent_step, + _a_guard_then_a_judge_of_the_guard, + _turns_in_threads_that_share_a_history, + _a_system_message_from_startup_after_a_chat_that_gets_it, + ], + ids=[ + "150-other-compiles-then-a-guard", + "a-guard-before-each-of-150-agent-steps", + "a-guard-then-a-judge-of-the-guard", + "turns-in-threads-that-share-a-history", + "a-system-message-from-startup-after-a-chat-that-gets-it", + ], +) +def test_a_request_keeps_its_variables_until_its_call( + make: MakeClient, run: Callable[[LivePrompts], LiveCalls] +) -> None: + make() + calls = run( + LivePrompts( + answer=chat_prompt(*RAG_MESSAGES), + guard=chat_prompt( + {"role": "system", "content": "Is this safe?"}, + {"role": "placeholder", "name": "conversation"}, + name="guard", + ), + judge=chat_prompt( + {"role": "system", "content": "Check the answer of the guard."}, + {"role": "placeholder", "name": "conversation"}, + name="judge", + ), + other=chat_prompt({"role": "user", "content": "{{q}}"}, name="other"), + chat=chat_prompt(*CHAT_MESSAGES), + system=chat_prompt( + {"role": "system", "content": "You support {{product}}."}, name="system" + ), + with_system=chat_prompt( + {"role": "placeholder", "name": "system"}, + {"role": "placeholder", "name": "history"}, + {"role": "user", "content": "{{question}}"}, + ), + ) + ) + + for recorded, compiled in calls: + assert recorded == compiled + + +def test_a_compile_keeps_its_variables_for_5_minutes_after_it_or_after_the_last_call_that_names_it( + make: MakeClient, clock: list[float] +) -> None: + make() + chat = chat_prompt(*CHAT_MESSAGES) + history: list[dict[str, Any]] = [] + for turn in range(5): + request = chat.compile(team="billing", history=history, question=f"q{turn}") + prompt_attributes_for(request) + history = [*history, request[-1], _answer(turn)] + # These calls send only messages of the chat, so they name the newest compile among them. + summary = [*history, {"role": "user", "content": "Summarize the chat."}] + earlier = history[:-2] + start = clock[0] + + assert '"question":"q4"' in str(_variables_of(summary)) + assert '"question":"q3"' in str(_variables_of(earlier)) + + clock[0] = start + 299.5 + assert '"question":"q4"' in str(_variables_of(summary)) + clock[0] = start + 300 + assert _variables_of(earlier) is None + # The call 0.5 seconds before moved the deadline of q4. + assert '"question":"q4"' in str(_variables_of(summary)) + clock[0] = start + 600 + assert prompt_attributes_for(summary) == { + "td.prompt.name": "support/reply", + "td.prompt.version": 12, + "td.prompt.label": "production", + } + + def test_a_fallback_prompt_records_its_name_and_label_but_no_version( server: CaptureServer, make: MakeClient ) -> None: