diff --git a/api/src/routes/ws/threads_and_messages/post_thread/post_thread.rs b/api/src/routes/ws/threads_and_messages/post_thread/post_thread.rs
index 5c52c2a77..e41c26d3a 100644
--- a/api/src/routes/ws/threads_and_messages/post_thread/post_thread.rs
+++ b/api/src/routes/ws/threads_and_messages/post_thread/post_thread.rs
@@ -1227,6 +1227,7 @@ async fn search_for_relevant_terms(
     terms_search
   where
     fts @@ websearch_to_tsquery('{prompt}')
+    and organization_id = '{organization_id}'
   order by rank_ix
   limit least(10, 30) * 2
 ),
@@ -1236,6 +1237,7 @@ semantic as (
     row_number() over (order by embedding <#> '{prompt_embedding}') as rank_ix
   from
     terms_search
+  where organization_id = '{organization_id}'
   order by rank_ix
   limit least(10, 30) * 2
 )
diff --git a/api/src/utils/agents/failed_to_fix_sql_agent.rs b/api/src/utils/agents/failed_to_fix_sql_agent.rs
index a75b3ec45..c512c8dcc 100644
--- a/api/src/utils/agents/failed_to_fix_sql_agent.rs
+++ b/api/src/utils/agents/failed_to_fix_sql_agent.rs
@@ -8,8 +8,8 @@ use crate::utils::{
         prompt_node::{prompt_node, PromptNodeMessage, PromptNodeSettings},
     },
     prompts::analyst_chat_prompts::failed_to_fix_sql_prompts::{
-            failed_to_fix_sql_system_prompt, failed_to_fix_sql_user_prompt,
-        },
+        failed_to_fix_sql_system_prompt, failed_to_fix_sql_user_prompt,
+    },
 };
 
 pub enum FailedToFixSqlAgent {
@@ -82,6 +82,7 @@ pub async fn failed_to_fix_sql_agent(
         stream: Some(options.output_sender),
         stream_name: Some("failed_to_fix_sql".to_string()),
         prompt_name: "failed_to_fix_sql".to_string(),
+        json_mode: true,
         ..Default::default()
     };
 
@@ -93,19 +94,24 @@ pub async fn failed_to_fix_sql_agent(
         }
     };
 
-    // Combine master response with first part of response
+    let sql = match response.get("sql") {
+        Some(sql) => sql.as_str().unwrap_or_default().to_string(),
+        None => "".to_string(),
+    };
+
+    // Combine SQL with first part of response
     let combined_response = match (
-        response.as_str(),
+        sql.as_str(),
         options.outputs.get("first_part_of_response"),
     ) {
-        (Some(master_str), Some(first_part)) => {
+        (sql_str, Some(first_part)) if !sql_str.is_empty() => {
             if let Some(first_part_str) = first_part.as_str() {
-                Value::String(format!("{}\n\n{}", first_part_str, master_str))
+                Value::String(format!("{}\n\n{}", first_part_str, sql_str))
             } else {
-                response
+                Value::String(sql)
             }
         }
-        _ => response,
+        _ => Value::String(sql),
     };
 
     Ok(combined_response)
diff --git a/api/src/utils/prompts/analyst_chat_prompts/failed_to_fix_sql_prompts.rs b/api/src/utils/prompts/analyst_chat_prompts/failed_to_fix_sql_prompts.rs
index 11bb95a2b..f148ae28a 100644
--- a/api/src/utils/prompts/analyst_chat_prompts/failed_to_fix_sql_prompts.rs
+++ b/api/src/utils/prompts/analyst_chat_prompts/failed_to_fix_sql_prompts.rs
@@ -8,11 +8,20 @@ You should talk about how you tried to fix the SQL query three times, and that y
 
 At the end make sure to apologize to the user
 
+PLEASE OUTPUT THE SQL QUERY IN THE FOLLOWING JSON FORMAT:
+```json
+{
+  \"sql\": \"SELECT...\"
+}
+```
+
 ### GENERAL GUIDELINES
 - Keep your response under 50 words
 - Do not output the failed SQL query in your response
 - Keep this response fairly non-technical
-- escape columns, datasets, tables, errors, etc. with backticks.".to_string()
+- escape columns, datasets, tables, errors, etc. with backticks.
+- You must output the sql query in the format specified above.
+".to_string()
 }
 
 pub fn failed_to_fix_sql_user_prompt(
@@ -44,5 +53,7 @@ pub fn failed_to_fix_sql_user_prompt(
         message.push_str(error);
     }
 
+    message.push_str("\n\nPlease output the SQL query in the format specified above. (```sql ... ```)");
+
     message
 }
diff --git a/api/src/utils/prompts/generate_sql_prompts/dataset_selector_prompt.rs b/api/src/utils/prompts/generate_sql_prompts/dataset_selector_prompt.rs
index 656249d72..a23c5a847 100644
--- a/api/src/utils/prompts/generate_sql_prompts/dataset_selector_prompt.rs
+++ b/api/src/utils/prompts/generate_sql_prompts/dataset_selector_prompt.rs
@@ -13,6 +13,7 @@ Your task is to pick all the datasets required to answer the user question/reque
 - Feel free to select multiple datasets that can be joined together to provide more complete answers
 - If the user requests advanced analysis like predictions, forecasts, correlation, impact analysis, etc., identify all datasets that could be combined for the analysis
 - Consider relationships between datasets and how they can be joined to provide comprehensive answers
+- Multiple dataset can be selected even while one completely answers the user request.
 "#,
         datasets
     )