このページを閲覧しているということは、皆さんが会社や自宅でのDIYプロジェクトにおいてRAG（Retriever-Augmented Generator）を活用している、または導入を検討していることと思います。カスタムローダーの作成に至るということは、既存のLlamaIndexやLangchainが標準で提供するRetriever機能だけでは要件を満たさない、あるいは私のケースのように、既存のRetrieverが十分ではなく、さらなるカスタマイズを求めているのでしょう。

インターネットで調べても、カスタムローダーの作成方法に関する情報はあまり見つかりません。私の検索方法が不適切なのか、それとも情報が少ないのかは分かりませんが、様々なキーワードで試してもあまり変わりませんでした。

そこで、この記事ではカスタムローダーの作り方について詳しく説明していきます。すぐに作成方法を知りたい方は、下記の出来上がったテンプレートまでスクロールして、コピーして使ってくださいね。

既存のローダーの観察

作り方どこにも記載がなかったので、既存のコードを参考にするしかありません。
今回は、WikipediaのRetrieverを参考に作っていきたいと思います。
Toolsを作ったときもWikipediaが一番簡素に作ってあったイメージですね。

まずは、観察ということでInputとOutputをどのような形でやればいいのか、見ていきたいと思います。（ここが肝）

from typing import List

from langchain.callbacks.manager import CallbackManagerForRetrieverRun
from langchain.schema import BaseRetriever, Document
from langchain.utilities.wikipedia import WikipediaAPIWrapper


class WikipediaRetriever(BaseRetriever, WikipediaAPIWrapper):
    """`Wikipedia API` retriever.

    It wraps load() to get_relevant_documents().
    It uses all WikipediaAPIWrapper arguments without any change.
    """

    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        return self.load(query=query)

こちらがRetriever側のコードです。ここでInputはqueryがあることがわかります。重要な部分は、Wikipediaのデータを取得するWikipediaAPIWrapperに含まれています。具体的には、_get_relevant_documentsという内部メソッドで、継承されているWikipediaAPIWrapperのloadメソッドが利用されています。

次に、このWikipediaAPIWrapperの詳細に焦点を当てて見ていきましょう。

    def load(self, query: str) -> List[Document]:
        """
        Run Wikipedia search and get the article text plus the meta information.
        See

        Returns: a list of documents.

        """
        page_titles = self.wiki_client.search(query[:WIKIPEDIA_MAX_QUERY_LENGTH])
        docs = []
        for page_title in page_titles[: self.top_k_results]:
            if wiki_page := self._fetch_page(page_title):
                if doc := self._page_to_document(page_title, wiki_page):
                    docs.append(doc)
        return docs

この部分のコードで__page_to_documentを呼んでいる事がわかります。
そのコードを見てみると以下のようになっています。

    def _page_to_document(self, page_title: str, wiki_page: Any) -> Document:
        main_meta = {
            "title": page_title,
            "summary": wiki_page.summary,
            "source": wiki_page.url,
        }
        add_meta = (
            {
                "categories": wiki_page.categories,
                "page_url": wiki_page.url,
                "image_urls": wiki_page.images,
                "related_titles": wiki_page.links,
                "parent_id": wiki_page.parent_id,
                "references": wiki_page.references,
                "revision_id": wiki_page.revision_id,
                "sections": wiki_page.sections,
            }
            if self.load_all_available_meta
            else {}
        )
        doc = Document(
            page_content=wiki_page.content[: self.doc_content_chars_max],
            metadata={
                **main_meta,
                **add_meta,
            },
        )
        return doc

ここの中の

        doc = Document(
            page_content=wiki_page.content[: self.doc_content_chars_max],
            metadata={
                **main_meta,
                **add_meta,
            },
        )

page_contentとmetadataを持ったDICTがOutputであることがわかります。

観察結果

上記を観察した結果下記のようにInput/Outputを書けば良いことがわかりました。

Input(入力)	Output（出力）
query: str 文字型のqueryというパラメータ	[Document(page_content="Retriveした内容",metadata={'source':'どこから取ってきたかなど'})]

※入力のqueryというパラメータ名は変更する事ができます。
※出力のmetadataのsource部分は任意に書き換えれます。

出来上がったテンプレート

下記のようにCustomLoaderを作るときのテンプレートを作りました。

注意
現在は固定値で返すようにしていますが、この部分を任意の場所からRetriveするように作り替えてください。

ドキュメントローダー(Retriver)

from typing import List

from langchain.callbacks.manager import CallbackManagerForRetrieverRun
from langchain.schema import BaseRetriever, Document

class CustomRetriever(BaseRetriever):
    """
    Custom retriever.
    """

    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
      doc  = [Document(page_content="島根県には出雲大社があります。",metadata={'summary':'test'})]
      return doc

使い方

import os
import openai
from langchain.llms import OpenAI
from langchain.chains import RetrievalQA
from langchain.chat_models import ChatOpenAI
from typing import List

from langchain.callbacks.manager import CallbackManagerForRetrieverRun
from langchain.schema import BaseRetriever, Document

class CustomRetriever(BaseRetriever):
    """
    Custom retriever.
    """

    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
      doc  = [Document(page_content="島根県には出雲大社があります。",metadata={'summary':'test'})]
      return doc

os.environ["OPENAI_API_KEY"] = #あなたのAPIKEY
retriever = CustomRetriever()
# retriever.get_relevant_documents(query='島根県には何がありますか？') #Retriverがどう拾って来るかみる

llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0)

qa = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=retriever,
    verbose=True
)

qa.run("島根県には何がありますか？")

こんな感じになります。Class部分は別のPythonファイルに分けるとスッキリすると思いますね。

実行結果

> Entering new RetrievalQA chain...

> Finished chain.
島根県には出雲大社があります。

まとめ

以上、実装方法とテンプレートになります。みなさまよきRAGライフを！