services:
  postgres:
    image: postgres:latest
    environment:
      POSTGRES_PASSWORD: password
    ports:
      - "5432:5432"
    volumes:
      - ./postgres_data:/var/lib/postgresql/data

  jupyterlab:
    image: quay.io/jupyter/minimal-notebook:latest
    environment:
      JUPYTER_TOKEN: "password"     # Jupyterlabにアクセスするためのトークン
      OPENAI_API_KEY: ${OPENAI_API_KEY}
    ports:
      - "8888:8888"
    volumes:
      - ./jupyter_work:/home/jovyan/work

  qdrant:
    image: qdrant/qdrant:latest
    ports:
      - "6333:6333"
      - "6334:6334"
    volumes:
      - ./qdrant_storage:/qdrant/storage:z

docker compose up -d

!pip install pandas psycopg2-binary

!wget https://tulip.kuee.kyoto-u.ac.jp/localgovfaq/localgovfaq.zip
!unzip localgovfaq.zip

import pandas as pd

def file2list(filename: str, prefix: str = "") -> tuple:
    """Q/AファイルをIDとコンテンツに分割、それぞれを配列で返す"""
    contents = []
    ids = []
    try:
        with open(filename, 'r') as file:
            for line in file:
                line = line.strip().replace(" ", "")
                id, content = line.split('\t')
                if prefix:
                    id = f"{prefix}_{id}"
                contents.append(content)
                ids.append(id)
    except Exception:
        raise

    return contents, ids

questions, ids = file2list("localgovfaq/qas/questions_in_Amagasaki.txt")
answers, _ = file2list("localgovfaq/qas/answers_in_Amagasaki.txt")

raw_df = pd.DataFrame({'qa_id': ids, 'question': questions, 'answer': answers})
df = raw_df.head(100)
df

import psycopg2
import pandas as pd

# 接続情報
host = "postgres"
port = "5432"
user = "postgres"
password = "password"

db_name = "faq"
table_name = "faq"
table_columns = ["qa_id", "question", "answer"]

try:
    # faqデータベース作成
    conn = psycopg2.connect(
        dbname="postgres",
        user=user,
        password=password,
        host=host,
        port=port
    )
    conn.autocommit = True
    cur = conn.cursor()
    cur.execute(f"CREATE DATABASE {db_name}")
    print(f"データベース作成: '{db_name}'")
    cur.close()
    conn.close()

    # faqデータベースに接続してテーブルとトリガーを作成
    with psycopg2.connect(
        dbname=db_name,
        user=user,
        password=password,
        host=host,
        port=port
    ) as conn:
        with conn.cursor() as cur:
            # テーブル作成
            create_table_query = f"""
            CREATE TABLE IF NOT EXISTS {table_name} (
                id BIGSERIAL PRIMARY KEY,
                qa_id INT,
                question TEXT,
                answer TEXT,
                updated_at TIMESTAMP DEFAULT NOW()
            )
            """
            cur.execute(create_table_query)

            # 注意:
            # 今回は使用しないが、新規登録・更新時にタイムスタンプを更新して
            # VectorETLである時点からの更新データのみを登録する、というような使い方
            # を想定してトリガーを設定した。
 
            # トリガー関数作成
            create_function_query = """
            CREATE OR REPLACE FUNCTION update_timestamp()
            RETURNS TRIGGER AS $$
            BEGIN
                NEW.updated_at = NOW();
                RETURN NEW;
            END;
            $$ LANGUAGE plpgsql;
            """
            cur.execute(create_function_query)

            # トリガー作成
            create_trigger_query = f"""
            CREATE TRIGGER set_timestamp
            BEFORE UPDATE ON {table_name}
            FOR EACH ROW
            EXECUTE FUNCTION update_timestamp();
            """
            cur.execute(create_trigger_query)

            conn.commit()
            print(f"テーブルとトリガー作成: '{table_name}'")

    # faqテーブルにデータ登録
    with psycopg2.connect(
        dbname=db_name,
        user=user,
        password=password,
        host=host,
        port=port
    ) as conn:
        with conn.cursor() as cur:
            insert_query = f"INSERT INTO {table_name} (qa_id, question, answer) VALUES (%s, %s, %s)"

            for _, row in df.iterrows():
                cur.execute(insert_query, tuple(row))
            
            conn.commit()
            print("データ登録完了")
            
except Exception as e:
    print(f"An error occurred: {e}")

finally:
    if conn:
        conn.close()

データベース作成: 'faq'
テーブルとトリガー作成: 'faq'
データ登録完了

with psycopg2.connect(
    dbname=db_name,
    user=user,
    password=password,
    host=host,
    port=port
) as conn:
    with conn.cursor() as cur:
        select_query = f"SELECT * FROM {table_name}"
        cur.execute(select_query)
        results = cur.fetchall()

for r in results[:5]:
    print("\n".join(map(str, r)))
    print("----")

!pip install qdrant-client

from qdrant_client import QdrantClient
from qdrant_client.http.models import Distance, VectorParams

client = QdrantClient(url="http://qdrant:6333")
collection_name = "faq"

if client.collection_exists(collection_name):
    client.delete_collection(collection_name=collection_name)

client.create_collection(
    collection_name=collection_name,
    vectors_config=VectorParams(size=1536, distance=Distance.COSINE),
)

True

!pip install --upgrade vector-etl

!pip freeze | grep -i vector-etl

vector-etl==0.1.7.1

from vector_etl import create_flow
import os

# データソースの定義
source = {
    "source_data_type": "database",
    "db_type": "postgres",
    "host": "postgres",
    "port": "5432",
    "database_name": "faq",
    "username": "postgres",
    "password": "password",
    "query": "select * from faq",
}

# Embeddingの定義
embedding = {
    "embedding_model": "OpenAI",
    "api_key": os.environ["OPENAI_API_KEY"],
    "model_name": "text-embedding-3-small"
}

# ターゲットの定義
target = {
    "target_database": "Qdrant",
    "qdrant_url": "http://qdrant:6333",
    "qdrant_api_key": "",
    "collection_name": "faq"
}

# ベクトル化対象のカラムの定義
embed_columns = ["answer"]

# パイプラインを作成
flow = create_flow()

# 各定義をパイプラインにセット
flow.set_source(source)
flow.set_embedding(embedding)
flow.set_target(target)
flow.set_embed_columns(embed_columns)

flow.execute()

INFO:vector_etl.orchestrator:ETL process completed successfully.

from openai import OpenAI
from qdrant_client import QdrantClient


def get_embedding(client, text, model):
    response = client.embeddings.create(
        input=text,
        model=model,
    )
    return response.data[0].embedding


openai_client = OpenAI()
qdrant_client = QdrantClient(url="http://qdrant:6333")

query = "地域総合センター今北への行き方を教えて。"
query_embed = get_embedding(openai_client, query, "text-embedding-3-small")

results = qdrant_client.query_points(
    collection_name="faq",
    query=query_embed,
    with_payload=True,
    limit=3
).points

for r in results:
    print(r.score)
    print(r.payload["answer"])
    print("----")

embedding:
  embedding_model: "OpenAI"
  api_key: ${OPENAI_API_KEY}
  model_name: "text-embedding-ada-002"