File size: 2,017 Bytes
28684d8 8bd064f 28684d8 8bd064f 28684d8 8bd064f 28684d8 8bd064f 28684d8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
from climateqa.engine.talk_to_data.myVanna import MyVanna
from climateqa.engine.talk_to_data.utils import loc2coords, detect_location_with_openai, detectTable, nearestNeighbourSQL, detect_relevant_tables, replace_coordonates
import sqlite3
import os
import pandas as pd
from climateqa.engine.llm import get_llm
import ast
llm = get_llm(provider="openai")
def ask_llm_to_add_table_names(sql_query, llm):
sql_with_table_names = llm.invoke(f"Make the following sql query display the source table in the rows {sql_query}. Just answer the query. The answer should not include ```sql\n").content
return sql_with_table_names
def ask_llm_column_names(sql_query, llm):
columns = llm.invoke(f"From the given sql query, list the columns that are being selected. The answer should only be a python list. Just answer the list. The SQL query : {sql_query}").content
columns_list = ast.literal_eval(columns.strip("```python\n").strip())
return columns_list
def ask_vanna(vn,db_vanna_path, query):
try :
location = detect_location_with_openai(query)
if location:
coords = loc2coords(location)
user_input = query.lower().replace(location.lower(), f"lat, long : {coords}")
relevant_tables = detect_relevant_tables(user_input, llm)
coords_tables = [nearestNeighbourSQL(db_vanna_path, coords, relevant_tables[i]) for i in range(len(relevant_tables))]
user_input_with_coords = replace_coordonates(coords, user_input, coords_tables)
sql_query, result_dataframe, figure = vn.ask(user_input_with_coords, print_results=False, allow_llm_to_see_data=True, auto_train=False)
return sql_query, result_dataframe, figure
else :
empty_df = pd.DataFrame()
empty_fig = None
return "", empty_df, empty_fig
except Exception as e:
print(f"Error: {e}")
empty_df = pd.DataFrame()
empty_fig = None
return "", empty_df, empty_fig |