TiDB 自動埋め込み機能と さくら AI Engine を使ったRAGシステム サンプル

https://zenn.dev/kameoncloud/articles/74239c7ca5057f
今日はこちらの記事の続きをやっていきます。

前回のサンプルはTiDBに対してベクトル検索を行った結果、類似の文字列を出力します。一般的なRAGシステムはそこからその結果をLLMに投げ最終的な回答を生成します。

今日はその部分を作っていきます。

そのまえに・・・

強制チャンク分割文字列、オーバーラップ

単純に与えられたPDFの中身を指定されたチャンク長で分割するだけでは、RAGの精度向上はなかなか難しく文章取り込み時のチューニングで精度は大きく向上します。よくつかわれる手法が強制チャンク分割文字列 と オーバーラップ です。
強制チャンク分割文字列とは引数として与えられたチャンク長より先に、指定された文字列が文章に現れた場合そこでチャンクを分割する仕組みです。例えば日本語の場合\n\nが現れると、別の段落に切り替わるため、そこでチャンクを分割してしまう、などで利用します。取り込み文章がマークダウンの場合###なども有効に機能します。

分割されたチャンクは直接的に関連を持たないのがRAGにおける文書取り込みの課題です。このためオーバーラップ文字長を指定して意図的に文書同士を重ね合わせます。例えば50と指定した場合、それぞれの文章の50チャンク文づつが重複して取り込まれるため、それぞれのチャンクに緩やかな関連性を持たせることができます。
これにより最大チャンク長（このサンプルでは512)を越えた長文でも関連した一つの文章として認識することができます。

更新版 index.js

以下がその対応版です。
node index.js test.pdf 512 --split="\n\n" --overlap=50
として実行可能です。

index.js

// TiDB Cloud Starter PDFデータ挿入サンプル
// 必要なパッケージ: npm install mysql2 dotenv pdfjs-dist gpt-tokenizer

import mysql from 'mysql2/promise';
import dotenv from 'dotenv';
import fs from 'fs/promises';
import { encode } from 'gpt-tokenizer';
import * as pdfjsLib from 'pdfjs-dist/legacy/build/pdf.mjs';

dotenv.config();

// TiDB接続設定
const config = {
  host: process.env.TIDB_HOST,
  port: process.env.TIDB_PORT || 4000,
  user: process.env.TIDB_USER,
  password: process.env.TIDB_PASSWORD,
  database: process.env.TIDB_DATABASE,
  ssl: {
    rejectUnauthorized: true
  }
};

// テーブルが存在しない場合は作成、または再作成
async function createTableIfNotExists(connection, recreate = false) {
  try {
    console.log('テーブルの存在を確認中...');
    
    const [tables] = await connection.execute(
      "SHOW TABLES LIKE 'documents'"
    );
    
    if (tables.length > 0 && recreate) {
      console.log('既存のテーブルを削除します...');
      await connection.execute('DROP TABLE documents');
      console.log('✓ テーブルを削除しました');
    }
    
    if (tables.length === 0 || recreate) {
      console.log('テーブルを作成します...');
      
      const createTableSQL = `
        CREATE TABLE documents (
          id INT PRIMARY KEY AUTO_INCREMENT,
          content TEXT,
          source VARCHAR(255),
          chunk_index INT,
          content_vector VECTOR(1024) GENERATED ALWAYS AS (
            EMBED_TEXT("tidbcloud_free/amazon/titan-embed-text-v2", content)
          ) STORED
        )
      `;
      
      await connection.execute(createTableSQL);
      console.log('✓ テーブル "documents" を作成しました (source, chunk_indexカラム付き)');
    } else {
      console.log('✓ テーブル "documents" は既に存在します');
      
      // テーブル構造を確認
      const [columns] = await connection.execute(
        "SHOW COLUMNS FROM documents"
      );
      const hasSource = columns.some(col => col.Field === 'source');
      const hasChunkIndex = columns.some(col => col.Field === 'chunk_index');
      
      if (!hasSource || !hasChunkIndex) {
        console.log('\n⚠️  警告: 既存テーブルにはsourceまたはchunk_indexカラムがありません');
        console.log('テーブルを再作成する場合は、以下のコマンドを実行してください:');
        console.log('node index.js <PDFファイル> <トークン数> --recreate');
        throw new Error('テーブル構造が不一致です。--recreateオプションを使用してテーブルを再作成してください。');
      }
    }
  } catch (error) {
    if (error.message.includes('テーブル構造が不一致')) {
      throw error;
    }
    console.error('テーブル作成エラー:', error.message);
    throw error;
  }
}

