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()

langchain.agents.mrkl.base.ZeroShotAgent.create_prompt langchain.agents.mrkl.base.ZeroShotAgent langchain.agents.agent.AgentExecutor.from_agent_and_tools langchain.chains.llm.LLMChain langchain_community.callbacks.get_openai_callback
[((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')]