import OpenAI from "openai";
import wiki from "wikipedia";
import { Client } from "langsmith";
import { traceable } from "langsmith/traceable";
import { wrapOpenAI } from "langsmith/wrappers";
const openai = wrapOpenAI(new OpenAI());
const generateWikiSearch = traceable(
async (input: { question: string }) => {
const messages = [
{
role: "system" as const,
content:
"Generate a search query to pass into Wikipedia to answer the user's question. Return only the search query and nothing more. This will be passed in directly to the Wikipedia search engine.",
},
{ role: "user" as const, content: input.question },
];
const chatCompletion = await openai.chat.completions.create({
model: "gpt-4.1-mini",
messages: messages,
temperature: 0,
});
return chatCompletion.choices[0].message.content ?? "";
},
{ name: "generateWikiSearch" }
);
const retrieve = traceable(
async (input: { query: string; numDocuments: number }) => {
const { results } = await wiki.search(input.query, { limit: 10 });
const finalResults: Array<{
page_content: string;
type: "Document";
metadata: { url: string };
}> = [];
for (const result of results) {
if (finalResults.length >= input.numDocuments) {
// Just return the top 2 pages for now
break;
}
const page = await wiki.page(result.title, { autoSuggest: false });
const summary = await page.summary();
finalResults.push({
page_content: summary.extract,
type: "Document",
metadata: { url: page.fullurl },
});
}
return finalResults;
},
{ name: "retrieve", run_type: "retriever" }
);
const generateAnswer = traceable(
async (input: { question: string; context: string }) => {
const messages = [
{
role: "system" as const,
content: `Answer the user's question based only on the content below:\n\n${input.context}`,
},
{ role: "user" as const, content: input.question },
];
const chatCompletion = await openai.chat.completions.create({
model: "gpt-4.1-mini",
messages: messages,
temperature: 0,
});
return chatCompletion.choices[0].message.content ?? "";
},
{ name: "generateAnswer" }
);
const ragPipeline = traceable(
async ({ question }: { question: string }, numDocuments: number = 2) => {
const query = await generateWikiSearch({ question });
const retrieverResults = await retrieve({ query, numDocuments });
const context = retrieverResults
.map((result) => result.page_content)
.join("\n\n");
const answer = await generateAnswer({ question, context });
return answer;
},
{ name: "ragPipeline" }
);