// PDFファイルを読み込んでテキストを抽出
async function extractTextFromPDF(pdfPath) {
  try {
    console.log(`\nPDFファイルを読み込み中: ${pdfPath}`);
    const dataBuffer = await fs.readFile(pdfPath);
    
    // PDFドキュメントを読み込む
    const loadingTask = pdfjsLib.getDocument({
      data: new Uint8Array(dataBuffer),
      useSystemFonts: true,
    });
    
    const pdfDocument = await loadingTask.promise;
    const numPages = pdfDocument.numPages;
    
    console.log(`✓ PDF読み込み完了`);
    console.log(`  - ページ数: ${numPages}`);
    
    // 全ページからテキストを抽出
    let fullText = '';
    for (let pageNum = 1; pageNum <= numPages; pageNum++) {
      const page = await pdfDocument.getPage(pageNum);
      const textContent = await page.getTextContent();
      const pageText = textContent.items.map(item => item.str).join(' ');
      fullText += pageText + '\n';
    }
    
    console.log(`  - テキスト長: ${fullText.length}文字`);
    
    return fullText;
  } catch (error) {
    console.error('PDF読み込みエラー:', error.message);
    throw error;
  }
}

// テキストを指定トークン数で分割（強制分割文字列とオーバーラップに対応）
function splitTextByTokens(text, maxTokens = 512, forceSplitMarker = null, overlapChars = 0) {
  console.log(`\nテキストを${maxTokens}トークンごとに分割中...`);
  if (forceSplitMarker) {
    console.log(`強制分割マーカー: "${forceSplitMarker}"`);
  }
  if (overlapChars > 0) {
    console.log(`オーバーラップ: ${overlapChars}文字`);
  }
  
  // 強制分割マーカーが指定されている場合、まずそれで分割
  let segments = [text];
  if (forceSplitMarker) {
    segments = text.split(forceSplitMarker);
    console.log(`✓ 強制分割マーカーにより${segments.length}個のセグメントに分割`);
  }
  
  const chunks = [];
  
  // 各セグメントを処理
  for (let segmentIndex = 0; segmentIndex < segments.length; segmentIndex++) {
    const segment = segments[segmentIndex].trim();
    if (!segment) continue;
    
    // セグメントをトークン数で更に分割（オーバーラップ対応）
    const segmentChunks = splitSegmentByTokens(segment, maxTokens, overlapChars);
    chunks.push(...segmentChunks);
  }
  
  console.log(`✓ 最終的に${chunks.length}個のチャンクに分割しました`);
  
  // 各チャンクのトークン数を表示
  chunks.forEach((chunk, index) => {
    const tokens = encode(chunk).length;
    console.log(`  チャンク${index + 1}: ${tokens}トークン, ${chunk.length}文字`);
  });
  
  return chunks;
}

