Updated wiki tool

package.json

@@ -1,6 +1,6 @@
 {
 	"name": "@ztimson/ai-utils",
-	"version": "1.1.0",
+	"version": "1.2.1",
 	"description": "AI Utility library",
 	"author": "Zak Timson",
 	"license": "MIT",

src/ai.ts

@@ -1,5 +1,5 @@
 import * as os from 'node:os';
-import LLM, {AnthropicConfig, OllamaConfig, OpenAiConfig, LLMRequest} from './llm';
+import LLM, {AnthropicConfig, OpenAiConfig, LLMRequest} from './llm';
 import { Audio } from './audio.ts';
 import {Vision} from './vision.ts';
 
@@ -18,7 +18,7 @@ export type AiOptions = {
 	embedder?: string;
 	/** Large language models, first is default */
 	llm?: Omit<LLMRequest, 'model'> & {
-		models: {[model: string]: AnthropicConfig | OllamaConfig | OpenAiConfig};
+		models: {[model: string]: AnthropicConfig | OpenAiConfig};
 	}
 	/** OCR model: eng, eng_best, eng_fast */
 	ocr?: string;

src/llm.ts

@@ -9,7 +9,6 @@ import {spawn} from 'node:child_process';
 import {Memory, MemoryManager} from './memory.ts';
 
 export type AnthropicConfig = {proto: 'anthropic', token: string};
-export type OllamaConfig = {proto: 'llama', host: string};
 export type OpenAiConfig = {proto: 'openai', host?: string, token: string};
 
 export type LLMMessage = {
@@ -95,7 +94,6 @@ class LLM {
 		Object.entries(ai.options.llm.models).forEach(([model, config]) => {
 			if(!this.defaultModel) this.defaultModel = model;
 			if(config.proto == 'anthropic') this.models[model] = new Anthropic(this.ai, config.token, model);
-			else if(config.proto == 'llama') this.models[model] = new OpenAi(this.ai, config.host, 'ignored', model, true);
 			else if(config.proto == 'openai') this.models[model] = new OpenAi(this.ai, config.host || null, config.token, model);
 		});
 		this.memoryManager = new MemoryManager(this);
@@ -429,9 +427,8 @@ ${relevant[0].content}
 		});
 	}
 
-	addModel(name: string, config: AnthropicConfig | OllamaConfig | OpenAiConfig, setDefault = false) {
+	addModel(name: string, config: AnthropicConfig | OpenAiConfig, setDefault = false) {
 		if(config.proto == 'anthropic') this.models[name] = new Anthropic(this.ai, config.token, name);
-		else if(config.proto == 'llama') this.models[name] = new OpenAi(this.ai, config.host, 'not-needed', name, true);
 		else if(config.proto == 'openai') this.models[name] = new OpenAi(this.ai, config.host || null, config.token, name);
 		if(setDefault || !this.defaultModel) this.defaultModel = name;
 	}
@@ -443,12 +440,11 @@ ${relevant[0].content}
 		}
 	}
 
