sql_tool.py

import os
from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain_google_genai import ChatGoogleGenerativeAI

def get_sql_toolkit():
    """Khởi tạo toolkit kết nối với Supabase PostgreSQL."""
    # Lấy chuỗi kết nối từ biến môi trường
    db_url = os.getenv("SUPABASE_DB_URL")
    
    if not db_url:
        print("CẢNH BÁO: Chưa cấu hình SUPABASE_DB_URL trong file .env")
        print("Vui lòng thêm SUPABASE_DB_URL=postgresql://postgres:[PASSWORD]@[HOST]:5432/postgres")
        return None

    # Khởi tạo kết nối db
    db = SQLDatabase.from_uri(db_url)
    
    # Khởi tạo LLM cho toolkit (để nó biết cách viết câu query)
    llm = ChatGoogleGenerativeAI(model="gemini-1.5-pro", temperature=0)
    
    # Tạo toolkit
    toolkit = SQLDatabaseToolkit(db=db, llm=llm)
    return toolkit