Microsoft AutoGen & Mistral Large로 PostgreSQL 데이터베이스 조회하기
Microsoft AutoGen & Mistral Large로 PostgreSQL 데이터베이스 조회하기
Microsoft AutoGen과 Mistral Large V2를 조합해 PostgreSQL 데이터베이스를 질의하는 방법을 보여주는 튜토리얼이에요. 함수 호출(function calling)을 이용해 AutoGen 에이전트가 안전하게(읽기 전용) 데이터베이스를 조회하고 답변하도록 구성합니다.
출처: 문서
본문
이 튜토리얼은 Microsoft AutoGen을 Mistral Large V2와 함께 사용해 Postgres 데이터베이스를 질의하는 방법의 예시로 작성됐어요. 짧고 간단한 개요를 제공하며, 다른 DB나 VectorDB를 질의하도록 추가 도구를 작성하고 붙이는 것도 가능합니다.
질의와 프롬프트 외에 좋은 문맥도 정말 중요합니다. 그래서 사용 가능한 모든 테이블의 문맥을 항상 제공하는 추가 함수도 넣었어요.
참고: 너무 많은 도구와 문맥을 함께 쓰기보다, 모든 것을 하는 단일 큰 모델 대신 서로 다른 '챗 모델'을 사용하도록 하세요.
준비 단계
처음 몇 단계에서는 PostgreSQL을 설치하고 데이터베이스 덤프를 가져옵니다. 여기서는 name.basics.tsv 데이터베이스 덤프를 사용합니다. (https://datasets.imdbws.com/, https://wiki.postgresql.org/wiki/Sample_Databases)
!sudo apt update
!sudo apt install dirmngr ca-certificates software-properties-common gnupg gnupg2 apt-transport-https curl -y
!curl -fSsL https://www.postgresql.org/media/keys/ACCC4CF8.asc | gpg --dearmor | sudo tee /usr/share/keyrings/postgresql.gpg > /dev/null
!echo 'deb [arch=amd64,arm64,ppc64el signed-by=/usr/share/keyrings/postgresql.gpg] http://apt.postgresql.org/pub/repos/apt/ jammy-pgdg main' | sudo tee /etc/apt/sources.list.d/pgdg.list
!sudo apt update
!sudo apt install postgresql-client-16 postgresql-16 -y
!sudo service postgresql start
!sudo -u postgres psql -U postgres -c "ALTER ROLE postgres WITH PASSWORD 'super_secret_postgres_password';"
# Download the dataset from IMDB
!cd /
!sudo mkdir data
!sudo wget https://datasets.imdbws.com/name.basics.tsv.gz
!sudo gunzip name.basics.tsv.gz
# Create a super user
!sudo -u postgres psql -U postgres -c "CREATE ROLE root WITH SUPERUSER;"
!sudo -u postgres psql -U postgres -c "ALTER ROLE root WITH LOGIN;"
!sudo -u postgres psql -U postgres -c "CREATE ROLE postgres WITH PASSWORD 'super_secret_postgres_password';"
# Import the dataset
!sudo -u postgres psql -U postgres -c "CREATE TABLE imdb ( nconst TEXT, primaryName TEXT, birthYear INT, deathYear INT, primaryProfession TEXT, knownForTitles TEXT);"
#It is possible you have to change the directory which containts this file
!sudo -u postgres psql -U postgres -c "COPY imdb FROM '/content/name.basics.tsv' WITH (HEADER true);"
!pip install pyautogen psycopg2
필요한 패키지를 모두 임포트합니다. 이 경우 autogen, postgresql 패키지와 기타 라이브러리가 필요해요.
from autogen import ConversableAgent, register_function
from typing import List, Optional, Union, Dict, Any
import psycopg2
import os
PostgreSQL 함수 정의하기
Postgres 함수는 유효한 json 입력을 받아 이 입력을 기반으로 쿼리를 실행합니다. 이후 함수는 결과를 LLM에 반환하고, LLM은 그에 대한 응답을 사용자에게 메시지로 전달합니다. 도구 사용법을 미리 정의했기 때문에 레코드를 삭제·수정·생성할 수 없고 읽기 전용 접근만 가능합니다.
def execute_postgres_query(
table_name: str,
columns: List[str],
filters: Optional[Dict[str, Any]] = None,
sort_column: Optional[str] = None,
sort_order: Optional[str] = None,
limit: Optional[int] = 150, # Default limit of 150 rows, you can edit this yourself if needed, the AI will also be able to change this.
):
# Validate input
if not table_name:
return "Error: table_name is required"
if not columns:
return "Error: columns is required"
if sort_column and not sort_order:
return "Error: sort_order is required when sort_column is specified"
# Generate SQL query
query = f"SELECT {', '.join(columns)} FROM {table_name}"
params = []
if filters:
filter_conditions = []
for column, value in filters.items():
if isinstance(value, str) and value.startswith('%') and value.endswith('%'):
filter_conditions.append(f"{column} LIKE %s")
params.append(value)
elif isinstance(value, list):
filter_conditions.append(f"{column} NOT IN %s")
params.append(tuple(value))
else:
filter_conditions.append(f"{column} = %s")
params.append(value)
query += " WHERE " + " AND ".join(filter_conditions)
if sort_column and sort_order:
query += f" ORDER BY {sort_column} {sort_order}"
if limit:
query += f" LIMIT {limit}"
# Execute SQL query
conn = psycopg2.connect(database="postgres", user="postgres", password="super_secret_postgres_password", host="localhost", port="5432")
cur = conn.cursor()
cur.execute(query, params)
results = cur.fetchall()
cur.close()
conn.close()
return results
Mistral API 키 입력
mistral-large-2407에 접근하려면 Mistral API 키를 입력합니다.
mistral_key = input("Enter your Mistral AI key: ")
모든 테이블의 문맥 제공
LLM에 모든 테이블의 문맥을 프롬프트하고 싶습니다. 이는 최신 상태여야 하므로, 모든 테이블을 질의하는 사전 정의된 입력으로 postgres 도구를 호출하는 별도 함수를 만듭니다.
def get_all_tables():
# Exclude default PostgreSQL schemas
excluded_schemas = ['information_schema', 'pg_catalog']
# Query to get all tables excluding the default schemas
table_columns = ['table_schema', 'table_name']
table_name = 'information_schema.tables'
filters = {'table_schema': excluded_schemas}
sort_column = 'table_schema'
sort_order = 'ASC'
# Execute the query to get all tables
tables_query_result = execute_postgres_query(
table_name,
table_columns,
filters,
sort_column,
sort_order
)
# Parse the results of the tables query
tables = [{'table_schema': row[0], 'table_name': row[1]} for row in tables_query_result]
# Prepare a list to store table information with columns
table_info = []
# Iterate over each table to get its columns
for table in tables:
schema_name = table['table_schema']
table_name = table['table_name']
# Query to get columns for the current table
columns_columns = ['column_name']
columns_table_name = 'information_schema.columns'
columns_filters = {'table_schema': schema_name, 'table_name': table_name}
columns_sort_column = 'ordinal_position'
columns_sort_order = 'ASC'
# Execute the query to get columns
columns_query_result = execute_postgres_query(
columns_table_name,
columns_columns,
columns_filters,
columns_sort_column,
columns_sort_order
)
# Parse the results of the columns query
columns = [row[0] for row in columns_query_result]
# Add table information with columns to the list
table_info.append({
'table_schema': schema_name,
'table_name': table_name,
'columns': columns
})
return table_info
챗봇 설정
이제 Mistral-Large-2407로 채팅하고 PostgreSQL 데이터베이스를 질의할 모든 준비가 됐습니다.
def chatbot(mistral_key):
config_list = [
{
'model': 'mistral-large-2407', # If the responses are very slow, change this model to open-mixtral-8x22b
'base_url': 'https://api.mistral.ai/v1',
"api_key": mistral_key,
"tool_choice": "auto",
},
]
llm_config={
"config_list": config_list,
"temperature": 0.1
}
user = ConversableAgent(
"user",
llm_config=False,
is_termination_msg=lambda msg: "tool_calls" not in msg,
human_input_mode="NEVER",
max_consecutive_auto_reply=3,
)
assistant = ConversableAgent(
name="assistant",
llm_config=llm_config,
system_message="You are an helpful AI assistant, you use your Postgres tool to query the database. Keep in mind the possibility of to long contexts lengths when using limits wrong."
)
assistant.register_for_llm(name="postgres_query", description="Useful for when you need query the postgres db")(execute_postgres_query)
user.register_for_execution(name="postgres_query")(execute_postgres_query)
LLM_CONTEXT = get_all_tables()
user.send(f"This are all the available tables; \n\n {LLM_CONTEXT} \n\n ", assistant, request_reply=False)
assistant.send("Thanks for the additonal context of all existing tables!", user, request_reply=False)
while True:
task = input("Enter the query for the LLM ('exit' to quit): ")
if task.lower() == 'exit':
break
context_task = f"{task}"
user.initiate_chat(assistant, message=context_task, clear_history=False)
chatbot(mistral_key)
이 모델에 물어볼 좋은 예시 질문은 다음과 같아요.
- Get me all people named Mistral
- Get me all actors born in 2000, limit them to 10
- Get me all actors from before 1950, limit them to 13
보시다시피 Mistral Large V2(또는 동급 모델)로 레코드 삭제 권한을 부여하지 않고도 Postgres DB를 질의하는 함수를 비교적 쉽게 만들 수 있습니다. 이는 사람들이 실수를 걱정하지 않고 서로 다른 데이터베이스에 더 안전하게 접근할 수 있는 방법을 제공합니다.
더 알아보기 (Learn more)
- Microsoft AutoGen 문서 — 멀티에이전트 프레임워크
ConversableAgent/register_for_llm— AutoGen 에이전트·도구 등록mistral-large-2407— 사용한 Mistral 모델 (느리면open-mixtral-8x22b로 교체 가능)- IMDB 데이터셋
name.basics.tsv— 예시 데이터베이스