gpt4o_infer.py

import os
import openai

client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

def call_gpt4o(prompt):
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": "You are a helpful assistant that converts natural language into SQL queries using the provided schema."},
            {"role": "user", "content": prompt}
        ],
        temperature=0.2,
        max_tokens=512
    )
    return response.choices[0].message.content