// セグメント内をトークン数で分割（オーバーラップ対応）
function splitSegmentByTokens(text, maxTokens, overlapChars = 0) {
  const sentences = text.split(/(?<=[.!?。！？])\s+|\n+/);
  const chunks = [];
  let currentChunk = '';
  let currentTokens = 0;
  let previousChunkEnd = ''; // 前のチャンクの末尾（オーバーラップ用）
  
  for (const sentence of sentences) {
    const sentenceTokens = encode(sentence).length;
    
    // 1文が最大トークン数を超える場合は強制分割
    if (sentenceTokens > maxTokens) {
      if (currentChunk) {
        // オーバーラップ用に末尾を保存
        previousChunkEnd = currentChunk.slice(-overlapChars);
        chunks.push(currentChunk.trim());
        currentChunk = previousChunkEnd; // 次のチャンクはオーバーラップで開始
        currentTokens = encode(currentChunk).length;
      }
      
      // 長い文を単語レベルで分割
      const words = sentence.split(/\s+/);
      let tempChunk = currentChunk;
      let tempTokens = currentTokens;
      
      for (const word of words) {
        const wordTokens = encode(word).length;
        if (tempTokens + wordTokens > maxTokens) {
          if (tempChunk.trim()) {
            // オーバーラップ用に末尾を保存
            previousChunkEnd = tempChunk.slice(-overlapChars);
            chunks.push(tempChunk.trim());
          }
          tempChunk = previousChunkEnd + word + ' '; // オーバーラップで開始
          tempTokens = encode(tempChunk).length;
        } else {
          tempChunk += word + ' ';
          tempTokens += wordTokens;
        }
      }
      
      if (tempChunk.trim()) {
        currentChunk = tempChunk;
        currentTokens = tempTokens;
      }
      continue;
    }
    
    // 現在のチャンクに追加できるか確認
    if (currentTokens + sentenceTokens <= maxTokens) {
      currentChunk += sentence + ' ';
      currentTokens += sentenceTokens;
    } else {
      // 現在のチャンクを保存
      if (currentChunk.trim()) {
        // オーバーラップ用に末尾を保存
        previousChunkEnd = currentChunk.slice(-overlapChars);
        chunks.push(currentChunk.trim());
      }
      // 新しいチャンクをオーバーラップで開始
      currentChunk = previousChunkEnd + sentence + ' ';
      currentTokens = encode(currentChunk).length;
    }
  }
  
  // 最後のチャンクを追加
  if (currentChunk.trim()) {
    chunks.push(currentChunk.trim());
  }
  
  return chunks;
}

// PDFから抽出したチャンクをデータベースに挿入
async function insertPDFChunks(pdfPath, maxTokens = 512, recreate = false, forceSplitMarker = null, overlapChars = 0) {
  let connection;
  
  try {
    // データベース接続
    console.log('\nTiDB Cloudに接続中...');
    connection = await mysql.createConnection(config);
    console.log('✓ 接続成功！');

    // テーブルが存在しない場合は作成
    await createTableIfNotExists(connection, recreate);

    // PDFからテキストを抽出
    const text = await extractTextFromPDF(pdfPath);
    
    // テキストをトークン数で分割（強制分割マーカーとオーバーラップ対応）
    const chunks = splitTextByTokens(text, maxTokens, forceSplitMarker, overlapChars);
    
    // チャンクをデータベースに挿入
    console.log('\nデータベースに挿入中...');
    const fileName = pdfPath.split(/[/\\]/).pop();
    
    for (let i = 0; i < chunks.length; i++) {
      const [result] = await connection.execute(
        'INSERT INTO documents (content, source, chunk_index) VALUES (?, ?, ?)',
        [chunks[i], fileName, i + 1]
      );
      console.log(`✓ チャンク${i + 1}/${chunks.length} 挿入完了 (ID: ${result.insertId})`);
    }

    // 挿入されたデータを確認
    console.log('\n挿入されたデータを確認:');
    const [rows] = await connection.execute(
      'SELECT id, source, chunk_index, LEFT(content, 100) as content_preview FROM documents WHERE source = ? ORDER BY chunk_index',
      [fileName]
    );
    console.table(rows);

    console.log(`\n✓ 処理完了: ${chunks.length}個のチャンクを挿入しました`);

  } catch (error) {
    console.error('エラーが発生しました:', error.message);
    throw error;
  } finally {
    if (connection) {
      await connection.end();
      console.log('\n接続を閉じました。');
    }
  }
}

