Added MCP, Hybrid Memories and Skill support
This commit is contained in:
243
src/llm.ts
243
src/llm.ts
@@ -1,5 +1,3 @@
|
||||
import {sum} from '@tensorflow/tfjs';
|
||||
import {JSONAttemptParse} from '@ztimson/utils';
|
||||
import {AbortablePromise, Ai} from './ai.ts';
|
||||
import {Anthropic} from './antrhopic.ts';
|
||||
import {OpenAi} from './open-ai.ts';
|
||||
@@ -7,7 +5,8 @@ import {LLMProvider} from './provider.ts';
|
||||
import {AiTool} from './tools.ts';
|
||||
import {fileURLToPath} from 'url';
|
||||
import {dirname, join} from 'path';
|
||||
import { spawn } from 'node:child_process';
|
||||
import {spawn} from 'node:child_process';
|
||||
import {Memory, MemoryManager} from './memory.ts';
|
||||
|
||||
export type AnthropicConfig = {proto: 'anthropic', token: string};
|
||||
export type OllamaConfig = {proto: 'ollama', host: string};
|
||||
@@ -37,16 +36,6 @@ export type LLMMessage = {
|
||||
timestamp?: number;
|
||||
}
|
||||
|
||||
/** Background information the AI will be fed */
|
||||
export type LLMMemory = {
|
||||
/** What entity is this fact about */
|
||||
owner: string;
|
||||
/** The information that will be remembered */
|
||||
fact: string;
|
||||
/** Owner and fact embedding vector */
|
||||
embeddings: [number[], number[]];
|
||||
}
|
||||
|
||||
export type LLMRequest = {
|
||||
/** System prompt */
|
||||
system?: string;
|
||||
@@ -63,17 +52,39 @@ export type LLMRequest = {
|
||||
/** Stream response */
|
||||
stream?: (chunk: {text?: string, tool?: string, done?: true}) => any;
|
||||
/** Compress old messages in the chat to free up context */
|
||||
compress?: {
|
||||
/** Trigger chat compression once context exceeds the token count */
|
||||
max: number;
|
||||
/** Compress chat until context size smaller than */
|
||||
min: number
|
||||
},
|
||||
/** Background information the AI will be fed */
|
||||
memory?: LLMMemory[],
|
||||
compress?: {max: number; min: number};
|
||||
/** User's memory documents - RAG injected automatically each turn */
|
||||
memory?: Memory[];
|
||||
/** Model to use for memory operations */
|
||||
memoryModel?: string;
|
||||
/** Skill documents the AI can browse and read on demand */
|
||||
skills?: Skill[];
|
||||
/** MCP servers to connect and expose as tools */
|
||||
mcp?: McpServer[];
|
||||
}
|
||||
|
||||
export type McpServer = {
|
||||
/** MCP server name for humans */
|
||||
name: string;
|
||||
/** Host URL */
|
||||
host: string;
|
||||
/** Server access token */
|
||||
token?: string;
|
||||
}
|
||||
|
||||
export type Skill = {
|
||||
/** Name of skill for humans */
|
||||
name: string;
|
||||
/** Description LLM will use to decide to learn a skill */
|
||||
description: string;
|
||||
/** Skill instructions */
|
||||
content: string;
|
||||
}
|
||||
|
||||
|
||||
class LLM {
|
||||
private memoryManager!: MemoryManager;
|
||||
|
||||
defaultModel!: string;
|
||||
models: {[model: string]: LLMProvider} = {};
|
||||
|
||||
@@ -85,14 +96,67 @@ class LLM {
|
||||
else if(config.proto == 'ollama') this.models[model] = new OpenAi(this.ai, config.host, 'not-needed', model);
|
||||
else if(config.proto == 'openai') this.models[model] = new OpenAi(this.ai, config.host || null, config.token, model);
|
||||
});
|
||||
this.memoryManager = new MemoryManager(this);
|
||||
}
|
||||
|
||||
private async setupMcp(servers: McpServer[] = []): Promise<{prompt: string, tools: AiTool[]}> {
|
||||
if(!servers?.length) return {prompt: '', tools: []};
|
||||
const allTools: AiTool[] = [];
|
||||
await Promise.all(servers.map(async server => {
|
||||
const res = await fetch(`${server.host}/tools`, {headers: server.token ? {Authorization: `Bearer ${server.token}`} : {}});
|
||||
const mcp: any = await res.json();
|
||||
if(!mcp?.tools) return;
|
||||
for(const t of mcp.tools) {
|
||||
const args: Record<string, any> = {};
|
||||
if(t.inputSchema?.properties) {
|
||||
for(const [key, val] of Object.entries<any>(t.inputSchema.properties)) {
|
||||
args[key] = {type: val.type || 'string', description: val.description || '', required: t.inputSchema.required?.includes(key)};
|
||||
}
|
||||
}
|
||||
allTools.push({
|
||||
name: `${server.name}_${t.name}`,
|
||||
description: t.description || '',
|
||||
args,
|
||||
fn: async (a: any) => {
|
||||
const r = await fetch(`${server.host}/tools/call`, {
|
||||
method: 'POST',
|
||||
headers: {'Content-Type': 'application/json', ...(server.token ? {Authorization: `Bearer ${server.token}`} : {})},
|
||||
body: JSON.stringify({name: t.name, arguments: a})
|
||||
});
|
||||
const data: any = await r.json();
|
||||
return data?.content?.[0]?.text ?? JSON.stringify(data);
|
||||
}
|
||||
});
|
||||
}
|
||||
}));
|
||||
|
||||
const list = allTools.map(t => `- ${t.name}: ${t.description}`).join('\n');
|
||||
return {
|
||||
prompt: `You have access to the following MCP tools:\n${list}`,
|
||||
tools: allTools
|
||||
};
|
||||
}
|
||||
|
||||
private setupSkills(skills: Skill[] = []): {prompt: string, tools: AiTool[]} {
|
||||
if(!skills?.length) return {prompt: '', tools: []};
|
||||
const list = skills.map(s => `- ${s.name}: ${s.description}`).join('\n');
|
||||
return {
|
||||
prompt: `You have access to the following skill documents, use \`read_skill\` to access them:\n${list}`,
|
||||
tools: [{
|
||||
name: 'read_skill',
|
||||
description: 'Read the full content of a skill/knowledge document',
|
||||
args: {
|
||||
name: {type: 'string', description: 'Exact skill name', required: true}
|
||||
},
|
||||
fn: (args: any) => {
|
||||
const skill = skills.find(s => s.name === args.name);
|
||||
if(!skill) return `Skill not found. Available:\n${list}`;
|
||||
return `# ${skill.name}\n${skill.content}`;
|
||||
}
|
||||
}]
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Chat with LLM
|
||||
* @param {string} message Question
|
||||
* @param {LLMRequest} options Configuration options and chat history
|
||||
* @returns {{abort: () => void, response: Promise<string>}} Function to abort response and chat history
|
||||
*/
|
||||
ask(message: string, options: LLMRequest = {}): AbortablePromise<string> {
|
||||
options = <any>{
|
||||
system: '',
|
||||
@@ -106,71 +170,51 @@ class LLM {
|
||||
if(!this.models[m]) throw new Error(`Model does not exist: ${m}`);
|
||||
let abort = () => {};
|
||||
return Object.assign(new Promise<string>(async res => {
|
||||
let tools: AiTool[] = options.tools || this.ai.options.llm?.tools || [];
|
||||
const prompts: string[] = [options.system || this.ai.options.llm?.system || ''];
|
||||
if(!options.history) options.history = [];
|
||||
// If memories were passed, find any relevant ones and add a tool for ADHOC lookups
|
||||
if(options.memory) {
|
||||
const search = async (query?: string | null, subject?: string | null, limit = 10) => {
|
||||
const [o, q] = await Promise.all([
|
||||
subject ? this.embedding(subject) : Promise.resolve(null),
|
||||
query ? this.embedding(query) : Promise.resolve(null),
|
||||
]);
|
||||
return (options.memory || []).map(m => {
|
||||
const score = (o ? this.cosineSimilarity(m.embeddings[0], o[0].embedding) : 0)
|
||||
+ (q ? this.cosineSimilarity(m.embeddings[1], q[0].embedding) : 0);
|
||||
return {...m, score};
|
||||
}).toSorted((a: any, b: any) => a.score - b.score).slice(0, limit)
|
||||
.map(m => `- ${m.owner}: ${m.fact}`).join('\n');
|
||||
}
|
||||
|
||||
options.system += '\nYou have RAG memory and will be given the top_k closest memories regarding the users query. Save anything new you have learned worth remembering from the user message using the remember tool and feel free to recall memories manually.\n';
|
||||
const relevant = await search(message);
|
||||
if(relevant.length) options.history.push({role: 'tool', name: 'recall', id: 'auto_recall_' + Math.random().toString(), args: {}, content: `Things I remembered:\n${relevant}`});
|
||||
options.tools = [{
|
||||
name: 'recall',
|
||||
description: 'Recall the closest memories you have regarding a query using RAG',
|
||||
args: {
|
||||
subject: {type: 'string', description: 'Find information by a subject topic, can be used with or without query argument'},
|
||||
query: {type: 'string', description: 'Search memory based on a query, can be used with or without subject argument'},
|
||||
topK: {type: 'number', description: 'Result limit, default 5'},
|
||||
},
|
||||
fn: (args) => {
|
||||
if(!args.subject && !args.query) throw new Error('Either a subject or query argument is required');
|
||||
return search(args.query, args.subject, args.topK);
|
||||
}
|
||||
}, {
|
||||
name: 'remember',
|
||||
description: 'Store important facts user shares for future recall',
|
||||
args: {
|
||||
owner: {type: 'string', description: 'Subject/person this fact is about'},
|
||||
fact: {type: 'string', description: 'The information to remember'}
|
||||
},
|
||||
fn: async (args) => {
|
||||
if(!options.memory) return;
|
||||
const e = await Promise.all([
|
||||
this.embedding(args.owner),
|
||||
this.embedding(`${args.owner}: ${args.fact}`)
|
||||
]);
|
||||
const newMem = {owner: args.owner, fact: args.fact, embeddings: <any>[e[0][0].embedding, e[1][0].embedding]};
|
||||
options.memory.splice(0, options.memory.length, ...[
|
||||
...options.memory.filter(m => {
|
||||
return !(this.cosineSimilarity(newMem.embeddings[0], m.embeddings[0]) >= 0.9 && this.cosineSimilarity(newMem.embeddings[1], m.embeddings[1]) >= 0.8);
|
||||
}),
|
||||
newMem
|
||||
]);
|
||||
return 'Remembered!';
|
||||
}
|
||||
}, ...options.tools || []];
|
||||
// MCP
|
||||
const mcp = options.mcp || this.ai.options?.llm?.mcp;
|
||||
if(mcp?.length) {
|
||||
const m = await this.setupMcp(mcp);
|
||||
prompts.unshift(m.prompt);
|
||||
tools.push(...m.tools);
|
||||
}
|
||||
|
||||
// Ask
|
||||
const resp = await this.models[m].ask(message, options);
|
||||
// Skills
|
||||
const skills = options.skills || this.ai.options?.llm?.skills;
|
||||
if(skills?.length) {
|
||||
const s = this.setupSkills(skills);
|
||||
prompts.unshift(s.prompt);
|
||||
tools.push(...s.tools);
|
||||
}
|
||||
|
||||
// Remove any memory calls from history
|
||||
if(options.memory) options.history.splice(0, options.history.length, ...options.history.filter(h => h.role != 'tool' || (h.name != 'recall' && h.name != 'remember')));
|
||||
// Memory
|
||||
if(options.memory) {
|
||||
const relevant = await this.memoryManager.recollect(message, options.memory);
|
||||
if(relevant.length) {
|
||||
const context = relevant.map(m => `### ${m.name}\n${m.content}`).join('\n\n');
|
||||
options.history.push({
|
||||
id: 'auto_recall_' + Math.random().toString(), role: 'tool', name: 'recall', args: {},
|
||||
content: `Knowledge Documents:\n\n${context}`
|
||||
});
|
||||
}
|
||||
prompts.unshift('You have access to a knowledge base. Relevant documents are injected automatically before each message. Use this knowledge to inform your responses.');
|
||||
}
|
||||
|
||||
// Compress message history
|
||||
const resp = await this.models[m].ask(message, {...options, tools, system: prompts.filter(Boolean).join('\n\n')});
|
||||
|
||||
// Trim memory injections from history
|
||||
if(options.memory) {
|
||||
options.history.splice(0, options.history.length, ...options.history.filter(h =>
|
||||
h.role !== 'tool' || h.name !== 'recall'));
|
||||
}
|
||||
|
||||
// Auto-memorize before compressing
|
||||
if(options.compress) {
|
||||
const compressed = await this.ai.language.compressHistory(options.history, options.compress.max, options.compress.min, options);
|
||||
if(options.memory) await this.memoryManager.memorize(options.history, options.memory, options);
|
||||
const compressed = await this.compressHistory(options.history, options.compress.max, options.compress.min, options);
|
||||
options.history.splice(0, options.history.length, ...compressed);
|
||||
}
|
||||
|
||||
@@ -178,13 +222,12 @@ class LLM {
|
||||
}), {abort});
|
||||
}
|
||||
|
||||
async code(message: string, options?: LLMRequest): Promise<any> {
|
||||
const resp = await this.ask(message, {...options, system: [
|
||||
options?.system,
|
||||
'Return your response in a code block'
|
||||
].filter(t => !!t).join(('\n'))});
|
||||
const codeBlock = /```(?:.+)?\s*([\s\S]*?)```/.exec(resp);
|
||||
return codeBlock ? codeBlock[1].trim() : null;
|
||||
/**
|
||||
* Digest full conversation history into memory documents.
|
||||
* Call on session end to persist the conversation.
|
||||
*/
|
||||
async updateMemory(history: LLMMessage[], memories: Memory[], options: LLMRequest = {}): Promise<void> {
|
||||
await this.memoryManager.memorize(history, memories, {model: this.defaultModel, ...options});
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -273,7 +316,7 @@ class LLM {
|
||||
* @param {maxTokens?: number, overlapTokens?: number} opts Options for embedding such as chunk sizes
|
||||
* @returns {Promise<Awaited<{index: number, embedding: number[], text: string, tokens: number}>[]>} Chunked embeddings
|
||||
*/
|
||||
embedding(target: object | string, opts: {maxTokens?: number, overlapTokens?: number} = {}): AbortablePromise<any[]> {
|
||||
embedding(target: object | string, opts: {maxTokens?: number, overlapTokens?: number} = {}): AbortablePromise<{index: number, embedding: number[], text: string, tokens: number}[]> {
|
||||
let {maxTokens = 500, overlapTokens = 50} = opts;
|
||||
let aborted = false;
|
||||
const abort = () => { aborted = true; };
|
||||
@@ -281,7 +324,6 @@ class LLM {
|
||||
const embed = (text: string): Promise<number[]> => {
|
||||
return new Promise((resolve, reject) => {
|
||||
if(aborted) return reject(new Error('Aborted'));
|
||||
|
||||
const args: string[] = [
|
||||
join(dirname(fileURLToPath(import.meta.url)), 'embedder.js'),
|
||||
<string>this.ai.options.path,
|
||||
@@ -290,7 +332,6 @@ class LLM {
|
||||
const proc = spawn('node', args, {stdio: ['pipe', 'pipe', 'ignore']});
|
||||
proc.stdin.write(text);
|
||||
proc.stdin.end();
|
||||
|
||||
let output = '';
|
||||
proc.stdout.on('data', (data: Buffer) => output += data.toString());
|
||||
proc.on('close', (code: number) => {
|
||||
@@ -300,7 +341,7 @@ class LLM {
|
||||
const result = JSON.parse(output);
|
||||
resolve(result.embedding);
|
||||
} catch(err) {
|
||||
reject(new Error('Failed to parse embedding output'));
|
||||
reject(err);
|
||||
}
|
||||
} else {
|
||||
reject(new Error(`Embedder process exited with code ${code}`));
|
||||
@@ -320,7 +361,7 @@ class LLM {
|
||||
}
|
||||
return results;
|
||||
})();
|
||||
return Object.assign(p, { abort });
|
||||
return <any>Object.assign(p, {abort});
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -346,8 +387,8 @@ class LLM {
|
||||
(char.charCodeAt(0) * (index + 1)) % dimensions / dimensions).slice(0, dimensions);
|
||||
}
|
||||
const v = vector(target);
|
||||
const similarities = searchTerms.map(t => vector(t)).map(refVector => this.cosineSimilarity(v, refVector))
|
||||
return {avg: similarities.reduce((acc, s) => acc + s, 0) / similarities.length, max: Math.max(...similarities), similarities}
|
||||
const similarities = searchTerms.map(t => vector(t)).map(refVector => this.cosineSimilarity(v, refVector));
|
||||
return {avg: similarities.reduce((acc, s) => acc + s, 0) / similarities.length, max: Math.max(...similarities), similarities};
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -358,7 +399,7 @@ class LLM {
|
||||
* @returns {Promise<{} | {} | RegExpExecArray | null>}
|
||||
*/
|
||||
async json(text: string, schema: string, options?: LLMRequest): Promise<any> {
|
||||
let system = `Your job is to convert input to JSON using tool calls. Call the \`submit\` tool at least once with JSON matching this schema:\n\`\`\`json\n${schema}\n\`\`\`\n\nResponses are ignored`;
|
||||
let system = `Your job is to convert input to JSON using tool calls. Call the \`submit\` tool at least once with JSON matching this schema:\n\`\`\`json\n${schema}\n\`\`\`\n\nResponses are ignored`;
|
||||
if(options?.system) system += '\n\n' + options.system;
|
||||
return new Promise(async (resolve, reject) => {
|
||||
let done = false;
|
||||
@@ -392,7 +433,7 @@ class LLM {
|
||||
* @returns {Promise<string>} Summary
|
||||
*/
|
||||
async summarize(text: string, length: number = 500, options?: LLMRequest): Promise<string | null> {
|
||||
let system = `Your job is to summarize the users message using tool calls. Call the \`submit\` tool at least once with the shortest summary possible that's <= ${length} words. The tool call will respond with the token count. Responses are ignored`;
|
||||
let system = `Your job is to summarize the users message using tool calls. Call the \`submit\` tool at least once with the shortest summary possible that's <= ${length} words. The tool call will respond with the token count. Responses are ignored`;
|
||||
if(options?.system) system += '\n\n' + options.system;
|
||||
return new Promise(async (resolve, reject) => {
|
||||
let done = false;
|
||||
|
||||
Reference in New Issue
Block a user