Yatheshr/DAP_Query_Agent

1import os2 3# Load .env locally; skip if not available (e.g., in HF)4try:5    from dotenv import load_dotenv6    load_dotenv()7except ImportError:8    pass9 10from langchain.agents import create_sql_agent11from langchain.agents.agent_toolkits import SQLDatabaseToolkit12from langchain.chat_models import ChatOpenAI13from langchain.sql_database import SQLDatabase14from langchain.agents.agent_types import AgentType15 16def create_dap_sql_agent():17    db_conn_str = os.environ.get("DB_CONN_STR")18    openai_api_key = os.environ.get("OPENAI_API_KEY")19 20    if not db_conn_str or not openai_api_key:21        raise ValueError("OPENAI_API_KEY or DB_CONN_STR is not set.")22 23    uri = f"mssql+pyodbc:///?odbc_connect={db_conn_str}"24    db = SQLDatabase.from_uri(uri)25 26    llm = ChatOpenAI(27        temperature=0,28        model_name="gpt-4",  # or "gpt-4o"29        openai_api_key=openai_api_key,30    )31 32    toolkit = SQLDatabaseToolkit(db=db, llm=llm)33 34    agent = create_sql_agent(35        llm=llm,36        toolkit=toolkit,37        verbose=True,38        agent_type=AgentType.OPENAI_FUNCTIONS,39    )40 41    return agent42