// メイン処理
(async () => {
  console.log('=== TiDB Auto Embedding PDF挿入サンプル ===');
  
  // コマンドライン引数からPDFパスを取得
  const args = process.argv.slice(2);
  const pdfPath = args[0];
  
  if (!pdfPath) {
    console.error('\nエラー: PDFファイルのパスを指定してください');
    console.log('使用方法: node index.js <PDFファイルのパス> [最大トークン数] [オプション]');
    console.log('\nオプション:');
    console.log('  --recreate              テーブルを再作成');
    console.log('  --split=<文字列>        指定した文字列で強制分割');
    console.log('  --overlap=<文字数>      チャンク間のオーバーラップ文字数');
    console.log('\n例:');
    console.log('  node index.js ./sample.pdf 512');
    console.log('  node index.js ./sample.pdf 512 --recreate');
    console.log('  node index.js ./sample.pdf 512 --split="---"');
    console.log('  node index.js ./sample.pdf 512 --overlap=100');
    console.log('  node index.js ./sample.pdf 512 --split="\\f" --overlap=50 --recreate');
    process.exit(1);
  }
  
  // 最大トークン数（オプション、デフォルトは512）
  const maxTokens = parseInt(args[1]) || 512;
  
  // --recreateオプションの確認
  const recreate = args.includes('--recreate');
  
  // --splitオプションの確認
  let forceSplitMarker = null;
  const splitArg = args.find(arg => arg.startsWith('--split='));
  if (splitArg) {
    forceSplitMarker = splitArg.replace('--split=', '');
    // エスケープシーケンスを実際の文字に変換
    forceSplitMarker = forceSplitMarker
      .replace(/\\n/g, '\n')
      .replace(/\\r/g, '\r')
      .replace(/\\t/g, '\t')
      .replace(/\\f/g, '\f');  // 改ページ文字
  }
  
  // --overlapオプションの確認
  let overlapChars = 0;
  const overlapArg = args.find(arg => arg.startsWith('--overlap='));
  if (overlapArg) {
    overlapChars = parseInt(overlapArg.replace('--overlap=', '')) || 0;
    if (overlapChars < 0) overlapChars = 0;
  }
  
  console.log(`\n設定:`);
  console.log(`  PDFファイル: ${pdfPath}`);
  console.log(`  最大トークン数: ${maxTokens}`);
  if (recreate) {
    console.log(`  テーブル再作成: はい`);
  }
  if (forceSplitMarker) {
    console.log(`  強制分割マーカー: "${forceSplitMarker.replace(/\n/g, '\\n').replace(/\f/g, '\\f')}"`);
  }
  if (overlapChars > 0) {
    console.log(`  オーバーラップ: ${overlapChars}文字`);
  }
  
  await insertPDFChunks(pdfPath, maxTokens, recreate, forceSplitMarker, overlapChars);
  
})().catch(error => {
  console.error('致命的なエラー:', error);
  process.exit(1);
});

RAG さくらの AI Engine版

.envに以下の行を追加します。

SAKURA_API_KEY=xxxc8ca1-870a-4be0-89fe-320f0acf73fd:xxxrFRv69Yz+ZP5r9jCSmBx9cfwa23QMPm04HmGj

その後search.jsを以下に置き換えて実行します。

search.js

// TiDB Cloud Starter RAGシステム (ベクトル検索 + さくらのAI Engine)
// 必要なパッケージ: npm install mysql2 dotenv

import mysql from 'mysql2/promise';
import dotenv from 'dotenv';

dotenv.config();

// TiDB接続設定
const config = {
  host: process.env.TIDB_HOST,
  port: process.env.TIDB_PORT || 4000,
  user: process.env.TIDB_USER,
  password: process.env.TIDB_PASSWORD,
  database: process.env.TIDB_DATABASE,
  ssl: {
    rejectUnauthorized: true
  }
};

// さくらのAI Engine設定
const SAKURA_API_URL = 'https://api.ai.sakura.ad.jp/v1/chat/completions';
const SAKURA_API_KEY = process.env.SAKURA_API_KEY || 'ad0c8ca1-870a-4be0-89fe-320f0acf73fd:qgRrFRv69Yz+ZP5r9jCSmBx9cfwa23QMPm04HmGj';

