langchain-awsのAmazonS3VectorsでLangGraphのRAGのベクトルストアをマネージドに差し替えてみる

Python 3.13
litellm 1.83.14
langgraph 1.1.10
langchain 1.1.0
langchain-litellm 0.6.4
langchain-aws 1.4.6
boto3 1.41.0

$ uv pip install litellm langgraph langchain langchain-core langchain-litellm langchain-aws boto3

# ベクトルバケットの作成（既に作成済みであれば不要）
$ aws s3vectors create-vector-bucket \
  --vector-bucket-name s3-vectors-482842011168

$ export ANTHROPIC_API_KEY="sk-ant-..."

import os

from langchain_aws import BedrockEmbeddings
from langchain_aws.vectorstores.s3_vectors import AmazonS3Vectors
from langchain_core.documents import Document

BUCKET_NAME = os.environ.get("S3_VECTORS_BUCKET", "s3-vectors-482842011168")
INDEX_NAME = os.environ.get("S3_VECTORS_INDEX", "company-docs")
EMBEDDING_MODEL_ID = "amazon.titan-embed-text-v2:0"
REGION = os.environ.get("AWS_REGION", "ap-northeast-1")

# 社内ドキュメント
documents = [
    Document(
        id="doc_1",
        page_content="有給休暇は入社6ヶ月後に10日付与されます。申請は3営業日前までに上長の承認が必要です。",
        metadata={"category": "leave"},
    ),
    Document(
        id="doc_2",
        page_content="リモートワークは週3日まで可能です。事前申請は不要ですが、勤怠システムへの記録が必要です。",
        metadata={"category": "workstyle"},
    ),
    Document(
        id="doc_3",
        page_content="経費精算は月末締め、翌月15日払いです。領収書の原本またはスキャンデータの提出が必要です。",
        metadata={"category": "expense"},
    ),
    Document(
        id="doc_4",
        page_content="育児休業は子が1歳になるまで取得可能です。延長の場合は最長2歳まで認められます。",
        metadata={"category": "leave"},
    ),
    Document(
        id="doc_5",
        page_content="健康診断は年1回、会社負担で受診できます。35歳以上は人間ドックも選択可能です。",
        metadata={"category": "health"},
    ),
]

# BedrockのTitan Text Embeddings V2で1024次元のベクトルを生成
embeddings = BedrockEmbeddings(model_id=EMBEDDING_MODEL_ID, region_name=REGION)

# S3 Vectorsをバックエンドにしたベクトルストア
# create_index_if_not_exist=True なのでインデックスが無ければ自動作成される
vector_store = AmazonS3Vectors(
    vector_bucket_name=BUCKET_NAME,
    index_name=INDEX_NAME,
    embedding=embeddings,
    region_name=REGION,
    distance_metric="cosine",
    create_index_if_not_exist=True,
)

def ingest() -> None:
    """サンプルドキュメントをS3 Vectorsに登録する。同じIDで上書きされる。"""
    vector_store.add_documents(documents)
    print(f"ドキュメント {len(documents)} 件を登録しました")

if __name__ == "__main__":
    ingest()

$ python common.py
ドキュメント 5 件を登録しました

from common import vector_store
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_litellm import ChatLiteLLM

COMPLETION_MODEL = "anthropic/claude-sonnet-4-6"

# S3 Vectorsをretrieverとして利用
retriever = vector_store.as_retriever(search_kwargs={"k": 2})

prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "以下の社内ドキュメントの情報のみをもとに質問に回答してください。"
            "ドキュメントに記載がない場合は「情報が見つかりませんでした」と回答してください。\n\n"
            "{context}",
        ),
        ("user", "{question}"),
    ]
)

llm = ChatLiteLLM(model=COMPLETION_MODEL)

def format_docs(docs):
    return "\n".join(d.page_content for d in docs)

# LCELでretrieve → prompt → LLM → 文字列抽出 のパイプラインを構築
chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

queries = [
    "リモートワークは週何日までできますか？",
    "健康診断の費用は自己負担ですか？",
]

for query in queries:
    print(f"\n{'='*50}")
    print(f"質問: {query}")
    print(f"{'='*50}")
    print(f"回答: {chain.invoke(query)}")

$ python simple_rag.py

==================================================
質問: リモートワークは週何日までできますか？
==================================================
回答: リモートワークは週3日まで可能です。事前申請は不要ですが、勤怠システムへの記録が必要です。

==================================================
質問: 健康診断の費用は自己負担ですか？
==================================================
回答: 自己負担ではありません。健康診断は年1回、会社負担で受診できます。35歳以上は人間ドックも選択可能です。

from typing_extensions import TypedDict

from common import vector_store
from langchain_core.documents import Document
from langchain_litellm import ChatLiteLLM
from langgraph.graph import END, START, StateGraph

COMPLETION_MODEL = "anthropic/claude-sonnet-4-6"

llm = ChatLiteLLM(model=COMPLETION_MODEL)
retriever = vector_store.as_retriever(search_kwargs={"k": 2})

class RAGState(TypedDict):
    question: str
    documents: list[Document]
    answer: str

