Press n or j to go to the next uncovered block, b, p or k for the previous block.
| 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x | import type {
GoogleEmbedResponse,
GoogleErrorResponse,
} from '@api/ai-providers/google/types';
import { generateInvalidProviderResponseError } from '@api/utils/ai-provider';
import type {
AIProviderFunctionConfig,
ResponseTransformFunction,
} from '@shared/types/ai-providers/config';
import type { ErrorResponseBody } from '@shared/types/api/response/body';
import type {
CreateEmbeddingsRequestBody,
CreateEmbeddingsResponseBody,
EmbeddingData,
} from '@shared/types/api/routes/embeddings-api';
import { AIProvider } from '@shared/types/constants';
import {
googleTransformEmbeddingInput,
googleTransformEmbeddingsDimension,
} from './transform-generation-config';
import { GoogleErrorResponseTransform } from './utils';
export const googleEmbedConfig: AIProviderFunctionConfig = {
input: {
param: 'instances',
required: true,
transform: (saRequestBody: CreateEmbeddingsRequestBody) =>
googleTransformEmbeddingInput(saRequestBody) as Record<string, unknown>[],
},
parameters: {
param: 'parameters',
required: false,
},
dimensions: {
param: 'parameters',
required: false,
transform: (saRequestBody: CreateEmbeddingsRequestBody) =>
googleTransformEmbeddingsDimension(saRequestBody),
},
};
export const vertexGoogleEmbedResponseTransform: ResponseTransformFunction = (
aiProviderResponseBody,
aiProviderResponseStatus,
_aiProviderResponseHeaders,
_strictOpenAiCompliance,
saRequestData,
) => {
const googleResponse = aiProviderResponseBody as unknown as
| GoogleEmbedResponse
| GoogleErrorResponse;
if (aiProviderResponseStatus !== 200) {
const errorResposne = GoogleErrorResponseTransform(
googleResponse as unknown as GoogleErrorResponse,
);
if (errorResposne) return errorResposne;
}
if ('predictions' in googleResponse) {
const data: EmbeddingData[] = [];
let tokenCount = 0;
googleResponse.predictions.forEach((prediction, index) => {
const item = {
object: 'embedding',
index: index,
...(prediction.imageEmbedding && {
image_embedding: prediction.imageEmbedding,
}),
...(prediction.videoEmbeddings && {
video_embeddings: prediction.videoEmbeddings.map(
(videoEmbedding, idx) => ({
object: 'embedding',
embedding: videoEmbedding.embedding,
index: idx,
start_offset: videoEmbedding.startOffsetSec,
end_offset: videoEmbedding.endOffsetSec,
}),
),
}),
...(prediction.textEmbedding && {
embedding: prediction.textEmbedding,
}),
...(prediction.embeddings && {
embedding: prediction.embeddings.values,
}),
};
tokenCount += prediction?.embeddings?.statistics?.token_count || 0;
data.push(item as unknown as EmbeddingData);
});
data.forEach((item, index) => {
item.index = index;
});
const embeddingsResponseBody: CreateEmbeddingsResponseBody = {
model: (saRequestData.requestBody as CreateEmbeddingsRequestBody).model,
object: 'list',
data: data,
usage: {
prompt_tokens: tokenCount,
total_tokens: tokenCount,
},
};
return embeddingsResponseBody;
}
return generateInvalidProviderResponseError(
googleResponse as unknown as Record<string, unknown>,
AIProvider.GOOGLE_VERTEX_AI,
) as ErrorResponseBody;
};
|