import { StartChatParams, CachedContent } from '@google-cloud/vertexai';
import { ApiClient } from '@google-cloud/vertexai/build/src/resources';
import { GenerateRequest, z } from 'genkit';
import { CacheConfigDetails } from './types.js';

/**
 * @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.
 */

/**
 * Handles context caching and transforms the chatRequest for Vertex AI.
 * @param apiKey
 * @param request
 * @param chatRequest
 * @param modelVersion
 * @returns
 */
declare function handleContextCache(apiClient: ApiClient, request: GenerateRequest<z.ZodTypeAny>, chatRequest: StartChatParams, modelVersion: string, cacheConfigDetails: CacheConfigDetails): Promise<{
    cache: CachedContent;
    newChatRequest: StartChatParams;
}>;
/**
 * Handles cache validation, creation, and usage, transforming the chatRequest if necessary.
 * @param apiClient The API client for Vertex AI.
 * @param options Plugin options containing project details and auth.
 * @param request The generate request passed to the model.
 * @param chatRequest The current chat request configuration.
 * @param modelVersion The version of the model being used.
 * @param cacheConfigDetails Configuration details for caching.
 * @returns A transformed chat request and cache data (if applicable).
 */
declare function handleCacheIfNeeded(apiClient: ApiClient, request: GenerateRequest<z.ZodTypeAny>, chatRequest: StartChatParams, modelVersion: string, cacheConfigDetails: CacheConfigDetails | null): Promise<{
    chatRequest: StartChatParams;
    cache: CachedContent | null;
}>;

export { handleCacheIfNeeded, handleContextCache };