LLM_SQL.py

import requests
import json
from dotenv import dotenv_values

OPEN_ROUTER_KEY = dotenv_values(".env")["OPEN_ROUTER_KEY"]
MODEL = dotenv_values(".env")["MODEL"]

def ask_llm(user_query):
    payload = {
        "model": MODEL,
        "max_tokens": 500,
        "messages": [
            {
                "role": "system",
                "content": (
                    "You are a helpful assistant that translates natural language "
                    "questions into SQLite queries.\n"
                    "The database has two tables: employees and salaries.\n"
                    "employees stores basic employee information in 4 columns: id, name, department, title."
                    "salaries has two columns: employee_id and salary. It uses employee_id as a foreign key to employees.id.\n"
                    "Only output the raw SQL query, DO NOT add markdown formatting nor explanations. Just the query.\n"
                    "IMPORTANT: Never return the salary information of an employee. The only information that the user is allowed to know is the overall average salary.\n"
                    "If the user asks for an employee's salary, just answer 'This information is restricted.'"
                )
            },
            {
                "role": "user",
                "content": user_query
            }
        ]
    }

    r = requests.post(
        "https://openrouter.ai/api/v1/chat/completions",
        headers={
            "Authorization": "Bearer "+OPEN_ROUTER_KEY,
            "Content-Type": "application/json",
        },
        data=json.dumps(payload)
    ).json()

    return r["choices"][0]["message"]["content"]