👏

AnthropicAI Tool で Retrieval-Augmented Generation を実装してみた

2024/04/26に公開

LangChain なんか使わなくてもシュッと作れたので記事にしておく。

RAG とは

生成AIに検索能力をもたせるやつ。

https://atmarkit.itmedia.co.jp/ait/articles/2403/13/news035.html

要は検索機能をこちらで提供してやって、AIにそれを読ませる。

AnthropicAI Tool

OpenAI でいう Function Calling

JSONSchema で関数シグネチャを与えると、それを使うDSLを生成する。実際の関数は自分で実装して、AI が生成した引数(JSONSchema に従う)を渡す。

https://docs.anthropic.com/claude/docs/tool-use

TypeScript の Mapped Types でツールの実装部分に型をつける簡単なラッパーを書いた。

https://jsr.io/@mizchi/anthropic-helper@0.0.3

RAG の CLI を作る

  • Google検索をするAPIを実装
    • Google Custom Engine API を使った
  • 本文要約をするAPIを実装
    • Mozilla の実装を使った
  • 与えられた URL を fetch して、その本文部分を抽出する Tool を定義
  • それらを使ってくれるようなシステムプロンプトを実装

AnthropicAI の API_KEY と Google Custom Search Engine のAPIの取得は略

こういう環境変数があればよい。

.env
ANTHROPIC_API_KEY=
GOOGLE_API_KEY=
GOOGLE_CSE_ID=
import AnthropicAI from "npm:@anthropic-ai/sdk@0.20.7";
import { createMessageHandler, createToolsHandler } from 'jsr:@mizchi/anthropic-helper@0.0.3';

// システムプロンプト
const system = `
あなたは優秀なアシスタントです。あなたはユーザーからの質問に可能な限り正確に答えることができます。

もしあなたが知らない場合は、Google で検索して情報を取得することができます。
URL について聞かれた場合は、その URL を開いて本文を取得して答えてください。
いずれもその先の URL に詳細な情報がある場合は、その URL を開いて詳細な情報を取得してください。
`;

import { Readability } from "npm:@mozilla/readability@0.5.0";
import { JSDOM } from "npm:jsdom@16.4.0";
import { google } from "npm:googleapis@61.0.0";

function extractMainContent(html: string) {
  const doc = new JSDOM(html).window.document;
  const reader = new Readability(doc);
  const article = reader.parse();
  return article?.content;
}

async function searchGoogle(query: string) {
  const res = await google.customsearch('v1').cse.list({
    key: Deno.env.get('GOOGLE_API_KEY')!,
    cx: Deno.env.get('GOOGLE_CSE_ID')!,
    q: query
  });
  return res.data.items?.slice(0, 5).map((item) => {
    // 不要なデータを捨てる
    return {
      title: item.title,
      link: item.link,
      snippet: item.snippet,
      kind: item.kind,
      labels: item.labels
    }
  });
}

async function runAnthropicAITools(
  options:
    | Partial<AnthropicAI.Beta.Tools.Messages.MessageCreateParamsNonStreaming>
    & Pick<AnthropicAI.Beta.Tools.Messages.MessageCreateParamsNonStreaming, "messages" | "tools">
): Promise<AnthropicAI.Beta.Tools.Messages.ToolsBetaMessage> {
  const apiKey = Deno.env.get("ANTHROPIC_API_KEY")!;
  const client = new AnthropicAI({ apiKey });
  const res = await client.beta.tools.messages.create({
    model: "claude-3-opus-20240229",
    max_tokens: 1024,
    system,
    ...options,
  });
  return res;
}

// ここでAIが呼び出せるツールを定義
const TOOLS = [
  {
    name: "ask_to_user",
    description: "Ask to user and get response. If you feel given message is not enough, you can use this tool.",
    input_schema: {
      type: "object",
      properties: {
        question: {
          type: "string",
          description: "question to ask"
        }
      },
      required: ["question"]
    }
  },
  {
    name: "open_url",
    description: "open url and extract main content",
    input_schema: {
      type: "object",
      properties: {
        url: {
          type: "string",
          description: "url to open"
        }
      },
      required: ["url"]
    }
  },
  {
    name: "search_google",
    description: "search google. If you want to search word, you can use this tool.",
    input_schema: {
      type: "object",
      properties: {
        query: {
          type: "string",
          description: "search query"
        }
      },
      required: ["query"]
    }
  },

] as const;

