Budowanie interfejsu bazy danych w języku naturalnym

Język naturalny na SQL: wizja

Przegląd architektury systemu

Definiowanie narzędzia Query Database

query_db_tool = {
    'type': 'function',
    'function': {
        'name': 'query_database',
        'description': '''Execute a read-only SQL query on the company database.
Use this to answer questions about customers, orders, and products.
Only SELECT statements are allowed. Never use DROP, DELETE, UPDATE, or INSERT.

Available tables:
- customers (id, name, email, created_at, country)
- orders (id, customer_id, total_amount, status, created_at)
- order_items (id, order_id, product_id, quantity, unit_price)
- products (id, name, category, price, stock_quantity)
''',
        'parameters': {
            'type': 'object',
            'properties': {
                'sql': {
                    'type': 'string',
                    'description': 'A valid PostgreSQL SELECT statement.'
                },
                'explanation': {
                    'type': 'string',
                    'description': 'One-sentence explanation of what this query does.'
                }
            },
            'required': ['sql', 'explanation']
        }
    }
}

Bezpieczne wykonywanie SQL

import re
import psycopg2

DANGEROUS_KEYWORDS = ['DROP', 'DELETE', 'UPDATE', 'INSERT', 'TRUNCATE', 'ALTER', 'CREATE', 'EXEC']

def execute_safe_query(sql: str, max_rows: int = 100) -> list:
    '''Execute a read-only SQL query with safety guards.'''
    sql_upper = sql.upper().strip()

    # Only allow SELECT
    if not sql_upper.startswith('SELECT'):
        raise ValueError('Only SELECT statements are allowed.')

    # Block dangerous keywords
    for keyword in DANGEROUS_KEYWORDS:
        if re.search(r'\b' + keyword + r'\b', sql_upper):
            raise ValueError(f'Forbidden keyword: {keyword}')

    conn = psycopg2.connect('postgresql://readonly_user:pass@localhost/appdb')
    with conn:
        with conn.cursor() as cur:
            # Enforce read-only transaction
            cur.execute('SET TRANSACTION READ ONLY')
            cur.execute(sql)
            columns = [desc[0] for desc in cur.description]
            rows = cur.fetchmany(max_rows)
    return [dict(zip(columns, row)) for row in rows]

Wstrzykiwanie kontekstu schematu do komunikatu systemowego

SYSTEM_PROMPT = '''You are a data analyst assistant with access to the company database.
When users ask data questions, use the query_database tool to look up the answer.
Always explain your query in plain English before executing it.

Database schema:

CREATE TABLE customers (
    id SERIAL PRIMARY KEY,
    name VARCHAR NOT NULL,
    email VARCHAR UNIQUE,
    created_at TIMESTAMPTZ DEFAULT NOW(),
    country VARCHAR(2)  -- ISO 2-letter code: 'US', 'UK', 'DE', etc.
);

CREATE TABLE orders (
    id SERIAL PRIMARY KEY,
    customer_id INTEGER REFERENCES customers(id),
    total_amount NUMERIC(10,2),
    status VARCHAR  -- 'pending', 'shipped', 'delivered', 'cancelled'
    created_at TIMESTAMPTZ DEFAULT NOW()
);

Only use columns that exist in the schema above.
'''

Formatowanie wyników zapytania dla modelu

import json

def format_results(rows: list, max_display: int = 20) -> str:
    if not rows:
        return 'The query returned no results.'

    total = len(rows)
    display = rows[:max_display]

    # Format as a simple table
    if display:
        columns = list(display[0].keys())
        lines = [' | '.join(columns)]
        lines.append('-' * len(lines[0]))
        for row in display:
            lines.append(' | '.join(str(row[col]) for col in columns))

    result = '\n'.join(lines)
    if total > max_display:
        result += f'\n... ({total - max_display} more rows not shown)'
    return result

Implementacja pełnego potoku

from openai import OpenAI
import json

client = OpenAI()

def answer_data_question(user_question: str) -> str:
    messages = [
        {'role': 'system', 'content': SYSTEM_PROMPT},
        {'role': 'user', 'content': user_question}
    ]

    # First call: get SQL from model
    resp = client.chat.completions.create(
        model='gpt-4o', messages=messages, tools=[query_db_tool]
    )
    assistant_msg = resp.choices[0].message
    messages.append(assistant_msg)

    if resp.choices[0].finish_reason == 'tool_calls':
        tc = assistant_msg.tool_calls[0]
        args = json.loads(tc.function.arguments)
        print(f'Executing: {args["explanation"]}')
        print(f'SQL: {args["sql"]}')

        try:
            rows = execute_safe_query(args['sql'])
            result_text = format_results(rows)
        except ValueError as e:
            result_text = f'Query blocked: {str(e)}'

        messages.append({'role': 'tool', 'tool_call_id': tc.id, 'content': result_text})

        # Second call: narrate results
        final = client.chat.completions.create(model='gpt-4o', messages=messages)
        return final.choices[0].message.content

    return assistant_msg.content

Obsługa wieloetapowych pytań dotyczących danych

def answer_complex_question(user_question: str) -> str:
    messages = [
        {'role': 'system', 'content': SYSTEM_PROMPT},
        {'role': 'user', 'content': user_question}
    ]

    for _ in range(5):  # Max 5 query rounds
        resp = client.chat.completions.create(
            model='gpt-4o', messages=messages, tools=[query_db_tool]
        )
        msg = resp.choices[0].message
        messages.append(msg)

        if resp.choices[0].finish_reason == 'stop':
            return msg.content  # Model is done

        # Process tool call and loop
        tc = msg.tool_calls[0]
        args = json.loads(tc.function.arguments)
        try:
            rows = execute_safe_query(args['sql'])
            result = format_results(rows)
        except Exception as e:
            result = f'Error: {str(e)}'

        messages.append({'role': 'tool', 'tool_call_id': tc.id, 'content': result})

    return 'Could not complete the analysis within the step limit.'

Zapobieganie ryzyku wstrzyknięcia SQL

ALLOWED_TABLES = {'customers', 'orders', 'order_items', 'products'}

def validate_tables_in_sql(sql: str) -> bool:
    '''Check that only whitelisted tables are referenced in the query.'''
    import sqlparse
    parsed = sqlparse.parse(sql)[0]
    table_names = set()
    from_seen = False
    for token in parsed.flatten():
        if token.ttype is sqlparse.tokens.Keyword and token.value.upper() in ('FROM', 'JOIN'):
            from_seen = True
        elif from_seen and token.ttype is sqlparse.tokens.Name:
            table_names.add(token.value.lower())
            from_seen = False
    unknown = table_names - ALLOWED_TABLES
    if unknown:
        raise ValueError(f'References unknown tables: {unknown}')
    return True

Buforowanie często używanych zapytań

import redis
import hashlib
import json

r = redis.Redis.from_url('redis://localhost:6379')

def cached_query(sql: str, ttl_seconds: int = 300) -> list:
    cache_key = 'nl_query:' + hashlib.sha256(sql.encode()).hexdigest()
    cached = r.get(cache_key)
    if cached:
        return json.loads(cached)
    rows = execute_safe_query(sql)
    r.setex(cache_key, ttl_seconds, json.dumps(rows, default=str))
    return rows

Wyjaśnianie zapytań użytkownikom

Szybkie sprawdzenie

Podsumowanie lekcji