123 lines
4.7 KiB
TypeScript
123 lines
4.7 KiB
TypeScript
import {findByProp, objectMap, JSONSanitize, JSONAttemptParse} from '@ztimson/utils';
|
|
import {Ai} from './ai.ts';
|
|
import {LLMMessage, LLMRequest} from './llm.ts';
|
|
import {AbortablePromise, LLMProvider} from './provider.ts';
|
|
import {Ollama as ollama} from 'ollama';
|
|
|
|
export class Ollama extends LLMProvider {
|
|
client!: ollama;
|
|
|
|
constructor(public readonly ai: Ai, public host: string, public model: string) {
|
|
super();
|
|
this.client = new ollama({host});
|
|
}
|
|
|
|
private toStandard(history: any[]): LLMMessage[] {
|
|
for(let i = 0; i < history.length; i++) {
|
|
if(history[i].role == 'assistant' && history[i].tool_calls) {
|
|
if(history[i].content) delete history[i].tool_calls;
|
|
else {
|
|
history.splice(i, 1);
|
|
i--;
|
|
}
|
|
} else if(history[i].role == 'tool') {
|
|
const error = history[i].content.startsWith('{"error":');
|
|
history[i] = {role: 'tool', name: history[i].tool_name, args: history[i].args, [error ? 'error' : 'content']: history[i].content, timestamp: history[i].timestamp};
|
|
}
|
|
if(!history[i]?.timestamp) history[i].timestamp = Date.now();
|
|
}
|
|
return history;
|
|
}
|
|
|
|
private fromStandard(history: LLMMessage[]): any[] {
|
|
return history.map((h: any) => {
|
|
const {timestamp, ...rest} = h;
|
|
if(h.role != 'tool') return rest;
|
|
return {role: 'tool', tool_name: h.name, content: h.error || h.content}
|
|
});
|
|
}
|
|
|
|
ask(message: string, options: LLMRequest = {}): AbortablePromise<LLMMessage[]> {
|
|
const controller = new AbortController();
|
|
const response = new Promise<any>(async (res, rej) => {
|
|
let system = options.system || this.ai.options.system;
|
|
let history = this.fromStandard([...options.history || [], {role: 'user', content: message, timestamp: Date.now()}]);
|
|
if(history[0].roll == 'system') {
|
|
if(!system) system = history.shift();
|
|
else history.shift();
|
|
}
|
|
if(options.compress) history = await this.ai.language.compressHistory(<any>history, options.compress.max, options.compress.min);
|
|
if(options.system) history.unshift({role: 'system', content: system})
|
|
|
|
const tools = options.tools || this.ai.options.tools || [];
|
|
const requestParams: any = {
|
|
model: options.model || this.model,
|
|
messages: history,
|
|
stream: !!options.stream,
|
|
signal: controller.signal,
|
|
options: {
|
|
temperature: options.temperature || this.ai.options.temperature || 0.7,
|
|
num_predict: options.max_tokens || this.ai.options.max_tokens || 4096,
|
|
},
|
|
tools: tools.map(t => ({
|
|
type: 'function',
|
|
function: {
|
|
name: t.name,
|
|
description: t.description,
|
|
parameters: {
|
|
type: 'object',
|
|
properties: t.args ? objectMap(t.args, (key, value) => ({...value, required: undefined})) : {},
|
|
required: t.args ? Object.entries(t.args).filter(t => t[1].required).map(t => t[0]) : []
|
|
}
|
|
}
|
|
}))
|
|
}
|
|
|
|
let resp: any, isFirstMessage = true;
|
|
do {
|
|
resp = await this.client.chat(requestParams).catch(err => {
|
|
err.message += `\n\nMessages:\n${JSON.stringify(history, null, 2)}`;
|
|
throw err;
|
|
});
|
|
|
|
if(options.stream) {
|
|
if(!isFirstMessage) options.stream({text: '\n\n'});
|
|
else isFirstMessage = false;
|
|
resp.message = {role: 'assistant', content: '', tool_calls: []};
|
|
for await (const chunk of resp) {
|
|
if(controller.signal.aborted) break;
|
|
if(chunk.message?.content) {
|
|
resp.message.content += chunk.message.content;
|
|
options.stream({text: chunk.message.content});
|
|
}
|
|
if(chunk.message?.tool_calls) resp.message.tool_calls = chunk.message.tool_calls;
|
|
if(chunk.done) break;
|
|
}
|
|
}
|
|
|
|
if(resp.message?.tool_calls?.length && !controller.signal.aborted) {
|
|
history.push(resp.message);
|
|
const results = await Promise.all(resp.message.tool_calls.map(async (toolCall: any) => {
|
|
const tool = tools.find(findByProp('name', toolCall.function.name));
|
|
if(!tool) return {role: 'tool', tool_name: toolCall.function.name, content: '{"error": "Tool not found"}'};
|
|
const args = typeof toolCall.function.arguments === 'string' ? JSONAttemptParse(toolCall.function.arguments, {}) : toolCall.function.arguments;
|
|
try {
|
|
const result = await tool.fn(args, this.ai);
|
|
return {role: 'tool', tool_name: toolCall.function.name, args, content: JSONSanitize(result)};
|
|
} catch (err: any) {
|
|
return {role: 'tool', tool_name: toolCall.function.name, args, content: JSONSanitize({error: err?.message || err?.toString() || 'Unknown'})};
|
|
}
|
|
}));
|
|
history.push(...results);
|
|
requestParams.messages = history;
|
|
}
|
|
} while (!controller.signal.aborted && resp.message?.tool_calls?.length);
|
|
|
|
if(options.stream) options.stream({done: true});
|
|
res(this.toStandard([...history, {role: 'assistant', content: resp.message?.content}]));
|
|
});
|
|
|
|
return Object.assign(response, {abort: () => controller.abort()});
|
|
}
|
|
}
|