// Tool の実装を TOOLS の定義から型を生成
const handleTool = createToolsHandler(TOOLS, {
  async ask_to_user(input, content) {
    const res = prompt(input.question);
    return {
      tool_use_id: content.id,
      type: 'tool_result',
      content: [
        { type: 'text', text: res ?? "no answer" }
      ],
      is_error: false
    };
  },
  async search_google(input, content) {
    try {
      const result = await searchGoogle(input.query);
      return {
        tool_use_id: content.id,
        type: 'tool_result',
        content: [
          { type: 'text', text: JSON.stringify(result, null, 2) }
        ],
        is_error: false
      };
    } catch (err) {
      return {
        tool_use_id: content.id,
        type: 'tool_result',
        content: [
          { type: 'text', text: 'failed to search' }
        ],
        is_error: true
      };
    }
  },
  async open_url(input, content) {
    try {
      const res = await fetch(input.url).then((res) => res.text());
      const main = extractMainContent(res) || res;
      return {
        tool_use_id: content.id,
        type: 'tool_result',
        content: [
          { type: 'text', text: main }
        ],
        is_error: false
      };
    } catch (err) {
      return {
        tool_use_id: content.id,
        type: 'tool_result',
        content: [
          { type: 'text', text: 'failed to fetch' }
        ],
        is_error: true
      };
    }
  },
});

const handler = createMessageHandler({
  tools: TOOLS as any,
  handleTool,
  messages: [
    {
      role: 'user',
      content: Deno.args.join(" ")
    }
  ]
});

while (!handler.isEnd()) {
  const res = await runAnthropicAITools({
    tools: TOOLS as any,
    messages: handler.current()
  });
  await handler.handleResponse(res);
}

できた。

AIが知らないであろう、最近発表されたゲームについて質問する。

$ deno run -A --env run.ts "Slay the Spire 2 について教えて"
Slay the Spire 2 について教えて
<thinking>
ユーザーから "Slay the Spire 2 について教えて" という質問を受けました。
Slay the Spire 2 についての情報を Google 検索することで答えられそうです。
検索に必要なパラメータは query だけなので、ユーザーから十分な情報が得られていると判断します。
</thinking>
[tool_use] search_google { query: "Slay the Spire 2" }
[tool_result] [{"type":"text","text":"[\n  {...
<search_quality_reflection>
検索結果より、Slay the Spire 2 は 2025 年リリース予定の続編ゲームであることがわかりました。Steam ページの情報や Reddit での議論を見ると、ゲームの概要についてもある程度の情報が得られました。ただし、ゲームの詳細な内容についてはまだ公開されていないようです。現時点ではこの程度の情報で十分でしょう。
</search_quality_reflection>

<search_quality_score>4</search_quality_score>

<result>
Slay the Spire 2 は、2025年にリリース予定のローグライクデッキビルディングゲーム「Slay the Spire」の続編です。

Steam のストアページによると、Slay the Spire 2 では以下のような特徴が予定されています。

- 前作よりも巨大になった塔を舞台に、新たな冒険が繰り広げられる
- さらに多様なデッキ構築要素が導入される
- 塔の内部構造がプレイのたびに変化する
- 新しいプレイモードが追加される 
- ゲームエンジンが一新され、グラフィックや操作性が向上する

リリースまではまだ時間があるため、ゲームの詳細な内容は今後徐々に明らかになっていくと思われます。ファンの間でも続編への期待が高まっているようです。オリジナルの Slay the Spire が非常に高い評価を得ているだけに、Slay the Spire 2 にも大きな注目が集まるでしょう。
</result>

うごいてそう。

Discussion