code: import datetime
  import difflib
  import logging
  import os
  from functools import wraps
  from queue import Queue
  from threading import Thread
  from typing import Any, Callable, Dict, List
  
  import numpy as np
  import openai
  import pandas as pd
  import sqlalchemy
  from google.api_core.exceptions import GoogleAPIError
  from langchain.agents.agent import AgentExecutor
  from langchain.agents.agent_toolkits.base import BaseToolkit
  from langchain.agents.mrkl.base import ZeroShotAgent
  from langchain.callbacks.base import BaseCallbackManager
  from langchain.callbacks.manager import (
      AsyncCallbackManagerForToolRun,
      CallbackManagerForToolRun,
  )
  from langchain.chains.llm import LLMChain
  from langchain.tools.base import BaseTool
  from langchain_community.callbacks import get_openai_callback
  from langchain_openai import OpenAIEmbeddings
  from overrides import override
  from pydantic import BaseModel, Field
  from sqlalchemy import MetaData
  from sqlalchemy.exc import SQLAlchemyError
  from sqlalchemy.sql import func
  
  from dataherald.context_store import ContextStore
  from dataherald.db import DB
  from dataherald.db_scanner.models.types import TableDescription, TableDescriptionStatus
  from dataherald.db_scanner.repository.base import TableDescriptionRepository
  from dataherald.repositories.sql_generations import (
      SQLGenerationRepository,
  )
  from dataherald.sql_database.base import SQLDatabase, SQLInjectionError
  from dataherald.sql_database.models.types import (
      DatabaseConnection,
  )
  from dataherald.sql_generator import EngineTimeOutORItemLimitError, SQLGenerator
  from dataherald.types import Prompt, SQLGeneration
  from dataherald.utils.agent_prompts import (
      AGENT_PREFIX,
      ERROR_PARSING_MESSAGE,
      FORMAT_INSTRUCTIONS,
      PLAN_BASE,
      PLAN_WITH_FEWSHOT_EXAMPLES,
      PLAN_WITH_FEWSHOT_EXAMPLES_AND_INSTRUCTIONS,
      PLAN_WITH_INSTRUCTIONS,
      SUFFIX_WITH_FEW_SHOT_SAMPLES,
      SUFFIX_WITHOUT_FEW_SHOT_SAMPLES,
  )
  from dataherald.utils.timeout_utils import run_with_timeout
  
  logger = logging.getLogger(__name__)
  
  
  TOP_K = SQLGenerator.get_upper_bound_limit()
  EMBEDDING_MODEL = "text-embedding-3-large"
  TOP_TABLES = 20
  
  
  def catch_exceptions():  # noqa: C901
      def decorator(fn: Callable[[str], str]) -> Callable[[str], str]:  # noqa: C901
          @wraps(fn)
          def wrapper(*args: Any, **kwargs: Any) -> Any:  # noqa: PLR0911
              try:
                  return fn(*args, **kwargs)
              except openai.AuthenticationError as e:
                  # Handle authentication error here
                  return f"OpenAI API authentication error: {e}"
              except openai.RateLimitError as e:
                  # Handle API error here, e.g. retry or log
                  return f"OpenAI API request exceeded rate limit: {e}"
              except openai.BadRequestError as e:
                  # Handle connection error here
                  return f"OpenAI API request timed out: {e}"
              except openai.APIResponseValidationError as e:
                  # Handle rate limit error (we recommend using exponential backoff)
                  return f"OpenAI API response is invalid: {e}"
              except openai.OpenAIError as e:
                  # Handle timeout error (we recommend using exponential backoff)
                  return f"OpenAI API returned an error: {e}"
              except GoogleAPIError as e:
                  return f"Google API returned an error: {e}"
              except SQLAlchemyError as e:
                  return f"Error: {e}"
  
          return wrapper
  
      return decorator
  
  
  def replace_unprocessable_characters(text: str) -> str:
      """Replace unprocessable characters with a space."""
      text = text.strip()
      return text.replace(r"\_", "_")
  
  
  # Classes needed for tools
  class BaseSQLDatabaseTool(BaseModel):
      """Base tool for interacting with the SQL database and the context information."""
  
      db: SQLDatabase = Field(exclude=True)
      context: List[dict] | None = Field(exclude=True, default=None)
  
      class Config(BaseTool.Config):
          """Configuration for this pydantic object."""
  
          arbitrary_types_allowed = True
          extra = "allow"
  
  
  class SystemTime(BaseSQLDatabaseTool, BaseTool):
      """Tool for finding the current data and time."""
  
      name = "SystemTime"
      description = """
      Input is an empty string, output is the current data and time.
      Always use this tool before generating a query if there is any time or date in the given question.
      """
  
      @catch_exceptions()
      def _run(
          self,
          tool_input: str = "",  # noqa: ARG002
          run_manager: CallbackManagerForToolRun | None = None,  # noqa: ARG002
      ) -> str:
          """Execute the query, return the results or an error message."""
          current_datetime = datetime.datetime.now()
          return f"Current Date and Time: {str(current_datetime)}"
  
      async def _arun(
          self,
          tool_input: str = "",
          run_manager: AsyncCallbackManagerForToolRun | None = None,
      ) -> str:
          raise NotImplementedError("GetCurrentTimeTool does not support async")
  
  
  class QuerySQLDataBaseTool(BaseSQLDatabaseTool, BaseTool):
      """Tool for querying a SQL database."""
  
      name = "SqlDbQuery"
      description = """
      Input: SQL query.
      Output: Result from the database or an error message if the query is incorrect.
      If an error occurs, rewrite the query and retry.
      Use this tool to execute SQL queries.
      """
  
      @catch_exceptions()
      def _run(
          self,
          query: str,
          top_k: int = TOP_K,
          run_manager: CallbackManagerForToolRun | None = None,  # noqa: ARG002
      ) -> str:
          """Execute the query, return the results or an error message."""
          query = replace_unprocessable_characters(query)
          if "```sql" in query:
              query = query.replace("```sql", "").replace("```", "")
  
          try:
              return run_with_timeout(
                  self.db.run_sql,
                  args=(query,),
                  kwargs={"top_k": top_k},
                  timeout_duration=int(os.getenv("SQL_EXECUTION_TIMEOUT", "60")),
              )[0]
          except TimeoutError:
              return "SQL query execution time exceeded, proceed without query execution"
  
      async def _arun(
          self,
          query: str,
          run_manager: AsyncCallbackManagerForToolRun | None = None,
      ) -> str:
          raise NotImplementedError("QuerySQLDataBaseTool does not support async")
  
  
  class GetUserInstructions(BaseSQLDatabaseTool, BaseTool):
      """Tool for retrieving the instructions from the user"""
  
      name = "GetAdminInstructions"
      description = """
      Input: is an empty string.
      Output: Database admin instructions before generating the SQL query.
      The generated SQL query MUST follow the admin instructions even it contradicts with the given question.
      """
      instructions: List[dict]
  
      @catch_exceptions()
      def _run(
          self,
          tool_input: str = "",  # noqa: ARG002
          run_manager: CallbackManagerForToolRun | None = None,  # noqa: ARG002
      ) -> str:
          response = "Admin: All of the generated SQL queries must follow the below instructions:\n"
          for instruction in self.instructions:
              response += f"{instruction['instruction']}\n"
          return response
  
      async def _arun(
          self,
          tool_input: str = "",  # noqa: ARG002
          run_manager: AsyncCallbackManagerForToolRun | None = None,
      ) -> str:
          raise NotImplementedError("GetUserInstructions does not support async")
  
  
  class TablesSQLDatabaseTool(BaseSQLDatabaseTool, BaseTool):
      """Tool which takes in the given question and returns a list of tables with their relevance score to the question"""
  
      name = "DbTablesWithRelevanceScores"
      description = """
      Input: Given question.
      Output: Comma-separated list of tables with their relevance scores, indicating their relevance to the question.
      Use this tool to identify the relevant tables for the given question.
      """
      db_scan: List[TableDescription]
      embedding: OpenAIEmbeddings
  
      def get_embedding(
          self,
          text: str,
      ) -> List[float]:
          text = text.replace("\n", " ")
          return self.embedding.embed_query(text)
  
      def get_docs_embedding(
          self,
          docs: List[str],
      ) -> List[List[float]]:
          return self.embedding.embed_documents(docs)
  
      def cosine_similarity(self, a: List[float], b: List[float]) -> float:
          return round(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)), 4)
  
      @catch_exceptions()
      def _run(
          self,
          user_question: str,
          run_manager: CallbackManagerForToolRun | None = None,  # noqa: ARG002
      ) -> str:
          """Use the concatenation of table name, columns names, and the description of the table as the table representation"""
          question_embedding = self.get_embedding(user_question)
          table_representations = []
          for table in self.db_scan:
              col_rep = ""
              for column in table.columns:
                  if column.description is not None:
                      col_rep += f"{column.name}: {column.description}, "
                  else:
                      col_rep += f"{column.name}, "
              if table.description is not None:
                  table_rep = f"Table {table.table_name} contain columns: [{col_rep}], this tables has: {table.description}"
              else:
                  table_rep = f"Table {table.table_name} contain columns: [{col_rep}]"
              table_representations.append([table.table_name, table_rep])
          df = pd.DataFrame(
              table_representations, columns=["table_name", "table_representation"]
          )
          df["table_embedding"] = self.get_docs_embedding(df.table_representation)
          df["similarities"] = df.table_embedding.apply(
              lambda x: self.cosine_similarity(x, question_embedding)
          )
          df = df.sort_values(by="similarities", ascending=True)
          df = df.tail(TOP_TABLES)
          table_relevance = ""
          for _, row in df.iterrows():
              table_relevance += (
                  f'Table: {row["table_name"]}, relevance score: {row["similarities"]}\n'
              )
          return table_relevance
  
      async def _arun(
          self,
          user_question: str = "",
          run_manager: AsyncCallbackManagerForToolRun | None = None,
      ) -> str:
          raise NotImplementedError("TablesSQLDatabaseTool does not support async")
  
  
  class ColumnEntityChecker(BaseSQLDatabaseTool, BaseTool):
      """Tool for checking the existance of an entity inside a column."""
  
      name = "DbColumnEntityChecker"
      description = """
      Input: Column name and its corresponding table, and an entity.
      Output: cell-values found in the column similar to the given entity.
      Use this tool to get cell values similar to the given entity in the given column.
  
      Example Input: table1 -> column2, entity
      """
  
      def find_similar_strings(
          self, input_list: List[tuple], target_string: str, threshold=0.4
      ):
          similar_strings = []
          for item in input_list:
              similarity = difflib.SequenceMatcher(
                  None, str(item[0]).strip().lower(), target_string.lower()
              ).ratio()
              if similarity >= threshold:
                  similar_strings.append((str(item[0]).strip(), similarity))
          similar_strings.sort(key=lambda x: x[1], reverse=True)
          return similar_strings[:25]
  
      @catch_exceptions()
      def _run(
          self,
          tool_input: str,
          run_manager: CallbackManagerForToolRun | None = None,  # noqa: ARG002
      ) -> str:
          try:
              schema, entity = tool_input.split(",")
              table_name, column_name = schema.split("->")
          except ValueError:
              return "Invalid input format, use following format: table_name -> column_name, entity (entity should be a string without ',')"
          search_pattern = f"%{entity.strip().lower()}%"
          meta = MetaData(bind=self.db.engine)
          table = sqlalchemy.Table(table_name.strip(), meta, autoload=True)
          try:
              search_query = sqlalchemy.select(
                  [func.distinct(table.c[column_name.strip()])]
              ).where(func.lower(table.c[column_name.strip()]).like(search_pattern))
              search_results = self.db.engine.execute(search_query).fetchall()
              search_results = search_results[:25]
          except SQLAlchemyError:
              search_results = []
          distinct_query = sqlalchemy.select(
              [func.distinct(table.c[column_name.strip()])]
          )
          results = self.db.engine.execute(distinct_query).fetchall()
          results = self.find_similar_strings(results, entity)
          similar_items = "Similar items:\n"
          already_added = {}
          for item in results:
              similar_items += f"{item[0]}\n"
              already_added[item[0]] = True
          if len(search_results) > 0:
              for item in search_results:
                  if item[0] not in already_added:
                      similar_items += f"{item[0]}\n"
          return similar_items
  
      async def _arun(
          self,
          tool_input: str,
          run_manager: AsyncCallbackManagerForToolRun | None = None,
      ) -> str:
          raise NotImplementedError("ColumnEntityChecker does not support async")
  
  
  class SchemaSQLDatabaseTool(BaseSQLDatabaseTool, BaseTool):
      """Tool for getting schema of relevant tables."""
  
      name = "DbRelevantTablesSchema"
      description = """
      Input: Comma-separated list of tables.
      Output: Schema of the specified tables.
      Use this tool to discover all columns of the relevant tables and identify potentially relevant columns.
  
      Example Input: table1, table2, table3
      """
      db_scan: List[TableDescription]
  
      @catch_exceptions()
      def _run(
          self,
          table_names: str,
          run_manager: CallbackManagerForToolRun | None = None,  # noqa: ARG002
      ) -> str:
          """Get the schema for tables in a comma-separated list."""
          table_names_list = table_names.split(", ")
          table_names_list = [
              replace_unprocessable_characters(table_name)
              for table_name in table_names_list
          ]
          tables_schema = ""
          for table in self.db_scan:
              if table.table_name in table_names_list:
                  tables_schema += table.table_schema + "\n"
                  descriptions = []
                  if table.description is not None:
                      descriptions.append(
                          f"Table `{table.table_name}`: {table.description}\n"
                      )
                      for column in table.columns:
                          if column.description is not None:
                              descriptions.append(
                                  f"Column `{column.name}`: {column.description}\n"
                              )
                  if len(descriptions) > 0:
                      tables_schema += f"/*\n{''.join(descriptions)}*/\n"
          if tables_schema == "":
              tables_schema += "Tables not found in the database"
          return tables_schema
  
      async def _arun(
          self,
          table_name: str,
          run_manager: AsyncCallbackManagerForToolRun | None = None,
      ) -> str:
          raise NotImplementedError("SchemaSQLDatabaseTool does not support async")
  
  
  class InfoRelevantColumns(BaseSQLDatabaseTool, BaseTool):
      """Tool for getting more information for potentially relevant columns"""
  
      name = "DbRelevantColumnsInfo"
      description = """
      Input: Comma-separated list of potentially relevant columns with their corresponding table.
      Output: Information about the values inside the columns and their descriptions.
      Use this tool to gather details about potentially relevant columns. then, filter them, and identify the relevant ones.
  
      Example Input: table1 -> column1, table1 -> column2, table2 -> column1
      """
      db_scan: List[TableDescription]
  
      @catch_exceptions()
      def _run(  # noqa: C901
          self,
          column_names: str,
          run_manager: CallbackManagerForToolRun | None = None,  # noqa: ARG002
      ) -> str:
          """Get the column level information."""
          items_list = column_names.split(", ")
          column_full_info = ""
          for item in items_list:
              if " -> " in item:
                  table_name, column_name = item.split(" -> ")
                  table_name = replace_unprocessable_characters(table_name)
                  column_name = replace_unprocessable_characters(column_name)
                  found = False
                  for table in self.db_scan:
                      if table_name == table.table_name:
                          col_info = ""
                          for column in table.columns:
                              if column_name == column.name:
                                  found = True
                                  col_info += f"Description: {column.description},"
                                  if column.low_cardinality:
                                      col_info += f" categories = {column.categories},"
                          col_info += " Sample rows: "
                          if found:
                              for row in table.examples:
                                  col_info += row[column_name] + ", "
                              col_info = col_info[:-2]
                              column_full_info += f"Table: {table_name}, column: {column_name}, additional info: {col_info}\n"
              else:
                  return "Malformed input, input should be in the following format Example Input: table1 -> column1, table1 -> column2, table2 -> column1"  # noqa: E501
              if not found:
                  column_full_info += f"Table: {table_name}, column: {column_name} not found in database\n"
          return column_full_info
  
      async def _arun(
          self,
          table_name: str,
          run_manager: AsyncCallbackManagerForToolRun | None = None,
      ) -> str:
          raise NotImplementedError("InfoRelevantColumnsTool does not support async")
  
  
  class GetFewShotExamples(BaseSQLDatabaseTool, BaseTool):
      """Tool to obtain few-shot examples from the pool of samples"""
  
      name = "FewshotExamplesRetriever"
      description = """
      Input: Number of required Question/SQL pairs.
      Output: List of similar Question/SQL pairs related to the given question.
      Use this tool to fetch previously asked Question/SQL pairs as examples for improving SQL query generation.
      For complex questions, request more examples to gain a better understanding of tables and columns and the SQL keywords to use.
      If the given question is very similar to one of the retrieved examples, it is recommended to use the same SQL query and modify it slightly to fit the given question.
      Always use this tool first and before any other tool!
      """  # noqa: E501
      few_shot_examples: List[dict]
  
      @catch_exceptions()
      def _run(
          self,
          number_of_samples: str,
          run_manager: CallbackManagerForToolRun | None = None,  # noqa: ARG002
      ) -> str:
          """Get the schema for tables in a comma-separated list."""
          if number_of_samples.strip().isdigit():
              number_of_samples = int(number_of_samples.strip())
          else:
              return "Action input for the fewshot_examples_retriever tool should be an integer"
          returned_output = ""
          for example in self.few_shot_examples[:number_of_samples]:
              returned_output += (
                  f"Question: {example['prompt_text']} -> SQL: {example['sql']}\n"
              )
          if returned_output == "":
              returned_output = "No previously asked Question/SQL pairs are available"
          return returned_output
  
      async def _arun(
          self,
          number_of_samples: str,
          run_manager: AsyncCallbackManagerForToolRun | None = None,
      ) -> str:
          raise NotImplementedError("GetFewShotExamplesTool does not support async")
  
  
  class SQLDatabaseToolkit(BaseToolkit):
      """Dataherald toolkit"""
  
      db: SQLDatabase = Field(exclude=True)
      context: List[dict] | None = Field(exclude=True, default=None)
      few_shot_examples: List[dict] | None = Field(exclude=True, default=None)
      instructions: List[dict] | None = Field(exclude=True, default=None)
      db_scan: List[TableDescription] = Field(exclude=True)
      embedding: OpenAIEmbeddings = Field(exclude=True)
  
      @property
      def dialect(self) -> str:
          """Return string representation of SQL dialect to use."""
          return self.db.dialect
  
      class Config:
          """Configuration for this pydantic object."""
  
          arbitrary_types_allowed = True
  
      def get_tools(self) -> List[BaseTool]:
          """Get the tools in the toolkit."""
          tools = []
          query_sql_db_tool = QuerySQLDataBaseTool(db=self.db, context=self.context)
          tools.append(query_sql_db_tool)
          if self.instructions is not None:
              tools.append(
                  GetUserInstructions(
                      db=self.db, context=self.context, instructions=self.instructions
                  )
              )
          get_current_datetime = SystemTime(db=self.db, context=self.context)
          tools.append(get_current_datetime)
          tables_sql_db_tool = TablesSQLDatabaseTool(
              db=self.db,
              context=self.context,
              db_scan=self.db_scan,
              embedding=self.embedding,
          )
          tools.append(tables_sql_db_tool)
          schema_sql_db_tool = SchemaSQLDatabaseTool(
              db=self.db, context=self.context, db_scan=self.db_scan
          )
          tools.append(schema_sql_db_tool)
          info_relevant_tool = InfoRelevantColumns(
              db=self.db, context=self.context, db_scan=self.db_scan
          )
          tools.append(info_relevant_tool)
          column_sample_tool = ColumnEntityChecker(db=self.db, context=self.context)
          tools.append(column_sample_tool)
          if self.few_shot_examples is not None:
              get_fewshot_examples_tool = GetFewShotExamples(
                  db=self.db,
                  context=self.context,
                  few_shot_examples=self.few_shot_examples,
              )
              tools.append(get_fewshot_examples_tool)
          return tools
  
  
  class DataheraldSQLAgent(SQLGenerator):
      """Dataherald SQL agent"""
  
      max_number_of_examples: int = 5  # maximum number of question/SQL pairs
      llm: Any = None
  
      def remove_duplicate_examples(self, fewshot_exmaples: List[dict]) -> List[dict]:
          returned_result = []
          seen_list = []
          for example in fewshot_exmaples:
              if example["prompt_text"] not in seen_list:
                  seen_list.append(example["prompt_text"])
                  returned_result.append(example)
          return returned_result
  
      def create_sql_agent(
          self,
          toolkit: SQLDatabaseToolkit,
          callback_manager: BaseCallbackManager | None = None,
          prefix: str = AGENT_PREFIX,
          suffix: str | None = None,
          format_instructions: str = FORMAT_INSTRUCTIONS,
          input_variables: List[str] | None = None,
          max_examples: int = 20,
          number_of_instructions: int = 1,
          max_iterations: int
          | None = int(os.getenv("AGENT_MAX_ITERATIONS", "15")),  # noqa: B008
          max_execution_time: float | None = None,
          early_stopping_method: str = "generate",
          verbose: bool = False,
          agent_executor_kwargs: Dict[str, Any] | None = None,
          **kwargs: Dict[str, Any],
      ) -> AgentExecutor:
          """Construct an SQL agent from an LLM and tools."""
          tools = toolkit.get_tools()
          if max_examples > 0 and number_of_instructions > 0:
              plan = PLAN_WITH_FEWSHOT_EXAMPLES_AND_INSTRUCTIONS
              suffix = SUFFIX_WITH_FEW_SHOT_SAMPLES
          elif max_examples > 0:
              plan = PLAN_WITH_FEWSHOT_EXAMPLES
              suffix = SUFFIX_WITH_FEW_SHOT_SAMPLES
          elif number_of_instructions > 0:
              plan = PLAN_WITH_INSTRUCTIONS
              suffix = SUFFIX_WITHOUT_FEW_SHOT_SAMPLES
          else:
              plan = PLAN_BASE
              suffix = SUFFIX_WITHOUT_FEW_SHOT_SAMPLES
          plan = plan.format(
              dialect=toolkit.dialect,
              max_examples=max_examples,
          )
          prefix = prefix.format(
              dialect=toolkit.dialect, max_examples=max_examples, agent_plan=plan
          )
          prompt = ZeroShotAgent.create_prompt(
              tools,
              prefix=prefix,
              suffix=suffix,
              format_instructions=format_instructions,
              input_variables=input_variables,
          )
          llm_chain = LLMChain(
              llm=self.llm,
              prompt=prompt,
              callback_manager=callback_manager,
          )
          tool_names = [tool.name for tool in tools]
          agent = ZeroShotAgent(llm_chain=llm_chain, allowed_tools=tool_names, **kwargs)
          return AgentExecutor.from_agent_and_tools(
              agent=agent,
              tools=tools,
              callback_manager=callback_manager,
              verbose=verbose,
              max_iterations=max_iterations,
              max_execution_time=max_execution_time,
              early_stopping_method=early_stopping_method,
              **(agent_executor_kwargs or {}),
          )
  
      @override
      def generate_response(
          self,
          user_prompt: Prompt,
          database_connection: DatabaseConnection,
          context: List[dict] = None,
      ) -> SQLGeneration:
          context_store = self.system.instance(ContextStore)
          storage = self.system.instance(DB)
          response = SQLGeneration(
              prompt_id=user_prompt.id,
              llm_config=self.llm_config,
              created_at=datetime.datetime.now(),
          )
          self.llm = self.model.get_model(
              database_connection=database_connection,
              temperature=0,
              model_name=self.llm_config.llm_name,
              api_base=self.llm_config.api_base,
          )
          repository = TableDescriptionRepository(storage)
          db_scan = repository.get_all_tables_by_db(
              {
                  "db_connection_id": str(database_connection.id),
                  "status": TableDescriptionStatus.SCANNED.value,
              }
          )
          if not db_scan:
              raise ValueError("No scanned tables found for database")
          few_shot_examples, instructions = context_store.retrieve_context_for_question(
              user_prompt, number_of_samples=self.max_number_of_examples
          )
          if few_shot_examples is not None:
              new_fewshot_examples = self.remove_duplicate_examples(few_shot_examples)
              number_of_samples = len(new_fewshot_examples)
          else:
              new_fewshot_examples = None
              number_of_samples = 0
          logger.info(f"Generating SQL response to question: {str(user_prompt.dict())}")
          self.database = SQLDatabase.get_sql_engine(database_connection)
          toolkit = SQLDatabaseToolkit(
              db=self.database,
              context=context,
              few_shot_examples=new_fewshot_examples,
              instructions=instructions,
              db_scan=db_scan,
              embedding=OpenAIEmbeddings(
                  openai_api_key=database_connection.decrypt_api_key(),
                  model=EMBEDDING_MODEL,
              ),
          )
          agent_executor = self.create_sql_agent(
              toolkit=toolkit,
              verbose=True,
              max_examples=number_of_samples,
              number_of_instructions=len(instructions) if instructions is not None else 0,
              max_execution_time=int(os.environ.get("DH_ENGINE_TIMEOUT", 150)),
          )
          agent_executor.return_intermediate_steps = True
          agent_executor.handle_parsing_errors = ERROR_PARSING_MESSAGE
          with get_openai_callback() as cb:
              try:
                  result = agent_executor.invoke({"input": user_prompt.text})
                  result = self.check_for_time_out_or_tool_limit(result)
              except SQLInjectionError as e:
                  raise SQLInjectionError(e) from e
              except EngineTimeOutORItemLimitError as e:
                  raise EngineTimeOutORItemLimitError(e) from e
              except Exception as e:
                  return SQLGeneration(
                      prompt_id=user_prompt.id,
                      tokens_used=cb.total_tokens,
                      completed_at=datetime.datetime.now(),
                      sql="",
                      status="INVALID",
                      error=str(e),
                  )
          sql_query = ""
          if "```sql" in result["output"]:
              sql_query = self.remove_markdown(result["output"])
          else:
              sql_query = self.extract_query_from_intermediate_steps(
                  result["intermediate_steps"]
              )
          logger.info(f"cost: {str(cb.total_cost)} tokens: {str(cb.total_tokens)}")
          response.sql = replace_unprocessable_characters(sql_query)
          response.tokens_used = cb.total_tokens
          response.completed_at = datetime.datetime.now()
          return self.create_sql_query_status(
              self.database,
              response.sql,
              response,
          )
  
      @override
      def stream_response(
          self,
          user_prompt: Prompt,
          database_connection: DatabaseConnection,
          response: SQLGeneration,
          queue: Queue,
      ):
          context_store = self.system.instance(ContextStore)
          storage = self.system.instance(DB)
          sql_generation_repository = SQLGenerationRepository(storage)
          self.llm = self.model.get_model(
              database_connection=database_connection,
              temperature=0,
              model_name=self.llm_config.llm_name,
              api_base=self.llm_config.api_base,
              streaming=True,
          )
          repository = TableDescriptionRepository(storage)
          db_scan = repository.get_all_tables_by_db(
              {
                  "db_connection_id": str(database_connection.id),
                  "status": TableDescriptionStatus.SCANNED.value,
              }
          )
          if not db_scan:
              raise ValueError("No scanned tables found for database")
          few_shot_examples, instructions = context_store.retrieve_context_for_question(
              user_prompt, number_of_samples=self.max_number_of_examples
          )
          if few_shot_examples is not None:
              new_fewshot_examples = self.remove_duplicate_examples(few_shot_examples)
              number_of_samples = len(new_fewshot_examples)
          else:
              new_fewshot_examples = None
              number_of_samples = 0
          self.database = SQLDatabase.get_sql_engine(database_connection)
          toolkit = SQLDatabaseToolkit(
              queuer=queue,
              db=self.database,
              context=[{}],
              few_shot_examples=new_fewshot_examples,
              instructions=instructions,
              db_scan=db_scan,
              embedding=OpenAIEmbeddings(
                  openai_api_key=database_connection.decrypt_api_key(),
                  model=EMBEDDING_MODEL,
              ),
          )
          agent_executor = self.create_sql_agent(
              toolkit=toolkit,
              verbose=True,
              max_examples=number_of_samples,
              number_of_instructions=len(instructions) if instructions is not None else 0,
              max_execution_time=int(os.environ.get("DH_ENGINE_TIMEOUT", 150)),
          )
          agent_executor.return_intermediate_steps = True
          agent_executor.handle_parsing_errors = ERROR_PARSING_MESSAGE
          thread = Thread(
              target=self.stream_agent_steps,
              args=(
                  user_prompt.text,
                  agent_executor,
                  response,
                  sql_generation_repository,
                  queue,
              ),
          )
          thread.start()
  
