Skip to content
Snippets Groups Projects
Commit dfcb9a47 authored by thucpn's avatar thucpn
Browse files

wiki -> wikipedia

parent 4e517718
No related branches found
No related tags found
No related merge requests found
import { tool } from "@llamaindex/core/tools"; import { tool } from "@llamaindex/core/tools";
import { default as wikipedia } from "wikipedia"; import { default as wiki } from "wikipedia";
import { z } from "zod"; import { z } from "zod";
export type WikiToolOutput = { export type WikiToolOutput = {
...@@ -7,7 +7,7 @@ export type WikiToolOutput = { ...@@ -7,7 +7,7 @@ export type WikiToolOutput = {
content: string; content: string;
}; };
export const wiki = () => { export const wikipedia = () => {
return tool({ return tool({
name: "wikipedia", name: "wikipedia",
description: "Use this function to search Wikipedia", description: "Use this function to search Wikipedia",
...@@ -16,11 +16,11 @@ export const wiki = () => { ...@@ -16,11 +16,11 @@ export const wiki = () => {
lang: z.string().describe("The language to search in").default("en"), lang: z.string().describe("The language to search in").default("en"),
}), }),
execute: async ({ query, lang }): Promise<WikiToolOutput> => { execute: async ({ query, lang }): Promise<WikiToolOutput> => {
wikipedia.setLang(lang); wiki.setLang(lang);
const searchResult = await wikipedia.search(query); const searchResult = await wiki.search(query);
const pageTitle = searchResult?.results[0]?.title; const pageTitle = searchResult?.results[0]?.title;
if (!pageTitle) return { title: "No search results.", content: "" }; if (!pageTitle) return { title: "No search results.", content: "" };
const result = await wikipedia.page(pageTitle, { autoSuggest: false }); const result = await wiki.page(pageTitle, { autoSuggest: false });
return { title: pageTitle, content: await result.content() }; return { title: pageTitle, content: await result.content() };
}, },
}); });
......
import { describe, expect, test } from "vitest"; import { describe, expect, test } from "vitest";
import { wiki } from "../src/tools/wiki"; import { wikipedia } from "../src/tools/wikipedia";
describe("Wikipedia Tool", () => { describe("Wikipedia Tool", () => {
test("wiki tool returns content for valid query", async () => { test("wiki tool returns content for valid query", async () => {
const wikipediaTool = wiki(); const wikipediaTool = wikipedia();
const result = await wikipediaTool.call({ const result = await wikipediaTool.call({
query: "Albert Einstein", query: "Albert Einstein",
lang: "en", lang: "en",
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment