@genkit-ai/vertexai
Version:
Genkit AI framework plugin for Google Cloud Vertex AI APIs including Gemini APIs, Imagen, and more.
90 lines • 2.82 kB
JavaScript
/**
* @license
* Copyright 2024 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
*/
import {
CachedContents
} from "@google-cloud/vertexai/build/src/resources";
import { GenkitError } from "genkit";
import { logger } from "genkit/logging";
import {
calculateTTL,
generateCacheKey,
getContentForCache,
lookupContextCache,
validateContextCacheRequest
} from "./utils.js";
async function handleContextCache(apiClient, request, chatRequest, modelVersion, cacheConfigDetails) {
const cachedContentsClient = new CachedContents(apiClient);
const { cachedContent, chatRequest: newChatRequest } = getContentForCache(
request,
chatRequest,
modelVersion,
cacheConfigDetails
);
cachedContent.model = modelVersion;
const cacheKey = generateCacheKey(cachedContent);
cachedContent.displayName = cacheKey;
let cache;
try {
cache = await lookupContextCache(cachedContentsClient, cacheKey);
logger.debug(`Cache hit: ${cache ? "true" : "false"}`);
} catch (error) {
logger.debug("No cache found, creating one.");
}
if (!cache) {
try {
const createParams = {
...cachedContent,
// TODO: make this neater - idk why they chose to stringify the ttl...
ttl: JSON.stringify(calculateTTL(cacheConfigDetails)) + "s"
};
cache = await cachedContentsClient.create(createParams);
logger.debug(`Created new cache entry with key: ${cacheKey}`);
} catch (cacheError) {
logger.error(
`Failed to create cache with key ${cacheKey}: ${cacheError}`
);
throw new GenkitError({
status: "INTERNAL",
message: `Failed to create cache: ${cacheError}`
});
}
}
if (!cache) {
throw new GenkitError({
status: "INTERNAL",
message: "Failed to use context cache feature"
});
}
newChatRequest.cachedContent = cache.name;
return { cache, newChatRequest };
}
async function handleCacheIfNeeded(apiClient, request, chatRequest, modelVersion, cacheConfigDetails) {
if (!cacheConfigDetails || !validateContextCacheRequest(request, modelVersion)) {
return { chatRequest, cache: null };
}
const { cache, newChatRequest } = await handleContextCache(
apiClient,
request,
chatRequest,
modelVersion,
cacheConfigDetails
);
return { chatRequest: newChatRequest, cache };
}
export {
handleCacheIfNeeded,
handleContextCache
};
//# sourceMappingURL=index.mjs.map