Понимание и внедрение схемы

import psycopg2

def get_schema(conn):
    query = '''
        SELECT table_name, column_name, data_type
        FROM information_schema.columns
        WHERE table_schema = 'public'
        ORDER BY table_name, ordinal_position
    '''
    with conn.cursor() as cur:
        cur.execute(query)
        return cur.fetchall()

from collections import defaultdict

def build_schema_dict(conn):
    rows = get_schema(conn)
    schema = defaultdict(list)
    for table_name, column_name, data_type in rows:
        schema[table_name].append({
            'name': column_name,
            'type': data_type
        })
    return dict(schema)

# Result:
# {
#   'users': [{'name': 'id', 'type': 'integer'}, {'name': 'email', 'type': 'character varying'}],
#   'orders': [{'name': 'id', 'type': 'integer'}, {'name': 'user_id', 'type': 'integer'}]
# }

def format_schema_for_prompt(schema_dict, pk_info=None, fk_info=None):
    lines = []
    for table, columns in schema_dict.items():
        col_parts = []
        for col in columns:
            label = col['name']
            if pk_info and (table, col['name']) in pk_info:
                label += ' PK'
            if fk_info and (table, col['name']) in fk_info:
                label += f' FK->{fk_info[(table, col["name"])]}'
            col_parts.append(f"{label} ({col['type']})")
        lines.append(f"Table {table}: {', '.join(col_parts)}")
    return '\n'.join(lines)

# Output:
# Table users: id PK (integer), email (varchar), created_at (timestamp)
# Table orders: id PK (integer), user_id FK->users.id (integer), total (float)

if __name__ == '__main__':
    demo_schema = {'users': [{'name': 'id', 'type': 'integer'}, {'name': 'email', 'type': 'varchar'}]}
    demo_pk = {('users', 'id')}
    print(format_schema_for_prompt(demo_schema, pk_info=demo_pk))

def get_foreign_keys(conn):
    query = '''
        SELECT
            kcu.table_name,
            kcu.column_name,
            ccu.table_name AS foreign_table,
            ccu.column_name AS foreign_column
        FROM information_schema.table_constraints AS tc
        JOIN information_schema.key_column_usage AS kcu
            ON tc.constraint_name = kcu.constraint_name
        JOIN information_schema.constraint_column_usage AS ccu
            ON ccu.constraint_name = tc.constraint_name
        WHERE tc.constraint_type = 'FOREIGN KEY'
    '''
    with conn.cursor() as cur:
        cur.execute(query)
        return {
            (row[0], row[1]): f'{row[2]}.{row[3]}'
            for row in cur.fetchall()
        }

if __name__ == '__main__':
    class FakeCursor:
        def __enter__(self): return self
        def __exit__(self, *a): return False
        def execute(self, query): pass
        def fetchall(self):
            return [('orders', 'user_id', 'users', 'id')]
    class FakeConn:
        def cursor(self): return FakeCursor()

    fks = get_foreign_keys(FakeConn())
    print('Foreign keys found:')
    for (table, col), ref in fks.items():
        print(f'  {table}.{col} -> {ref}')

def estimate_schema_tokens(schema_dict):
    text = format_schema_for_prompt(schema_dict)
    # Rough estimate: 1 token per 4 characters
    estimated_tokens = len(text) // 4
    print(f'Tables: {len(schema_dict)}')
    print(f'Estimated schema tokens: {estimated_tokens}')
    return estimated_tokens

# 200 tables * 15 columns * 25 chars/col = 75,000 chars = ~18,750 tokens
# Plus user question + system prompt = easily over context limit

def select_relevant_tables(question, all_table_names, n=5):
    table_list = ', '.join(all_table_names)
    prompt = f'''Database tables: {table_list}

Question: {question}

List the {n} most relevant table names as a JSON array.
Example: ["users", "orders", "products"]'''

    response = llm_call(prompt)
    import json
    return json.loads(response)

def compressed_schema(question, conn):
    all_tables = list(build_schema_dict(conn).keys())
    relevant = select_relevant_tables(question, all_tables)
    full_schema = build_schema_dict(conn)
    return {t: full_schema[t] for t in relevant if t in full_schema}

AUDIT_COLUMNS = {
    'created_at', 'updated_at', 'deleted_at', 'created_by',
    'updated_by', 'version', 'is_deleted', 'modified_at'
}

def compress_schema(schema_dict, exclude_audit=True):
    compressed = {}
    for table, columns in schema_dict.items():
        # Skip internal/system tables
        if table.startswith('_') or table.startswith('pg_'):
            continue
        if exclude_audit:
            columns = [c for c in columns if c['name'] not in AUDIT_COLUMNS]
        if columns:  # only include if columns remain
            compressed[table] = columns
    return compressed

if __name__ == '__main__':
    demo_schema = {
        'users': [{'name': 'id', 'type': 'INT'}, {'name': 'email', 'type': 'VARCHAR'}, {'name': 'created_at', 'type': 'TIMESTAMP'}],
        'pg_stat': [{'name': 'x', 'type': 'INT'}],
    }
    compressed = compress_schema(demo_schema)
    print('Tables kept:', list(compressed.keys()))
    print('users columns after compression:', [c['name'] for c in compressed['users']])

TABLE_DESCRIPTIONS = {
    'users': 'Registered app users with authentication info',
    'orders': 'Customer purchase orders',
    'order_items': 'Individual line items within an order',
    'products': 'Product catalog with pricing',
    'payments': 'Payment transactions linked to orders'
}

def format_schema_with_descriptions(schema_dict):
    lines = []
    for table, columns in schema_dict.items():
        desc = TABLE_DESCRIPTIONS.get(table, '')
        col_str = ', '.join(f"{c['name']} ({c['type']})" for c in columns)
        if desc:
            lines.append(f"Table {table} ({desc}): {col_str}")
        else:
            lines.append(f"Table {table}: {col_str}")
    return '\n'.join(lines)

if __name__ == '__main__':
    demo_schema = {'users': [{'name': 'id', 'type': 'INT'}], 'orders': [{'name': 'id', 'type': 'INT'}]}
    print(format_schema_with_descriptions(demo_schema))

import time

class SchemaCache:
    def __init__(self, ttl_seconds=300):
        self._cache = None
        self._timestamp = 0
        self.ttl = ttl_seconds

    def get(self, conn):
        now = time.time()
        if self._cache is None or (now - self._timestamp) > self.ttl:
            print('Refreshing schema cache...')
            schema_dict = build_schema_dict(conn)
            fk_info = get_foreign_keys(conn)
            self._cache = format_schema_for_prompt(schema_dict, fk_info=fk_info)
            self._timestamp = now
        return self._cache

schema_cache = SchemaCache(ttl_seconds=300)

def build_sql_agent_prompt(question, conn, large_db=False):
    if large_db:
        schema = compressed_schema(question, conn)
        schema_text = format_schema_with_descriptions(schema)
    else:
        schema_text = schema_cache.get(conn)

    system = f'''You are a PostgreSQL expert.
Return ONLY a valid SELECT query based on this schema:

{schema_text}

Rules:
- Use only SELECT statements
- Use table aliases for clarity
- Limit results to 100 rows unless asked for all
'''
    return system