// ベクトル検索を実行
async function searchDocuments(searchQuery, limit = 3) {
  let connection;
  
  try {
    console.log('TiDB Cloudに接続中...');
    connection = await mysql.createConnection(config);
    console.log('✓ 接続成功！\n');

    console.log(`検索クエリ: "${searchQuery}"`);
    console.log(`取得件数: 上位${limit}件\n`);
    
    const sql = `
      SELECT id, content, source, chunk_index 
      FROM documents
      ORDER BY VEC_EMBED_COSINE_DISTANCE(
        content_vector,
        ?
      )
      LIMIT ${limit}
    `;

    const [rows] = await connection.execute(sql, [searchQuery]);

    if (rows.length === 0) {
      console.log('検索結果が見つかりませんでした。');
      return [];
    }

    console.log('=== ベクトル検索結果 ===\n');
    rows.forEach((row, index) => {
      console.log(`[${index + 1}] ID: ${row.id} | Source: ${row.source || 'N/A'} | Chunk: ${row.chunk_index || 'N/A'}`);
      console.log(`    ${row.content.substring(0, 150)}...`);
      console.log('');
    });

    return rows;

  } catch (error) {
    console.error('ベクトル検索エラー:', error.message);
    throw error;
  } finally {
    if (connection) {
      await connection.end();
    }
  }
}

// さくらのAI Engineで回答を生成
async function generateAnswer(query, contexts) {
  try {
    console.log('=== さくらのAI Engineで回答生成中... ===\n');

    // コンテキストを結合
    const contextText = contexts
      .map((ctx, idx) => `[文書${idx + 1}]\n${ctx.content}`)
      .join('\n\n');

    // プロンプトを構築
    const systemPrompt = `あなたは親切なアシスタントです。以下の文書を参考にして、ユーザーの質問に正確に答えてください。
文書に記載されていない情報については「文書には記載がありません」と答えてください。`;

    const userPrompt = `【参考文書】
${contextText}

【質問】
${query}

【回答】`;

    // APIリクエスト
    const response = await fetch(SAKURA_API_URL, {
      method: 'POST',
      headers: {
        'Accept': 'application/json',
        'Authorization': `Bearer ${SAKURA_API_KEY}`,
        'Content-Type': 'application/json'
      },
      body: JSON.stringify({
        model: 'gpt-oss-120b',
        messages: [
          { role: 'system', content: systemPrompt },
          { role: 'user', content: userPrompt }
        ],
        temperature: 0.7,
        max_tokens: 1000,
        stream: false
      })
    });

    if (!response.ok) {
      const errorText = await response.text();
      throw new Error(`API Error: ${response.status} - ${errorText}`);
    }

    const data = await response.json();
    
    if (!data.choices || data.choices.length === 0) {
      throw new Error('APIから有効な応答が返されませんでした');
    }

    const answer = data.choices[0].message.content;
    
    console.log('=== 生成された回答 ===\n');
    console.log(answer);
    console.log('\n');
    
    // 使用トークン数を表示
    if (data.usage) {
      console.log('--- トークン使用量 ---');
      console.log(`入力トークン: ${data.usage.prompt_tokens}`);
      console.log(`出力トークン: ${data.usage.completion_tokens}`);
      console.log(`合計トークン: ${data.usage.total_tokens}`);
    }

    return answer;

  } catch (error) {
    console.error('回答生成エラー:', error.message);
    throw error;
  }
}

// RAGパイプライン: 検索 → 回答生成
async function ragSearch(query, limit = 3) {
  console.log('=== TiDB + さくらのAI Engine RAGシステム ===\n');
  
  try {
    // 1. ベクトル検索で関連文書を取得
    const contexts = await searchDocuments(query, limit);
    
    if (contexts.length === 0) {
      console.log('関連する文書が見つからなかったため、回答を生成できません。');
      return;
    }

    console.log('─'.repeat(60) + '\n');

    // 2. さくらのAI Engineで回答を生成
    await generateAnswer(query, contexts);

  } catch (error) {
    console.error('RAG処理エラー:', error.message);
    throw error;
  }
}

