LangChain.jsでCloudflare Workers AI用のカスタムLLMクラスを作る

LangChain.jsのLLMクラスを継承して、Cloudflare Workers AIのgpt-oss-20bを呼び出すカスタムモデルを作る手順を解説します。

LangChain.jsに用意されていないLLMプロバイダーやAPIを使いたいときは、@langchain/coreのLLMクラスを継承したカスタムモデルを作ります。Cloudflare Workers AIのgpt-oss-20bを呼び出すクラスを例に、invokeやチェインから使えるようになるまでの手順を解説します。

LLMクラスを継承する

継承元は@langchain/core/language_models/llmsのLLMです。実装が必須なのは、モデルの種類を返す_llmType()と、プロンプト文字列を受け取って回答の文字列を返す_call()の2つだけです。

import { LLM, type BaseLLMParams } from '@langchain/core/language_models/llms';
import { CallbackManagerForLLMRun } from '@langchain/core/callbacks/manager';
import { GenerationChunk } from '@langchain/core/outputs';

export class CloudflareResponsesModel extends LLM {
  // 以降のステップで実装する
}

パラメータを定義する

モデルが受け取る設定は、BaseLLMParamsを継承したインタフェースで定義します。Workers AIのenv.AIバインディングとモデル名は必須にして、それ以外は省略できるようにしました。

export interface ReasoningConfig {
  effort?: 'low' | 'medium' | 'high';
  summary?: 'auto' | 'concise' | 'detailed';
}

export interface CloudflareResponsesModelParams extends BaseLLMParams {
  modelName: keyof AiModels;
  aiBinding: Ai;
  jsonSchema?: {
    schema: Record<string, unknown>;
    name: string;
  };
  maxTokens?: number;
  reasoning?: ReasoningConfig;
}

ReasoningConfigの値をリテラル型にしておくと、effort: 'hgih'のような綴り間違いをコンパイル時に検出できます。AiModelsとAiは@cloudflare/workers-typesの型です。

コンストラクタで初期化する

クラスのフィールドを宣言し、コンストラクタで初期化します。super(fields ?? {})で親クラスを初期化する点と、maxTokensにデフォルト値を持たせる点がポイントです。

modelName: keyof AiModels;
aiBinding: Ai;
jsonSchema?: { schema: Record<string, unknown>; name: string };
maxTokens: number;
reasoning?: ReasoningConfig;

constructor(
  modelName: keyof AiModels,
  aiBinding: Ai,
  fields?: Partial<CloudflareResponsesModelParams>
) {
  super(fields ?? {});
  this.modelName = modelName;
  this.aiBinding = aiBinding;
  this.jsonSchema = fields?.jsonSchema;
  this.maxTokens = fields?.maxTokens ?? 8192;
  this.reasoning = fields?.reasoning;
}

必須メソッドを実装する

_llmType()はモデルの種類を識別する文字列を返すだけです。

_llmType(): string {
  return 'cloudflare_responses';
}

_call()が本体です。プロンプトからリクエストボディを組み立て、Workers AIを呼び出し、回答の文字列を取り出して返します。JSON Schemaと推論の設定は、指定されたときだけリクエストに足します。

async _call(
  prompt: string,
  _options?: this['ParsedCallOptions'],
  _runManager?: CallbackManagerForLLMRun
): Promise<string> {
  const messages = [{ role: 'user', content: prompt }];

  const requestBody: Record<string, unknown> = {
    messages,
    max_tokens: this.maxTokens,
  };

  if (this.jsonSchema) {
    requestBody.response_format = {
      type: 'json_schema',
      json_schema: { name: this.jsonSchema.name, schema: this.jsonSchema.schema },
    };
  }

  if (this.reasoning?.effort) {
    requestBody.reasoning_effort = this.reasoning.effort;
  }

  const response = await this.aiBinding.run(this.modelName, requestBody);
  return this.extractContent(response);
}

レスポンスの形式の違いを吸収する

_call()は文字列を返す必要があります。一方、aiBinding.run()の戻り値はモデルや指定したオプションによって形が異なります。そこで、形式ごとに分岐して文字列を取り出すextractContent()を用意しました。

private extractContent(response: unknown): string {
  // 文字列の場合はそのまま返す
  if (typeof response === 'string') {
    return response;
  }

  if (typeof response === 'object' && response !== null) {
    const resp = response as Record<string, unknown>;

    // JSON Schemaを指定したときの形式
    if ('response' in resp && typeof resp.response === 'object') {
      return JSON.stringify(resp.response);
    }

    // OpenAI形式: choices[0].message.content
    if ('choices' in resp && Array.isArray(resp.choices)) {
      const choice = resp.choices[0] as Record<string, unknown>;
      if ('message' in choice) {
        const message = choice.message as Record<string, unknown>;
        if ('content' in message && typeof message.content === 'string') {
          let content = message.content;

          // Markdownのコードブロックで囲まれていたら中身だけ取り出す
          const match = content.match(/```(?:json)?\s*\n([\s\S]*?)\n```/);
          if (match) {
            content = match[1].trim();
          }

          return content;
        }
      }
    }

    // どれにも当てはまらなければJSON文字列にして返す
    return JSON.stringify(response);
  }

  return String(response);
}

