Updated wiki tool

package.json

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

src/tools.ts

@@ -306,93 +306,90 @@ export const WebSearchTool: AiTool = {
 	}
 }
 
+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 {
-	private async get(url: string): Promise<any> {
+			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 {
+			clean(text: string) {
-		return text.replace(/\n{3,}/g, '\n\n').replace(/ {2,}/g, ' ').replace(/\[\d+\]/g, '').trim();
+				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 {
+				return text
-		if(text.length <= max) return text;
+					.replace(/^={4}\s*(.+?)\s*={4}$/gm, '#### $1')
-		const cut = text.slice(0, max);
+					.replace(/^={3}\s*(.+?)\s*={3}$/gm, '### $1')
-		const lastPara = cut.lastIndexOf('\n\n');
+					.replace(/^={2}\s*(.+?)\s*={2}$/gm, '## $1')
-		return lastPara > max * 0.7 ? cut.slice(0, lastPara) : cut;
+					.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 introOnly = detail !== 'full';
-		const text = this.truncate(content, detail === 'intro' ? 2000 : 8000);
+				const content = await this.fetchExtract(title, introOnly);
-		return `## ${title}\n🔗 ${url}\n\n${text}`;
+				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');
+		if(args.mode == 'search') return wiki.search(args.query);
-	}
+		return wiki.lookup(args.query, args.mode || 'summary');
-};
-
-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);
 	}
 };