Bygg ett databasingränssnitt på naturligt språk

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']
        }
    }
}

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]

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.
'''

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

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

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.'

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

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