// メイン処理
(async () => {
  // コマンドライン引数から検索クエリを取得
  const args = process.argv.slice(2);
  const query = args.join(' ');
  
  if (!query) {
    console.error('エラー: 検索クエリを指定してください\n');
    console.log('使用方法: node search.js <質問> [取得文書数]');
    console.log('\n例:');
    console.log('  node search.js "再生可能エネルギーとは何ですか"');
    console.log('  node search.js "AIの医療応用について教えて" 5');
    console.log('\n環境変数:');
    console.log('  SAKURA_API_KEY - さくらのAI Engine APIキー (オプション)');
    process.exit(1);
  }
  
  // 最後の引数が数値なら取得文書数として扱う
  let limit = 3;
  const lastArg = args[args.length - 1];
  if (/^\d+$/.test(lastArg)) {
    limit = parseInt(lastArg);
    const queryWithoutLimit = args.slice(0, -1).join(' ');
    if (queryWithoutLimit) {
      await ragSearch(queryWithoutLimit, limit);
    } else {
      console.error('エラー: 検索クエリを指定してください');
      process.exit(1);
    }
  } else {
    await ragSearch(query, limit);
  }

})().catch(error => {
  console.error('\n致命的なエラー:', error);
  process.exit(1);
});

node search.js 個人情報とは？で実行できます。

ode search.js 個人情報とは？
[dotenv@17.2.3] injecting env (5) from .env -- tip: ✅ audit secrets and track compliance: https://dotenvx.com/ops
=== TiDB + さくらのAI Engine RAGシステム ===

TiDB Cloudに接続中...
✓ 接続成功！

検索クエリ: "個人情報とは？"
取得件数: 上位3件

=== ベクトル検索結果 ===

[1] ID: 4 | Source: test.pdf | Chunk: 4
    ものであることその他の個  人情報の有用性に配慮しつつ、個人の権利利益を保護することを目的とする。 （平二七法六五・令三法三七・一部改正）  （定義）  第二条   この法律に おいて「個人情報」とは、生存する個人に関する情報であって、次の各号のいずれかに該当するものをいう。 一   当該情報に含ま...

[2] ID: 64 | Source: test.pdf | Chunk: 64
    章において「個人情報ファイル」とは、保有個人情報を含む情報の集合物であって、次に掲げるものをいう。 一   一定の事務の目的を達成するために特定の保有個人情報を電子計算機を用いて検索することができるように体系的に構成したもの  二   前号に掲げるもののほか、一定の事務の目的を達成するために氏名、生...

[3] ID: 15 | Source: test.pdf | Chunk: 15
    用方法からみて個人の権利利  益を害するおそれが少ないものとして政令で定めるものを除く。）をいう。 一   特定の個人情報を電子計算機を用いて検索することができるように体系的に構成したもの  二   前号に掲げるもののほか、特定の個人情報を容易に検索することができるように体系 的に構成したものとして...

────────────────────────────────────────────────────────────

=== さくらのAI Engineで回答生成中... ===

=== 生成された回答 ===

**個人情報の定義（文書1）**

この法律において「個人情報」とは、**生存する個人に関する情報**であって、次のいずれかに該当するものをいいます。

1. **氏名、生年月日その他の記述等**（文書、図画、電磁的記録、音声、動作など）に記載・記録され、またはそれらの方法で表された一切の事項（個人識別符号を除く）で、**特定の個人を識別できるもの**。
   - 他の情報と照合すれば容易に個人を特定できる情報も含む。

2. **個人識別符号が含まれるもの**。

（※「個人識別符号」とは、政令で定められた文字・番号・記号等のうち、個人を識別できる符号を指します。）


--- トークン使用量 ---
入力トークン: 1315
出力トークン: 397
合計トークン: 1712