apis: langchain.chains.llm.LLMChain langchain.agents.agent.AgentExecutor.from_agent_and_tools langchain.agents.mrkl.base.ZeroShotAgent langchain_community.callbacks.get_openai_callback langchain.agents.mrkl.base.ZeroShotAgent.create_prompt
extract_api: [((2000, 2027), 'logging.getLogger', 'logging.getLogger', (['__name__'], {}), '(__name__)\n', (2017, 2027), False, 'import logging\n'), ((2038, 2074), 'dataherald.sql_generator.SQLGenerator.get_upper_bound_limit', 'SQLGenerator.get_upper_bound_limit', ([], {}), '()\n', (2072, 2074), False, 'from dataherald.sql_generator import EngineTimeOutORItemLimitError, SQLGenerator\n'), ((3869, 3888), 'pydantic.Field', 'Field', ([], {'exclude': '(True)'}), '(exclude=True)\n', (3874, 3888), False, 'from pydantic import BaseModel, Field\n'), ((3922, 3955), 'pydantic.Field', 'Field', ([], {'exclude': '(True)', 'default': 'None'}), '(exclude=True, default=None)\n', (3927, 3955), False, 'from pydantic import BaseModel, Field\n'), ((20062, 20081), 'pydantic.Field', 'Field', ([], {'exclude': '(True)'}), '(exclude=True)\n', (20067, 20081), False, 'from pydantic import BaseModel, Field\n'), ((20115, 20148), 'pydantic.Field', 'Field', ([], {'exclude': '(True)', 'default': 'None'}), '(exclude=True, default=None)\n', (20120, 20148), False, 'from pydantic import BaseModel, Field\n'), ((20192, 20225), 'pydantic.Field', 'Field', ([], {'exclude': '(True)', 'default': 'None'}), '(exclude=True, default=None)\n', (20197, 20225), False, 'from pydantic import BaseModel, Field\n'), ((20264, 20297), 'pydantic.Field', 'Field', ([], {'exclude': '(True)', 'default': 'None'}), '(exclude=True, default=None)\n', (20269, 20297), False, 'from pydantic import BaseModel, Field\n'), ((20336, 20355), 'pydantic.Field', 'Field', ([], {'exclude': '(True)'}), '(exclude=True)\n', (20341, 20355), False, 'from pydantic import BaseModel, Field\n'), ((20390, 20409), 'pydantic.Field', 'Field', ([], {'exclude': '(True)'}), '(exclude=True)\n', (20395, 20409), False, 'from pydantic import BaseModel, Field\n'), ((2266, 2275), 'functools.wraps', 'wraps', (['fn'], {}), '(fn)\n', (2271, 2275), False, 'from functools import wraps\n'), ((4731, 4754), 'datetime.datetime.now', 'datetime.datetime.now', ([], {}), '()\n', (4752, 4754), False, 'import datetime\n'), ((9562, 9649), 'pandas.DataFrame', 'pd.DataFrame', (['table_representations'], {'columns': "['table_name', 'table_representation']"}), "(table_representations, columns=['table_name',\n    'table_representation'])\n", (9574, 9649), True, 'import pandas as pd\n'), ((11971, 12000), 'sqlalchemy.MetaData', 'MetaData', ([], {'bind': 'self.db.engine'}), '(bind=self.db.engine)\n', (11979, 12000), False, 'from sqlalchemy import MetaData\n'), ((24413, 24555), 'langchain.agents.mrkl.base.ZeroShotAgent.create_prompt', 'ZeroShotAgent.create_prompt', (['tools'], {'prefix': 'prefix', 'suffix': 'suffix', 'format_instructions': 'format_instructions', 'input_variables': 'input_variables'}), '(tools, prefix=prefix, suffix=suffix,\n    format_instructions=format_instructions, input_variables=input_variables)\n', (24440, 24555), False, 'from langchain.agents.mrkl.base import ZeroShotAgent\n'), ((24643, 24715), 'langchain.chains.llm.LLMChain', 'LLMChain', ([], {'llm': 'self.llm', 'prompt': 'prompt', 'callback_manager': 'callback_manager'}), '(llm=self.llm, prompt=prompt, callback_manager=callback_manager)\n', (24651, 24715), False, 'from langchain.chains.llm import LLMChain\n'), ((24830, 24900), 'langchain.agents.mrkl.base.ZeroShotAgent', 'ZeroShotAgent', ([], {'llm_chain': 'llm_chain', 'allowed_tools': 'tool_names'}), '(llm_chain=llm_chain, allowed_tools=tool_names, **kwargs)\n', (24843, 24900), False, 'from langchain.agents.mrkl.base import ZeroShotAgent\n'), ((24916, 25187), 'langchain.agents.agent.AgentExecutor.from_agent_and_tools', 'AgentExecutor.from_agent_and_tools', ([], {'agent': 'agent', 'tools': 'tools', 'callback_manager': 'callback_manager', 'verbose': 'verbose', 'max_iterations': 'max_iterations', 'max_execution_time': 'max_execution_time', 'early_stopping_method': 'early_stopping_method'}), '(agent=agent, tools=tools,\n    callback_manager=callback_manager, verbose=verbose, max_iterations=\n    max_iterations, max_execution_time=max_execution_time,\n    early_stopping_method=early_stopping_method, **agent_executor_kwargs or {})\n', (24950, 25187), False, 'from langchain.agents.agent import AgentExecutor\n'), ((25998, 26033), 'dataherald.db_scanner.repository.base.TableDescriptionRepository', 'TableDescriptionRepository', (['storage'], {}), '(storage)\n', (26024, 26033), False, 'from dataherald.db_scanner.repository.base import TableDescriptionRepository\n'), ((26897, 26944), 'dataherald.sql_database.base.SQLDatabase.get_sql_engine', 'SQLDatabase.get_sql_engine', (['database_connection'], {}), '(database_connection)\n', (26923, 26944), False, 'from dataherald.sql_database.base import SQLDatabase, SQLInjectionError\n'), ((29041, 29064), 'datetime.datetime.now', 'datetime.datetime.now', ([], {}), '()\n', (29062, 29064), False, 'import datetime\n'), ((29527, 29559), 'dataherald.repositories.sql_generations.SQLGenerationRepository', 'SQLGenerationRepository', (['storage'], {}), '(storage)\n', (29550, 29559), False, 'from dataherald.repositories.sql_generations import SQLGenerationRepository\n'), ((29836, 29871), 'dataherald.db_scanner.repository.base.TableDescriptionRepository', 'TableDescriptionRepository', (['storage'], {}), '(storage)\n', (29862, 29871), False, 'from dataherald.db_scanner.repository.base import TableDescriptionRepository\n'), ((30648, 30695), 'dataherald.sql_database.base.SQLDatabase.get_sql_engine', 'SQLDatabase.get_sql_engine', (['database_connection'], {}), '(database_connection)\n', (30674, 30695), False, 'from dataherald.sql_database.base import SQLDatabase, SQLInjectionError\n'), ((31576, 31703), 'threading.Thread', 'Thread', ([], {'target': 'self.stream_agent_steps', 'args': '(user_prompt.text, agent_executor, response, sql_generation_repository, queue)'}), '(target=self.stream_agent_steps, args=(user_prompt.text,\n    agent_executor, response, sql_generation_repository, queue))\n', (31582, 31703), False, 'from threading import Thread\n'), ((23228, 23267), 'os.getenv', 'os.getenv', (['"""AGENT_MAX_ITERATIONS"""', '"""15"""'], {}), "('AGENT_MAX_ITERATIONS', '15')\n", (23237, 23267), False, 'import os\n'), ((27798, 27819), 'langchain_community.callbacks.get_openai_callback', 'get_openai_callback', ([], {}), '()\n', (27817, 27819), False, 'from langchain_community.callbacks import get_openai_callback\n'), ((8452, 8464), 'numpy.dot', 'np.dot', (['a', 'b'], {}), '(a, b)\n', (8458, 8464), True, 'import numpy as np\n'), ((25715, 25738), 'datetime.datetime.now', 'datetime.datetime.now', ([], {}), '()\n', (25736, 25738), False, 'import datetime\n'), ((8468, 8485), 'numpy.linalg.norm', 'np.linalg.norm', (['a'], {}), '(a)\n', (8482, 8485), True, 'import numpy as np\n'), ((8488, 8505), 'numpy.linalg.norm', 'np.linalg.norm', (['b'], {}), '(b)\n', (8502, 8505), True, 'import numpy as np\n'), ((27607, 27647), 'os.environ.get', 'os.environ.get', (['"""DH_ENGINE_TIMEOUT"""', '(150)'], {}), "('DH_ENGINE_TIMEOUT', 150)\n", (27621, 27647), False, 'import os\n'), ((28056, 28076), 'dataherald.sql_database.base.SQLInjectionError', 'SQLInjectionError', (['e'], {}), '(e)\n', (28073, 28076), False, 'from dataherald.sql_database.base import SQLDatabase, SQLInjectionError\n'), ((28161, 28193), 'dataherald.sql_generator.EngineTimeOutORItemLimitError', 'EngineTimeOutORItemLimitError', (['e'], {}), '(e)\n', (28190, 28193), False, 'from dataherald.sql_generator import EngineTimeOutORItemLimitError, SQLGenerator\n'), ((31381, 31421), 'os.environ.get', 'os.environ.get', (['"""DH_ENGINE_TIMEOUT"""', '(150)'], {}), "('DH_ENGINE_TIMEOUT', 150)\n", (31395, 31421), False, 'import os\n'), ((6019, 6059), 'os.getenv', 'os.getenv', (['"""SQL_EXECUTION_TIMEOUT"""', '"""60"""'], {}), "('SQL_EXECUTION_TIMEOUT', '60')\n", (6028, 6059), False, 'import os\n'), ((28402, 28425), 'datetime.datetime.now', 'datetime.datetime.now', ([], {}), '()\n', (28423, 28425), False, 'import datetime\n')]