cookbook / db /postgres_init.py
JasonFinley0821's picture
feat : upgrade embedding
a4137be
import psycopg2
from langchain_huggingface import HuggingFaceEmbeddings
# 💡 將 Chroma 替換為 PostgreSQL 的向量儲存類
# 這裡我們使用 LangChain 提供的 PostgresVectorStore
from langchain_community.vectorstores.pgvector import PGVector
import os
from dotenv import load_dotenv
load_dotenv()
# ⚠️ 請替換為您的實際配置
# 假設您的資料庫中已經創建了 'recipes_vectors' 表格
# --- 1. 資料庫連線設定 ---
# --- 資料庫連線設定 ---
DB_HOST=os.getenv("DB_HOST", "localhost")
DB_PORT=os.getenv("DB_PORT")
DB_DATABASE=os.getenv("DB_DATABASE")
DB_USER=os.getenv("DB_USER")
DB_PASSWORD=os.getenv("DB_PASSWORD")
DB_CONFIG = {
"host": DB_HOST, # 如果找不到,提供預設值
"database": DB_DATABASE,
"user": DB_USER,
"password": DB_PASSWORD,
"port": DB_PORT
}
TABLE_NAME = "recipes"
VECTOR_DIM = 384 # 確保這個維度和您的 'embedding' 欄位定義一致
# 🚨 選擇一個多語言模型。例如:BGE-M3 (維度 1024) 或 paraphrase-multilingual-mpnet-base-v2 (維度 768)
MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
# 這是 PGVector 實例使用的連接字串格式
CONNECTION_STRING = PGVector.connection_string_from_db_params(
driver="psycopg2",
**DB_CONFIG
)
# --- 將此函數定義為 get_vectordb ---
def get_vectordb():
"""
獲取已存在的 PostgreSQL 向量儲存實例 (PGVector)。
此函數取代了原先獲取 Chroma 實例的功能。
"""
embeddings = HuggingFaceEmbeddings(
model_name=MODEL_NAME
)
# 使用您的 PostgreSQL 連接字串和集合名稱
db = PGVector(
connection_string=CONNECTION_STRING,
embedding_function=embeddings,
collection_name=TABLE_NAME, # 您的 recipe table 名稱
use_jsonb=True
)
return db