rag-toolkit

FAISS ベースの RAG（Retrieval-Augmented Generation）ユーティリティライブラリです。

インストール

pip install -e .

# PDF から DB を作る場合は追加で必要
pip install pdfminer.six

ディレクトリ構成

rag_toolkit/
├── __init__.py
├── builders/
│   └── db_builder.py        # ベクトル DB 作成クラス
├── retrievers/
│   └── faiss_retriever.py   # 検索クラス（ライブラリの中心）
└── tools/
    └── tool.py              # LLM ツール基底クラス

クラス概要

DBBuilder が出力するファイル形式

DB の作成（DBBuilder）

基本的な流れ

from rag_toolkit import DBBuilder

builder = DBBuilder(output_dir="database/my_system/")

# PDF からマニュアル DB を作成
builder.build_from_pdf(["manual.pdf", "release_notes.pdf"])

# CSV（Q/A）から QA DB を作成
builder.build_from_csv("qa.csv")

# JSON からキャラクター知識 DB を作成
builder.build_from_json("knowledge.json")

例1: PDF からマニュアル DB

from rag_toolkit import DBBuilder

builder = DBBuilder(
    output_dir="database/sample_XYZ_system/",
    chunk_size=1000,    # 1チャンクあたりの文字数
    chunk_overlap=200,  # チャンク間のオーバーラップ文字数
)
builder.build_from_pdf(
    pdf_files=[
        "database/sample_XYZ_system/XYZシステムリリースノート.pdf",
        "database/sample_XYZ_system/XYZシステム統合ユーザーマニュアル.pdf",
    ],
    index_filename="vector_index.faiss",
    meta_filename="metadata.json",
)

例2: CSV（Q/A形式）から QA DB

builder.build_from_csv(
    csv_file="database/sample_XYZ_system/XYZ_system_QA.csv",
    question_col="Q",   # 質問列のヘッダー名
    answer_col="A",     # 回答列のヘッダー名
    index_filename="qa_index.faiss",
    meta_filename="qa_metadata.json",
)

例3: JSON からキャラクター知識 DB

from rag_toolkit import DBBuilder

knowledge_builder = DBBuilder(output_dir="database/DB_ikazuchi_mal/")
knowledge_builder.build_from_json(
    json_file="database/documents_mal_knowledge.json",
    index_filename="knowledge_index.faiss",
    meta_filename="knowledge_metadata.json",
    category_value="いかづちマル",  # metadata["カテゴリ"] を上書きしたい場合
)

複数 DB を作る際のモデル共有（メモリ節約）

from rag_toolkit import DBBuilder

# builder を 1 つ作り、同じモデルで複数の DB を続けて作成できる
builder = DBBuilder(output_dir="database/sample_XYZ_system/")
builder.build_from_pdf([...])
builder.build_from_csv(...)

# 別の出力先に作る場合もモデルを渡して使い回せる
knowledge_builder = DBBuilder(
    output_dir="database/DB_ikazuchi_mal/",
    model=builder.model,  # ← ロード済みモデルを渡す
)
knowledge_builder.build_from_json(...)

DB の検索（FaissRetriever）

1. 単体検索

from rag_toolkit import FaissRetriever

# マニュアル検索
manual = FaissRetriever(
    name="search_xyz_manual",
    description="XYZシステムの仕様書・マニュアルを検索する",
    database_dir="database/sample_XYZ_system/",
    meta_page_key="page",  # ページ番号を表示したい場合に指定
)

# QA 検索（"question+answer" で Q/A 結合テキストを参照）
qa = FaissRetriever(
    name="search_xyz_qa",
    description="XYZシステムの過去QAを検索する",
    database_dir="database/sample_XYZ_system/",
    index_filename="qa_index.faiss",
    meta_filename="qa_metadata.json",
    meta_text_keys=["text", "question+answer"],
)

print(manual.run_query("パスワードのリセット手順"))
print(qa.run_query("ログインできない場合の対処法"))

2. 複数 DB 横断検索

# スコア上位順でまとめて返す
results = FaissRetriever.search_multi([manual, qa], "ログインエラーの対処法")
for r in results:
    print(r["source"], r["total_score"], r["text"])

# 文字列として取得
text = FaissRetriever.search_multi_as_str([manual, qa], "ログインエラーの対処法")
print(text)

3. LLM エージェントのツールとして使う

retriever = FaissRetriever(
    name="search_xyz_manual",
    description="XYZシステムの仕様書・マニュアルを検索する",
    database_dir="database/sample_XYZ_system/",
    prompt="以下の参考情報をもとに回答してください。",
)

context = {"query_text": "エラーコード E-501 の意味"}
result = retriever.run(context)
# result["status"]  -> "success" or "ignore"
# result["message"] -> prompt + 検索結果（LLM に渡す文字列）
print(result["message"])

依存ライブラリ