Skip to content

Commit

Permalink
feat(vertex): support token counting
Browse files Browse the repository at this point in the history
  • Loading branch information
RobertCraigie authored and stainless-app[bot] committed Dec 17, 2024
1 parent 7fd4830 commit a4a3729
Showing 1 changed file with 32 additions and 1 deletion.
33 changes: 32 additions & 1 deletion packages/vertex-sdk/src/client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,7 @@ export class AnthropicVertex extends Core.APIClient {
this._authClientPromise = this._auth.getClient();
}

messages: Resources.Messages = new Resources.Messages(this);
messages: MessagesResource = makeMessagesResource(this);
beta: BetaResource = makeBetaResource(this);

protected override defaultQuery(): Core.DefaultQuery | undefined {
Expand Down Expand Up @@ -147,10 +147,41 @@ export class AnthropicVertex extends Core.APIClient {
options.path = `/projects/${this.projectId}/locations/${this.region}/publishers/anthropic/models/${model}:${specifier}`;
}

if (
options.path === '/v1/messages/count_tokens' ||
(options.path == '/v1/messages/count_tokens?beta=true' && options.method === 'post')
) {
if (!this.projectId) {
throw new Error(
'No projectId was given and it could not be resolved from credentials. The client should be instantiated with the `projectId` option or the `ANTHROPIC_VERTEX_PROJECT_ID` environment variable should be set.',
);
}

if (Core.isObj(options.body)) {
options.body['anthropic_version'] = undefined;
}

options.path = `/projects/${this.projectId}/locations/${this.region}/publishers/anthropic/models/count-tokens:rawPredict`;
}

return super.buildRequest(options);
}
}

/**
* The Vertex SDK does not currently support the Batch API.
*/
type MessagesResource = Omit<Resources.Messages, 'batches'>;

function makeMessagesResource(client: AnthropicVertex): MessagesResource {
const resource = new Resources.Messages(client);

// @ts-expect-error we're deleting non-optional properties
delete resource.batches;

return resource;
}

/**
* The Vertex API does not currently support prompt caching, token counting or the Batch API.
*/
Expand Down

0 comments on commit a4a3729

Please sign in to comment.