LangChainでLlama Indexを使う

LangChainで、Llama Indexで作成したIndexを使う方法が、調べてもいまいち出てこなかったので備忘録です.

llama-index==0.8.55
langchain==0.0.325
langchain-experimental==0.0.36

LangChainとは

LangChain は LLM を幅広くアプリケーションへ実装できるようにするためのフレームワーク. ReActを実現する Agent の機能がキー. javascript/typescript による API も提供しているので Web アプリとの相性も良いらしい..

LlamaIndexとは

LlamaIndex は独自のドキュメントを LLM に知識と与えて拡張するためにプロンプトに対して与えるインデックスを作ることに中心としたフレームワーク. Llama Hubという様々なカスタムローダーを提供するサービスもありインデックス作成にただならぬ熱意を感じる.

※参考: LLM 対応自然言語処理フレームワークをいろいろ触ってテーブルデータで試してみたかった

LlamaIndexでQA

ここまで見たように, Pipelineとして優秀なLangChainに, 秀逸なLlamaIndexを使えればかなりの精度のRAGが作れるのでは? ということです.

もちろんLlamaIndexでもQuestion & Answerを実現する機構は整っています. LlamaIndexでは以下のようにインデックスを作ります. まずはディレクトリ以下のドキュメントを読み込んで, Vector Indexを作ります.

from llama_index import SimpleDirectoryReader
from llama_index import StorageContext, ServiceContext
from llama_index import VectorStoreIndex

# ドキュメントのロード
directory_reader = SimpleDirectoryReader(
        input_dir="path/to/your/document", recursive=True,
        filename_as_id = True,
    )
documents = directory_reader.load_data()

# index作成
storage_context = StorageContext.from_defaults()
service_context = ServiceContext.from_defaults()
index = VectorStoreIndex.from_documents(
        documents, storage_context=storage_context, service_context=service_context, show_progress=True
    )

filename_as_id=Trueにすると, Answerに対する根拠(source)を得ることが簡単になります. (contextの引数は省略してます)LlamaIndexにおいて回答を得る方法は以下の通りです.

query_engine = index.as_query_engine()
response = query_engine.query("How are you?")

これだけです.

LangChainでLlamaIndex

さてLangChainではどうするかというと, 結論は以下のretrieverを作成するだけです.

from langchain.pydantic_v1 import Field
from typing import Any, Dict, List, cast

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

class LlamaIndexRetriever(BaseRetriever):
    """`LlamaIndex` retriever.

    It is used for the question-answering with sources over
    an LlamaIndex data structure."""
    index: Any
    query_kwargs: Dict = Field(default_factory=dict)

    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        """Get documents relevant for a query."""
        try:
            from llama_index.indices.base import BaseGPTIndex
            from llama_index.response.schema import Response
        except ImportError:
            raise ImportError(
                "You need to install `pip install llama-index` to use this retriever."
            )
        index = cast(BaseGPTIndex, self.index)
        query_engine = index.as_query_engine(
            response_mode="no_text",
            similarity_top_k=3,
            **self.query_kwargs)
        response = query_engine.query(query)
        response = cast(Response, response)
        # parse source nodes
        docs = []
        for source_node in response.source_nodes:
            docs.append(
                Document(page_content=source_node.node.text,
                         metadata={"source": source_node.node.relationships[NodeRelationship.SOURCE].node_id})
            )
        return docs

あとはこんな感じで使えます.

from langchain.chains import RetrievalQA
from langchain.llms import OpenAI

qa = RetrievalQA.from_chain_type(llm=OpenAI(), chain_type="stuff",retriever=LlamaIndexRetriever(index=index))
query = "What did the president say about Ketanji Brown Jackson"
qa.run(query)

公式で, 実装するクラスもありますがなぜか動かないので自前でやる必要があるみたいですね. (2023.10.26現在)

感想

もし他の方法でできたひといたら教えてください!