ストリーミングに対応する

_streamResponseChunks()を実装すると、stream()で呼び出せるようになります。ここでは_call()の結果を1つのチャンクとして返すだけにしました。Workers AIのストリーミング出力は使っていないため、回答は生成が終わってからまとめて届きます。

async *_streamResponseChunks(
  prompt: string,
  options: this['ParsedCallOptions'],
  runManager?: CallbackManagerForLLMRun
): AsyncGenerator<GenerationChunk> {
  const result = await this._call(prompt, options, runManager);
  yield new GenerationChunk({
    text: result,
    generationInfo: {},
  });
}

ビルダーメソッドで設定違いのインスタンスを作る

JSON Schemaや推論の設定を変えたインスタンスを作れるように、withJsonSchema()とwithReasoning()を用意します。元のインスタンスは変更せず、新しいインスタンスを返します。

withJsonSchema(schema: Record<string, unknown>, name: string): CloudflareResponsesModel {
  return new CloudflareResponsesModel(this.modelName, this.aiBinding, {
    jsonSchema: { schema, name },
    maxTokens: this.maxTokens,
    reasoning: this.reasoning,
  });
}

withReasoning(config: ReasoningConfig): CloudflareResponsesModel {
  return new CloudflareResponsesModel(this.modelName, this.aiBinding, {
    jsonSchema: this.jsonSchema,
    maxTokens: this.maxTokens,
    reasoning: config,
  });
}

メソッドチェインでつなげて使います。

const model = new CloudflareResponsesModel('@cf/openai/gpt-oss-20b', env.AI)
  .withJsonSchema(mySchema, 'MyResponse')
  .withReasoning({ effort: 'high', summary: 'detailed' });

独自のメソッドを足す

LangChainの標準インタフェースにない機能は、独自のメソッドとして足せます。callWithReasoning()はResponses API形式（inputとreasoning）でモデルを呼び出し、回答と推論の要約を分けて返します。推論の要約は、レスポンスのoutput配列のreasoning要素から取り出します。

async callWithReasoning(prompt: string): Promise<{ content: string; reasoning?: string }> {
  const requestBody: Record<string, unknown> = {
    input: prompt,
    max_output_tokens: this.maxTokens,
    reasoning: {
      effort: this.reasoning?.effort ?? 'medium',
      summary: this.reasoning?.summary ?? 'auto',
    },
  };
  const response = await this.aiBinding.run(this.modelName, requestBody);
  const resp = (response ?? {}) as Record<string, unknown>;
  const output = Array.isArray(resp.output) ? (resp.output as Record<string, unknown>[]) : [];

  // output配列から、指定した種類の要素のテキストを連結する
  const texts = (type: string, key: string) =>
    output
      .filter((item) => item.type === type && Array.isArray(item[key]))
      .flatMap((item) => (item[key] as { text?: string }[]).map((part) => part.text ?? ''))
      .join('\n');

  return {
    content: typeof resp.output_text === 'string' ? resp.output_text : texts('message', 'content'),
    reasoning: texts('reasoning', 'summary') || undefined,
  };
}

使い方

invoke()で呼び出します。

import { CloudflareResponsesModel } from './cloudflare-responses-model';

const model = new CloudflareResponsesModel('@cf/openai/gpt-oss-20b', env.AI, {
  maxTokens: 4096,
});

const response = await model.invoke('こんにちは、元気ですか？');
console.log(response);

withJsonSchema()を使うと、構造化した出力を指定できます。戻り値は文字列なので、JSON.parseで変換します。

const schema = {
  type: 'object',
  properties: {
    name: { type: 'string' },
    age: { type: 'number' },
  },
  required: ['name', 'age'],
};

const model = new CloudflareResponsesModel('@cf/openai/gpt-oss-20b', env.AI).withJsonSchema(
  schema,
  'Person'
);

const response = await model.invoke('田中太郎さんは30歳です。この情報を構造化してください。');
const person = JSON.parse(response);

LLMクラスを継承しているので、PromptTemplateとつなげたチェインにも組み込めます。

import { PromptTemplate } from '@langchain/core/prompts';

const prompt = PromptTemplate.fromTemplate('次の質問に答えてください: {question}');
const chain = prompt.pipe(model);

const result = await chain.invoke({ question: 'TypeScriptとは何ですか？' });
console.log(result);

まとめ

LangChain.jsのカスタムモデルはLLMクラスを継承して_llmType()と_call()を実装すれば動きます。今回のコードでは、次の3点を足しました。

extractContent()で、レスポンスの形式の違いを文字列に揃える

withJsonSchema()とwithReasoning()で、設定違いのインスタンスを作る

callWithReasoning()で、推論の要約を回答と分けて受け取る

別のAPIやローカルモデルを使う場合も、_call()の中身を差し替えれば同じ形で組み込めます。型定義は@langchain/coreと@cloudflare/workers-typesのパッケージで確認できます。