VannaAIにおけるカスタムLLMの統合と高度な構成

概要

VannaはデフォルトでOpenAIのGPTモデルを利用しますが、アーキテクチャ上はOpenAI互換のAPIエンドポイントを持つ任意の大規模言語モデル(LLM)と統合可能です。本記事では、Vannaの拡張性を活かし、セルフホスト型やプロキシサーバー経由のカスタムLLMを接続してText-to-SQL環境を構築する手順について解説します。

カスタムLLMクライアントの実装

標準的な構成ではVannaが提供するデフォルトの設定を利用しますが、特定のLLM(例: Kimi、Llama等)を使用する場合は、OpenAIクライアントを初期化し、エンドポイントやAPIキーをカスタマイズする必要があります。

以下の例では、`VannaDB_VectorStore`と`OpenAI_Chat`を継承したカスタムクラスを作成し、独自のクライアント設定を注入します。

import os
from openai import OpenAI
from vanna.openai import OpenAI_Chat
from vanna.vannadb import VannaDB_VectorStore

# 環境変数や設定ファイルから読み込むことを推奨
VANNA_MODEL_NAME = "custom_vanna_model"
VANNA_API_KEY = "your_vanna_api_key"
CUSTOM_LLM_ENDPOINT = "https://api.your-custom-server.com/v1"
CUSTOM_LLM_MODEL = "kimi-large"
AUTH_TOKEN = "your_llm_service_token"

# OpenAI互換のクライアントを初期化
llm_client = OpenAI(
    api_key=AUTH_TOKEN,
    base_url=CUSTOM_LLM_ENDPOINT,
    default_headers={
        "Authorization": f"Bearer {AUTH_TOKEN}",
        "Content-Type": "application/json"
    }
)

class CustomVanna(VannaDB_VectorStore, OpenAI_Chat):
    def __init__(self, client=None, config=None):
        # ベクターストアの初期化(VannaDB利用時)
        VannaDB_VectorStore.__init__(
            self,
            vanna_model=VANNA_MODEL_NAME,
            vanna_api_key=VANNA_API_KEY,
            config=config
        )
        # LLMチャット機能の初期化(カスタムクライアントを渡す)
        OpenAI_Chat.__init__(
            self,
            client=client,
            config=config
        )

# カスタム設定でインスタンス化
vn = CustomVanna(
    client=llm_client,
    config={"model": CUSTOM_LLM_MODEL}
)

データベースへの接続

インスタンスの作成後、対象のデータベースに接続します。VannaはMySQLをはじめ、PostgreSQL、Snowflakeなどの主要なDBに対応しています。ここではMySQLへの接続例を示します。

vn.connect_to_mysql(
    host="localhost",
    user="db_user",
    password="secure_password",
    dbname="analytics_db",
    port=3306
)

クエリ生成とデータの取得

設定が完了したら、自然言語での質問(Text-to-SQL)を試行します。初回はシステムの応答を確認するためにシンプルな問い合わせを行います。

# システムの挙動確認
vn.ask(question="自己紹介をしてください")

続いて、具体的なデータ検索を要求します。適切なトレーニングデータが与えられていない状態でも、LLMの汎用的な知識に基づいてSQLを生成しようとします。

# 特定の条件を持つデータの検索
response = vn.ask(question="student_infoテーブルからscoreが92を超えるレコードを全て取得してください")

応答には生成されたSQL、実行結果のデータ、および可視化用のプロットコードが含まれます。データが正しく取得できれば、システムは期待通りに機能しています。

メタデータの学習(Training)

より精度の高いSQLを生成させるためには、データベースのスキーマ情報(DDL)や、質問とSQLのペアを学習させる必要があります。`train`メソッドを使用してベクターデータベースに情報を追加します。

# DDLによるテーブル構造の学習
vn.train(ddl="""
CREATE TABLE student_info (
    id INT NOT NULL AUTO_INCREMENT PRIMARY KEY COMMENT '学生ID',
    name VARCHAR(100) NOT NULL COMMENT '氏名',
    age INT NOT NULL COMMENT '年齢',
    province VARCHAR(100) NOT NULL COMMENT '居住都道府県',
    score DECIMAL(5,2) NOT NULL COMMENT 'テストの点数'
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
""")

# 質問とSQLのペアを学習させる
vn.train(
    question="90点以上の学生をリストアップしてください",
    sql="SELECT * FROM student_info WHERE score >= 90;"
)

vn.train(
    question="85点未満の学生が居住している都道府県を表示してください",
    sql="SELECT DISTINCT province FROM student_info WHERE score < 85;"
)

学習後、より複雑な集計クエリを試します。

vn.ask(question="都道府県ごとの平均点を算出してください")

プロンプト構造の分析

内部では、Vannaは以下の要素を含んだ詳細なシステムプロンプトをLLMに送信しています。これにより、LLMは「SQL専門家」として振る舞い、提供されたDDL(テーブル定義)に基づいたクエリを生成します。

  • 役割定義: SQLエキスパートとして、与えられたコンテキストに基づきクエリを生成するよう指示されます。
  • コンテキスト: 学習させたDDL文が `===Tables` セクションとして挿入されます。
  • ガイドライン:
    • 十分な情報がある場合は、説明を加えずにSQLのみを出力する。
    • 情報が不足している場合は、中間SQL(`intermediate_sql`)を実行して候補値を確認する。
    • 過去に同一の質問があれば、その回答を再利用する。

このプロンプトエンジニアリングにより、自然言語の質問を正確なSQLクエリへと変換する精度が向上しています。

タグ: VannaAI Python OpenAI API Text-to-SQL RAG

6月25日 00:09 投稿