UNPKG

@genkit-ai/vertexai

Version:

Genkit AI framework plugin for Google Cloud Vertex AI APIs including Gemini APIs, Imagen, and more.

1 lines 5.67 kB
{"version":3,"sources":["../../src/context-caching/index.ts"],"sourcesContent":["/**\n * @license\n * Copyright 2024 Google LLC\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n * https://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n */\n\nimport { CachedContent, StartChatParams } from '@google-cloud/vertexai';\nimport {\n ApiClient,\n CachedContents,\n} from '@google-cloud/vertexai/build/src/resources';\nimport { GenerateRequest, GenkitError, z } from 'genkit';\nimport { logger } from 'genkit/logging';\nimport type { CacheConfigDetails } from './types.js';\nimport {\n calculateTTL,\n generateCacheKey,\n getContentForCache,\n lookupContextCache,\n validateContextCacheRequest,\n} from './utils.js';\n\n/**\n * Handles context caching and transforms the chatRequest for Vertex AI.\n * @param apiKey\n * @param request\n * @param chatRequest\n * @param modelVersion\n * @returns\n */\nexport async function handleContextCache(\n apiClient: ApiClient,\n request: GenerateRequest<z.ZodTypeAny>,\n chatRequest: StartChatParams,\n modelVersion: string,\n cacheConfigDetails: CacheConfigDetails\n): Promise<{ cache: CachedContent; newChatRequest: StartChatParams }> {\n const cachedContentsClient = new CachedContents(apiClient);\n\n const { cachedContent, chatRequest: newChatRequest } = getContentForCache(\n request,\n chatRequest,\n modelVersion,\n cacheConfigDetails\n );\n cachedContent.model = modelVersion;\n const cacheKey = generateCacheKey(cachedContent);\n\n cachedContent.displayName = cacheKey;\n\n let cache;\n try {\n cache = await lookupContextCache(cachedContentsClient, cacheKey);\n logger.debug(`Cache hit: ${cache ? 'true' : 'false'}`);\n } catch (error) {\n logger.debug('No cache found, creating one.');\n }\n\n if (!cache) {\n try {\n const createParams: CachedContent = {\n ...cachedContent,\n // TODO: make this neater - idk why they chose to stringify the ttl...\n ttl: JSON.stringify(calculateTTL(cacheConfigDetails)) + 's',\n };\n cache = await cachedContentsClient.create(createParams);\n logger.debug(`Created new cache entry with key: ${cacheKey}`);\n } catch (cacheError) {\n logger.error(\n `Failed to create cache with key ${cacheKey}: ${cacheError}`\n );\n throw new GenkitError({\n status: 'INTERNAL',\n message: `Failed to create cache: ${cacheError}`,\n });\n }\n }\n\n if (!cache) {\n throw new GenkitError({\n status: 'INTERNAL',\n message: 'Failed to use context cache feature',\n });\n }\n // This isn't necessary, but it's nice to have for debugging purposes.\n newChatRequest.cachedContent = cache.name;\n\n return { cache, newChatRequest };\n}\n\n/**\n * Handles cache validation, creation, and usage, transforming the chatRequest if necessary.\n * @param apiClient The API client for Vertex AI.\n * @param options Plugin options containing project details and auth.\n * @param request The generate request passed to the model.\n * @param chatRequest The current chat request configuration.\n * @param modelVersion The version of the model being used.\n * @param cacheConfigDetails Configuration details for caching.\n * @returns A transformed chat request and cache data (if applicable).\n */\nexport async function handleCacheIfNeeded(\n apiClient: ApiClient,\n request: GenerateRequest<z.ZodTypeAny>,\n chatRequest: StartChatParams,\n modelVersion: string,\n cacheConfigDetails: CacheConfigDetails | null\n): Promise<{ chatRequest: StartChatParams; cache: CachedContent | null }> {\n if (\n !cacheConfigDetails ||\n !validateContextCacheRequest(request, modelVersion)\n ) {\n return { chatRequest, cache: null };\n }\n\n const { cache, newChatRequest } = await handleContextCache(\n apiClient,\n request,\n chatRequest,\n modelVersion,\n cacheConfigDetails\n );\n return { chatRequest: newChatRequest, cache };\n}\n"],"mappings":"AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA;AAgBA;AAAA,EAEE;AAAA,OACK;AACP,SAA0B,mBAAsB;AAChD,SAAS,cAAc;AAEvB;AAAA,EACE;AAAA,EACA;AAAA,EACA;AAAA,EACA;AAAA,EACA;AAAA,OACK;AAUP,eAAsB,mBACpB,WACA,SACA,aACA,cACA,oBACoE;AACpE,QAAM,uBAAuB,IAAI,eAAe,SAAS;AAEzD,QAAM,EAAE,eAAe,aAAa,eAAe,IAAI;AAAA,IACrD;AAAA,IACA;AAAA,IACA;AAAA,IACA;AAAA,EACF;AACA,gBAAc,QAAQ;AACtB,QAAM,WAAW,iBAAiB,aAAa;AAE/C,gBAAc,cAAc;AAE5B,MAAI;AACJ,MAAI;AACF,YAAQ,MAAM,mBAAmB,sBAAsB,QAAQ;AAC/D,WAAO,MAAM,cAAc,QAAQ,SAAS,OAAO,EAAE;AAAA,EACvD,SAAS,OAAO;AACd,WAAO,MAAM,+BAA+B;AAAA,EAC9C;AAEA,MAAI,CAAC,OAAO;AACV,QAAI;AACF,YAAM,eAA8B;AAAA,QAClC,GAAG;AAAA;AAAA,QAEH,KAAK,KAAK,UAAU,aAAa,kBAAkB,CAAC,IAAI;AAAA,MAC1D;AACA,cAAQ,MAAM,qBAAqB,OAAO,YAAY;AACtD,aAAO,MAAM,qCAAqC,QAAQ,EAAE;AAAA,IAC9D,SAAS,YAAY;AACnB,aAAO;AAAA,QACL,mCAAmC,QAAQ,KAAK,UAAU;AAAA,MAC5D;AACA,YAAM,IAAI,YAAY;AAAA,QACpB,QAAQ;AAAA,QACR,SAAS,2BAA2B,UAAU;AAAA,MAChD,CAAC;AAAA,IACH;AAAA,EACF;AAEA,MAAI,CAAC,OAAO;AACV,UAAM,IAAI,YAAY;AAAA,MACpB,QAAQ;AAAA,MACR,SAAS;AAAA,IACX,CAAC;AAAA,EACH;AAEA,iBAAe,gBAAgB,MAAM;AAErC,SAAO,EAAE,OAAO,eAAe;AACjC;AAYA,eAAsB,oBACpB,WACA,SACA,aACA,cACA,oBACwE;AACxE,MACE,CAAC,sBACD,CAAC,4BAA4B,SAAS,YAAY,GAClD;AACA,WAAO,EAAE,aAAa,OAAO,KAAK;AAAA,EACpC;AAEA,QAAM,EAAE,OAAO,eAAe,IAAI,MAAM;AAAA,IACtC;AAAA,IACA;AAAA,IACA;AAAA,IACA;AAAA,IACA;AAAA,EACF;AACA,SAAO,EAAE,aAAa,gBAAgB,MAAM;AAC9C;","names":[]}