Google GeminiとMyScaleDBを使用したRAGパワードチャットボットの構築方法

from langchain_community.document_loaders import PyPDFLoader
loader = PyPDFLoader("Employee_Handbook.pdf")
pages = loader.load_and_split()
pages = pages[4:]  # 最初の数ページは不要なのでスキップします
text = "\n".join([doc.page_content for doc in pages])
 
        Copied!
    

from langchain_text_splitters import RecursiveCharacterTextSplitter
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,
    chunk_overlap=150,
    length_function=len,
    is_separator_regex=False,
)
docs = text_splitter.create_documents([text])
for i, d in enumerate(docs):
    d.metadata = {"doc_id": i}
 
        Copied!
    

import os
import google.generativeai as genai
import pandas as pd

os.environ["GEMINI_API_KEY"] = "your_key_here"

# この関数は文を引数として受け取り、その埋め込みを返します
def get_embeddings(text):
    # 埋め込みモデルを定義します
    model = 'models/embedding-001'
    # 埋め込みを取得します
    embedding = genai.embed_content(model=model,
                                    content=text,
                                    task_type="retrieval_document")
    return embedding['embedding']

# ドキュメントからpage_contentを取得し、新しいリストを作成します
content_list = [doc.page_content for doc in docs]
# 1つのpage_contentを送信します
embeddings = [get_embeddings(content) for content in content_list]

# データフレームを作成してデータベースにインジェストします
dataframe = pd.DataFrame({
    'page_content': content_list,
    'embeddings': embeddings
})
 
        Copied!
    

import clickhouse_connect

client = clickhouse_connect.get_client(
    host='your_host_name',
    port="port_number,
    username='your_username',
    password='yiur_password_hhere'
)

 
        Copied!
    

# テーブル名を 'handbook' としてテーブルを作成します
client.command("""
    CREATE TABLE default.handbook (
        id Int64,
        page_content String,
        embeddings Array(Float32),
        CONSTRAINT check_data_length CHECK length(embeddings) = 768
    ) ENGINE = MergeTree()
    ORDER BY id
""")

# CONSTRAINTは、各埋め込みベクトルの長さが768であることを確認します

# データをバッチで挿入します
batch_size = 10
num_batches = len(dataframe) // batch_size

for i in range(num_batches):
    start_idx = i * batch_size
    end_idx = start_idx + batch_size
    batch_data = dataframe[start_idx:end_idx]
    # データを挿入します
    client.insert("default.handbook", batch_data.to_records(index=False).tolist(), column_names=batch_data.columns.tolist())
    print(f"Batch {i+1}/{num_batches} inserted.")

# クイックなデータの検索のためにベクトルインデックスを作成します
client.command("""
ALTER TABLE default.handbook
    ADD VECTOR INDEX vector_index embeddings
    TYPE MSTG
""")
 
        Copied!
    

def get_relevant_docs(user_query):
    # get_embeddings関数を再度呼び出して、ユーザークエリをベクトル埋め込みに変換します
    query_embeddings = get_embeddings(user_query)
    # クエリを実行します
    results = client.query(f"""
        SELECT page_content,
        distance(embeddings, {query_embeddings}) as dist FROM default.handbook ORDER BY dist LIMIT 3
    """)
    relevant_docs = []
    for row in results.named_results():
        relevant_docs.append(row['page_content'])
    return relevant_docs
 
        Copied!
    

def make_rag_prompt(query, relevant_passage):
    relevant_passage = ' '.join(relevant_passage)
    prompt = (
        f"You are a helpful and informative chatbot that answers questions using text from the reference passage included below. "
        f"Respond in a complete sentence and make sure that your response is easy to understand for everyone. "
        f"Maintain a friendly and conversational tone. If the passage is irrelevant, feel free to ignore it.\n\n"
        f"QUESTION: '{query}'\n"
        f"PASSAGE: '{relevant_passage}'\n\n"
        f"ANSWER:"
    )
    return prompt

import google.generativeai as genai

def generate_response(user_prompt):
    model = genai.GenerativeModel('gemini-pro')
    answer = model.generate_content(user_prompt)
    return answer.text

def generate_answer(query):
    relevant_text = get_relevant_docs(query)
    text = " ".join(relevant_text)
    prompt = make_rag_prompt(query, relevant_passage=relevant_text)
    answer = generate_response(prompt)
    return answer

answer = generate_answer(query="What is the Work Dress Code?")
print(answer)
 
        Copied!