diff --git a/.changeset/quiet-portkey-tracing.md b/.changeset/quiet-portkey-tracing.md new file mode 100644 index 000000000..dee22ef5e --- /dev/null +++ b/.changeset/quiet-portkey-tracing.md @@ -0,0 +1,5 @@ +--- +"braintrust": minor +--- + +Add automatic instrumentation and `wrapPortkey` for Portkey AI v3 chat completions, including streaming responses, token usage, tool calls, and errors. diff --git a/js/package.json b/js/package.json index a72758597..61ac882be 100644 --- a/js/package.json +++ b/js/package.json @@ -212,6 +212,7 @@ "eslint-plugin-n": "^18.2.1", "eslint-plugin-node-import": "^1.0.5", "openai": "6.25.0", + "portkey-ai": "3.1.0", "rollup": "^4.60.3", "tar": "^7.5.22", "tinybench": "^4.0.1", diff --git a/js/src/auto-instrumentations/configs/all.ts b/js/src/auto-instrumentations/configs/all.ts index deb7aad8f..777a76d75 100644 --- a/js/src/auto-instrumentations/configs/all.ts +++ b/js/src/auto-instrumentations/configs/all.ts @@ -20,6 +20,7 @@ import { googleADKConfigs } from "./google-adk"; import { googleGenerativeAIConfigs } from "./google-generative-ai"; import { googleGenAIConfigs } from "./google-genai"; import { groqConfigs } from "./groq"; +import { portkeyConfigs } from "./portkey"; import { huggingFaceConfigs } from "./huggingface"; import { huggingFaceTransformersConfigs } from "./huggingface-transformers"; import { langchainConfigs } from "./langchain"; @@ -108,6 +109,7 @@ const defaultInstrumentationConfigGroups: readonly InstrumentationConfigGroup[] { integrations: ["googleADK"], configs: googleADKConfigs }, { integrations: ["cohere"], configs: cohereConfigs }, { integrations: ["groq"], configs: groqConfigs }, + { integrations: ["portkey"], configs: portkeyConfigs }, { integrations: ["genkit"], configs: genkitConfigs, diff --git a/js/src/auto-instrumentations/configs/portkey.ts b/js/src/auto-instrumentations/configs/portkey.ts new file mode 100644 index 000000000..00263188d --- /dev/null +++ b/js/src/auto-instrumentations/configs/portkey.ts @@ -0,0 +1,18 @@ +import type { InstrumentationConfig } from "../orchestrion-js"; +import { portkeyChannels } from "../../instrumentation/plugins/portkey-channels"; + +export const portkeyConfigs: InstrumentationConfig[] = [ + { + channelName: portkeyChannels.chatCompletionsCreate.channelName, + module: { + name: "portkey-ai", + versionRange: ">=3.0.0 <4.0.0", + filePath: "dist/src/apis/chatCompletions.js", + }, + functionQuery: { + className: "ChatCompletions", + methodName: "create", + kind: "Async", + }, + }, +]; diff --git a/js/src/exports.ts b/js/src/exports.ts index 98c512d01..9769544f7 100644 --- a/js/src/exports.ts +++ b/js/src/exports.ts @@ -240,6 +240,7 @@ export { wrapCohere } from "./wrappers/cohere"; export { wrapVoyageAI } from "./wrappers/voyageai"; export { wrapTypeSafe } from "./wrappers/typesafe"; export { wrapGroq } from "./wrappers/groq"; +export { wrapPortkey } from "./wrappers/portkey"; export { wrapBedrockRuntime } from "./wrappers/bedrock-runtime"; export { wrapCopilotClient } from "./wrappers/github-copilot"; export { diff --git a/js/src/instrumentation/braintrust-plugin.test.ts b/js/src/instrumentation/braintrust-plugin.test.ts index f6510cbe2..d051d6283 100644 --- a/js/src/instrumentation/braintrust-plugin.test.ts +++ b/js/src/instrumentation/braintrust-plugin.test.ts @@ -17,6 +17,7 @@ import { MistralPlugin } from "./plugins/mistral-plugin"; import { OllamaPlugin } from "./plugins/ollama-plugin"; import { CoherePlugin } from "./plugins/cohere-plugin"; import { GroqPlugin } from "./plugins/groq-plugin"; +import { PortkeyPlugin } from "./plugins/portkey-plugin"; import { GitHubCopilotPlugin } from "./plugins/github-copilot-plugin"; import { LangChainPlugin } from "./plugins/langchain-plugin"; import { LangSmithPlugin } from "./plugins/langsmith-plugin"; @@ -37,6 +38,10 @@ function createPluginClassMock() { }); } +vi.mock("./plugins/portkey-plugin", () => ({ + PortkeyPlugin: createPluginClassMock(), +})); + // Preserve the re-exported utility functions (parseMetricsFromUsage, etc.) // while mocking out the OpenAIPlugin class — those utilities are also // imported and tested at the bottom of this file via braintrust-plugin's @@ -673,6 +678,22 @@ describe("BraintrustPlugin", () => { expect(MistralPlugin).toHaveBeenCalledTimes(1); }); + it("enables and tears down Portkey by default", () => { + const plugin = new BraintrustPlugin(); + plugin.enable(); + const portkey = vi.mocked(PortkeyPlugin).mock.results[0].value; + expect(portkey.enable).toHaveBeenCalledTimes(1); + plugin.disable(); + expect(portkey.disable).toHaveBeenCalledTimes(1); + }); + + it("does not enable Portkey when portkey: false", () => { + const plugin = new BraintrustPlugin({ integrations: { portkey: false } }); + plugin.enable(); + expect(PortkeyPlugin).not.toHaveBeenCalled(); + plugin.disable(); + }); + it("should not create GitHubCopilot plugin when gitHubCopilot: false", () => { const plugin = new BraintrustPlugin({ integrations: { gitHubCopilot: false }, @@ -751,6 +772,7 @@ describe("BraintrustPlugin", () => { ollama: false, cohere: false, groq: false, + portkey: false, gitHubCopilot: false, langchain: false, langsmith: false, @@ -777,6 +799,7 @@ describe("BraintrustPlugin", () => { expect(OllamaPlugin).not.toHaveBeenCalled(); expect(CoherePlugin).not.toHaveBeenCalled(); expect(GroqPlugin).not.toHaveBeenCalled(); + expect(PortkeyPlugin).not.toHaveBeenCalled(); expect(GitHubCopilotPlugin).not.toHaveBeenCalled(); expect(LangChainPlugin).not.toHaveBeenCalled(); expect(LangSmithPlugin).not.toHaveBeenCalled(); diff --git a/js/src/instrumentation/braintrust-plugin.ts b/js/src/instrumentation/braintrust-plugin.ts index 93015d351..adfd294d2 100644 --- a/js/src/instrumentation/braintrust-plugin.ts +++ b/js/src/instrumentation/braintrust-plugin.ts @@ -19,6 +19,7 @@ import { OllamaPlugin } from "./plugins/ollama-plugin"; import { GoogleADKPlugin } from "./plugins/google-adk-plugin"; import { CoherePlugin } from "./plugins/cohere-plugin"; import { GroqPlugin } from "./plugins/groq-plugin"; +import { PortkeyPlugin } from "./plugins/portkey-plugin"; import { BedrockRuntimePlugin } from "./plugins/bedrock-runtime-plugin"; import { GenkitPlugin } from "./plugins/genkit-plugin"; import { GitHubCopilotPlugin } from "./plugins/github-copilot-plugin"; @@ -80,6 +81,7 @@ export class BraintrustPlugin extends BasePlugin { private googleADKPlugin: GoogleADKPlugin | null = null; private coherePlugin: CoherePlugin | null = null; private groqPlugin: GroqPlugin | null = null; + private portkeyPlugin: PortkeyPlugin | null = null; private bedrockRuntimePlugin: BedrockRuntimePlugin | null = null; private genkitPlugin: GenkitPlugin | null = null; private gitHubCopilotPlugin: GitHubCopilotPlugin | null = null; @@ -223,6 +225,11 @@ export class BraintrustPlugin extends BasePlugin { this.groqPlugin.enable(); } + if (integrations.portkey !== false) { + this.portkeyPlugin = new PortkeyPlugin(); + this.portkeyPlugin.enable(); + } + if ( integrations.bedrock !== false && integrations.awsBedrock !== false && @@ -401,6 +408,11 @@ export class BraintrustPlugin extends BasePlugin { this.groqPlugin = null; } + if (this.portkeyPlugin) { + this.portkeyPlugin.disable(); + this.portkeyPlugin = null; + } + if (this.bedrockRuntimePlugin) { this.bedrockRuntimePlugin.disable(); this.bedrockRuntimePlugin = null; diff --git a/js/src/instrumentation/config.ts b/js/src/instrumentation/config.ts index 057191b68..b4349dda8 100644 --- a/js/src/instrumentation/config.ts +++ b/js/src/instrumentation/config.ts @@ -42,6 +42,7 @@ export interface InstrumentationIntegrationsConfig { ollama?: boolean; cohere?: boolean; groq?: boolean; + portkey?: boolean; bedrock?: boolean; awsBedrock?: boolean; awsBedrockRuntime?: boolean; @@ -138,6 +139,8 @@ const envIntegrationAliases: Record< cohere: "cohere", groq: "groq", "groq-sdk": "groq", + portkey: "portkey", + "portkey-ai": "portkey", bedrock: "bedrock", "aws-bedrock": "awsBedrock", awsbedrock: "awsBedrock", @@ -196,6 +199,7 @@ export function getDefaultInstrumentationIntegrations(): Record< ollama: true, cohere: true, groq: true, + portkey: true, bedrock: true, awsBedrock: true, awsBedrockRuntime: true, diff --git a/js/src/instrumentation/plugins/instrumentation-names.test.ts b/js/src/instrumentation/plugins/instrumentation-names.test.ts index fef4d8bfb..e89f5adc9 100644 --- a/js/src/instrumentation/plugins/instrumentation-names.test.ts +++ b/js/src/instrumentation/plugins/instrumentation-names.test.ts @@ -21,6 +21,7 @@ import { googleADKChannels } from "./google-adk-channels"; import { googleGenerativeAIChannels } from "./google-generative-ai-channels"; import { googleGenAIChannels } from "./google-genai-channels"; import { groqChannels } from "./groq-channels"; +import { portkeyChannels } from "./portkey-channels"; import { huggingFaceChannels } from "./huggingface-channels"; import { langGraphSDKChannels } from "./langgraph-sdk-channels"; import { langChainChannels } from "./langchain-channels"; @@ -69,6 +70,7 @@ describe("built-in instrumentation provenance names", () => { ], [googleGenAIChannels.generateContent, INSTRUMENTATION_NAMES.GOOGLE_GENAI], [groqChannels.chatCompletionsCreate, INSTRUMENTATION_NAMES.GROQ], + [portkeyChannels.chatCompletionsCreate, INSTRUMENTATION_NAMES.PORTKEY], [huggingFaceChannels.chatCompletion, INSTRUMENTATION_NAMES.HUGGINGFACE], [langGraphSDKChannels.wait, INSTRUMENTATION_NAMES.LANGGRAPH_SDK], [langChainChannels.configure, INSTRUMENTATION_NAMES.LANGCHAIN], diff --git a/js/src/instrumentation/plugins/portkey-channels.ts b/js/src/instrumentation/plugins/portkey-channels.ts new file mode 100644 index 000000000..0260afc81 --- /dev/null +++ b/js/src/instrumentation/plugins/portkey-channels.ts @@ -0,0 +1,23 @@ +import { channel, defineChannels } from "../core/channel-definitions"; +import { INSTRUMENTATION_NAMES } from "../../span-origin"; +import type { + PortkeyChatCompletion, + PortkeyChatCompletionChunk, + PortkeyChatCreateParams, + PortkeyChatStream, +} from "../../vendor-sdk-types/portkey"; + +// Portkey's chat API uses OpenAI-compatible requests and responses, but takes +// gateway configuration and HTTP request options as separate arguments. +export const portkeyChannels = defineChannels( + "portkey-ai", + { + chatCompletionsCreate: channel< + [PortkeyChatCreateParams, unknown?, unknown?], + PortkeyChatCompletion | PortkeyChatStream, + Record, + PortkeyChatCompletionChunk + >({ channelName: "chat.completions.create", kind: "async" }), + }, + { instrumentationName: INSTRUMENTATION_NAMES.PORTKEY }, +); diff --git a/js/src/instrumentation/plugins/portkey-plugin.test.ts b/js/src/instrumentation/plugins/portkey-plugin.test.ts new file mode 100644 index 000000000..e4ef200d5 --- /dev/null +++ b/js/src/instrumentation/plugins/portkey-plugin.test.ts @@ -0,0 +1,342 @@ +import { readFileSync } from "node:fs"; +import { createRequire } from "node:module"; +import { runInNewContext } from "node:vm"; +import { + afterEach, + beforeAll, + beforeEach, + describe, + expect, + it, + vi, +} from "vitest"; +import { Portkey } from "portkey-ai"; +import { _exportsForTestingOnly, initLogger, traced } from "../../logger"; +import { configureNode } from "../../node/config"; +import { create } from "../../auto-instrumentations/orchestrion-js"; +import { getDefaultInstrumentationConfigs } from "../../auto-instrumentations/configs/all"; +import { readDisabledInstrumentationEnvConfig } from "../config"; +import { wrapPortkey } from "../../wrappers/portkey"; + +configureNode(); + +const require = createRequire(import.meta.url); +const filename = require.resolve("portkey-ai/dist/src/apis/chatCompletions.js"); +const source = readFileSync(filename, "utf8"); +const { version } = require("portkey-ai/package.json") as { version: string }; +const request = { + model: "test-model", + messages: [{ role: "user" as const, content: "Say hello" }], + max_tokens: 16, +}; +const usage = { prompt_tokens: 8, completion_tokens: 2, total_tokens: 10 }; +const completion = { + id: "test-completion", + object: "chat.completion", + model: "test-model", + choices: [ + { + index: 0, + message: { role: "assistant", content: "Hello!" }, + finish_reason: "stop", + }, + ], + usage, +}; + +// Execute the installed SDK's actual resource through the same transformer used +// by the loader/bundlers. Only HTTP is stubbed: these are hermetic unit tests, +// not provider recordings. +function autoClient(disabledList?: string) { + const client = new Portkey({ + apiKey: "test-key", + baseURL: "https://portkey.test/v1", + }); + const transformer = create( + getDefaultInstrumentationConfigs({ + disabledIntegrationConfig: + readDisabledInstrumentationEnvConfig(disabledList).integrations, + }), + ).getTransformer("portkey-ai", version, "dist/src/apis/chatCompletions.js"); + const code = transformer?.transform(source, "cjs").code ?? source; + const exports: { Chat?: new (client: Portkey) => Portkey["chat"] } = {}; + runInNewContext(code, { + exports, + require: createRequire(filename), + globalThis, + Map, + }); + client.chat = new exports.Chat!(client); + return client; +} + +describe("Portkey instrumentation", () => { + let backgroundLogger: ReturnType< + typeof _exportsForTestingOnly.useTestBackgroundLogger + >; + const fetchMock = vi.fn(); + + beforeAll(async () => { + await _exportsForTestingOnly.simulateLoginForTests(); + }); + beforeEach(() => { + backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); + initLogger({ + projectName: "portkey-plugin.test.ts", + projectId: "test-project-id", + }); + fetchMock.mockReset(); + fetchMock.mockImplementation(async () => Response.json(completion)); + vi.stubGlobal("fetch", fetchMock); + }); + afterEach(() => { + vi.unstubAllGlobals(); + _exportsForTestingOnly.clearTestBackgroundLogger(); + }); + + for (const mode of ["wrapped", "auto", "wrapped-auto"] as const) { + describe(mode, () => { + function makeClient() { + const client = + mode === "wrapped" + ? new Portkey({ + apiKey: "test-key", + baseURL: "https://portkey.test/v1", + }) + : autoClient(); + return mode === "auto" ? client : wrapPortkey(client); + } + + it("records one child span with input, output, usage and Portkey provenance", async () => { + const client = makeClient(); + await traced( + async () => { + const response = await client.chat.completions.create(request); + expect(response.choices).toEqual(completion.choices); + expect(typeof response.getHeaders).toBe("function"); + }, + { name: "parent" }, + ); + const spans = await backgroundLogger.drain(); + const calls = spans.filter( + (span) => + "span_attributes" in span && + span.span_attributes?.name === "portkey.chat.completions.create", + ); + expect(calls).toHaveLength(1); + const parent = spans.find( + (span) => + "span_attributes" in span && + span.span_attributes?.name === "parent", + ); + expect(calls[0]).toMatchObject({ + input: request.messages, + output: completion.choices, + metadata: { + provider: "portkey", + model: request.model, + max_tokens: 16, + }, + metrics: { prompt_tokens: 8, completion_tokens: 2, tokens: 10 }, + span_parents: [ + parent && "span_id" in parent ? parent.span_id : undefined, + ], + span_attributes: { type: "llm" }, + context: { span_origin: { instrumentation: { name: "portkey" } } }, + }); + expect(fetchMock).toHaveBeenCalledTimes(1); + }); + + it("forwards gateway configuration and request options without logging credentials", async () => { + await makeClient().chat.completions.create( + request, + { + apiKey: "private-gateway-key", + virtualKey: "private-provider-key", + provider: "openai", + }, + { extraHeaders: { "x-test-request": "forwarded" } }, + ); + const [url, options] = fetchMock.mock.calls[0]; + expect(String(url)).toBe("https://portkey.test/v1/chat/completions"); + expect(new Headers(options?.headers).get("x-test-request")).toBe( + "forwarded", + ); + const spans = await backgroundLogger.drain(); + expect(spans).toHaveLength(1); + expect(JSON.stringify(spans)).not.toContain("private-gateway-key"); + expect(JSON.stringify(spans)).not.toContain("private-provider-key"); + }); + + it("aggregates streaming text, tool arguments and final usage without consuming chunks", async () => { + const chunks = [ + { + choices: [ + { + index: 0, + delta: { + role: "assistant", + content: "Hello", + tool_calls: [ + { + index: 0, + id: "call-1", + type: "function", + function: { name: "weather", arguments: '{"city":' }, + }, + ], + }, + }, + ], + }, + { + choices: [ + { + index: 0, + delta: { + content: "!", + tool_calls: [ + { index: 0, function: { arguments: '"Paris"}' } }, + ], + }, + finish_reason: "tool_calls", + }, + ], + }, + { choices: [], usage }, + ]; + fetchMock.mockImplementation( + async () => + new Response( + chunks + .map((chunk) => `data: ${JSON.stringify(chunk)}\n\n`) + .join("") + "data: [DONE]\n\n", + { headers: { "content-type": "text/event-stream" } }, + ), + ); + const stream = await makeClient().chat.completions.create({ + ...request, + stream: true, + }); + const received = []; + for await (const chunk of stream) received.push(chunk); + expect(received).toEqual(chunks); + const spans = await backgroundLogger.drain(); + expect(spans).toHaveLength(1); + expect(spans[0]).toMatchObject({ + output: [ + { + index: 0, + message: { + role: "assistant", + content: "Hello!", + tool_calls: [ + { + id: "call-1", + type: "function", + function: { + name: "weather", + arguments: '{"city":"Paris"}', + }, + }, + ], + }, + finish_reason: "tool_calls", + }, + ], + metrics: { + prompt_tokens: 8, + completion_tokens: 2, + tokens: 10, + time_to_first_token: expect.any(Number), + }, + }); + }); + + it("preserves SDK errors and finishes the error span", async () => { + fetchMock.mockImplementation(async () => + Response.json( + { error: { message: "invalid model" } }, + { status: 400 }, + ), + ); + await expect( + makeClient().chat.completions.create(request), + ).rejects.toThrow("invalid model"); + const spans = await backgroundLogger.drain(); + expect(spans).toHaveLength(1); + expect(spans[0]).toMatchObject({ + error: expect.stringContaining("invalid model"), + metrics: { end: expect.any(Number) }, + }); + }); + + it("finishes a partially consumed stream", async () => { + fetchMock.mockImplementation( + async () => + new Response( + 'data: {"choices":[{"index":0,"delta":{"role":"assistant","content":"Hello"}}]}\n\ndata: [DONE]\n\n', + { headers: { "content-type": "text/event-stream" } }, + ), + ); + const stream = await makeClient().chat.completions.create({ + ...request, + stream: true, + }); + for await (const chunk of stream) { + expect(chunk.choices).toHaveLength(1); + break; + } + const spans = await backgroundLogger.drain(); + expect(spans).toHaveLength(1); + expect(spans[0]).toMatchObject({ + metrics: { end: expect.any(Number) }, + }); + }); + + it("propagates stream errors and closes the span", async () => { + fetchMock.mockImplementation( + async () => + new Response( + 'event: error\ndata: {"message":"stream failed"}\n\n', + { headers: { "content-type": "text/event-stream" } }, + ), + ); + const stream = await makeClient().chat.completions.create({ + ...request, + stream: true, + }); + await expect(async () => { + for await (const chunk of stream) void chunk; + }).rejects.toThrow("stream failed"); + const spans = await backgroundLogger.drain(); + expect(spans).toHaveLength(1); + expect(spans[0]).toMatchObject({ + error: expect.stringContaining("stream failed"), + metrics: { end: expect.any(Number) }, + }); + }); + }); + } + + it.each(["portkey", "portkey-ai"])( + "honors the %s disable alias", + async (alias) => { + const response = await autoClient(alias).chat.completions.create(request); + expect(response.choices).toEqual(completion.choices); + expect(await backgroundLogger.drain()).toEqual([]); + }, + ); + + it("keeps wrapping idempotent and preserves the client API", async () => { + const original = new Portkey({ apiKey: "test-key" }); + const client = wrapPortkey(original); + expect(wrapPortkey(original)).toBe(client); + expect(wrapPortkey(client)).toBe(client); + expect(client).toBeInstanceOf(Portkey); + expect(client.chat.completions.create).toBe(client.chat.completions.create); + expect(client.models).toBe(original.models); + await client.chat.completions.create(request); + expect(await backgroundLogger.drain()).toHaveLength(1); + }); +}); diff --git a/js/src/instrumentation/plugins/portkey-plugin.ts b/js/src/instrumentation/plugins/portkey-plugin.ts new file mode 100644 index 000000000..4a47f39f5 --- /dev/null +++ b/js/src/instrumentation/plugins/portkey-plugin.ts @@ -0,0 +1,36 @@ +import { BasePlugin } from "../core"; +import { traceStreamingChannel, unsubscribeAll } from "../core/channel-tracing"; +import { SpanTypeAttribute } from "../../../util/index"; +import { parseMetricsFromUsage } from "../../openai-utils"; +import { getCurrentUnixTimestamp } from "../../util"; +import { aggregateChatCompletionChunks } from "./openai-plugin"; +import { extractOpenAIChatInput } from "./openai-span-data"; +import { portkeyChannels } from "./portkey-channels"; + +export class PortkeyPlugin extends BasePlugin { + protected onEnable(): void { + this.unsubscribers.push( + traceStreamingChannel(portkeyChannels.chatCompletionsCreate, { + name: "portkey.chat.completions.create", + type: SpanTypeAttribute.LLM, + extractInput: ([params]) => { + const { input, metadata } = extractOpenAIChatInput(params); + return { input, metadata: { ...metadata, provider: "portkey" } }; + }, + extractOutput: (result) => result?.choices, + extractMetrics: (result, startTime) => { + const metrics = parseMetricsFromUsage(result?.usage); + if (startTime) { + metrics.time_to_first_token = getCurrentUnixTimestamp() - startTime; + } + return metrics; + }, + aggregateChunks: aggregateChatCompletionChunks, + }), + ); + } + + protected onDisable(): void { + unsubscribeAll(this.unsubscribers); + } +} diff --git a/js/src/span-origin.ts b/js/src/span-origin.ts index 755fc43f5..b9ec4762e 100644 --- a/js/src/span-origin.ts +++ b/js/src/span-origin.ts @@ -40,6 +40,7 @@ export const INSTRUMENTATION_NAMES = { OPENROUTER: "openrouter", OPENROUTER_AGENT: "openrouter-agent", PI_CODING_AGENT: "pi-coding-agent", + PORTKEY: "portkey", STRANDS_AGENT_SDK: "strands-agent-sdk", TEMPORAL: "temporal", TYPESAFE: "typesafe", diff --git a/js/src/vendor-sdk-types/portkey.ts b/js/src/vendor-sdk-types/portkey.ts new file mode 100644 index 000000000..ea02a02f2 --- /dev/null +++ b/js/src/vendor-sdk-types/portkey.ts @@ -0,0 +1,24 @@ +import type { + OpenAIChatCompletion, + OpenAIChatCompletionChunk, + OpenAIChatCreateParams, + OpenAIChatStream, +} from "./openai"; + +// Only the OpenAI-compatible portion of Portkey consumed by instrumentation. +export type PortkeyChatCreateParams = OpenAIChatCreateParams; +export type PortkeyChatCompletion = OpenAIChatCompletion; +export type PortkeyChatCompletionChunk = OpenAIChatCompletionChunk; +export type PortkeyChatStream = OpenAIChatStream; + +export interface PortkeyClient { + chat: { + completions: { + create( + body: PortkeyChatCreateParams, + params?: unknown, + options?: unknown, + ): Promise; + }; + }; +} diff --git a/js/src/wrappers/portkey.ts b/js/src/wrappers/portkey.ts new file mode 100644 index 000000000..86529e101 --- /dev/null +++ b/js/src/wrappers/portkey.ts @@ -0,0 +1,69 @@ +import { portkeyChannels } from "../instrumentation/plugins/portkey-channels"; +import { runWithAutoInstrumentationSuppressed } from "../instrumentation/auto-instrumentation-suppression"; +import type { ArgsOf } from "../instrumentation/core/channel-definitions"; +import type { PortkeyClient } from "../vendor-sdk-types/portkey"; +import { isObject } from "../../util/index"; +import { debugLogger } from "../debug-logger"; + +type Create = PortkeyClient["chat"]["completions"]["create"]; + +const proxyCache = new WeakMap(); + +/** + * Trace chat completions from a client created with `new Portkey(...)`. + * Supports both regular responses and `stream: true`. + */ +export function wrapPortkey(client: T): T { + if (!isPortkeyClient(client)) { + debugLogger.warn("Unsupported Portkey client. Not wrapping."); + return client; + } + const cached = proxyCache.get(client); + if (cached) return cached as T; + + const proxy = proxyPath(client, ["chat", "completions", "create"]); + proxyCache.set(client, proxy); + proxyCache.set(proxy, proxy); + return proxy; +} + +function isPortkeyClient(value: unknown): value is PortkeyClient { + return ( + isObject(value) && + isObject(value.chat) && + isObject(value.chat.completions) && + typeof value.chat.completions.create === "function" + ); +} + +function proxyPath(target: T, path: string[]): T { + const cache = new Map(); + return new Proxy(target, { + get(object, prop, receiver) { + const value = Reflect.get(object, prop, receiver); + if (prop !== path[0]) return value; + if (cache.has(prop)) return cache.get(prop); + + if (path.length === 1 && typeof value === "function") { + const wrapped = ( + ...args: ArgsOf + ) => + portkeyChannels.chatCompletionsCreate.tracePromise( + () => + runWithAutoInstrumentationSuppressed(() => + Reflect.apply(value as Create, object, args), + ), + { arguments: args }, + ); + cache.set(prop, wrapped); + return wrapped; + } + if (value !== null && typeof value === "object") { + const wrapped = proxyPath(value, path.slice(1)); + cache.set(prop, wrapped); + return wrapped; + } + return value; + }, + }); +} diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 1e69fe489..4fe6bac02 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -515,6 +515,9 @@ importers: openai: specifier: 6.25.0 version: 6.25.0(ws@8.21.3)(zod@3.25.76) + portkey-ai: + specifier: 3.1.0 + version: 3.1.0(zod@3.25.76) rollup: specifier: ^4.60.3 version: 4.60.3 @@ -2232,6 +2235,10 @@ packages: engines: {node: '>=0.4.0'} hasBin: true + agentkeepalive@4.6.0: + resolution: {integrity: sha512-kja8j7PjmncONqaTsB8fQ+wE2mSU2DJ9D4XKoJ5PFWIdRMa6SLSN1ff4mOr4jCbfRSsxR4keIiySJU0N9T5hIQ==} + engines: {node: '>= 8.0.0'} + ai@6.0.37: resolution: {integrity: sha512-GUMBYx0TXKxXRcWy6DY3ryHJZ7cATxc99WPT7WCD8hGCYkdhFVItKPTtINeTlK+FlUomDqzjPGtiDcVftcw//w==} engines: {node: '>=18'} @@ -2578,6 +2585,10 @@ packages: resolution: {integrity: sha512-ZmdL2rui+eB2YwhsWzjInR8LldtZHGDoQ1ugH85ppHKwpUHL7j7rN0Ti9NCnGiQbhaZ11FpR+7ao1dNsmduNUg==} engines: {node: '>=12'} + dotenv@16.6.1: + resolution: {integrity: sha512-uBq4egWHTcTt33a72vpSG0z3HnPuIl6NqYcTrKEg2azoEyl2hpW0zqlxysq2pK9HlDIHyHyakeYaYnSAwd8bow==} + engines: {node: '>=12'} + dotenv@17.3.1: resolution: {integrity: sha512-IO8C/dzEb6O3F9/twg6ZLXz164a2fhTnEWb95H23Dm4OuN+92NmEAlTrupP9VW6Jm3sO26tQlqyvyi4CsnY9GA==} engines: {node: '>=12'} @@ -3025,6 +3036,9 @@ packages: resolution: {integrity: sha512-tsYlhAYpjCKa//8rXZ9DqKEawhPoSytweBC2eNvcaDK+57RZLHGqNs3PZTQO6yekLFSuvA6AlnAfrw1uBvtb+Q==} hasBin: true + humanize-ms@1.2.1: + resolution: {integrity: sha512-Fl70vYtsAFb/C06PTS9dZBo7ihau+Tu/DNCk/OyHhea07S+aeMWpFFkUaXRa8fI+ScZbEI8dfSxwY7gxZ9SAVQ==} + husky@9.1.7: resolution: {integrity: sha512-5gs5ytaNjBrh5Ow3zrvdUUY+0VxIuWVL4i9irt6friV+BqdCfmV11CQTWMiBYWHbXhco+J1kHfTOUkePhCDvMA==} engines: {node: '>=18'} @@ -3563,6 +3577,18 @@ packages: zod: optional: true + openai@6.8.1: + resolution: {integrity: sha512-ACifslrVgf+maMz9vqwMP4+v9qvx5Yzssydizks8n+YUJ6YwUoxj51sKRQ8HYMfR6wgKLSIlaI108ZwCk+8yig==} + hasBin: true + peerDependencies: + ws: ^8.18.0 + zod: ^3.25 || ^4.0 + peerDependenciesMeta: + ws: + optional: true + zod: + optional: true + optionator@0.9.4: resolution: {integrity: sha512-6IpQ7mKUxRcZNLIObR0hz7lxsapSSIYNZJwXPGeF0mTVqGKFIXj1DQcMoT22S3ROcLyY/rz0PWaWZ9ayWmad9g==} engines: {node: '>= 0.8.0'} @@ -3702,6 +3728,9 @@ packages: resolution: {integrity: sha512-Nc3IT5yHzflTfbjgqWcCPpo7DaKy4FnpB0l/zCAW0Tc7jxAiuqSxHasntB3D7887LSrA93kDJ9IXovxJYxyLCA==} engines: {node: '>=4'} + portkey-ai@3.1.0: + resolution: {integrity: sha512-ER6N6TBVAvil/BLBLOWzBEzWeBn9N1St5snG+m2mDxpP1lKdvUbbsog3c6xoLML5AxZkMlt3oDv3nUqeT5x62g==} + postcss-load-config@6.0.1: resolution: {integrity: sha512-oPtTM4oerL+UXmx+93ytZVN82RrlY/wPUV8IeDxFrzIjXOLF1pN+EmKPLbubvKHT2HC20xXsCAH2Z+CKV6Oz/g==} engines: {node: '>= 18'} @@ -6252,6 +6281,10 @@ snapshots: acorn@8.16.0: {} + agentkeepalive@4.6.0: + dependencies: + humanize-ms: 1.2.1 + ai@6.0.37(zod@3.25.76): dependencies: '@ai-sdk/gateway': 3.0.15(zod@3.25.76) @@ -6537,6 +6570,8 @@ snapshots: dotenv@16.4.5: {} + dotenv@16.6.1: {} + dotenv@17.3.1: {} dunder-proto@1.0.1: @@ -7042,6 +7077,10 @@ snapshots: human-id@4.1.3: {} + humanize-ms@1.2.1: + dependencies: + ms: 2.1.3 + husky@9.1.7: {} hyperdyperid@1.2.0: {} @@ -7594,6 +7633,11 @@ snapshots: zod: 3.25.76 optional: true + openai@6.8.1(ws@8.21.3)(zod@3.25.76): + optionalDependencies: + ws: 8.21.3 + zod: 3.25.76 + optionator@0.9.4: dependencies: deep-is: 0.1.4 @@ -7729,6 +7773,17 @@ snapshots: pluralize@8.0.0: {} + portkey-ai@3.1.0(zod@3.25.76): + dependencies: + agentkeepalive: 4.6.0 + dotenv: 16.6.1 + openai: 6.8.1(ws@8.21.3)(zod@3.25.76) + ws: 8.21.3 + transitivePeerDependencies: + - bufferutil + - utf-8-validate + - zod + postcss-load-config@6.0.1(jiti@2.6.1)(postcss@8.5.26)(tsx@4.21.0)(yaml@2.9.0): dependencies: lilconfig: 3.1.3