{"version":3,"sources":["../../src/context-caching/utils.ts"],"sourcesContent":["/**\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 * http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\nimport type { CachedContent, StartChatParams } from '@google-cloud/vertexai';\nimport type { CachedContents } from '@google-cloud/vertexai/build/src/resources';\nimport crypto from 'crypto';\nimport { GenkitError, type MessageData, type z } from 'genkit';\nimport type { GenerateRequest } from 'genkit/model';\nimport {\n CONTEXT_CACHE_SUPPORTED_MODELS,\n DEFAULT_TTL,\n INVALID_ARGUMENT_MESSAGES,\n} from './constants';\nimport {\n cacheConfigSchema,\n type CacheConfig,\n type CacheConfigDetails,\n} from './types';\n\n/**\n * Generates a SHA-256 hash to use as a cache key.\n * @param request CachedContent - request object to hash\n * @returns string - the generated cache key\n */\nexport function generateCacheKey(request: CachedContent): string {\n return crypto\n .createHash('sha256')\n .update(JSON.stringify(request))\n .digest('hex');\n}\n\n/**\n * Retrieves the content needed for the cache based on the chat history and config details.\n */\nexport function getContentForCache(\n request: GenerateRequest,\n chatRequest: StartChatParams,\n modelVersion: string,\n cacheConfigDetails: CacheConfigDetails\n): {\n cachedContent: CachedContent;\n chatRequest: StartChatParams;\n cacheConfig?: CacheConfig;\n} {\n if (!chatRequest.history?.length) {\n throw new Error('No history provided for context caching');\n }\n\n validateHistoryLength(request, chatRequest);\n\n const { endOfCachedContents, cacheConfig } = cacheConfigDetails;\n const cachedContent: CachedContent = {\n model: modelVersion,\n contents: chatRequest.history.slice(0, endOfCachedContents + 1),\n };\n chatRequest.history = chatRequest.history.slice(endOfCachedContents + 1);\n\n return { cachedContent, chatRequest, cacheConfig };\n}\n\n/**\n * Validates that the request and chat request history lengths align.\n * @throws GenkitError if lengths are mismatched\n */\nfunction validateHistoryLength(\n request: GenerateRequest,\n chatRequest: StartChatParams\n) {\n if (chatRequest.history?.length !== request.messages.length - 1) {\n throw new GenkitError({\n status: 'INTERNAL',\n message:\n 'Genkit request history and Gemini chat request history length do not match',\n });\n }\n}\n\n/**\n * Looks up context cache using a cache manager and returns the found item, if any.\n */\nexport async function lookupContextCache(\n cacheManager: CachedContents,\n cacheKey: string,\n maxPages = 100,\n pageSize = 100\n) {\n let currentPage = 0;\n let pageToken: string | undefined;\n\n while (currentPage < maxPages) {\n const { cachedContents, nextPageToken } = await cacheManager.list(\n pageSize,\n pageToken\n );\n const found = cachedContents?.find(\n (content) => content.displayName === cacheKey\n );\n\n if (found) return found;\n if (!nextPageToken) break;\n\n pageToken = nextPageToken;\n currentPage++;\n }\n return null;\n}\n\n/**\n * Extracts the cache configuration from the request if available.\n */\nexport const extractCacheConfig = (\n request: GenerateRequest\n): {\n cacheConfig: { ttlSeconds?: number } | boolean;\n endOfCachedContents: number;\n} | null => {\n const endOfCachedContents = findLastIndex(\n request.messages,\n (message) => !!message.metadata?.cache\n );\n\n return endOfCachedContents === -1\n ? null\n : {\n endOfCachedContents,\n cacheConfig: cacheConfigSchema.parse(\n request.messages[endOfCachedContents].metadata?.cache\n ),\n };\n};\n\n/**\n * Validates context caching request for compatibility with model and request configurations.\n */\nexport function validateContextCacheRequest(\n request: GenerateRequest,\n modelVersion: string\n): boolean {\n if (!modelVersion || !CONTEXT_CACHE_SUPPORTED_MODELS.includes(modelVersion)) {\n throw new GenkitError({\n status: 'INVALID_ARGUMENT',\n message: INVALID_ARGUMENT_MESSAGES.modelVersion,\n });\n }\n if (request.tools?.length)\n throw new GenkitError({\n status: 'INVALID_ARGUMENT',\n message: INVALID_ARGUMENT_MESSAGES.tools,\n });\n if (request.config?.codeExecution)\n throw new GenkitError({\n status: 'INVALID_ARGUMENT',\n message: INVALID_ARGUMENT_MESSAGES.codeExecution,\n });\n\n return true;\n}\n\n/**\n * Polyfill function for Array.prototype.findLastIndex for ES2015 compatibility.\n */\nfunction findLastIndex(\n array: T[],\n callback: (element: T, index: number, array: T[]) => boolean\n): number {\n for (let i = array.length - 1; i >= 0; i--) {\n if (callback(array[i], i, array)) return i;\n }\n return -1;\n}\n\n/**\n * Calculates the TTL (Time-To-Live) for the cache based on cacheConfigDetails.\n * @param cacheConfig - The caching configuration details.\n * @returns The TTL in seconds.\n */\nexport function calculateTTL(cacheConfig: CacheConfigDetails): number {\n if (cacheConfig.cacheConfig === true) {\n return DEFAULT_TTL;\n }\n if (cacheConfig.cacheConfig === false) {\n return 0;\n }\n return cacheConfig.cacheConfig.ttlSeconds || DEFAULT_TTL;\n}\n"],"mappings":"AAkBA,OAAO,YAAY;AACnB,SAAS,mBAA6C;AAEtD;AAAA,EACE;AAAA,EACA;AAAA,EACA;AAAA,OACK;AACP;AAAA,EACE;AAAA,OAGK;AAOA,SAAS,iBAAiB,SAAgC;AAC/D,SAAO,OACJ,WAAW,QAAQ,EACnB,OAAO,KAAK,UAAU,OAAO,CAAC,EAC9B,OAAO,KAAK;AACjB;AAKO,SAAS,mBACd,SACA,aACA,cACA,oBAKA;AACA,MAAI,CAAC,YAAY,SAAS,QAAQ;AAChC,UAAM,IAAI,MAAM,yCAAyC;AAAA,EAC3D;AAEA,wBAAsB,SAAS,WAAW;AAE1C,QAAM,EAAE,qBAAqB,YAAY,IAAI;AAC7C,QAAM,gBAA+B;AAAA,IACnC,OAAO;AAAA,IACP,UAAU,YAAY,QAAQ,MAAM,GAAG,sBAAsB,CAAC;AAAA,EAChE;AACA,cAAY,UAAU,YAAY,QAAQ,MAAM,sBAAsB,CAAC;AAEvE,SAAO,EAAE,eAAe,aAAa,YAAY;AACnD;AAMA,SAAS,sBACP,SACA,aACA;AACA,MAAI,YAAY,SAAS,WAAW,QAAQ,SAAS,SAAS,GAAG;AAC/D,UAAM,IAAI,YAAY;AAAA,MACpB,QAAQ;AAAA,MACR,SACE;AAAA,IACJ,CAAC;AAAA,EACH;AACF;AAKA,eAAsB,mBACpB,cACA,UACA,WAAW,KACX,WAAW,KACX;AACA,MAAI,cAAc;AAClB,MAAI;AAEJ,SAAO,cAAc,UAAU;AAC7B,UAAM,EAAE,gBAAgB,cAAc,IAAI,MAAM,aAAa;AAAA,MAC3D;AAAA,MACA;AAAA,IACF;AACA,UAAM,QAAQ,gBAAgB;AAAA,MAC5B,CAAC,YAAY,QAAQ,gBAAgB;AAAA,IACvC;AAEA,QAAI,MAAO,QAAO;AAClB,QAAI,CAAC,cAAe;AAEpB,gBAAY;AACZ;AAAA,EACF;AACA,SAAO;AACT;AAKO,MAAM,qBAAqB,CAChC,YAIU;AACV,QAAM,sBAAsB;AAAA,IAC1B,QAAQ;AAAA,IACR,CAAC,YAAY,CAAC,CAAC,QAAQ,UAAU;AAAA,EACnC;AAEA,SAAO,wBAAwB,KAC3B,OACA;AAAA,IACE;AAAA,IACA,aAAa,kBAAkB;AAAA,MAC7B,QAAQ,SAAS,mBAAmB,EAAE,UAAU;AAAA,IAClD;AAAA,EACF;AACN;AAKO,SAAS,4BACd,SACA,cACS;AACT,MAAI,CAAC,gBAAgB,CAAC,+BAA+B,SAAS,YAAY,GAAG;AAC3E,UAAM,IAAI,YAAY;AAAA,MACpB,QAAQ;AAAA,MACR,SAAS,0BAA0B;AAAA,IACrC,CAAC;AAAA,EACH;AACA,MAAI,QAAQ,OAAO;AACjB,UAAM,IAAI,YAAY;AAAA,MACpB,QAAQ;AAAA,MACR,SAAS,0BAA0B;AAAA,IACrC,CAAC;AACH,MAAI,QAAQ,QAAQ;AAClB,UAAM,IAAI,YAAY;AAAA,MACpB,QAAQ;AAAA,MACR,SAAS,0BAA0B;AAAA,IACrC,CAAC;AAEH,SAAO;AACT;AAKA,SAAS,cACP,OACA,UACQ;AACR,WAAS,IAAI,MAAM,SAAS,GAAG,KAAK,GAAG,KAAK;AAC1C,QAAI,SAAS,MAAM,CAAC,GAAG,GAAG,KAAK,EAAG,QAAO;AAAA,EAC3C;AACA,SAAO;AACT;AAOO,SAAS,aAAa,aAAyC;AACpE,MAAI,YAAY,gBAAgB,MAAM;AACpC,WAAO;AAAA,EACT;AACA,MAAI,YAAY,gBAAgB,OAAO;AACrC,WAAO;AAAA,EACT;AACA,SAAO,YAAY,YAAY,cAAc;AAC/C;","names":[]}