Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 41 additions & 11 deletions js/plugins/anthropic/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -44,24 +44,38 @@ export type { AnthropicCacheControl, AnthropicCitation } from './types.js';
export { cacheControl } from './utils.js';

/**
* Gets or creates an Anthropic client instance.
* Supports test client injection for internal testing.
* Gets the test client if injected (for testing only).
* @internal
*/
function getAnthropicClient(options?: PluginOptions): Anthropic {
// Check for test client injection first (internal use only)
function getTestClient(options?: PluginOptions): Anthropic | undefined {
const internalOptions = options as InternalPluginOptions | undefined;
if (internalOptions?.[__testClient]) {
return internalOptions[__testClient];
return internalOptions?.[__testClient];
}

/**
* Validates the API key configuration at plugin initialization.
* When apiKey is false, validation is deferred to request time.
*
* @throws Error if API key is required but not available
*/
function validateApiKey(options?: PluginOptions): void {
// If apiKey is explicitly false, defer validation to request time
if (options?.apiKey === false) {
return;
}

// Production path: create real client
// Check for test client injection (for testing only)
if (getTestClient(options)) {
return;
}

// Validate that we have an API key available
const apiKey = options?.apiKey || process.env.ANTHROPIC_API_KEY;
if (!apiKey) {
throw new Error(
'Please pass in the API key or set the ANTHROPIC_API_KEY environment variable'
);
}
return new Anthropic({ apiKey });
}

/**
Expand Down Expand Up @@ -94,8 +108,12 @@ function getAnthropicClient(options?: PluginOptions): Anthropic {
* ```
*/
function anthropicPlugin(options?: PluginOptions): GenkitPluginV2 {
const client = getAnthropicClient(options);
// Validate API key at plugin init (unless deferred with apiKey: false)
validateApiKey(options);

const pluginApiKey = options?.apiKey;
const defaultApiVersion = options?.apiVersion;
const testClient = getTestClient(options);

let listActionsCache: ActionMetadata[] | null = null;

Expand All @@ -106,8 +124,9 @@ function anthropicPlugin(options?: PluginOptions): GenkitPluginV2 {
for (const name of Object.keys(KNOWN_CLAUDE_MODELS)) {
const action = claudeModel({
name,
client,
pluginApiKey,
defaultApiVersion,
testClient,
});
actions.push(action);
}
Expand All @@ -119,14 +138,25 @@ function anthropicPlugin(options?: PluginOptions): GenkitPluginV2 {
const modelName = name.startsWith('anthropic/') ? name.slice(10) : name;
return claudeModel({
name: modelName,
client,
pluginApiKey,
defaultApiVersion,
testClient,
});
}
return undefined;
},
list: async () => {
if (listActionsCache) return listActionsCache;
// For listing, we need a client. Create one if we have an API key available.
const listApiKey =
pluginApiKey !== false
? pluginApiKey || process.env.ANTHROPIC_API_KEY
: undefined;
if (!listApiKey && !testClient) {
// Can't list models without an API key
return [];
}
const client = testClient ?? new Anthropic({ apiKey: listApiKey });
listActionsCache = await listActions(client);
return listActionsCache;
},
Expand Down
44 changes: 30 additions & 14 deletions js/plugins/anthropic/src/models.ts
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
* limitations under the License.
*/

import Anthropic from '@anthropic-ai/sdk';
import type {
GenerateRequest,
GenerateResponseData,
Expand All @@ -33,6 +34,7 @@ import {
AnthropicBaseConfigSchemaType,
AnthropicConfigSchema,
AnthropicThinkingConfigSchema,
calculateApiKey,
resolveBetaEnabled,
type ClaudeModelParams,
type ClaudeRunnerParams,
Expand Down Expand Up @@ -198,14 +200,7 @@ export function claudeRunner<TConfigSchema extends z.ZodTypeAny>(
params: ClaudeRunnerParams,
configSchema: TConfigSchema
) {
const { defaultApiVersion, ...runnerParams } = params;

if (!runnerParams.client) {
throw new Error('Anthropic client is required to create a runner');
}

let stableRunner: Runner | null = null;
let betaRunner: BetaRunner | null = null;
const { defaultApiVersion, pluginApiKey, testClient, name } = params;

return async (
request: GenerateRequest<TConfigSchema>,
Expand All @@ -223,13 +218,28 @@ export function claudeRunner<TConfigSchema extends z.ZodTypeAny>(
const normalizedRequest = request as unknown as GenerateRequest<
typeof AnthropicConfigSchema
>;

// Determine the client to use
// Test client takes precedence, otherwise calculate API key at request time
const client = testClient
? testClient
: new Anthropic({
apiKey: calculateApiKey(
pluginApiKey,
normalizedRequest.config?.apiKey
),
});
Comment on lines +224 to +231
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Creating a new Anthropic client instance for every request can lead to performance issues. Each new instance may establish a new TCP connection and perform a new TLS handshake, adding significant latency to each API call. This prevents connection reuse (HTTP keep-alive).

Consider caching Anthropic client instances based on their API key. The cache could be maintained at the plugin level and passed down to the runner. This would allow reusing clients for subsequent requests that use the same API key, improving performance by reusing existing connections.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My take is that the bottleneck would be inference, optimising here for TLS handshake is adding complexity for marginal gain. Not caching the client is a pattern in other plugins.


const isBeta = resolveBetaEnabled(
normalizedRequest.config,
defaultApiVersion
);

// Create runner with the client
const runner = isBeta
? (betaRunner ??= new BetaRunner(runnerParams))
: (stableRunner ??= new Runner(runnerParams));
? new BetaRunner({ name, client })
: new Runner({ name, client });

return runner.run(normalizedRequest, {
streamingRequested,
sendChunk,
Expand Down Expand Up @@ -265,17 +275,22 @@ export function claudeModelReference(
}

/**
* Defines a Claude model with the given name and Anthropic client.
* Defines a Claude model with the given name and API key configuration.
* Accepts any model name and lets the API validate it. If the model is in KNOWN_CLAUDE_MODELS, uses that modelRef
* for better defaults; otherwise creates a generic model reference.
*/
export function claudeModel(
params: ClaudeModelParams
): ModelAction<z.ZodTypeAny> {
const { name, client: runnerClient, defaultApiVersion: apiVersion } = params;
const {
name,
pluginApiKey,
defaultApiVersion: apiVersion,
testClient,
} = params;
// Use supported model ref if available, otherwise create generic model ref
const knownModelRef = KNOWN_CLAUDE_MODELS[name];
let modelInfo = knownModelRef
const modelInfo = knownModelRef
? knownModelRef.info
: GENERIC_CLAUDE_MODEL_INFO;
const configSchema = knownModelRef?.configSchema ?? AnthropicConfigSchema;
Expand All @@ -291,8 +306,9 @@ export function claudeModel(
claudeRunner(
{
name,
client: runnerClient,
pluginApiKey,
defaultApiVersion: apiVersion,
testClient,
},
configSchema
)
Expand Down
4 changes: 2 additions & 2 deletions js/plugins/anthropic/src/runner/base.ts
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ import {
MediaSchema,
MediaType,
MediaTypeSchema,
type ClaudeRunnerParams,
type RunnerConstructorParams,
type ThinkingConfig,
} from '../types.js';

Expand Down Expand Up @@ -66,7 +66,7 @@ export abstract class BaseRunner<ApiTypes extends RunnerTypes> {
*/
protected readonly DEFAULT_MAX_OUTPUT_TOKENS = 4096;

constructor(params: ClaudeRunnerParams) {
constructor(params: RunnerConstructorParams) {
this.name = params.name;
this.client = params.client;
}
Expand Down
10 changes: 6 additions & 4 deletions js/plugins/anthropic/src/runner/beta.ts
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ import { KNOWN_CLAUDE_MODELS, extractVersion } from '../models.js';
import {
AnthropicConfigSchema,
type AnthropicDocumentOptions,
type ClaudeRunnerParams,
type RunnerConstructorParams,
} from '../types.js';
import { removeUndefinedProperties } from '../utils.js';
import { BaseRunner } from './base.js';
Expand Down Expand Up @@ -142,7 +142,7 @@ interface BetaRunnerTypes extends RunnerTypes {
* Runner for the Anthropic Beta API.
*/
export class BetaRunner extends BaseRunner<BetaRunnerTypes> {
constructor(params: ClaudeRunnerParams) {
constructor(params: RunnerConstructorParams) {
super(params);
}

Expand Down Expand Up @@ -318,12 +318,13 @@ export class BetaRunner extends BaseRunner<BetaRunnerTypes> {
// Need to extract topP and topK from request.config to avoid duplicate properties being added to the body
// This happens because topP and topK have different property names (top_p and top_k) in the Anthropic API.
// Thinking is extracted separately to avoid type issues.
// ApiVersion is extracted separately as it's not a valid property for the Anthropic API.
// apiVersion and apiKey are extracted separately as they're not valid properties for the Anthropic API.
const {
topP,
topK,
apiVersion: _1,
thinking: _2,
apiKey: _3,
...restConfig
} = request.config ?? {};

Expand Down Expand Up @@ -375,12 +376,13 @@ export class BetaRunner extends BaseRunner<BetaRunnerTypes> {
// Need to extract topP and topK from request.config to avoid duplicate properties being added to the body
// This happens because topP and topK have different property names (top_p and top_k) in the Anthropic API.
// Thinking is extracted separately to avoid type issues.
// ApiVersion is extracted separately as it's not a valid property for the Anthropic API.
// apiVersion and apiKey are extracted separately as they're not valid properties for the Anthropic API.
const {
topP,
topK,
apiVersion: _1,
thinking: _2,
apiKey: _3,
...restConfig
} = request.config ?? {};

Expand Down
10 changes: 6 additions & 4 deletions js/plugins/anthropic/src/runner/stable.ts
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ import { KNOWN_CLAUDE_MODELS, extractVersion } from '../models.js';
import {
AnthropicConfigSchema,
type AnthropicDocumentOptions,
type ClaudeRunnerParams,
type RunnerConstructorParams,
} from '../types.js';
import { removeUndefinedProperties } from '../utils.js';
import { BaseRunner } from './base.js';
Expand Down Expand Up @@ -85,7 +85,7 @@ interface RunnerTypes extends BaseRunnerTypes {
}

export class Runner extends BaseRunner<RunnerTypes> {
constructor(params: ClaudeRunnerParams) {
constructor(params: RunnerConstructorParams) {
super(params);
}

Expand Down Expand Up @@ -237,12 +237,13 @@ export class Runner extends BaseRunner<RunnerTypes> {
// Need to extract topP and topK from request.config to avoid duplicate properties being added to the body
// This happens because topP and topK have different property names (top_p and top_k) in the Anthropic API.
// Thinking is extracted separately to avoid type issues.
// ApiVersion is extracted separately as it's not a valid property for the Anthropic API.
// apiVersion and apiKey are extracted separately as they're not valid properties for the Anthropic API.
const {
topP,
topK,
apiVersion: _1,
thinking: _2,
apiKey: _3,
...restConfig
} = request.config ?? {};

Expand Down Expand Up @@ -288,12 +289,13 @@ export class Runner extends BaseRunner<RunnerTypes> {
// Need to extract topP and topK from request.config to avoid duplicate properties being added to the body
// This happens because topP and topK have different property names (top_p and top_k) in the Anthropic API.
// Thinking is extracted separately to avoid type issues.
// ApiVersion is extracted separately as it's not a valid property for the Anthropic API.
// apiVersion and apiKey are extracted separately as they're not valid properties for the Anthropic API.
const {
topP,
topK,
apiVersion: _1,
thinking: _2,
apiKey: _3,
...restConfig
} = request.config ?? {};

Expand Down
Loading