DH-5001/removing the markdowns from queries for gpt-4-turbo (#247)

* DH-5001/removing the markdowns from queries for gpt-4-turbo

* DH-5001/reformat with black

* DH-5001/add logs

dataherald/sql_generator/dataherald_sqlagent.py

@@ -155,6 +155,10 @@ def _run(
         run_manager: CallbackManagerForToolRun | None = None,  # noqa: ARG002
     ) -> str:
         """Execute the query, return the results or an error message."""
+        if "```sql" in query:
+            logger.info("**** Removing markdown formatting from the query\n")
+            query = query.replace("```sql", "").replace("```", "")
+            logger.info(f"**** Query after removing markdown formatting: {query}\n")
         return self.db.run_sql(query, top_k=top_k)[0]
 
     async def _arun(
@@ -690,7 +694,14 @@ def generate_response(
         for step in result["intermediate_steps"]:
             action = step[0]
             if type(action) == AgentAction and action.tool == "sql_db_query":
-                sql_query_list.append(self.format_sql_query(action.tool_input))
+                query = self.format_sql_query(action.tool_input)
+                if "```sql" in query:
+                    logger.info("**** Removing markdown formatting from the query\n")
+                    query = query.replace("```sql", "").replace("```", "")
+                    logger.info(
+                        f"**** Query after removing markdown formatting: {query}\n"
+                    )
+                sql_query_list.append(query)
         intermediate_steps = self.format_intermediate_representations(
             result["intermediate_steps"]
         )

dataherald/utils/models_context_window.py

@@ -11,4 +11,6 @@
     "gpt-3.5-turbo-0613": 4000,
     "gpt-3.5-turbo-16k-0613": 16000,
     "gpt-3.5-turbo-0301": 4000,
+    "gpt-4-1106-preview": 128000,
+    "gpt-3.5-turbo-1106": 16000,
 }