mirror of https://github.com/buster-so/buster.git
87 lines
2.8 KiB
TypeScript
87 lines
2.8 KiB
TypeScript
import { createVertexAnthropic } from '@ai-sdk/google-vertex/anthropic';
|
|
import type { LanguageModelV1 } from '@ai-sdk/provider';
|
|
import { wrapAISDKModel } from 'braintrust';
|
|
|
|
export const vertexModel = (modelId: string): LanguageModelV1 => {
|
|
// Create a proxy that validates credentials on first use
|
|
let actualModel: LanguageModelV1 | null = null;
|
|
|
|
const getActualModel = () => {
|
|
if (!actualModel) {
|
|
const clientEmail = process.env.VERTEX_CLIENT_EMAIL;
|
|
let privateKey = process.env.VERTEX_PRIVATE_KEY;
|
|
const project = process.env.VERTEX_PROJECT;
|
|
|
|
if (!clientEmail || !privateKey || !project) {
|
|
throw new Error(
|
|
'Missing required environment variables: VERTEX_CLIENT_EMAIL or VERTEX_PRIVATE_KEY'
|
|
);
|
|
}
|
|
|
|
// Handle escaped newlines in private key
|
|
privateKey = privateKey.replace(/\\n/g, '\n');
|
|
|
|
const vertex = createVertexAnthropic({
|
|
baseURL: `https://aiplatform.googleapis.com/v1/projects/${project}/locations/global/publishers/anthropic/models`,
|
|
location: 'global',
|
|
project,
|
|
googleAuthOptions: {
|
|
credentials: {
|
|
client_email: clientEmail,
|
|
private_key: privateKey,
|
|
},
|
|
},
|
|
headers: {
|
|
'anthropic-beta': 'fine-grained-tool-streaming-2025-05-14',
|
|
},
|
|
fetch: ((url, options) => {
|
|
if (options?.body) {
|
|
try {
|
|
// Parse existing body if it's a string
|
|
const existingBody =
|
|
typeof options.body === 'string' ? JSON.parse(options.body) : options.body;
|
|
|
|
// Append disable_parallel_tool_use if tool_choice is present
|
|
const modifiedBody = {
|
|
...existingBody,
|
|
};
|
|
|
|
if (modifiedBody.tool_choice) {
|
|
modifiedBody.tool_choice = {
|
|
...modifiedBody.tool_choice,
|
|
disable_parallel_tool_use: true,
|
|
};
|
|
}
|
|
|
|
// Return modified options
|
|
return fetch(url, {
|
|
...options,
|
|
body: JSON.stringify(modifiedBody),
|
|
});
|
|
} catch (error) {
|
|
console.error('Failed to parse request body:', error);
|
|
// If body parsing fails, fall back to original request
|
|
return fetch(url, options);
|
|
}
|
|
}
|
|
|
|
// For requests without body, pass through unchanged
|
|
return fetch(url, options);
|
|
}) as typeof fetch,
|
|
});
|
|
|
|
// Wrap the model with Braintrust tracing
|
|
actualModel = wrapAISDKModel(vertex(modelId));
|
|
}
|
|
return actualModel;
|
|
};
|
|
|
|
// Create a proxy that delegates all calls to the actual model
|
|
return new Proxy({} as LanguageModelV1, {
|
|
get(_target, prop) {
|
|
const model = getActualModel();
|
|
return Reflect.get(model, prop);
|
|
},
|
|
});
|
|
};
|