improvement: Add vertex embeddings support #622
Conversation
There was a problem hiding this comment.
Tried making a request to this provider in local, the request failed when the request is openai compliant, suggested changes in comments
Failing request
curl --location 'http://localhost:8787/v1/embeddings' \
--header 'x-portkey-provider: vertex-ai' \
--header 'x-portkey-vertex-region: us-central1' \
--header 'Content-Type: application/json' \
--header 'x-portkey-api-key: jd-' \
--header 'Authorization: ya29.c....' \
--header 'x-portkey-vertex-project-id: {{YOUR_PROJECT_ID}}' \
--data-raw '{
"model": "textembedding-gecko@001",
"input": "Hello this is a test",
}'
| @@ -68,3 70,49 @@ export interface VertexLlamaChatCompleteStreamChunk { | |||
| created?: number; | |||
| provider?: string; | |||
| } | |||
|
|
|||
| export const GoogleErrorResponseTransform: ( | |||
There was a problem hiding this comment.
this method can be moved to a utils.ts file (preferred), or just the embed.ts file if it's only being used there.
| import { GOOGLE_VERTEX_AI } from '../../globals'; | ||
| import { generateInvalidProviderResponseError } from '../utils'; | ||
|
|
||
| export const GoogleEmbedConfig: ProviderConfig = { |
There was a problem hiding this comment.
The request structure is incorrect, you've typed params as having type VertexEmbedParams, this wouldn't be OpenAI compliant. The gateway transforms an OpenAI embeddings request to a Vertex embeddings request, so the code should be as below
export interface EmbedInstancesData {
content: string;
}
export const GoogleEmbedConfig: ProviderConfig = {
input: {
param: 'instances',
required: true,
transform: (params: EmbedParams): Array<EmbedInstancesData> => {
const instances = Array<EmbedInstancesData>();
if (Array.isArray(params.input)) {
params.input.forEach((text) => {
instances.push({
content: text,
});
});
} else {
instances.push({
content: params.input,
});
}
return instances;
},
},
parameters: {
param: 'parameters',
required: false,
},
};
There was a problem hiding this comment.
Better implementation, with support for task_type
export interface EmbedInstancesData {
content: string;
}
enum TASK_TYPE {...}
interface GoogleEmbedParams extends EmbedParams {
task_type: TASK_TYPE | string
}
export const GoogleEmbedConfig: ProviderConfig = {
input: {
param: 'instances',
required: true,
transform: (params: EmbedParams): Array<EmbedInstancesData> => {
const instances = Array<EmbedInstancesData>();
if (Array.isArray(params.input)) {
params.input.forEach((text) => {
instances.push({
content: text,
task_type: params.task_type
});
});
} else {
instances.push({
content: params.input,
});
}
return instances;
},
},
parameters: {
param: 'parameters',
required: false,
},
};
There was a problem hiding this comment.
Modified according to suggestions. Please check if now works ok with task type and the SDK
|
Thanks for the PR! We will merge this today. |
|
Hey @elentaure - I noticed one difference while going through the vertex embeddings documentation. The gateway is currently picking up the Documentation for the tokenCount parameter: https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/text-embeddings-api#response_body I will raise a quick PR to make this change. |
Title:
Description:
Based on the embeddings code for google ai studio and modified to adapt to the vertex api differences.