Fixed Reason, Type extraction

app.py

@@ -139,7 +139,7 @@ Recommend a visualization:
                 ])
 
     final_prompt = prompt.format_prompt(question=text_query,
-                                        sql_query=sql_query, results=sql_result)
+                                        sql_query=sql_query, results=sql_result[:20])
     response = run_llm(final_prompt)
     response = response.replace('```', '')
     lines = response.strip().split('\n')
@@ -157,7 +157,7 @@ def format_data(text_query, sql_query, sql_result, visualization_type):
 
     template = ChatPromptTemplate.from_messages([
     ("system", "You are a Data expert who formats data according to the required needs. You are given the question asked by the user, it's sql query, the result of the query and the format you need to format it in."),
-    ("human", "For the given question: {question}\n\nSQL query: {sql_query}\n\Result: {results}\n\nUse the following example to structure the data: {instructions}. If there is None in Result please change it to '0'. Just give the json string. Do not format it. Do not give backticks."),
+    ("human", "For the given question: {question}\n\nSQL query: {sql_query}\n\Result: {results}\n\nUse the following example to structure the data: {instructions}. If there is None in Result please change it to '0'. Just give the json string. Do not format it."),
     ])
 
 
@@ -165,6 +165,8 @@ def format_data(text_query, sql_query, sql_result, visualization_type):
                                     results=sql_result, instructions=instruction)
     print(prompt)
     formatted_data = run_llm(prompt)
+    formatted_data = formatted_data.reaplce('```', '')
+    formatted_data = formatted_data.reaplce('json', '')
     print(f'Formatted Data {formatted_data}')
     return json.loads(formatted_data.replace('.', '').strip())
 
@@ -226,6 +228,11 @@ def main(table, text_query):
     except Exception as e:
         return generate_output(schema, prompt, generated_sql_query, fig, pd.DataFrame([{"error": f"❌ Unable to execute the SQL query. {e}"}]))
 
+    if len(sql_query_result) >= 100:
+        gr.Warning(f"⚠️ Data is too large for visualization. Please refine your query.")
+        return generate_output(schema, prompt, generated_sql_query,
+                               fig, sql_query_df)
+    
     try:
       visualization_type, reason = get_visualization_type(text_query=text_query,
                                                   sql_query=generated_sql_query, sql_result=sql_query_result)
@@ -237,9 +244,9 @@ def main(table, text_query):
         return generate_output(schema, prompt, generated_sql_query,
                                fig, sql_query_df)
     
+    
     if visualization_type != 'none':
         try:
-
             plot = visualize_result(text_query=text_query, sql_query=generated_sql_query,
                         sql_result=sql_query_result, visualization_type=visualization_type)