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 = {
 	}
 }
 
-class WikipediaClient {
+export const WikipediaTool: AiTool = {
-	private async get(url: string): Promise<any> {
-		const resp = await fetch(url, {headers: {'User-Agent': UA}});
-		return resp.json();
-	}
-
-	private api(params: Record<string, any>): Promise<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();
-	}
-
-	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;
-	}
-
-	private async searchTitles(query: string, limit = 6): Promise<any[]> {
-		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> {
-		const params: any = {action: 'query', prop: 'extracts', titles: title, explaintext: 1, redirects: 1};
-		if(intro) params.exintro = 1;
-		const data = await this.api(params);
-		const page = Object.values(data.query?.pages || {})[0] as any;
-		return this.clean(page?.extract || '');
-	}
-
-	private pageUrl(title: string): string {
-		return `https://en.wikipedia.org/wiki/${encodeURIComponent(title.replace(/ /g, '_'))}`;
-	}
-
-	private stripHtml(text: string): string {
-		return text.replace(/<[^>]+>/g, '');
-	}
-
-	async lookup(query: string, detail: 'intro' | 'full' = 'intro'): Promise<string> {
-		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}`;
-	}
-
-	async search(query: string): Promise<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);
-			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}
+		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}) => {
+	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();
+			}
+
+			api(params: any) {
+				const qs = new URLSearchParams({...params, format: 'json', utf8: '1'}).toString();
+				return this.get(`https://en.wikipedia.org/w/api.php?${qs}`);
+			}
+
+			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);
+				}
+
+				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();
+			}
+
+			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 || [];
+			}
+
+			async fetchExtract(title: string, introOnly = false) {
+				const params: any = {action: 'query', prop: 'extracts', titles: title, explaintext: 1, redirects: 1};
+				if(introOnly) params.exintro = 1;
+				const data = await this.api(params);
+				const page: any = Object.values(data.query?.pages || {})[0];
+				return this.clean(page?.extract || '');
+			}
+
+			pageUrl(title: string) {
+				return `https://en.wikipedia.org/wiki/${encodeURIComponent(title.replace(/ /g, '_'))}`;
+			}
+
+			stripHtml(text: string) {
+				return text.replace(/<[^>]+>/g, '');
+			}
+
+			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 introOnly = detail !== 'full';
+				const content = await this.fetchExtract(title, introOnly);
+				return `## ${title}\n🔗 ${url}\n\n${content}`;
+			}
+
+			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.stripHtml(r.snippet || '').trim();
+					lines.push(`**${i + 1}. ${r.title}**\n${snippet}\n${this.pageUrl(r.title)}`);
+				}
+				return lines.join('\n\n');
+			}
+		}
+
 		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');
 	}
 };