{"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 type { CachedContent, StartChatParams } from '@google-cloud/vertexai';\nimport {\n  CachedContents,\n  type ApiClient,\n} from '@google-cloud/vertexai/build/src/resources';\nimport { GenkitError, type GenerateRequest, type 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,EACE;AAAA,OAEK;AACP,SAAS,mBAAiD;AAC1D,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":[]}