Cloudflare AI Gatewayを利用して、Amazon Bedrockを使ったAPIの利用状況の可視化、分析、ロギングできるか試してみた
2024.05.10どうも！オペレーション部の西村祐二です。
以前、Hono + Cloudflare Workers + AWS SDK for JavaScript v3 + Amazon Bedrockの構成でAIモデルの動作確認用のAPIを作りました。
Hono + Cloudflare Workers + AWS SDK for JavaScript v3の構成でAmazon Bedrockが利用できるか試してみた | DevelopersIO
今回このAPIに対して、AIアプリケーションの利用状況を可視化、分析、キャッシュ、生成された回答のロギングなどの環境を提供するCloudflare AI Gatewayを試してみたいと思います。
Cloudflare AI Gatewayについては下記、公式ブログを参照ください。
AI Gatewayの発表：AIアプリケーションの可観測性、信頼性、スケーラビリティを高める
※2024/05/10時点でCloudflare AI Gatewayはベータ版の状態となります。
作成した検証用APIはどんなAPIか
pathパラメータにモデルID、クエリパラメータにプロンプトを指定すると、AWS SDKからAmazon Bedrock Runtime API経由でモデルを実行し、回答をレスポンスするAPIです。
以下に例を示します。
例
- 利用モデルID：anthropic.claude-v2:1
http://localhost:8787/api/bedrock/anthropic.claude-v2:1?p=hello
- 利用モデルID：anthropic.claude-3-sonnet-20240229-v1:0
http://localhost:8787/api/bedrock/anthropic.claude-3-sonnet-20240229-v1:0?p=hello
- 利用モデルID：amazon.titan-text-express-v1
http://localhost:8787/api/bedrock/amazon.titan-text-express-v1?p=hello
主要なコード
AWS SDKを使用してAmazon Bedrockのサービスを利用する実装をしてます。
import { Hono } from 'hono';
import { BedrockRuntimeClient, InvokeModelCommand } from '@aws-sdk/client-bedrock-runtime';
type Bindings = {
AWS_ACCESS_KEY_ID: string;
AWS_SECRET_ACCESS_KEY: string;
AWS_REGION: string;
};
const bedrock = new Hono<{ Bindings: Bindings }>();
bedrock.get('/:modelId', async (c) => {
const modelId = c.req.param('modelId');
let prompt = c.req.query('p');
let payload = {};
if (modelId.startsWith('anthropic.claude-3')) {
payload = {
anthropic_version: 'bedrock-2023-05-31',
max_tokens: 1000,
messages: [
{
role: 'user',
content: [{ type: 'text', text: prompt }],
},
],
};
} else if (modelId.startsWith('anthropic')) {
// promptを変更
prompt = `Human: ${prompt}\n\nAssistant:`;
payload = {
prompt,
max_tokens_to_sample: 500,
temperature: 0.5,
stop_sequences: ['\n\nHuman:'],
};
} else {
const textGenerationConfig = {
maxTokenCount: 4096,
stopSequences: [],
temperature: 0,
topP: 1,
};
payload = {
inputText: prompt,
textGenerationConfig,
};
}
const { AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_REGION } = c.env;
const client = new BedrockRuntimeClient({
region: AWS_REGION,
credentials: {
accessKeyId: AWS_ACCESS_KEY_ID,
secretAccessKey: AWS_SECRET_ACCESS_KEY,
},
});
const command = new InvokeModelCommand({
body: JSON.stringify(payload),
contentType: 'application/json',
accept: 'application/json',
modelId,
});
const response = await client.send(command);
const decodedResponseBody = new TextDecoder().decode(response.body);
const responseBody = JSON.parse(decodedResponseBody);
const res = responseBody;
return c.json({ modelId, prompt, res });
});
export { bedrock };
Cloudflare AI Gatewayの環境を作成
まずは、Cloudflare上に環境を作成します。
今回、「nishimura-test-ai-app」という名前で環境を作成しました。
API Endpointsをクリックすると、払い出されたエンドポイントを確認できるのでメモしておきましょう。
どうやってCloudflare AI Gatewayと連携するのか
AIを利用するAPIのエンドポイントをCloudflare AI Gatewayが払い出したエンドポイントに差し替えることで、Cloudflare側にデータが連携されてAI Gatewayのコンソール上でAPIの利用状況の可視化、分析、ロギングなど行うことができます。
実装
手っ取り早く結論の実装を知りたい人向け
AWS SDK for JavaScript v3にはミドルウェアスタックという機能があり、AWSのAPIを実行する任意のタイミングで処理を挟むことができます。このミドルウェアの機能を利用します。
Introducing Middleware Stack in Modular AWS SDK for JavaScript | AWS Developer Tools Blog
下記のように修正することで、うまくCloudflare AI Gatewayと連携することができました。
AI_GATEWAY_HOSTにはgateway.ai.cloudflare.com、AI_GATEWAY_PATHには/v1/xxxxxxxxxxxxx/nishimura-test-ai-appを設定しています。
import { Hono } from 'hono';
import { BedrockRuntimeClient, InvokeModelCommand } from '@aws-sdk/client-bedrock-runtime';
type Bindings = {
AWS_ACCESS_KEY_ID: string;
AWS_SECRET_ACCESS_KEY: string;
AWS_REGION: string;
AI_GATEWAY_HOST: string;
AI_GATEWAY_PATH: string;
};
const bedrock = new Hono<{ Bindings: Bindings }>();
bedrock.get('/:modelId', async (c) => {
const modelId = c.req.param('modelId');
let prompt = c.req.query('p');
let payload = {};
if (modelId.startsWith('anthropic.claude-3')) {
payload = {
anthropic_version: 'bedrock-2023-05-31',
max_tokens: 1000,
messages: [
{
role: 'user',
content: [{ type: 'text', text: prompt }],
},
],
};
} else if (modelId.startsWith('anthropic')) {
// promptを変更
prompt = `Human: ${prompt}\n\nAssistant:`;
payload = {
prompt,
max_tokens_to_sample: 500,
temperature: 0.5,
stop_sequences: ['\n\nHuman:'],
};
} else {
const textGenerationConfig = {
maxTokenCount: 4096,
stopSequences: [],
temperature: 0,
topP: 1,
};
payload = {
inputText: prompt,
textGenerationConfig,
};
}
const { AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_REGION, AI_GATEWAY_HOST, AI_GATEWAY_PATH } = c.env;
const client = new BedrockRuntimeClient({
region: AWS_REGION,
credentials: {
accessKeyId: AWS_ACCESS_KEY_ID,
secretAccessKey: AWS_SECRET_ACCESS_KEY,
},
});
client.middlewareStack.add(
(next, context) => async (args: any) => {
args.request.headers['host'] = AI_GATEWAY_HOST;
args.request.hostname = AI_GATEWAY_HOST;
args.request.path = `${AI_GATEWAY_PATH}/aws-bedrock/bedrock-runtime/${AWS_REGION}${args.request.path}`;
const result = await next(args);
return result;
},
{
step: 'finalizeRequest',
}
);
const command = new InvokeModelCommand({
body: JSON.stringify(payload),
contentType: 'application/json',
accept: 'application/json',
modelId,
});
const response = await client.send(command);
const decodedResponseBody = new TextDecoder().decode(response.body);
const responseBody = JSON.parse(decodedResponseBody);
const res = responseBody;
return c.json({ modelId, prompt, res });
});
export { bedrock };
ハマったところ
この実装を見つけるまでにハマり散らかしました。ので、私向けにメモしておきます。
まず、Cloudflareから提供されているチュートリアルを参考にしました。
CloudflareからAmazon Bedrockを使ったサンプルコードがドキュメントとして提供されています。
Amazon Bedrock · Cloudflare AI Gateway docs
内容はAWS SDKは使わずに、 aws4fetchというライブラリを使いSigv4に署名しAPIを直接実行するもので、その流れで、Cloudflare AI Gatewayのエンドポイントに変更していました。
Workersなどのエッジ環境でAWS SDKが利用できない場合は、ありかもと思いますが、なかなかAPI直接実行するのは現実的ではないかなぁ + 現状、Workers上でAWS SDKが利用できているので違う方法を探すことにしました。
（ちなみに私の環境ではこのチュートリアル通りに実施してもクレデンシャルのエラーが解消できずでうまく動作できませんでした。。）
次にAWS SDKのエンドポイントの差し替えでよくやる方法として、クライアントを作成する際にendpointを指定することで変更できますが、これもうまく動作せずクレデンシャルエラーとなってしまいました。
const client = new BedrockRuntimeClient({
endpoint:
"https://gateway.ai.cloudflare.com/v1/xxxx/nishimura-test-ai-app",
region: "ap-northeast-1",
});
Cloudflareのサンプルコードを眺めていると、はじめはデフォルトのエンドポイントで署名、API実行直前にCloudflare AI Gatewayのエンドポイントに差し替える実装をしていました。
それを参考に下記のように処理実行直前にエンドポイントを変更する処理を入れてみたけどだめでした。
... client.config.endpoint = "https://gateway.ai.cloudflare.com/v1/xxxx/nishimura-test-ai-app" const response = await client.send(command); ...
AIまわりはPythonのほうが活発なので、Python指定でGitHubを回遊していたら下記PRをみつけ https://github.com/BerriAI/litellm/pull/3467/files
これはBoto3のイベントシステムという仕組みを利用して対応していました。
AWS SDK for JavaScript v3にも同様の仕組みがないか探したところ、ミドルウェアの機能を紹介しているブログを見つけこれを参考に実装し解決したという流れになります。
このミドルウェアの機能については別途ブログ化したいと思います。
Cloudflare AI Gatewayで確認できること
※2024/05/10時点でCloudflare AI Gatewayはベータ版の状態です。内容は変更される可能性があります。
分析のためのダッシュボード画面が提供されています。ここではリクエスト数、キャッシュ、トークン数、エラー数を確認することができます。
今回Bedrockしか利用していませんが、複数のプロバイダーを利用している場合はプロバイダー事に切り替えることができます。
Bedrockで利用しているモデルごとに利用状況を切り替えることはできないみたいでした。今後に期待したいと思います。
- ログ
リクエストの内容や、レスポンスを確認することができます。
リクエストが成功したか、キャッシュした値を返したか、どのプロバイダーを利用したログなのかわかるようになっているのは、複数のAIプロバイダーうを利用している人にとってはとても便利なのではないでしょうか
- リアルタイムログ
リアルタイムでログを確認することができます。
下記はログが流れる様子です。
同じ内容のリクエストのとき、即レスポンスが返されるのも確認できます。
さいごに
Cloudflare AI Gatewayの連携にハマりはしたのですが、モデルの利用状況をモニタリングできるダッシュボードであったり、ログ、キャッシュの機能もありとても便利だなと感じました。
今後、DLPの機能や使用状況によるアラートの機能など多くの便利な機能が計画されているということでとても期待しています。
誰かの参考になれば幸いです。
