💡
TiDB 自動埋め込み機能をスクリプトから操作する+PDF取り込み可変チャンク長サンプル
過去以下の記事でTiDBが新しくリリースした自動埋め込み機能を触ってみました 複雑なAIの知識なしでも、与えた文章を自動で埋め込み処理（ベクトル化）を行ってベクトル検索用テーブルに格納してくれる機能です。
使い方はシンプルで以下の様なテーブルをまずは作成します。
CREATE TABLE documents (
id INT PRIMARY KEY AUTO_INCREMENT,
content TEXT,
content_vector VECTOR(1024) GENERATED ALWAYS AS (
EMBED_TEXT("tidbcloud_free/amazon/titan-embed-text-v2", content)
) STORED
);
その後以下のSQLを実行すれば自動でベクトル化が行われ、オリジナルのテキストとベクトルがペアでテーブルに保存されます。
INSERT INTO documents (content) VALUES
("Electric vehicles reduce air pollution in cities."),
("Solar panels convert sunlight into renewable energy."),
("Plant-based diets lower carbon footprints significantly."),
("Deep learning algorithms improve medical diagnosis accuracy."),
("Blockchain technology enhances data security systems.");
今日はまずこれをスクリプト化してみます。
Node.js 用サンプル
まず以下の.envを作成します。
TIDB_HOST=gateway01.ap-northeast-1.prod.aws.tidbcloud.com
TIDB_PORT=4000
TIDB_USER=xxxw2jrpASrr9fT.root
TIDB_PASSWORD=xxxsVstxlAxXtiXx
TIDB_DATABASE=test
次に2つのスクリプトをファイルを作成します。
index.js:与えられた文字列をベクトル化しテーブルに保存を行います
search.js:検索用文字列をベクトル化し、上記で保存したベクトルに対して検索を行い結果を出力します。
index.js
// TiDB Cloud Starterへのデータ挿入サンプル
// 必要なパッケージ: 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: {
// TiDB Cloudは通常SSL接続が必要
rejectUnauthorized: true
}
};
// 挿入するデータ
const documents = [
"Electric vehicles reduce air pollution in cities.",
"Solar panels convert sunlight into renewable energy.",
"Plant-based diets lower carbon footprints significantly.",
"Deep learning algorithms improve medical diagnosis accuracy.",
"Blockchain technology enhances data security systems."
];
// テーブルが存在しない場合は作成
async function createTableIfNotExists(connection) {
try {
console.log('テーブルの存在を確認中...');
// テーブルの存在確認
const [tables] = await connection.execute(
"SHOW TABLES LIKE 'documents'"
);
if (tables.length === 0) {
console.log('テーブルが存在しないため、作成します...');
const createTableSQL = `
CREATE TABLE documents (
id INT PRIMARY KEY AUTO_INCREMENT,
content TEXT,
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" を作成しました');
} else {
console.log('✓ テーブル "documents" は既に存在します');
}
} catch (error) {
console.error('テーブル作成エラー:', error.message);
throw error;
}
}
async function insertDocuments() {
let connection;
try {
// データベース接続
console.log('TiDB Cloudに接続中...');
connection = await mysql.createConnection(config);
console.log('接続成功！');
// テーブルが存在しない場合は作成
await createTableIfNotExists(connection);
// データを1件ずつ挿入
console.log('\nデータを挿入中...');
for (const content of documents) {
const [result] = await connection.execute(
'INSERT INTO documents (content) VALUES (?)',
[content]
);
console.log(`✓ 挿入完了 (ID: ${result.insertId}): ${content.substring(0, 50)}...`);
}
// 挿入されたデータを確認
console.log('\n挿入されたデータを確認:');
const [rows] = await connection.execute(
'SELECT id, content FROM documents ORDER BY id DESC LIMIT 5'
);
console.table(rows);
console.log('\n処理が完了しました！');
} catch (error) {
console.error('エラーが発生しました:', error.message);
throw error;
} finally {
// 接続を閉じる
if (connection) {
await connection.end();
console.log('\n接続を閉じました。');
}
}
}
// 一括挿入バージョン（より効率的）
async function insertDocumentsBatch() {
let connection;
try {
console.log('TiDB Cloudに接続中...');
connection = await mysql.createConnection(config);
console.log('接続成功！');
// テーブルが存在しない場合は作成
await createTableIfNotExists(connection);
// プレースホルダーを生成
const placeholders = documents.map(() => '(?)').join(',');
const sql = `INSERT INTO documents (content) VALUES ${placeholders}`;
console.log('\nデータを一括挿入中...');
const [result] = await connection.execute(sql, documents);
console.log(`✓ ${result.affectedRows}件のデータを挿入しました`);
// 挿入されたデータを確認
console.log('\n挿入されたデータを確認:');
const [rows] = await connection.execute(
'SELECT id, content FROM documents ORDER BY id DESC LIMIT 5'
);
console.table(rows);
console.log('\n処理が完了しました！');
} catch (error) {
console.error('エラーが発生しました:', error.message);
throw error;
} finally {
if (connection) {
await connection.end();
console.log('\n接続を閉じました。');
}
}
}
// メイン処理
(async () => {
console.log('=== TiDB Auto Embedding データ挿入サンプル ===\n');
// 1件ずつ挿入する場合
await insertDocuments();
// 一括挿入する場合（コメントを外して使用）
// await insertDocumentsBatch();
})().catch(error => {
console.error('致命的なエラー:', error);
process.exit(1);
});
search.js
// TiDB Cloud Starter ベクトル検索サンプル
// 必要なパッケージ: 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
}
};
// ベクトル検索を実行
async function searchDocuments(searchQuery, limit = 3) {
let connection;
try {
// データベース接続
console.log('TiDB Cloudに接続中...');
connection = await mysql.createConnection(config);
console.log('接続成功！');
// ベクトル検索を実行
console.log(`\n検索クエリ: "${searchQuery}"`);
console.log(`結果件数: 上位${limit}件\n`);
const sql = `
SELECT id, content
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('検索結果が見つかりませんでした。');
} else {
console.log('=== 検索結果 ===\n');
rows.forEach((row, index) => {
console.log(`[${index + 1}] ID: ${row.id}`);
console.log(` ${row.content}`);
console.log('');
});
}
return rows;
} catch (error) {
console.error('エラーが発生しました:', error.message);
throw error;
} finally {
// 接続を閉じる
if (connection) {
await connection.end();
console.log('接続を閉じました。');
}
}
}
// 複数のクエリで検索を試す
async function searchMultipleQueries() {
const queries = [
"Renewable energy solutions for environmental protection",
"Artificial intelligence in healthcare",
"Climate change and sustainability"
];
console.log('=== 複数クエリでのベクトル検索デモ ===\n');
for (const query of queries) {
await searchDocuments(query, 3);
console.log('\n' + '='.repeat(60) + '\n');
}
}
// メイン処理
(async () => {
console.log('=== TiDB Auto Embedding ベクトル検索サンプル ===\n');
// 単一クエリでの検索
await searchDocuments("Renewable energy solutions for environmental protection", 3);
// 複数クエリでの検索（コメントを外して使用）
// await searchMultipleQueries();
})().catch(error => {
console.error('致命的なエラー:', error);
process.exit(1);
});
ファイルが出来上がったら必要なライブラリをインストールします。
npm install mysql2 dotenv
次にindex.js,search.jsの順番に実行します。
node index.js
[dotenv@17.2.3] injecting env (5) from .env -- tip: ⚙️ specify custom .env file path with { path: '/custom/path/.env' }
=== TiDB Auto Embedding データ挿入サンプル ===
TiDB Cloudに接続中...
接続成功！
テーブルの存在を確認中...
テーブルが存在しないため、作成します...
✓ テーブル "documents" を作成しました
データを挿入中...
✓ 挿入完了 (ID: 1): Electric vehicles reduce air pollution in cities....
✓ 挿入完了 (ID: 2): Solar panels convert sunlight into renewable energ...
✓ 挿入完了 (ID: 3): Plant-based diets lower carbon footprints signific...
✓ 挿入完了 (ID: 4): Deep learning algorithms improve medical diagnosis...
✓ 挿入完了 (ID: 5): Blockchain technology enhances data security syste...
挿入されたデータを確認:
┌─────────┬────┬────────────────────────────────────────────────────────────────┐
│ (index) │ id │ content │
├─────────┼────┼────────────────────────────────────────────────────────────────┤
│ 0 │ 5 │ 'Blockchain technology enhances data security systems.' │
│ 1 │ 4 │ 'Deep learning algorithms improve medical diagnosis accuracy.' │
│ 2 │ 3 │ 'Plant-based diets lower carbon footprints significantly.' │
│ 3 │ 2 │ 'Solar panels convert sunlight into renewable energy.' │
│ 4 │ 1 │ 'Electric vehicles reduce air pollution in cities.' │
└─────────┴────┴────────────────────────────────────────────────────────────────┘
処理が完了しました！
接続を閉じました。
node .\search.js
[dotenv@17.2.3] injecting env (5) from .env -- tip: 🔄 add secrets lifecycle management: https://dotenvx.com/ops
=== TiDB Auto Embedding ベクトル検索サンプル ===
TiDB Cloudに接続中...
接続成功！
検索クエリ: "Renewable energy solutions for environmental protection"
結果件数: 上位3件
=== 検索結果 ===
[1] ID: 2
Solar panels convert sunlight into renewable energy.
[2] ID: 1
Electric vehicles reduce air pollution in cities.
[3] ID: 4
Deep learning algorithms improve medical diagnosis accuracy.
接続を閉じました。
PDF の埋め込み用改造
次にindex.jsを改造してPDFを読み込めるようにします。
まず必要なライブラリをインストールします。
npm install pdfjs-dist gpt-tokenizer
次にindex.jsを以下に置き換えます。
index.js
// TiDB Cloud Starter PDFデータ挿入サンプル
// 必要なパッケージ: npm install mysql2 dotenv pdf-parse gpt-tokenizer
import mysql from 'mysql2/promise';
import dotenv from 'dotenv';
import fs from 'fs/promises';
import pdfParse from 'pdf-parse';
import { encode } from 'gpt-tokenizer';
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) {
try {
console.log('テーブルの存在を確認中...');
const [tables] = await connection.execute(
"SHOW TABLES LIKE 'documents'"
);
if (tables.length === 0) {
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" を作成しました');
} else {
console.log('✓ テーブル "documents" は既に存在します');
}
} catch (error) {
console.error('テーブル作成エラー:', error.message);
throw error;
}
}
// PDFファイルを読み込んでテキストを抽出
async function extractTextFromPDF(pdfPath) {
try {
console.log(`\nPDFファイルを読み込み中: ${pdfPath}`);
const dataBuffer = await fs.readFile(pdfPath);
const pdfData = await pdfParse(dataBuffer);
console.log(`✓ PDF読み込み完了`);
console.log(` - ページ数: ${pdfData.numpages}`);
console.log(` - テキスト長: ${pdfData.text.length}文字`);
return pdfData.text;
} catch (error) {
console.error('PDF読み込みエラー:', error.message);
throw error;
}
}
// テキストを指定トークン数で分割
function splitTextByTokens(text, maxTokens = 512) {
console.log(`\nテキストを${maxTokens}トークンごとに分割中...`);
// テキストを段落や文で分割（改行を保持）
const sentences = text.split(/(?<=[.!?。！？])\s+|\n+/);
const chunks = [];
let currentChunk = '';
let currentTokens = 0;
for (const sentence of sentences) {
const sentenceTokens = encode(sentence).length;
// 1文が最大トークン数を超える場合は強制分割
if (sentenceTokens > maxTokens) {
if (currentChunk) {
chunks.push(currentChunk.trim());
currentChunk = '';
currentTokens = 0;
}
// 長い文を単語レベルで分割
const words = sentence.split(/\s+/);
let tempChunk = '';
let tempTokens = 0;
for (const word of words) {
const wordTokens = encode(word).length;
if (tempTokens + wordTokens > maxTokens) {
chunks.push(tempChunk.trim());
tempChunk = word + ' ';
tempTokens = wordTokens;
} 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()) {
chunks.push(currentChunk.trim());
}
currentChunk = sentence + ' ';
currentTokens = sentenceTokens;
}
}
// 最後のチャンクを追加
if (currentChunk.trim()) {
chunks.push(currentChunk.trim());
}
console.log(`✓ ${chunks.length}個のチャンクに分割しました`);
// 各チャンクのトークン数を表示
chunks.forEach((chunk, index) => {
const tokens = encode(chunk).length;
console.log(` チャンク${index + 1}: ${tokens}トークン, ${chunk.length}文字`);
});
return chunks;
}
// PDFから抽出したチャンクをデータベースに挿入
async function insertPDFChunks(pdfPath, maxTokens = 512) {
let connection;
try {
// データベース接続
console.log('\nTiDB Cloudに接続中...');
connection = await mysql.createConnection(config);
console.log('✓ 接続成功！');
// テーブルが存在しない場合は作成
await createTableIfNotExists(connection);
// PDFからテキストを抽出
const text = await extractTextFromPDF(pdfPath);
// テキストをトークン数で分割
const chunks = splitTextByTokens(text, maxTokens);
// チャンクをデータベースに挿入
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 pdfPath = process.argv[2];
if (!pdfPath) {
console.error('\nエラー: PDFファイルのパスを指定してください');
console.log('使用方法: node index.js <PDFファイルのパス> [最大トークン数]');
console.log('例: node index.js ./sample.pdf 512');
process.exit(1);
}
// 最大トークン数（オプション、デフォルトは512）
const maxTokens = parseInt(process.argv[3]) || 512;
console.log(`\n設定:`);
console.log(` PDFファイル: ${pdfPath}`);
console.log(` 最大トークン数: ${maxTokens}`);
await insertPDFChunks(pdfPath, maxTokens);
})().catch(error => {
console.error('致命的なエラー:', error);
process.exit(1);
});
以下の書式で実行できます。
node index.js <PDFファイルのパス> [最大トークン数] --recreate
--recreateはテーブルを再作成するオプションです。初回実行のみ付与します。
node index.js test.pdf 512で実行します。
処理経過がコンソールに大量に出力されますが、最後に✓ 処理完了: x個のチャンクを挿入しましたが出力されれば完了です。
search.jsの質問文を埋め込んだPDFの中身用に書き換えて、検索を実行してみます。
node .\search.js
[dotenv@17.2.3] injecting env (5) from .env -- tip: 🔄 add secrets lifecycle management: https://dotenvx.com/ops
=== TiDB Auto Embedding ベクトル検索サンプル ===
TiDB Cloudに接続中...
接続成功！
検索クエリ: "個人情報とは"
結果件数: 上位3件
=== 検索結果 ===
[1] ID: 4
（平二七法六五・令三法三七・一部改正） （定義） 第二条 この法律に おいて「個人情報」とは、生存する個人に関する情報であって、次の各号のいずれかに該当するものをいう。 一 当該情報に含まれる氏名、生年月日その他の記述等（文書、図画若しくは電磁的記録（電磁的方式（電子的方式、磁気的方式その他人の知覚によっては認識 することができない方式をいう。次項第二号において同じ。）で作られる記録をいう。以下同じ。）に記載され、若しくは記録され、又は音声、動作その他の 方法を用いて表された一切の事項（個人識別符号を除く。）をいう。以下同じ。）により特定の個人を識別することができるもの（他の情報 と容易に照合する ことができ、それにより特定の個 人を識別することができることとなるものを含む。） 二 個人識別符号が含まれるもの ２ この法律において「個人識別符号」とは、次の各号のいずれかに該当する文字、 番号、記号その他の符号のうち、政令で定めるものをいう。
[2] ID: 5
一 特定の個人の身体の一部の特徴を電子計算機の用に供するために変換した文字、番号、記号その他の符号であって、当該特定の個人を識別することができる もの 二 個人に提供される役務の利用若しくは個人に販売される商品の購入に関し割り当てられ、又は個人に発行され るカードその他の書類に記載され、若しくは電 磁的方式により 記録された文字、番号、記号その他の符号であって、その利用者若しくは購入者又は発行を受ける者ごとに異なるものとなるように割り当てら れ、又は記載され、若しくは記録されることにより、特定の利用者若しくは購入者又は発行を受ける者を識別することができるもの ３ この法律において「要配慮個人情報」とは、本人の人種、信条、社会的 身分、病歴、犯罪の経歴、犯罪により害を被った事実その他本人に対する不当な差別、 偏見その他の不利益が生じないようにその取扱いに特に配慮を要する ものとして政令で 定める記述等が含まれる個人情報をいう。 ４ この法律において個人情報について「本人」とは、個人情報によって識別される特定の個人をいう。 ５ この法律において「仮名加工情報」とは、次の各号に掲げる個人情報の区分に応じて当該各号に定める措置を講じて他の情報と照合しない限り特定の個人を識 別することができないように個人情報を加工して得られる個人に関する情報をいう。
[3] ID: 13
（令三法 三七・旧第十四条繰下） 第四章 個人情報取扱事業者等の義務等 （令三法三七・改称） 第一節 総則 （令三法三七・追加） （定義） 第十六条 この章及び第八章において「個人情報データベース等」とは、個人情報を含む情報の集合物であって、次に掲げるもの（利用方法からみて個人の権利利 益を害するおそれが少ないものとして政令で定めるものを除く。）をいう。 一 特定の個人情報を電子計算機を用いて検索することができるように体系的に構成したもの 二 前号に掲げるもののほか、 特定の個人情報を容易に検索することができるように体系 的に構成したものとして政令で定めるもの ２ この章及び第六章から第八章までにおいて「個人情報取扱事業者」とは、個人情報データベース等を事業の用に供している者をいう。ただし、次に掲げる者を 除く。 一 国の機関 二 地方公共団体 三 独立行政法人等 四 地方独立行政法人 ３ この章において「個人データ」とは、個人情報データベース等を構成する個人情報をいう。 ４ この章において「保有個人データ」とは、個人情報取扱事業者が、 開示、内容の訂正、追加又は削除、利用の停止、消去及び第三者への提供の停止を行うこと のできる権 限を有する個人データであって、その存否が明らかになることにより公 益その他の利益が害されるものとして政令で定めるもの以外のものをいう。
接続を閉じました。
次回予告：LLM連携
TiDBは自動での埋め込み（ベクトル化）とベクトル検索は提供してくれますが、最終的な回答を生成するLLMは提供していません。次回の記事では、RAGと言われるベクトル検索の結果をLLMが受け取り最終回答を生成するサンプルを作ります。
Discussion