def retrieve_node(state: RAGState) -> dict:
    """S3 Vectorsから関連ドキュメントを取得する。"""
    docs = retriever.invoke(state["question"])
    return {"documents": docs}

def generate_node(state: RAGState) -> dict:
    """取得したドキュメントをコンテキストにしてLLMで回答を生成する。"""
    context = "\n".join(d.page_content for d in state["documents"])
    messages = [
        {
            "role": "system",
            "content": (
                "以下の社内ドキュメントの情報のみをもとに質問に回答してください。"
                "ドキュメントに記載がない場合は「情報が見つかりませんでした」と回答してください。\n\n"
                f"{context}"
            ),
        },
        {"role": "user", "content": state["question"]},
    ]
    response = llm.invoke(messages)
    return {"answer": response.content}

# retrieve → generate の2ノードグラフを構築
graph = StateGraph(RAGState)
graph.add_node("retrieve", retrieve_node)
graph.add_node("generate", generate_node)
graph.add_edge(START, "retrieve")
graph.add_edge("retrieve", "generate")
graph.add_edge("generate", END)

app = graph.compile()

queries = [
    "リモートワークは週何日までできますか？",
    "健康診断の費用は自己負担ですか？",
]

for query in queries:
    print(f"\n{'='*50}")
    print(f"質問: {query}")
    print(f"{'='*50}")
    result = app.invoke({"question": query})
    print("\n[検索ヒット]")
    for d in result["documents"]:
        print(f"- {d.page_content}")
    print(f"\n[回答] {result['answer']}")

$ python langgraph_rag.py

==================================================
質問: リモートワークは週何日までできますか？
==================================================

[検索ヒット]
- リモートワークは週3日まで可能です。事前申請は不要ですが、勤怠システムへの記録が必要です。
- 経費精算は月末締め、翌月15日払いです。領収書の原本またはスキャンデータの提出が必要です。

[回答] リモートワークは週3日まで可能です。

==================================================
質問: 健康診断の費用は自己負担ですか？
==================================================

[検索ヒット]
- 健康診断は年1回、会社負担で受診できます。35歳以上は人間ドックも選択可能です。
- 有給休暇は入社6ヶ月後に10日付与されます。申請は3営業日前までに上長の承認が必要です。

[回答] いいえ、自己負担ではありません。健康診断は年1回、会社負担で受診できます。

from common import vector_store
from langchain.agents import create_agent
from langchain_core.tools import tool
from langchain_litellm import ChatLiteLLM

COMPLETION_MODEL = "anthropic/claude-sonnet-4-6"

retriever = vector_store.as_retriever(search_kwargs={"k": 3})

@tool
def search_company_docs(query: str) -> str:
    """社内規定や制度に関するドキュメントをS3 Vectorsから検索します。"""
    docs = retriever.invoke(query)
    if not docs:
        return "該当するドキュメントが見つかりませんでした"
    return "\n\n".join(d.page_content for d in docs)

system_prompt = (
    "あなたは社内ドキュメントに基づいて質問に回答するアシスタントです。\n"
    "以下のルールに従ってください:\n"
    "1. 社内制度に関する質問には、必ず search_company_docs ツールで検索してから回答してください\n"
    "2. 検索結果に情報がない場合は「社内ドキュメントに該当する情報が見つかりませんでした」と回答してください\n"
    "3. 一般的な知識で回答できる質問（挨拶など）には、ツールを使わず直接回答してください\n"
    "4. 回答は簡潔にまとめてください"
)

llm = ChatLiteLLM(model=COMPLETION_MODEL)
agent = create_agent(llm, tools=[search_company_docs], system_prompt=system_prompt)

queries = [
    "リモートワークのルールを教えてください",
    "こんにちは！",
    "育児休業はいつまで取れますか？",
    "Pythonのリスト内包表記について教えてください",
]

for query in queries:
    print(f"\n{'='*60}")
    print(f"質問: {query}")
    print(f"{'='*60}")
    result = agent.invoke({"messages": [{"role": "user", "content": query}]})
    print(f"回答: {result['messages'][-1].content}")

$ python agentic_rag.py

============================================================
質問: リモートワークのルールを教えてください
============================================================
回答: リモートワークは週3日まで可能です。事前申請は不要ですが、勤怠システムへの記録が必要です。

============================================================
質問: こんにちは！
============================================================
回答: こんにちは！社内制度に関するご質問があればお気軽にどうぞ。

============================================================
質問: 育児休業はいつまで取れますか？
============================================================
回答: 育児休業は子が1歳になるまで取得可能です。延長の場合は最長2歳まで認められます。

============================================================
質問: Pythonのリスト内包表記について教えてください
============================================================
回答: 申し訳ありませんが、こちらは社内ドキュメントに基づいて回答するアシスタントのため、Pythonに関する技術的な質問にはお答えできません。

leave_retriever = vector_store.as_retriever(
    search_kwargs={
        "k": 3,
        "filter": {"category": {"$eq": "leave"}},
    }
)

docs = leave_retriever.invoke("休みについて教えて")
for d in docs:
    print(d.page_content)