GoogleColabでGeminiを使ってCSVを元データにした簡易RAGを試してみた

Google ColabでLangChainのドキュメントに従い、CSVファイルを元にしたRAGを実装してみたので、メモ。 ※実際は、SQLiteにCSVデータを入れて、SQLでデータを読み出している

!pip install -qU langchain  langchain-community langchain-experimental pandas
!wget https://web.stanford.edu/class/archive/cs/cs109/cs109.1166/stuff/titanic.csv -O titanic.csv

import pandas as pd

df = pd.read_csv("titanic.csv")

from langchain_community.utilities import SQLDatabase
from sqlalchemy import create_engine

engine = create_engine("sqlite:///titanic.db")

# SQLiteにデータを書き込む
df.to_sql("titanic", engine, index=False)


db = SQLDatabase(engine=engine)

from langchain_community.agent_toolkits import create_sql_agent
import os
from google.colab import userdata
os.environ["GOOGLE_API_KEY"] = userdata.get("GOOGLE_API_KEY")

from langchain_google_genai import ChatGoogleGenerativeAI
# LLMの準備
model  = ChatGoogleGenerativeAI(model="gemini-1.5-flash-latest", temperature=0.9)
agent_executor = create_sql_agent(model, db=db, agent_type="tool-calling", verbose=True)

agent_executor.invoke({"input": "生き残りやすかったのはどのような人？"})

参考ページ

python.langchain.com