import { GenerateRequest, z } from 'genkit';
import { EmbedRequest } from 'ollama';

/**
 * 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
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

type ApiType = 'chat' | 'generate';
interface ModelDefinition {
    name: string;
    type?: ApiType;
    supports?: {
        tools?: boolean;
    };
}
interface EmbeddingModelDefinition {
    name: string;
    dimensions: number;
}
declare const OllamaEmbeddingPredictionSchema: z.ZodObject<{
    embedding: z.ZodArray<z.ZodNumber, "many">;
}, "strip", z.ZodTypeAny, {
    embedding: number[];
}, {
    embedding: number[];
}>;
type OllamaEmbeddingPrediction = z.infer<typeof OllamaEmbeddingPredictionSchema>;
interface DefineOllamaEmbeddingParams {
    name: string;
    modelName: string;
    dimensions: number;
    options: OllamaPluginParams;
}
/**
 * Parameters for the Ollama plugin configuration.
 */
interface OllamaPluginParams {
    /**
     * Array of models to be defined.
     *
     * ```ts
     * const ai = genkit({
     *   plugins: [
     *     ollama({
     *       models: [{ name: 'gemma' }],
     *       serverAddress: 'http://127.0.0.1:11434', // default local address
     *     }),
     *   ],
     * });
     * ```
     */
    models?: ModelDefinition[];
    /**
     * Array of embedding models to be defined.
     *
     * ```ts
     * const ai = genkit({
     *   plugins: [
     *     ollama({
     *       serverAddress: 'http://localhost:11434',
     *       embedders: [{ name: 'nomic-embed-text', dimensions: 768 }],
     *     }),
     *   ],
     * });
     * ```
     */
    embedders?: EmbeddingModelDefinition[];
    /**
     * The address of the Ollama server. Default: http://localhost:11434
     */
    serverAddress?: string;
    /**
     * Optional request headers, which can be either static or dynamically generated.
     *
     * ```ts
     * const ai = genkit({
     *   plugins: [
     *     ollama({
     *       models: [...],
     *       serverAddress: 'https://my-deployment.server.address',
     *       requestHeaders: async (params) => {
     *         const headers = await fetchAuthHeaders(params.serverAddress);
     *         return { Authorization: headers['Authorization'] };
     *       },
     *     }),
     *   ],
     * });
     * ```
     */
    requestHeaders?: RequestHeaders;
}
type RequestHeaderFunction = (params: {
    serverAddress: string;
    model?: ModelDefinition | EmbeddingModelDefinition;
    modelRequest?: GenerateRequest;
    embedRequest?: EmbedRequest;
}, modelRequest?: GenerateRequest) => Promise<Record<string, string> | void>;
type RequestHeaders = Record<string, string> | RequestHeaderFunction;
type OllamaRole = 'assistant' | 'tool' | 'system' | 'user';
interface OllamaTool {
    type: string;
    function: {
        name: string;
        description: string;
        parameters: Record<string, any>;
    };
}
interface OllamaToolCall {
    function: {
        index?: number;
        name: string;
        arguments: Record<string, any>;
    };
}
interface Message {
    role: string;
    content: string;
    images?: string[];
    tool_calls?: any[];
}
interface LocalModel {
    name: string;
    model: string;
    modified_at: string;
    size: number;
    digest: string;
    details?: {
        parent_model?: string;
        format?: string;
        family?: string;
        families?: string[];
        parameter_size?: string;
        quantization_level?: string;
    };
}
interface ListLocalModelsResponse {
    models: LocalModel[];
}

export { type ApiType, type DefineOllamaEmbeddingParams, type EmbeddingModelDefinition, type ListLocalModelsResponse, type LocalModel, type Message, type ModelDefinition, type OllamaEmbeddingPrediction, OllamaEmbeddingPredictionSchema, type OllamaPluginParams, type OllamaRole, type OllamaTool, type OllamaToolCall, type RequestHeaderFunction, type RequestHeaders };