@agenite/ollama
Version:
Ollama provider for Agenite
356 lines (351 loc) • 9.52 kB
JavaScript
;
var ollama = require('ollama');
var llm = require('@agenite/llm');
// src/provider.ts
// src/utils.ts
function createTokenInfo(model, inputTokens = 0, outputTokens = 0) {
return {
inputTokens,
outputTokens,
model,
inputCost: 0,
outputCost: 0
};
}
function createResponse(content, model, inputTokens = 0, outputTokens = 0, stopReason) {
return {
content,
tokenUsage: createTokenInfo(model, inputTokens, outputTokens),
stopReason
};
}
function createTextContent(text) {
return { type: "text", text };
}
function createToolUseContent(name, input) {
return {
type: "toolUse",
id: crypto.randomUUID(),
name,
input
};
}
function parseToolArguments(args) {
return typeof args === "string" ? JSON.parse(args) : args;
}
function convertFunctionCallsToToolUses(toolCalls) {
return toolCalls.map(
(toolCall) => createToolUseContent(
toolCall.function.name,
parseToolArguments(toolCall.function.arguments)
)
);
}
function mapStopReason(finishReason) {
if (!finishReason) return void 0;
const stopReasonMap = {
stop: "endTurn",
length: "maxTokens",
tool_calls: "toolUse"
};
return stopReasonMap[finishReason] || void 0;
}
function extractTextContent(blocks) {
return blocks.map((block) => block.type === "text" ? block.text : "").join("\n").trim();
}
function extractImages(blocks) {
return blocks.map((block) => {
if (block.type === "image" && block.source.type === "base64") {
return block.source.data;
}
return null;
}).filter((img) => img !== null);
}
function createToolResultMessage(toolResult, toolUse) {
return {
role: "tool",
content: JSON.stringify({
output: toolResult.content,
isError: toolResult.isError,
function_call: {
name: toolUse.name,
arguments: toolUse.input
}
}),
name: toolResult.toolName
};
}
function createRegularMessage(msg) {
const images = extractImages(msg.content);
const content = extractTextContent(msg.content);
return {
role: msg.role,
content,
...images.length > 0 && { images }
};
}
function findToolResult(toolUse, nextMsg) {
return nextMsg?.content.find(
(block) => block.type === "toolResult" && block.toolUseId === toolUse.id
);
}
function processMessagePair(msg, nextMsg) {
const toolUses = msg.content.filter(
(block) => block.type === "toolUse"
);
if (!toolUses.length) {
return {
messages: [createRegularMessage(msg)],
skipNext: false
};
}
const toolCallsMessage = {
role: "assistant",
content: "",
tool_calls: toolUses.map((toolUse) => ({
function: {
name: toolUse.name,
arguments: toolUse.input
}
}))
};
const messages = [toolCallsMessage];
for (const toolUse of toolUses) {
const toolResult = findToolResult(toolUse, nextMsg);
if (toolResult) {
messages.push(createToolResultMessage(toolResult, toolUse));
}
}
const hasToolResults = messages.some((msg2) => msg2.role === "tool");
return { messages, skipNext: hasToolResults };
}
function convertMessages(messages) {
const ollamaMessages = [];
for (let i = 0; i < messages.length; i++) {
const msg = messages[i];
if (!msg) continue;
const { messages: newMessages, skipNext } = processMessagePair(
msg,
messages[i + 1]
);
ollamaMessages.push(...newMessages);
if (skipNext) {
i++;
}
}
return ollamaMessages;
}
function extractParameterType(value) {
return typeof value === "object" && value && "type" in value ? String(value.type) : "string";
}
function extractParameterDescription(value, key) {
return typeof value === "object" && value && "description" in value ? String(value.description) : key;
}
function extractParameterEnum(value) {
return typeof value === "object" && value && "enum" in value ? { enum: value.enum } : {};
}
function convertParameterDefinition(key, value) {
return [
key,
{
type: extractParameterType(value),
description: extractParameterDescription(value, key),
...extractParameterEnum(value)
}
];
}
function convertToolParameters(tool) {
return {
type: "object",
properties: Object.fromEntries(
Object.entries(tool.inputSchema.properties || {}).map(
([key, value]) => convertParameterDefinition(key, value)
)
),
required: tool.inputSchema.required ?? []
};
}
function convertToolDefinition(tool) {
return {
type: "function",
function: {
name: tool.name,
description: tool.description,
parameters: convertToolParameters(tool)
}
};
}
function convertToolDefinitions(tools) {
if (!tools?.length) return void 0;
return tools.map(convertToolDefinition);
}
function createError(error, context) {
console.error(`Ollama ${context} failed:`, error);
return error instanceof Error ? new Error(`Ollama ${context} failed: ${error.message}`) : new Error(`Ollama ${context} failed with unknown error`);
}
// src/provider.ts
var OllamaProvider = class extends llm.BaseLLMProvider {
client;
config;
name = "Ollama";
version = "1.0";
constructor(config) {
super();
this.config = config;
this.client = new ollama.Ollama({
host: config.host
});
}
/**
* Creates base chat request parameters
*/
createBaseRequest(messages, options) {
if (options?.systemPrompt) {
messages.unshift({
role: "system",
content: options.systemPrompt
});
}
return {
model: this.config.model,
messages,
options: {
temperature: options?.temperature,
num_predict: options?.maxTokens,
stop: options?.stopSequences,
...this.config.parameters
},
tools: convertToolDefinitions(options?.tools)
};
}
/**
* Prepares messages for chat request
*/
prepareMessages(input) {
const messageArray = llm.convertStringToMessages(input);
return convertMessages(messageArray);
}
/**
* Combines text and tool calls into a single response content
*/
combineResponseContent(text, toolCalls) {
const content = [];
if (text) {
content.push(createTextContent(text));
}
if (toolCalls?.length) {
content.push(...convertFunctionCallsToToolUses(toolCalls));
}
return content;
}
async generate(input, options) {
try {
const ollamaMessages = this.prepareMessages(input);
const response = await this.client.chat({
...this.createBaseRequest(ollamaMessages, options),
stream: false
});
return createResponse(
this.combineResponseContent(
response.message.content,
response.message.tool_calls
),
this.config.model,
response.prompt_eval_count,
response.eval_count,
response.message.tool_calls?.length ? "toolUse" : mapStopReason("stop")
);
} catch (error) {
throw createError(error, "generation");
}
}
async *stream(input, options) {
try {
const ollamaMessages = this.prepareMessages(input);
let buffer = "";
let finalResponse = void 0;
let textAccumulator = "";
const toolCalls = [];
const response = await this.client.chat({
...this.createBaseRequest(ollamaMessages, options),
stream: true
});
let hasTextStart = false;
let hasTextEnd = false;
for await (const chunk of response) {
const content = chunk.message?.content;
if (content) {
if (!hasTextStart) {
yield {
type: "text",
text: "",
isStart: true
};
hasTextStart = true;
}
buffer += content;
textAccumulator += content;
if (buffer.length > 4) {
yield {
type: "text",
text: buffer
};
buffer = "";
}
}
if (chunk.message?.tool_calls?.length) {
if (buffer.length > 0) {
yield {
type: "text",
text: buffer,
isEnd: true
};
hasTextEnd = true;
buffer = "";
}
for (const toolCall of chunk.message.tool_calls) {
const tool = {
type: "toolUse",
toolUse: createToolUseContent(
toolCall.function.name,
toolCall.function.arguments
),
isEnd: true
};
toolCalls.push(toolCall);
yield tool;
}
}
finalResponse = chunk;
}
if (buffer.length > 0) {
yield {
type: "text",
text: buffer,
isEnd: true
};
} else if (!hasTextEnd && hasTextStart) {
yield {
type: "text",
text: "",
isEnd: true
};
}
if (!finalResponse) {
throw new Error("No final response received");
}
return createResponse(
this.combineResponseContent(textAccumulator, toolCalls),
this.config.model,
finalResponse.prompt_eval_count,
finalResponse.eval_count,
toolCalls?.length ? "toolUse" : mapStopReason(finalResponse.done ? "stop" : null)
);
} catch (error) {
throw createError(error, "stream");
}
}
};
exports.OllamaProvider = OllamaProvider;
//# sourceMappingURL=index.cjs.map
//# sourceMappingURL=index.cjs.map