import os
from langchain_community.utilities import SQLDatabase
from langchain_openai import ChatOpenAI # Replace with your preferred provider
from langchain_community.agent_toolkits import create_sql_agent
# 🔐 Credentials - Make sure these are securely stored
FIREBOLT_CLIENT_ID = os.getenv("FIREBOLT_CLIENT_ID")
FIREBOLT_CLIENT_SECRET = os.getenv("FIREBOLT_CLIENT_SECRET")
FIREBOLT_ACCOUNT_NAME = os.getenv("FIREBOLT_ACCOUNT_NAME")
FIREBOLT_DATABASE = os.getenv("FIREBOLT_DATABASE")
FIREBOLT_ENGINE_NAME = os.getenv("FIREBOLT_ENGINE_NAME")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
# 🔗 SQLAlchemy connection string for Firebolt
connection_url = (f"firebolt://{FIREBOLT_CLIENT_ID}:{FIREBOLT_CLIENT_SECRET}@{FIREBOLT_DATABASE}/"
f"{FIREBOLT_ENGINE_NAME}?account_name={FIREBOLT_ACCOUNT_NAME}")
llm = ChatOpenAI(temperature=0, openai_api_key=OPENAI_API_KEY, model_name="gpt-4o")
# For other providers, use the appropriate class:
# from langchain_anthropic import ChatAnthropic
# llm = ChatAnthropic(model="claude-3-sonnet-20240229", anthropic_api_key=os.getenv("ANTHROPIC_API_KEY"))
#
# from langchain_google_genai import ChatGoogleGenerativeAI
# llm = ChatGoogleGenerativeAI(model="gemini-pro", google_api_key=os.getenv("GOOGLE_API_KEY"))
db = SQLDatabase.from_uri(connection_url, schema="public")
agent_executor = create_sql_agent(llm, db=db, agent_type="openai-tools", verbose=True)