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