SQL Generation
SQL Generation
Pydantic AI로 사용자 입력에 맞는 SQL 쿼리를 생성하는 예시예요. 여기서 재미있는 부분은, 생성된 SQL이 단순히 텍스트로 끝나지 않고 실제 PostgreSQL에서 EXPLAIN 쿼리로 실행되어 검증된다는 점이에요. LLM이 만들어낸 쿼리가 문법적으로 유효한지까지 에이전트가 직접 확인하는 셈이죠.
출처: 공식문서
예시가 보여주는 것
예시 실행하기
검증을 위해 PostgreSQL이 필요해요. 예를 들어 Docker로 실행할 수 있어요.
docker run --rm -e POSTGRES_PASSWORD=postgres -p 54320:5432 postgres
(다른 postgres 인스턴스와의 충돌을 피하려고 포트 54320을 사용해요)
의존성을 설치하고 환경 변수를 설정했다면, 실행해요.
python -m pydantic_ai_examples.sql_gen
uv run -m pydantic_ai_examples.sql_gen
혹은 커스텀 프롬프트로 실행하려면:
python -m pydantic_ai_examples.sql_gen "find me errors"
uv run -m pydantic_ai_examples.sql_gen "find me errors"
이 모델은 Gemini가 이런 종류의 단발성 쿼리에 강하기 때문에 기본값으로 gemini-3-flash-preview를 사용해요.
예시 코드
sql_gen.py
import asyncio
import sys
from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
from dataclasses import dataclass
from datetime import date
from typing import Annotated, Any, TypeAlias
import asyncpg
import logfire
from annotated_types import MinLen
from devtools import debug
from pydantic import BaseModel, Field
from pydantic_ai import Agent, ModelRetry, RunContext, format_as_xml
# 'if-token-present' means nothing will be sent (and the example will work) if you don't have logfire configured
logfire.configure(send_to_logfire='if-token-present')
logfire.instrument_asyncpg()
logfire.instrument_pydantic_ai()
DB_SCHEMA = """
CREATE TABLE records (
created_at timestamptz,
start_timestamp timestamptz,
end_timestamp timestamptz,
trace_id text,
span_id text,
parent_span_id text,
level log_level,
span_name text,
message text,
attributes_json_schema text,
attributes jsonb,
tags text[],
is_exception boolean,
otel_status_message text,
service_name text
);
"""
SQL_EXAMPLES = [
{
'request': 'show me records where foobar is false',
'response': "SELECT * FROM records WHERE attributes->>'foobar' = false",
},
{
'request': 'show me records where attributes include the key "foobar"',
'response': "SELECT * FROM records WHERE attributes ? 'foobar'",
},
{
'request': 'show me records from yesterday',
'response': "SELECT * FROM records WHERE start_timestamp::date > CURRENT_TIMESTAMP - INTERVAL '1 day'",
},
{
'request': 'show me error records with the tag "foobar"',
'response': "SELECT * FROM records WHERE level = 'error' and 'foobar' = ANY(tags)",
},
]
@dataclass
class Deps:
conn: asyncpg.Connection
class Success(BaseModel):
"""Response when SQL could be successfully generated."""
sql_query: Annotated[str, MinLen(1)]
explanation: str = Field(
'', description='Explanation of the SQL query, as markdown'
)
class InvalidRequest(BaseModel):
"""Response the user input didn't include enough information to generate SQL."""
error_message: str
Response: TypeAlias = Success | InvalidRequest
agent = Agent[Deps, Response](
'google:gemini-3-flash-preview',
# Pass the union members directly: a `Response` type alias isn't yet accepted as a `TypeForm` value (PEP-747)
output_type=Success | InvalidRequest,
deps_type=Deps,
)
@agent.system_prompt
async def system_prompt() -> str:
return f"""\
Given the following PostgreSQL table of records, your job is to
write a SQL query that suits the user's request.
Database schema:
{DB_SCHEMA}
today's date = {date.today()}
{format_as_xml(SQL_EXAMPLES)}
"""
@agent.output_validator
async def validate_output(ctx: RunContext[Deps], output: Response) -> Response:
if isinstance(output, InvalidRequest):
return output
# gemini often adds extraneous backslashes to SQL
output.sql_query = output.sql_query.replace('\\', '')
if not output.sql_query.upper().startswith('SELECT'):
raise ModelRetry('Please create a SELECT query')
try:
await ctx.deps.conn.execute(f'EXPLAIN {output.sql_query}')
except asyncpg.exceptions.PostgresError as e:
raise ModelRetry(f'Invalid query: {e}') from e
else:
return output
async def main():
if len(sys.argv) == 1:
prompt = 'show me logs from yesterday, with level "error"'
else:
prompt = sys.argv[1]
async with database_connect(
'postgresql://postgres:***@localhost:54320', 'pydantic_ai_sql_gen'
) as conn:
deps = Deps(conn)
result = await agent.run(prompt, deps=deps)
debug(result.output)
# pyright: reportUnknownMemberType=false
# pyright: reportUnknownVariableType=false
@asynccontextmanager
async def database_connect(server_dsn: str, database: str) -> AsyncGenerator[Any, None]:
with logfire.span('check and create DB'):
conn = await asyncpg.connect(server_dsn)
try:
db_exists = await conn.fetchval(
'SELECT 1 FROM pg_database WHERE datname = $1', database
)
if not db_exists:
await conn.execute(f'CREATE DATABASE {database}')
finally:
await conn.close()
conn = await asyncpg.connect(f'{server_dsn}/{database}')
try:
with logfire.span('create schema'):
async with conn.transaction():
if not db_exists:
await conn.execute(
"CREATE TYPE log_level AS ENUM ('debug', 'info', 'warning', 'error', 'critical')"
)
await conn.execute(DB_SCHEMA)
yield conn
finally:
await conn.close()
if __name__ == '__main__':
asyncio.run(main())
여기서 핵심 구조를 짚어볼게요.
다이나믹 시스템 프롬프트 — @agent.system_prompt 데코레이터가 붙은 함수가 DB 스키마, 오늘 날짜, 그리고 few-shot 예시(SQL_EXAMPLES)를 XML로 묶어 모델에 전달해요. format_as_xml(SQL_EXAMPLES)가 예시 목록을 구조화된 XML로 바꿔주는 거죠.
구조화된 출력 + 검증 — output_type=Success | InvalidRequest로 응답을 유니온으로 선언해요. 성공이면 sql_query와 explanation, 요청이 모호하면 InvalidRequest의 error_message를 반환하죠. @agent.output_validator는 두 갈래 중 Success인 경우에만 검증을 진행해요.
검증 로직 — gemini가 자주 추가하는 불필요한 백슬래시를 제거하고, SELECT로 시작하는지 확인한 뒤, 실제로 EXPLAIN <query>를 PostgreSQL에 실행해요. 쿼리가 유효하지 않으면(PostgresError) ModelRetry를 던져 모델이 다시 시도하게 해요. 이 패턴이 "생성된 SQL이 진짜 동작하는 SQL"임을 보장해요.