LangChain.jsで「関連性の高い記事」を検索する

OpenAIのGPT-3で遊ぶためにLangChain.jsを触っていたのですが、「どうも関連記事検索ができそうだ」と感じた部分があったので、ひとまずまとめてみました。

やろうとしていること

• 特定のキーワードやテキストに対して、関連性の高いデータを表示する
• データはREST APIなどで事前に取得し、LangChain.jsに判断をまかせる
• すべての記事を対象にはせず、「最近の記事50件」などある程度フィルタする

ライブラリのインストール

まずはライブラリを追加します。

% yarn add langchain hnswlib-node

langchainと一緒にhnswlib-nodeも追加しましょう。

Langchain.jsのサンプルコードを動かす

まずはサンプルコードを動かしましょう。

元コード: https://github.com/hwchase17/langchainjs/blob/main/examples/src/indexes/vector_stores/hnswlib.ts

import { HNSWLib } from "langchain/vectorstores";
import { OpenAIEmbeddings } from "langchain/embeddings";

export const run = async () => {
  const vectorStore = await HNSWLib.fromTexts(
    ["Hello world", "Bye bye", "hello nice world"],
    [{ id: 2 }, { id: 1 }, { id: 3 }],
    new OpenAIEmbeddings()
  );

  const resultOne = await vectorStore.similaritySearch("hello world", 1);
  console.log(resultOne);
};

このスクリプトを動かすと、hello worldと関連の高いデータを1件取得できます。

[ Document { pageContent: 'Hello world', metadata: { id: 2 } } ]

OpenAIのAPIも使わずに判定していますので、APIキーなどは不要です。

WordPressの記事データを使ってみる

サンプルコードのおかげで、渡すデータのフォーマットが大体把握できました。

早速WP APIをfetchで呼び出してデータ元データにしてみましょう。

import fetch from "node-fetch";
import { HNSWLib } from "langchain/vectorstores";
import { OpenAIEmbeddings } from "langchain/embeddings";

export const run = async () => {
  const result = await fetch(`https://wp-api.example.com/wp-json/wp/v2/posts?per_page=50`)
  const resposne = await result.json()
  const vectorStore = await HNSWLib.fromTexts(
    (resposne as any[]).map((post) => post.title.rendered),
    (resposne as any[]).map((post) => ({id: post.id})),
    new OpenAIEmbeddings()
  );

  const resultOne = await vectorStore.similaritySearch("ReactからAstroへの移行", 1);
  console.log(resultOne);
};
run();

動かすの優先でas anyを使っていますが、実際に使う場合は型定義をつけることをお勧めします。

これを実行すると、ReactからAstroへの移行と関連性の高い記事を、取得した50記事の中から1つ選んでくれます。

[
  Document {
    pageContent: 'ReactコンポーネントをAstroコンポーネントに書き換える際の、クラス名の取り扱いについて',
    metadata: { id: 11960 }
  }
]

「今表示している記事」を除外する

WP APIで取得した記事一覧の中には、「今見ている記事」も含まれることがあります。

その場合は、filterを使って結果から除外するとよいでしょう。

  const results = await vectorStore.similaritySearch("ReactからAstroへの移行", 5);
  const searchTargets = (results as any[]).filter(post => post.id === resultOne[0].metadata.id)

フリーワード検索として利用する

検索した結果をそのままページ表示などに使いたい場合、例えばmapを使ってIDから記事データを取得すると良いかもしれません。

export const run = async () => {
  const result = await fetch(`https://wp-api.example/wp-json/wp/v2/posts?per_page=50`)
  const resposne = await result.json()
  const vectorStore = await HNSWLib.fromTexts(
    (resposne as any[]).map((post) => post.title.rendered),
    (resposne as any[]).map((post) => ({id: post.id})),
    new OpenAIEmbeddings()
  );

  const searchResult = await vectorStore.similaritySearch("AngularとReact", 3);
  return searchResult.map(item => {
    const post = (resposne as any[]).find(post => post.id === item.metadata.id)
    return post
  })
};
run().then(posts => console.log(posts.map(post => ({id: post.id, title: post.title.rendered}))))

上のスクリプトは、このような結果がでます。

[
  { id: 11371, title: 'Angular + nxrxにRedux Toolkitを追加してみた覚書' },
  { id: 11804, title: 'Nxで、Ionic(React)アプリを開発する方法' },
  { id: 11378, title: 'Ionic Angularで汎用的に利用したいカスタムComponentを作る方法の覚書' }
]

関連記事表示機能として使う

もしくは表示中の記事のタイトルを使って、関連記事を表示させることも可能そうです。

const POST_TITLE = 'ReactコンポーネントをAstroコンポーネントに書き換える際の、クラス名の取り扱いについて'
export const run = async () => {
  const result = await fetch(`https://wp-api.example/wp-json/wp/v2/posts?per_page=50`)
  const resposne = await result.json()
  const vectorStore = await HNSWLib.fromTexts(
    (resposne as any[]).map((post) => post.title.rendered),
    (resposne as any[]).map((post) => ({id: post.id})),
    new OpenAIEmbeddings()
  );

  const searchResult = await vectorStore.similaritySearch(POST_TITLE, 4);
  return searchResult.map(item => {
    const post = (resposne as any[]).find(post => post.id === item.metadata.id)
    return post
  }).filter(post => post.title.rendered !== POST_TITLE)
};
run().then(posts => console.log(posts.map(post => ({id: post.id, title: post.title.rendered}))))

この場合は、表示中の記事が含まれないようにfilterなどで除外してやりましょう。