genkitx-ollama
Version:
Genkit AI framework plugin for Ollama APIs.
153 lines (149 loc) • 4.4 kB
TypeScript
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 };