-	setModels(models: {[model: string]: AnthropicConfig | OllamaConfig | OpenAiConfig}, replace = true) {
+	setModels(models: {[model: string]: AnthropicConfig | OpenAiConfig}, replace = true) {
 		if(replace) this.models = {};
 		Object.entries(models).forEach(([model, config]) => {
 			if(!this.defaultModel) this.defaultModel = model;
 			if(config.proto == 'anthropic') this.models[model] = new Anthropic(this.ai, config.token, model);
-			else if(config.proto == 'llama') this.models[model] = new OpenAi(this.ai, config.host, 'not-needed', model, true);
 			else if(config.proto == 'openai') this.models[model] = new OpenAi(this.ai, config.host || null, config.token, model);
 		});
 		this.defaultModel = Object.keys(this.models)[0] ?? '';

src/open-ai.ts

@@ -8,7 +8,7 @@ import {convertSchema} from './tools.ts';
 export class OpenAi extends LLMProvider {
 	client!: openAI;
 
-	constructor(public readonly ai: Ai, public readonly host: string | null, public readonly token: string, public model: string, public llama?: boolean) {
+	constructor(public readonly ai: Ai, public readonly host: string | null, public readonly token: string, public model: string) {
 		super();
 		this.client = new openAI(clean({
 			baseURL: host,
@@ -96,13 +96,6 @@ export class OpenAi extends LLMProvider {
 
 			if(options.schema) {
 				const schema = convertSchema(options.schema);
-				if(this.llama) {
-					delete requestParams.tools;
-					requestParams.response_format = {
-						type: 'json_schema',
-						json_schema: {name: 'json', schema}
-					}
-				} else {
 				requestParams.response_format = {
 					type: 'json_schema',
 					json_schema: {
@@ -112,7 +105,6 @@ export class OpenAi extends LLMProvider {
 					}
 				};
 			}
-			}
 
 			let resp: any, isFirstMessage = true;
 			do {

src/tools.ts

@@ -306,93 +306,90 @@ export const WebSearchTool: AiTool = {
 	}
 }
 
-class WikipediaClient {
-	private async get(url: string): Promise<any> {
+export const WikipediaTool: AiTool = {
+	name: 'wikipedia_search',
+	description: 'Search Wikipedia for matching articles',
+	args: {
+		query: {type: 'string', description: 'Search term or article title', required: true},
+		mode: {type: 'string', description: 'search - look for articles, summary - intro of first found article (default), full - complete first found article', enum: ['search', 'summary', 'full'], default: 'summary'}
+	},
+	fn: async (args: {query: string, mode: 'search' | 'summary' | 'full'}) => {
+		const UA = 'Mozilla/5.0 (Windows NT 10.0; Win64; x64)';
+
+		class WikipediaClient {
+			async get(url: string) {
 				const resp = await fetch(url, {headers: {'User-Agent': UA}});
 				return resp.json();
 			}
 
-	private api(params: Record<string, any>): Promise<any> {
+			api(params: any) {
 				const qs = new URLSearchParams({...params, format: 'json', utf8: '1'}).toString();
 				return this.get(`https://en.wikipedia.org/w/api.php?${qs}`);
 			}
 
-	private clean(text: string): string {
-		return text.replace(/\n{3,}/g, '\n\n').replace(/ {2,}/g, ' ').replace(/\[\d+\]/g, '').trim();
+			clean(text: string) {
+				const cutoffs = ['== See also ==', '== References ==', '== Bibliography ==', '== External links =='];
+				for (const marker of cutoffs) {
+					const idx = text.indexOf(marker);
+					if (idx !== -1) text = text.slice(0, idx);
 				}
 
-	private truncate(text: string, max: number): string {
-		if(text.length <= max) return text;
-		const cut = text.slice(0, max);
-		const lastPara = cut.lastIndexOf('\n\n');
-		return lastPara > max * 0.7 ? cut.slice(0, lastPara) : cut;
+				return text
+					.replace(/^={4}\s*(.+?)\s*={4}$/gm, '#### $1')
+					.replace(/^={3}\s*(.+?)\s*={3}$/gm, '### $1')
+					.replace(/^={2}\s*(.+?)\s*={2}$/gm, '## $1')
+					.replace(/\n{3,}/g, '\n\n')
+					.replace(/ {2,}/g, ' ')
+					.replace(/\[\d+\]/g, '')
+					.trim();
 			}
 
-	private async searchTitles(query: string, limit = 6): Promise<any[]> {
+			async searchTitles(query: string, limit = 6) {
 				const data = await this.api({action: 'query', list: 'search', srsearch: query, srlimit: limit, srprop: 'snippet'});
 				return data.query?.search || [];
 			}
 
-	private async fetchExtract(title: string, intro = false): Promise<string> {
+			async fetchExtract(title: string, introOnly = false) {
 				const params: any = {action: 'query', prop: 'extracts', titles: title, explaintext: 1, redirects: 1};
-		if(intro) params.exintro = 1;
+				if(introOnly) params.exintro = 1;
 				const data = await this.api(params);
-		const page = Object.values(data.query?.pages || {})[0] as any;
+				const page: any = Object.values(data.query?.pages || {})[0];
 				return this.clean(page?.extract || '');
 			}
 
-	private pageUrl(title: string): string {
+			pageUrl(title: string) {
 				return `https://en.wikipedia.org/wiki/${encodeURIComponent(title.replace(/ /g, '_'))}`;
 			}
 
-	private stripHtml(text: string): string {
+			stripHtml(text: string) {
 				return text.replace(/<[^>]+>/g, '');
 			}
 
-	async lookup(query: string, detail: 'intro' | 'full' = 'intro'): Promise<string> {
+			async lookup(query: string, detail = 'summary') {
 				const results = await this.searchTitles(query, 6);
 				if(!results.length) return `❌ No Wikipedia articles found for "${query}"`;
 				const title = results[0].title;
 				const url = this.pageUrl(title);
-		const content = await this.fetchExtract(title, detail === 'intro');
-		const text = this.truncate(content, detail === 'intro' ? 2000 : 8000);
-		return `## ${title}\n🔗 ${url}\n\n${text}`;
+				const introOnly = detail !== 'full';
+				const content = await this.fetchExtract(title, introOnly);
+				return `## ${title}\n🔗 ${url}\n\n${content}`;
 			}
 
-	async search(query: string): Promise<string> {
+			async search(query: string) {
 				const results = await this.searchTitles(query, 8);
 				if(!results.length) return `❌ No results for "${query}"`;
 				const lines = [`### Search results for "${query}"\n`];
 				for(let i = 0; i < results.length; i++) {
 					const r = results[i];
-			const snippet = this.truncate(this.stripHtml(r.snippet || ''), 150);
+					const snippet = this.stripHtml(r.snippet || '').trim();
 					lines.push(`**${i + 1}. ${r.title}**\n${snippet}\n${this.pageUrl(r.title)}`);
 				}
 				return lines.join('\n\n');
 			}
-}
+		}
 
-export const WikipediaLookupTool: AiTool = {
-	name: 'wikipedia_lookup',
-	description: 'Get Wikipedia article content',
-	args: {
-		query: {type: 'string', description: 'Topic or article title', required: true},
-		detail: {type: 'string', description: 'Content level: "intro" (summary, default) or "full" (complete article)', enum: ['intro', 'full'], default: 'intro'}
-	},
-	fn: async (args: {query: string; detail?: 'intro' | 'full'}) => {
 		const wiki = new WikipediaClient();
-		return wiki.lookup(args.query, args.detail || 'intro');
-	}
-};
-
-export const WikipediaSearchTool: AiTool = {
-	name: 'wikipedia_search',
-	description: 'Search Wikipedia for matching articles',
-	args: {
-		query: {type: 'string', description: 'Search terms', required: true}
-	},
-	fn: async (args: {query: string}) => {
-		const wiki = new WikipediaClient();
-		return wiki.search(args.query);
+		if(args.mode == 'search') return wiki.search(args.query);
+		return wiki.lookup(args.query, args.mode || 'summary');
 	}
 };