diff --git a/packages/ai/src/protocols/gemini.ts b/packages/ai/src/protocols/gemini.ts index 7dd3ff4be89..23a02b8c6a8 100644 --- a/packages/ai/src/protocols/gemini.ts +++ b/packages/ai/src/protocols/gemini.ts @@ -166,6 +166,7 @@ const GeminiGenerationConfig = Schema.Struct({ const GeminiBodyFields = { cachedContent: Schema.optional(Schema.String), contents: Schema.Array(GeminiContent), + labels: Schema.optional(Schema.Record(Schema.String, Schema.String)), safetySettings: optionalArray(GeminiSafetySetting), serviceTier: Schema.optional(Schema.String), systemInstruction: Schema.optional(GeminiSystemInstruction), diff --git a/packages/ai/src/providers/google-vertex.ts b/packages/ai/src/providers/google-vertex.ts index 4180412d860..da0f96b82bc 100644 --- a/packages/ai/src/providers/google-vertex.ts +++ b/packages/ai/src/providers/google-vertex.ts @@ -1,14 +1,21 @@ +import { Effect } from "effect" import type { ProviderPackage } from "../provider-package.js" import { Gemini } from "../protocols/gemini.js" +import { ProviderShared } from "../protocols/shared.js" import { Auth } from "../route/auth.js" import { Route, type RouteDefaultsInput } from "../route/client.js" import { Endpoint } from "../route/endpoint.js" import { Framing } from "../route/framing.js" -import { ProviderID, type ModelID } from "../schema/index.js" +import { ProviderID, type LLMRequest, type ModelID, type ProviderOptions } from "../schema/index.js" import { GoogleVertexShared } from "./google-vertex-shared.js" -export type GeminiOptionsInput = Gemini.OptionsInput -export type GeminiProviderOptionsInput = Gemini.ProviderOptionsInput +export interface GeminiOptionsInput extends Gemini.OptionsInput { + readonly labels?: Readonly> +} + +export type GeminiProviderOptionsInput = ProviderOptions & { + readonly gemini?: GeminiOptionsInput +} export const id = ProviderID.make("google-vertex") @@ -17,7 +24,7 @@ export type Config = RouteDefaultsInput & readonly baseURL?: string readonly location?: string readonly project?: string - readonly providerOptions?: Gemini.ProviderOptionsInput + readonly providerOptions?: GeminiProviderOptionsInput } export type Settings = ProviderPackage.Settings & @@ -28,14 +35,31 @@ export type Settings = ProviderPackage.Settings & readonly baseURL?: string readonly location?: string readonly project?: string - readonly providerOptions?: Gemini.ProviderOptionsInput + readonly providerOptions?: GeminiProviderOptionsInput } +const fromRequest = Effect.fn("GoogleVertex.fromRequest")(function* (request: LLMRequest) { + const body = yield* Gemini.protocol.body.from(request) + const value = request.providerOptions?.gemini?.labels + const labels = ProviderShared.isRecord(value) + ? Object.fromEntries(Object.entries(value).filter((entry): entry is [string, string] => typeof entry[1] === "string")) + : undefined + return { ...body, labels } +}) + +const protocol = { + ...Gemini.protocol, + body: { + ...Gemini.protocol.body, + from: fromRequest, + }, +} + const route = Route.make({ id: "google-vertex-gemini", provider: id, providerMetadataKey: "google", - protocol: Gemini.protocol, + protocol, endpoint: Endpoint.path(({ request }) => { const model = String(request.model.id) return `/${model.startsWith("endpoints/") ? model : `models/${model}`}:streamGenerateContent?alt=sse` @@ -78,7 +102,7 @@ export const configure = (input: Config = {}) => { return { id, model: (modelID: string | ModelID) => - configuredRoute(input, modelID).model({ id: modelID }), + configuredRoute(input, modelID).model({ id: modelID }), configure, } } @@ -87,7 +111,7 @@ export const provider = { id, configure, } -export const model: ProviderPackage.Definition["model"] = ( +export const model: ProviderPackage.Definition["model"] = ( modelID, settings, ) => { diff --git a/packages/ai/test/provider/google-vertex.test.ts b/packages/ai/test/provider/google-vertex.test.ts index 5deacdab4d2..c0b310ffd12 100644 --- a/packages/ai/test/provider/google-vertex.test.ts +++ b/packages/ai/test/provider/google-vertex.test.ts @@ -54,6 +54,27 @@ describe("Google Vertex providers", () => { }), ) + it.effect("adds billing labels to Vertex Gemini requests", () => + Effect.gen(function* () { + const prepared = yield* compileRequest( + LLM.request({ + model: GoogleVertex.configure({ + accessToken: "vertex-token", + project: "vertex-project", + providerOptions: { + gemini: { labels: { component: "opencode", environment: "test" } }, + }, + }).model("gemini-3.5-flash"), + prompt: "Say hello.", + }), + ) + + expect(prepared.body).toMatchObject({ + labels: { component: "opencode", environment: "test" }, + }) + }), + ) + it.effect("projects Anthropic Messages onto the Vertex raw-predict API", () => Effect.gen(function* () { const model = GoogleVertexMessages.configure({