全端 LLM 應用開發-Day15-Python SDK 整合 pgvector

昨天都使用 SQL 語法來處理向量搜尋，今天我們來用 pgvector 的 Python SDK 來處理，並且整合 text-embedding-ada-002。

環境架設

1. 先安裝 pgvector 的 python SDK 以及連線到 PostgreSQL 的 psycopg2，使用指令 poetry add pgvector psycopg2-binary。

2. 在 FastAPI 裡，常常會使用 sqlalchemy 來做為 ORM，因此我們也需要安裝 sqlalchemy，使用指令 poetry add sqlalchemy

3. alembic 是常常和 sqlalchemy 搭配使用的 db migration 工具，這會讓我們更好管理 DB 的 migration，也一起安裝起來吧。使用指令 poetry add alembic

4. 使用指令 alembic init alembic，就可以初始化你的 alembic。你會看到你的專案裡多了一個 alembic 的資料夾，這裡就是可以管理你的 migration 的地方。專案裡也會多一個 alembic.ini。

5. 在 alembic 資料夾裡，打開 alembic.ini 檔，找到 sqlalchemy.url 這個屬性，在這裡放上你的資料庫連線字串。如果你是用昨天那個 docker-composer 建起來的話，那麼就是 postgresql://user:password@localhost:5432/vector

6. 接著我們建立一個檔案，叫 models.py。貼上下面的程式碼。注意 1536 是指 OpenAI text-embedding-ada-002 的 1536 維，這是要建立一個資料庫的 table 叫做 Embedding。

from sqlalchemy import Column, BigInteger
from sqlalchemy.ext.declarative import DeclarativeMeta, declarative_base
from pgvector.sqlalchemy import Vector

Base: DeclarativeMeta = declarative_base()


class Embedding(Base):
    __tablename__ = 'embeddings'
    id = Column(BigInteger, primary_key=True, index=True)
    vector = Column(Vector(1536)) 

7. 接著我們使用指令 alembic revision --autogenerate -m "Created embeddings table"

8. 這時候後該會跳出這段錯誤： SAWarning: Did not recognize type 'vector' of column 'embedding'。這是因為我們的 vector 欄位型態不是 PostgreSQL 的原始欄位，是靠著 extension 取得的。因此我們要再做一些修改，才會讓 alembic 知道。

9. 我們去到 alembic 創建的 versions 資料夾，然後找到 f665c637196d_created_embeddings_table.py 這個檔案，注意前面的亂數可能會不一樣。然後再這個 python 檔做以下的修改。

from pgvector.sqlalchemy import Vector

# 原本是這一行，import 後做修改
# sa.Column('vector', pgvector.sqlalchemy.Vector(dim=1536), nullable=True),
sa.Column('vector', Vector(dim=1536), nullable=True),

11. 接著我們再使用 alembic upgrade head，就會看到下面資料庫的 embedding table 創建成功的訊息了。你也可以進到 pgAdmin 裡看。

INFO  [alembic.runtime.migration] Context impl PostgresqlImpl.
INFO  [alembic.runtime.migration] Will assume transactional DDL.
INFO  [alembic.runtime.migration] Running upgrade  -> f665c637196d, Created embeddings table

開始寫程式

我們把環境架好之後，就來開始寫程式吧！

1. 先建了一個檔案叫 postgre_vecoter.py ，先建立資料庫連線的 session ，然後建立 Day 13 的 function 來用 OpenAI 來做 embedding。

def get_embedding(text, model_name):
    response = openai.Embedding.create(
        input=text,
        engine=model_name
    )
    return response['data'][0]['embedding']

2. 然後我們建立和資料庫的連線，以及把 embedding 的向量值存進去資料庫。

from sqlalchemy import create_engine, select
from sqlalchemy.orm import Session
from sqlalchemy.orm import sessionmaker

from models import Embedding

engine = create_engine("postgresql://user:password@localhost:5432/vector")
Session = sessionmaker(bind=engine)
session = Session()

def add_to_pg_vector(session, embeddings):
    embeddings_obj = [Embedding(vector=vector) for vector in embeddings]
    session.add_all(embeddings_obj)
    session.commit()

3. 接著我們再建立一個 function，是搜尋 pgvector 的相似向量。這裡我們用 cosine_distance。

def search_from_pg_vector(session, text_array, query_embedding, k=1):
    results = session.scalars(select(Embedding).order_by(
        Embedding.vector.cosine_distance(query_embedding)).limit(k))
    distance = session.scalars(
        select((Embedding.vector.cosine_distance(query_embedding))))
    
    # zip distance 和 text_array，並且把 cosine_distance 轉成 similarity。
    # results.id - 1 是因為 id 從 1 開始，而我們的 text_array 從 0 開始
    return [(text_array[result.id - 1], 1 - float(dist)) for dist, result in zip(distance, results)]

4. 然後再複製貼上下面的主程式吧，記得要把 session close 才是好習慣哦！

def main():
    EMBEDDING_MODEL_NAME = "embedding-ada-002"  # 你前幾天在 Azure OpenAI 上建立的模型名稱
    openai.api_base = "https://japanopenai2023ironman.openai.azure.com/"
    openai.api_key = "yourkey"
    openai.api_type = "azure"
    openai.api_version = "2023-03-15-preview"

    text_array = ["我會披星戴月的想你，我會奮不顧身的前進，遠方煙火越來越唏噓，凝視前方身後的距離",
                  "而我，在這座城市遺失了你，順便遺失了自己，以為荒唐到底會有捷徑。而我，在這座城市失去了你，輸給慾望高漲的自己，不是你，過分的感情"]

    embedding_array = [get_embedding(
        text, EMBEDDING_MODEL_NAME) for text in text_array]

    pg_vector = add_to_pg_vector(session, embedding_array)

    # Example search
    query_text = "工程師寫城市"
    query_embedding = get_embedding(query_text, EMBEDDING_MODEL_NAME)
    search_results = search_from_pg_vector(session, text_array, query_embedding)
    print(f"尋找 {query_text}:", search_results)


if __name__ == '__main__':
    try:
        main()
    finally:
        session.close()

跑起來後，應該就會看到結果如下：

尋找 工程師寫城市: [('而我，在這座城市遺失了你，順便遺失了自己，以為荒唐到底會有捷徑。而我，在這座城市失去了你，輸給慾望高漲的自己，不是你，過分的感情' 0.779329981883772)]

這相似度也太高了吧！

Pgvector 算是還滿不錯的一個向量資料庫 extension ，尤其是你們公司就在使用 PostgreSQL 的話，那就再適合不過了。但是 pgvector 最為人詬病的就是資料量一大，效能就會變差。如果真的碰到這一天，就用專業的向量資料庫吧！明天開始我們就來講這一塊。