PDFファイルを使ったRAGに挑戦（３）LangChain

from langchain.chains import RetrievalQA
from langchain.chains.question_answering import load_qa_chain
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.text_splitter import CharacterTextSplitter
from langchain.text_splitter import MarkdownHeaderTextSplitter, RecursiveCharacterTextSplitter
from langchain.document_loaders import TextLoader
from langchain.document_loaders import UnstructuredMarkdownLoader
from langchain.document_loaders import PDFPlumberLoader
from langchain.document_loaders import UnstructuredPDFLoader
from langchain.chat_models import ChatOpenAI
from langchain_anthropic import ChatAnthropic
from langchain.prompts import HumanMessagePromptTemplate, SystemMessagePromptTemplate, ChatPromptTemplate

import os


class QASystem:

    def __init__(self, persist_directory="/content/drive/MyDrive/work/chroma_db"):
        # Anthropic APIキーの設定
        os.environ["ANTHROPIC_API_KEY"] = claude_api_key
        os.environ["OPENAI_API_KEY"] = openai_api_key

        self.embeddings = OpenAIEmbeddings(
            model="text-embedding-3-large",
            openai_api_key=openai_api_key,
            chunk_size=1000,
            max_retries=3
        )

        self.llm = ChatAnthropic(
            model="claude-3-5-sonnet-20241022",
            temperature=0.0,
            anthropic_api_key=claude_api_key,
            max_tokens = 1000
        )

        self.persist_directory = persist_directory

        # 既存のDBがあれば読み込み、なければ新規作成
        if os.path.exists(persist_directory):
            self.vectorstore = Chroma(
                persist_directory=persist_directory,
                embedding_function=self.embeddings
            )
        else:
            self.vectorstore = None

    def add_documents(self, file_path):
        loader = UnstructuredMarkdownLoader(file_path)
        documents = loader.load()

        headers_to_split_on = [
            ("#", "header1"),
            ("##", "header2"),
            ("###", "header3"),
        ]
        markdown_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on)

        filename = os.path.basename(file_path)

        # メタデータを保持しながらドキュメントを分割
        md_splits = markdown_splitter.split_text(documents[0].page_content)
        for split in md_splits:
            split.metadata.update({
                'filename': filename,
                'header1': split.metadata.get('header1', ''),
                'header2': split.metadata.get('header2', ''),
                'header3': split.metadata.get('header3', '')
            })

        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000, 
            chunk_overlap=100,  
            separators=["\n\n", "\n", "。", "、", " ", ""],
            keep_separator=True
        )

        texts = text_splitter.split_documents(md_splits)

        if self.vectorstore is None:
            self.vectorstore = Chroma.from_documents(
                documents=texts,
                embedding=self.embeddings,
                persist_directory=self.persist_directory
            )
        else:
            self.vectorstore.add_documents(texts)


    def ask_question(self, filename, question: str) -> str:


        if self.vectorstore is None:
            return "ドキュメントが登録されていません。"

        # システムメッセージとヒューマンメッセージを作成
        system_template = """あなたは優秀なアナリストです。与えられた企業レポートを読みユーザからの質問に回答します。
        企業レポートはmarkdown形式で記述されています。セクションや表の構造を意識しながら注意深く分析してください。
        
        <<システムプロンプト　手順と制約事項>>
     
        企業レポート:
        {context}"""

        human_template = "{question}"

        system_message_prompt = SystemMessagePromptTemplate.from_template(system_template)
        human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)

        # チャットプロンプトテンプレートの作成
        chat_prompt = ChatPromptTemplate.from_messages([
            system_message_prompt,
            human_message_prompt
        ])


        retriever = self.vectorstore.as_retriever(
            search_type="mmr",  # Maximum Marginal Relevance
            search_kwargs={
                "k": 5,
                "filter": {"filename": filename},
                "fetch_k": 30,  # より多くの候補から選択
                "lambda_mult": 0.7  # 関連性と多様性のバランス
            }
        )

        qa_chain = RetrievalQA.from_chain_type(
            llm=self.llm,
            chain_type="stuff",
            retriever=retriever,
            return_source_documents=True,
            chain_type_kwargs={"prompt": chat_prompt}
        )

        result = qa_chain.invoke({"query": question})
        answer = result['result']
        sources = [doc.page_content[:100] + "..." for doc in result['source_documents']]


        return f"回答: {answer}\n"
        #return f"回答: {answer}\n\参照元:\n" + "\n".join(sources)