Slik fungerer NL-to-SQL-agenter

# High-level pipeline
def nl_to_sql_agent(user_question, db_connection):
    schema = get_schema(db_connection)
    sql = llm_generate_sql(user_question, schema)
    rows = execute_query(db_connection, sql)
    answer = format_results(rows, user_question)
    return answer

def build_schema_prompt(schema_info):
    lines = []
    for table in schema_info:
        cols = ', '.join(
            f"{c['name']} ({c['type']})"
            for c in table['columns']
        )
        lines.append(f"Table {table['name']}: {cols}")
    return '\n'.join(lines)

# Output:
# Table users: id (INT), email (VARCHAR), created_at (TIMESTAMP)
# Table orders: id (INT), user_id (INT), status (VARCHAR), total (FLOAT)

if __name__ == '__main__':
    demo_schema = [
        {'name': 'users', 'columns': [{'name': 'id', 'type': 'INT'}, {'name': 'email', 'type': 'VARCHAR'}]},
        {'name': 'orders', 'columns': [{'name': 'id', 'type': 'INT'}, {'name': 'user_id', 'type': 'INT'}]},
    ]
    print(build_schema_prompt(demo_schema))

SYSTEM_PROMPT = '''You are a SQL expert. Given a database schema and a question,
generate a valid {dialect} SELECT query. Return ONLY the SQL query, no explanation.
Do not use INSERT, UPDATE, DELETE, or DROP.

Schema:
{schema}
'''

def llm_generate_sql(question, schema, dialect='PostgreSQL'):
    prompt = SYSTEM_PROMPT.format(schema=schema, dialect=dialect)
    response = client.chat.completions.create(
        model='gpt-4o',
        messages=[
            {'role': 'system', 'content': prompt},
            {'role': 'user', 'content': question}
        ]
    )
    return response.choices[0].message.content.strip()

import psycopg2

def execute_query(conn, sql):
    try:
        with conn.cursor() as cur:
            cur.execute(sql)
            columns = [desc[0] for desc in cur.description]
            rows = cur.fetchmany(100)  # limit rows
            return {'columns': columns, 'rows': rows}
    except psycopg2.Error as e:
        return {'error': str(e), 'sql': sql}

def format_results(result, original_question):
    if 'error' in result:
        return f'Query failed: {result["error"]}'

    rows = result['rows']
    columns = result['columns']

    if not rows:
        return 'No results found.'

    # For simple counts/aggregates — just return the value
    if len(columns) == 1 and len(rows) == 1:
        return f'Result: {rows[0][0]}'

    # For multi-row results — summarize
    summary = f'Found {len(rows)} rows.\n'
    for row in rows[:5]:  # show first 5
        summary += ', '.join(f'{columns[i]}: {row[i]}' for i in range(len(columns))) + '\n'
    return summary

if __name__ == '__main__':
    demo_result = {'rows': [[42]], 'columns': ['count']}
    print(format_results(demo_result, 'How many users signed up?'))
    demo_result2 = {'rows': [], 'columns': ['id']}
    print(format_results(demo_result2, 'Any orders today?'))

AMBIGUITY_PROMPT = '''If the question is ambiguous, respond with JSON:
{"needs_clarification": true, "question": "your clarifying question"}

If clear, respond with the SQL query directly.

User question: {question}
'''

def generate_or_clarify(question, schema):
    response = llm_call(AMBIGUITY_PROMPT.format(
        question=question, schema=schema
    ))
    if '"needs_clarification"' in response:
        import json
        return json.loads(response)
    return {'sql': response}

# Two-phase approach for large schemas
def get_relevant_tables(question, all_tables):
    prompt = f'''Given these tables: {all_tables}
Which 3-5 tables are most relevant to answer: "{question}"?
Return a JSON list of table names only.'''
    response = llm_call(prompt)
    import json
    return json.loads(response)

def nl_to_sql_large_db(question, conn):
    all_tables = list_all_tables(conn)  # just names
    relevant = get_relevant_tables(question, all_tables)
    schema = get_schema_for_tables(conn, relevant)
    return llm_generate_sql(question, schema)

DIALECT_EXAMPLES = {
    'postgresql': 'Use LIMIT for row limits. Use NOW() for current time.',
    'mysql': 'Use LIMIT for row limits. Use NOW() for current time.',
    'sqlite': 'Use LIMIT. Use datetime("now") for current time.',
    'mssql': 'Use TOP N for row limits. Use GETDATE() for current time.',
    'bigquery': 'Use LIMIT. Use CURRENT_TIMESTAMP() for current time. Use backtick for table names.'
}

def get_dialect_hint(dialect):
    return DIALECT_EXAMPLES.get(dialect.lower(), '')

if __name__ == '__main__':
    for dialect in ['postgresql', 'sqlite', 'mssql']:
        print(f'{dialect}: {get_dialect_hint(dialect)}')

def nl_to_sql_with_retry(question, schema, conn, max_retries=3):
    sql = llm_generate_sql(question, schema)
    for attempt in range(max_retries):
        result = execute_query(conn, sql)
        if 'error' not in result:
            return format_results(result, question)
        # Ask LLM to fix the error
        fix_prompt = f'The SQL query failed with error: {result["error"]}\n'\
                     f'Original SQL: {sql}\n'\
                     f'Please fix the SQL query.'
        sql = llm_call(fix_prompt)
        print(f'Retry {attempt + 1} with fixed SQL')
    return 'Could not generate a valid query after retries.'

import hashlib

query_cache = {}

def cached_nl_to_sql(question, schema_hash, conn):
    cache_key = hashlib.md5((question + schema_hash).encode()).hexdigest()
    if cache_key in query_cache:
        print('Cache hit!')
        sql = query_cache[cache_key]
    else:
        schema = get_schema(conn)
        sql = llm_generate_sql(question, schema)
        query_cache[cache_key] = sql

    result = execute_query(conn, sql)
    return format_results(result, question)