|
| 1 | +import type { JSONSchemaType } from "ajv"; |
| 2 | +import type { BaseTool, ToolMetadata } from "llamaindex"; |
| 3 | +import { default as wiki } from "wikipedia"; |
| 4 | + |
| 5 | +type WikipediaParameter = { |
| 6 | + query: string; |
| 7 | + lang?: string; |
| 8 | +}; |
| 9 | + |
| 10 | +export type WikipediaToolParams = { |
| 11 | + metadata?: ToolMetadata<JSONSchemaType<WikipediaParameter>>; |
| 12 | +}; |
| 13 | + |
| 14 | +const DEFAULT_META_DATA: ToolMetadata<JSONSchemaType<WikipediaParameter>> = { |
| 15 | + name: "wikipedia_tool", |
| 16 | + description: "A tool that uses a query engine to search Wikipedia.", |
| 17 | + parameters: { |
| 18 | + type: "object", |
| 19 | + properties: { |
| 20 | + query: { |
| 21 | + type: "string", |
| 22 | + description: "The query to search for", |
| 23 | + }, |
| 24 | + lang: { |
| 25 | + type: "string", |
| 26 | + description: "The language to search in", |
| 27 | + nullable: true, |
| 28 | + }, |
| 29 | + }, |
| 30 | + required: ["query"], |
| 31 | + }, |
| 32 | +}; |
| 33 | + |
| 34 | +export class WikipediaTool implements BaseTool<WikipediaParameter> { |
| 35 | + private readonly DEFAULT_LANG = "en"; |
| 36 | + metadata: ToolMetadata<JSONSchemaType<WikipediaParameter>>; |
| 37 | + |
| 38 | + constructor(params?: WikipediaToolParams) { |
| 39 | + this.metadata = params?.metadata || DEFAULT_META_DATA; |
| 40 | + } |
| 41 | + |
| 42 | + async loadData( |
| 43 | + page: string, |
| 44 | + lang: string = this.DEFAULT_LANG, |
| 45 | + ): Promise<string> { |
| 46 | + wiki.setLang(lang); |
| 47 | + const pageResult = await wiki.page(page, { autoSuggest: false }); |
| 48 | + const content = await pageResult.content(); |
| 49 | + return content; |
| 50 | + } |
| 51 | + |
| 52 | + async call({ |
| 53 | + query, |
| 54 | + lang = this.DEFAULT_LANG, |
| 55 | + }: WikipediaParameter): Promise<string> { |
| 56 | + const searchResult = await wiki.search(query); |
| 57 | + if (searchResult.results.length === 0) return "No search results."; |
| 58 | + return await this.loadData(searchResult.results[0].title, lang); |
| 59 | + } |
| 60 | +} |
0 commit comments