Package: llama-index-tools-waii
Waii ToolThis tool connects to database connections managed by Waii, which allows generic SQL queries, do performance analyze, describe a SQL query, and more.UsageFirst you need to create a waii.ai account, you request an account fromhere.Initialize the tool with your account credentials:fromllama_index.tools.waiiimportWaiiToolSpecwaii_tool=WaiiToolSpec(url="https://tweakit.waii.ai/api/",# API Key of Waii (not OpenAI API key)api_key="...",# Connection key of WAII connected database, see https://github.com/waii-ai/waii-sdk-py#get-connectionsdatabase_key="...",)ToolsThe tools available are:get_answer: Get answer to natural language question (which generate a SQL query, run it, explain the result)describe_query: Describe a SQL queryperformance_analyze: Analyze performance of a SQL query (by query_id)diff_query: Compare two SQL queriesdescribe_dataset: Describe dataset, such as table, schema, etc.transcode: Transcode SQL query to another SQL dialectget_semantic_contexts: Get semantic contexts of a SQL querygenerate_query_only: Generate SQL query only (not run it)run_query: Run a SQL queryYou can also load the data directly callload_dataExamplesLoad datadocuments=waii_tool.load_data("Get all tables with their number of columns")index=VectorStoreIndex.from_documents(documents).as_query_engine()print(index.query("Which table contains most columns?"))Use as a ToolInitialize the agent:fromllama_index.agentimportOpenAIAgentfromllama_index.llmsimportOpenAIagent=OpenAIAgent.from_tools(waii_tool.to_tool_list(),llm=OpenAI(model="gpt-4-1106-preview"))Ask simple questionagent.chat("Give me top 3 countries with the most number of car factory")agent.chat("What are the car factories of these countries")Do performance analyzeagent.chat("Give me top 3 longest running queries, and their duration.")agent.chat("analyze the 2nd-longest running query")Diff two queriesprevious_query="""SELECTemployee_id,department,salary,AVG(salary) OVER (PARTITION BY department) AS department_avg_salary,salary - AVG(salary) OVER (PARTITION BY department) AS diff_from_avgFROMemployees;"""current_query="""SELECTemployee_id,department,salary,MAX(salary) OVER (PARTITION BY department) AS department_max_salary,salary - AVG(salary) OVER (PARTITION BY department) AS diff_from_avgFROMemployees;LIMIT 100;"""agent.chat(f"tell me difference between{previous_query}and{current_query}")Describe datasetagent.chat("Summarize the dataset")agent.chat("Give me questions which I can ask about this dataset")Describe a queryq="""SELECTemployee_id,department,salary,AVG(salary) OVER (PARTITION BY department) AS department_avg_salary,salary - AVG(salary) OVER (PARTITION BY department) AS diff_from_avgFROMemployees;"""agent.chat(f"what this query can do?{q}")Migrate query to another dialectq="""from pyspark.sql import SparkSessionfrom pyspark.sql.functions import avg, colfrom pyspark.sql.window import Window# Initialize Spark sessionspark = SparkSession.builder.appName("example").getOrCreate()# Assuming you have a DataFrame called 'employees'# If not, you need to read your data into a DataFrame first# Define window specificationwindowSpec = Window.partitionBy("department")# Perform the queryresult = (employees.select(col("employee_id"),col("department"),col("salary"),avg("salary").over(windowSpec).alias("department_avg_salary"),(col("salary") - avg("salary").over(windowSpec)).alias("diff_from_avg")))# Show the resultresult.show()"""agent.chat(f"translate this pyspark query{q}, to Snowflake")Use Waii API directlyYou can also use Waii API directly, seehere