Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

improvement: Add vertex embeddings support #622

Merged
merged 7 commits into from
Oct 1, 2024
Merged

improvement: Add vertex embeddings support #622

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift click to select a range
d6778f3
Add vertex embeddings support
Sep 22, 2024
557fcfc
Merge branch 'main' into feat/vertex-embed
elentaure Sep 23, 2024
0c18b6e
Merge branch 'main' into feat/vertex-embed
elentaure Sep 26, 2024
3ac8217
move GoogleErrorResponseTransform to utils
Sep 27, 2024
41baf0e
Modify task_type param
Sep 27, 2024
71cc978
Remove custom vertex params
Sep 28, 2024
0f6a66d
Merge branch 'main' into feat/vertex-embed
VisargD Oct 1, 2024
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
22 changes: 15 additions & 7 deletions src/providers/google-vertex-ai/api.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -50,14 50,24 @@ export const GoogleApiConfig: ProviderAPIConfig = {

const { provider, model } = getModelAndProvider(inputModel as string);
const projectRoute = getProjectRoute(providerOptions, inputModel as string);
const googleUrlMap = new Map<string, string>([
[
'chatComplete',
`${projectRoute}/publishers/${provider}/models/${model}:generateContent`,
],
[
'stream-chatComplete',
`${projectRoute}/publishers/${provider}/models/${model}:streamGenerateContent?alt=sse`,
],
[
'embed',
`${projectRoute}/publishers/${provider}/models/${model}:predict`,
],
]);

switch (provider) {
case 'google': {
if (mappedFn === 'chatComplete') {
return `${projectRoute}/publishers/${provider}/models/${model}:generateContent`;
} else if (mappedFn === 'stream-chatComplete') {
return `${projectRoute}/publishers/${provider}/models/${model}:streamGenerateContent?alt=sse`;
}
return googleUrlMap.get(mappedFn) || `${projectRoute}`;
}

case 'anthropic': {
Expand All @@ -72,8 82,6 @@ export const GoogleApiConfig: ProviderAPIConfig = {
return `${projectRoute}/endpoints/openapi/chat/completions`;
}

// Embed API is not yet implemented in the gateway
// This may be as easy as copy-paste from Google provider, but needs to be tested
default:
return `${projectRoute}`;
}
Expand Down
93 changes: 93 additions & 0 deletions src/providers/google-vertex-ai/embed.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 1,93 @@
import { ErrorResponse, ProviderConfig } from '../types';
import {
EmbedResponse,
EmbedResponseData,
EmbedParams,
} from '../../types/embedRequestBody';
import {
GoogleErrorResponse,
EmbedInstancesData,
GoogleEmbedResponse,
} from './types';
import { GOOGLE_VERTEX_AI } from '../../globals';
import { generateInvalidProviderResponseError } from '../utils';
import { GoogleErrorResponseTransform } from './utils';

enum TASK_TYPE {
RETRIEVAL_QUERY = 'RETRIEVAL_QUERY',
RETRIEVAL_DOCUMENT = 'RETRIEVAL_DOCUMENT',
SEMANTIC_SIMILARITY = 'SEMANTIC_SIMILARITY',
CLASSIFICATION = 'CLASSIFICATION',
CLUSTERING = 'CLUSTERING',
QUESTION_ANSWERING = 'QUESTION_ANSWERING',
FACT_VERIFICATION = 'FACT_VERIFICATION',
CODE_RETRIEVAL_QUERY = 'CODE_RETRIEVAL_QUERY',
}

interface GoogleEmbedParams extends EmbedParams {
task_type: TASK_TYPE | string;
}

export const GoogleEmbedConfig: ProviderConfig = {
Copy link
Contributor

Choose a reason for hiding this comment

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

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,
  },
};

Copy link
Contributor

Choose a reason for hiding this comment

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

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,
  },
};

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Modified according to suggestions. Please check if now works ok with task type and the SDK

input: {
param: 'instances',
required: true,
transform: (params: GoogleEmbedParams): 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,
task_type: params.task_type,
});
}
return instances;
},
},
parameters: {
param: 'parameters',
required: false,
},
};

export const GoogleEmbedResponseTransform: (
response: GoogleEmbedResponse | GoogleErrorResponse,
responseStatus: number
) => EmbedResponse | ErrorResponse = (response, responseStatus) => {
if (responseStatus !== 200) {
const errorResposne = GoogleErrorResponseTransform(
response as GoogleErrorResponse
);
if (errorResposne) return errorResposne;
}

if ('predictions' in response) {
const data: EmbedResponseData[] = [];
let tokenCount = 0;
response.predictions.forEach((prediction, index) => {
data.push({
object: 'embedding',
embedding: prediction.embeddings.values,
index: index,
});
tokenCount = prediction.embeddings.statistics.token_count;
});
return {
object: 'list',
data: data,
model: '', // Todo: find a way to send the google embedding model name back
usage: {
prompt_tokens: tokenCount,
total_tokens: response.metadata.billableCharacterCount,
},
};
}

return generateInvalidProviderResponseError(response, GOOGLE_VERTEX_AI);
};
3 changes: 3 additions & 0 deletions src/providers/google-vertex-ai/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 12,7 @@ import {
VertexLlamaChatCompleteStreamChunkTransform,
} from './chatComplete';
import { getModelAndProvider } from './utils';
import { GoogleEmbedConfig, GoogleEmbedResponseTransform } from './embed';

const VertexConfig: ProviderConfigs = {
api: VertexApiConfig,
Expand All @@ -24,9 25,11 @@ const VertexConfig: ProviderConfigs = {
return {
chatComplete: VertexGoogleChatCompleteConfig,
api: GoogleApiConfig,
embed: GoogleEmbedConfig,
responseTransforms: {
'stream-chatComplete': GoogleChatCompleteStreamChunkTransform,
chatComplete: GoogleChatCompleteResponseTransform,
embed: GoogleEmbedResponseTransform,
},
};
case 'anthropic':
Expand Down
22 changes: 22 additions & 0 deletions src/providers/google-vertex-ai/types.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,3 68,25 @@ export interface VertexLlamaChatCompleteStreamChunk {
created?: number;
provider?: string;
}

export interface EmbedInstancesData {
task_type: string;
content: string;
}

interface EmbedPredictionsResponse {
embeddings: {
values: number[];
statistics: {
truncated: string;
token_count: number;
};
};
}

export interface GoogleEmbedResponse {
predictions: EmbedPredictionsResponse[];
metadata: {
billableCharacterCount: number;
};
}
24 changes: 24 additions & 0 deletions src/providers/google-vertex-ai/utils.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 1,8 @@
import { GoogleErrorResponse } from './types';
import { generateErrorResponse } from '../utils';
import { GOOGLE_VERTEX_AI } from '../../globals';
import { ErrorResponse } from '../types';

/**
* Encodes an object as a Base64 URL-encoded string.
* @param obj The object to encode.
Expand Down Expand Up @@ -164,3 169,22 @@ export const getMimeType = (url: string) => {
] as keyof typeof fileExtensionMimeTypeMap;
return fileExtensionMimeTypeMap[extension] || 'image/jpeg';
};

export const GoogleErrorResponseTransform: (
response: GoogleErrorResponse,
provider?: string
) => ErrorResponse | undefined = (response, provider = GOOGLE_VERTEX_AI) => {
if ('error' in response) {
return generateErrorResponse(
{
message: response.error.message ?? '',
type: response.error.status ?? null,
param: null,
code: response.error.status ?? null,
},
provider
);
}

return undefined;
};
Loading