update OpenAI usage from Vanna

app.py

@@ -13,6 +13,7 @@ from climateqa.engine.graph import make_graph_agent,make_graph_agent_poc
 from climateqa.engine.chains.retrieve_papers import find_papers
 from climateqa.chat import start_chat, chat_stream, finish_chat
 from climateqa.engine.talk_to_data.main import ask_vanna
+from climateqa.engine.talk_to_data.myVanna import MyVanna
 
 from front.tabs import (create_config_modal, create_examples_tab, create_papers_tab, create_figures_tab, create_chat_interface, create_about_tab)
 from front.utils import process_figures
@@ -77,6 +78,14 @@ else :
 agent = make_graph_agent(llm=llm, vectorstore_ipcc=vectorstore, vectorstore_graphs=vectorstore_graphs, vectorstore_region = vectorstore_region, reranker=reranker, threshold_docs=0.2)
 agent_poc = make_graph_agent_poc(llm=llm, vectorstore_ipcc=vectorstore, vectorstore_graphs=vectorstore_graphs, vectorstore_region = vectorstore_region, reranker=reranker, threshold_docs=0, version="v4")#TODO put back default 0.2
 
+#Vanna object
+
+vn = MyVanna(config = {"temperature": 0, "api_key": os.getenv('THEO_API_KEY'), 'model': os.getenv('VANNA_MODEL'), 'pc_api_key': os.getenv('VANNA_PINECONE_API_KEY'), 'index_name': os.getenv('VANNA_INDEX_NAME'), "top_k" : 4})
+db_vanna_path = os.path.join(os.getcwd(), "data/drias/drias.db")
+vn.connect_to_sqlite(db_vanna_path)
+
+def ask_vanna_query(query):
+    return ask_vanna(vn, db_vanna_path, query)
 
 async def chat(query, history, audience, sources, reports, relevant_content_sources_selection, search_only):
     print("chat cqa - message received")
@@ -126,7 +135,7 @@ def create_drias_tab():
             show_vanna_table.click(lambda: Modal(visible=True),None ,[vanna_table_modal])
 
         vanna_display = gr.Plot()
-        vanna_direct_question.submit(ask_vanna, [vanna_direct_question], [vanna_sql_query ,vanna_table, vanna_display])
+        vanna_direct_question.submit(ask_vanna_query, [vanna_direct_question], [vanna_sql_query ,vanna_table, vanna_display])
     
 # # UI Layout Components
 def cqa_tab(tab_name):

climateqa/engine/talk_to_data/main.py

@@ -4,24 +4,10 @@ import sqlite3
 import os
 import pandas as pd
 from climateqa.engine.llm import get_llm
-
-from dotenv import load_dotenv
 import ast
 
-load_dotenv()
-
-
-OPENAI_API_KEY = os.getenv('THEO_API_KEY')
-PC_API_KEY = os.getenv('VANNA_PINECONE_API_KEY')
-INDEX_NAME = os.getenv('VANNA_INDEX_NAME')
-VANNA_MODEL = os.getenv('VANNA_MODEL')
 
 
-#Vanna object
-vn = MyVanna(config = {"temperature": 0, "api_key": OPENAI_API_KEY, 'model': VANNA_MODEL, 'pc_api_key': PC_API_KEY, 'index_name': INDEX_NAME, "top_k" : 4})
-db_vanna_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))), "data/drias/drias.db")
-vn.connect_to_sqlite(db_vanna_path)
-
 llm = get_llm(provider="openai")
 
 def ask_llm_to_add_table_names(sql_query, llm):
@@ -33,9 +19,10 @@ def ask_llm_column_names(sql_query, llm):
     columns_list = ast.literal_eval(columns.strip("```python\n").strip())
     return columns_list
 
-def ask_vanna(query):
+def ask_vanna(vn,db_vanna_path, query):
+
     try :
-        location = detect_location_with_openai(OPENAI_API_KEY, query)
+        location = detect_location_with_openai(query)
         if location:
 
             coords = loc2coords(location)
@@ -51,10 +38,10 @@ def ask_vanna(query):
 
         else : 
             empty_df = pd.DataFrame()
-            empty_fig = {}
+            empty_fig = None
             return "", empty_df, empty_fig
     except Exception as e:
         print(f"Error: {e}")
         empty_df = pd.DataFrame()
-        empty_fig = {}
+        empty_fig = None
         return "", empty_df, empty_fig

climateqa/engine/talk_to_data/utils.py

@@ -4,13 +4,13 @@ import pandas as pd
 from geopy.geocoders import Nominatim
 import sqlite3
 import ast
+from climateqa.engine.llm import get_llm
 
-
-def detect_location_with_openai(api_key, sentence):
+def detect_location_with_openai(sentence):
     """
-    Detects locations in a sentence using OpenAI's API.
+    Detects locations in a sentence using OpenAI's API via LangChain.
     """
-    openai.api_key = api_key
+    llm = get_llm()
 
     prompt = f"""
     Extract all locations (cities, countries, states, or geographical areas) mentioned in the following sentence.
@@ -19,18 +19,12 @@ def detect_location_with_openai(api_key, sentence):
     Sentence: "{sentence}"
     """
 
-    response = openai.chat.completions.create(
-        model="gpt-4o-mini",  
-        messages=[
-            {"role": "system", "content": "You are a helpful assistant skilled in identifying locations in text."},
-            {"role": "user", "content": prompt}
-        ],
-        max_tokens=100,
-        temperature=0
-    )
-
-    return response.choices[0].message.content.split("\n")[1][2:-2]
-
+    response = llm.invoke(prompt)
+    location_list = ast.literal_eval(response.content.strip("```python\n").strip())
+    if location_list:
+        return location_list[0]
+    else:
+        return ""
 
 def detectTable(sql_query):
     pattern = r'(?i)\bFROM\s+((?:`[^`]+`|"[^"]+"|\'[^\']+\'|\w+)(?:\.(?:`[^`]+`|"[^"]+"|\'[^\']+\'|\w+))*)'