import os import dotenv import gradio as gr # LLMs import openai from langchain.chat_models import AzureChatOpenAI from langchain.llms import AzureOpenAI from langchain.schema import AIMessage, HumanMessage from langchain_experimental.sql import SQLDatabaseChain from langchain.agents import create_sql_agent from langchain.agents.agent_toolkits import SQLDatabaseToolkit from langchain.agents.agent_types import AgentType from langchain.agents.mrkl.prompt import FORMAT_INSTRUCTIONS # Databases import adal import struct from sqlalchemy.engine import URL from sqlalchemy import create_engine, event from langchain import SQLDatabase # Load environment variables dotenv.load_dotenv() # Set OpenAI API settings from the .env file openai.api_type = "azure" openai.api_base = os.getenv("OPENAI_API_BASE") openai.api_key = os.getenv("OPENAI_API_KEY") openai.api_version = os.getenv('OPENAI_API_VERSION') def get_token(): tenantId = os.getenv("DB_TENANT_ID") clientId = os.getenv("DB_CLIENT_ID") clientSecret = os.getenv("DB_CLIENT_SECRET") authorityHostUrl = "https://login.microsoftonline.com" authorityUrl = authorityHostUrl + "/" + tenantId context = adal.AuthenticationContext(authorityUrl, api_version=None) token = context.acquire_token_with_client_credentials("https://database.windows.net/", clientId, clientSecret) tokenb = bytes(token["accessToken"], "UTF-8") exptoken = b'' for i in tokenb: exptoken += bytes({i}) exptoken += bytes(1) return struct.pack("=i", len(exptoken)) + exptoken def get_conn_url(): server = "sql-ae-dsdev-dna-hack01-pnpdebzuohv7y.database.windows.net" database = 'dna-hack-db01' return f"mssql+pyodbc://@{server}/{database}?driver=ODBC+Driver+17+for+SQL+Server" def get_database(): # connection_string = "mssql+pyodbc://@my-server.database.windows.net/myDb?driver=ODBC+Driver+17+for+SQL+Server" conn_url = get_conn_url() print(conn_url) engine = create_engine(conn_url) @event.listens_for(engine, "do_connect") def provide_token(dialect, conn_rec, cargs, cparams): # remove the "Trusted_Connection" parameter that SQLAlchemy adds cargs[0] = cargs[0].replace(";Trusted_Connection=Yes", "") # create token credential token_struct = get_token() # apply it to keyword arguments SQL_COPT_SS_ACCESS_TOKEN = 1256 cparams["attrs_before"] = {SQL_COPT_SS_ACCESS_TOKEN: token_struct} return SQLDatabase(engine = engine, schema="[dntmatrix]") def get_llm(): return AzureChatOpenAI(temperature=1.0, model_name='gpt-4', deployment_name='gpt-4', model_kwargs={ "engine": "gpt-4", "api_key": openai.api_key, "api_base": openai.api_base, "api_type": openai.api_type, "api_version": openai.api_version } ) def get_sql_chain(): db = get_database() llm = get_llm() return SQLDatabaseChain.from_llm(llm, db, verbose=True) def get_sql_agent(): db = get_database() llm = get_llm() format_instruction = """ """ + FORMAT_INSTRUCTIONS toolkit = SQLDatabaseToolkit(db=db, llm=llm) return create_sql_agent( llm=llm, toolkit=toolkit, verbose=True, agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION, format_instructions=format_instruction ) sql_agent = get_sql_agent() def predict(message, history): gpt_response = sql_agent.run(message) return gpt_response with gr.Blocks() as demo: gr.ChatInterface(predict) if __name__ == "__main__": demo.launch(server_name="0.0.0.0")