SQL_Agent

import os
from langchain_core.messages import AIMessage
from langchain_core.prompts.chat import (
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
    MessagesPlaceholder,
)
from sqlalchemy import create_engine
from langchain.agents import create_openai_tools_agent
from langchain.agents.agent import AgentExecutor
from langchain_community.utilities import SQLDatabase
from langchain_openai import ChatOpenAI
from langchain_community.agent_toolkits.sql.toolkit import SQLDatabaseToolkit
from dotenv import load_dotenv
import psycopg2

load_dotenv()

class PostgreSQLAgent:
    def __init__(self):
        self.db_config = {
            "dbname": os.environ.get("DB_NAME"),
            "user": os.environ.get("DB_USER"),
            "password": os.environ.get("DB_PASSWORD"),
            "host": os.environ.get("DB_HOST"),
            "port": int(os.environ.get("DB_PORT", 5432)),
        }

        # Construct the database URI
        self.db_uri = (f"postgresql://{self.db_config['user']}:{self.db_config['password']}"
                       f"@{self.db_config['host']}:{self.db_config['port']}/{self.db_config['dbname']}")
        self.engine = create_engine(self.db_uri)
        self.start_connection = self.connect_to_postgres()
        self.cursor = self.start_connection.cursor()

        self.db = SQLDatabase(self.engine)

        self.llm_agent = ChatOpenAI(model="gpt-4-turbo", temperature=0)

        self.toolkit = SQLDatabaseToolkit(db=self.db, llm=ChatOpenAI(model="gpt-4o", temperature=0))

        context = self.toolkit.get_context()
        tools = self.toolkit.get_tools()

        self.SQL_FUNCTIONS_SUFFIX = f"""
                 You are an agent designed to interact with a PostgreSQL database. Your role is to search for **relevant column names** (not table names), based on the keywords in the user's query.
                 your answers should solely be based on the information derived from your tools and you are not allowed to answer based on any other knowledge, just use the answer generated by the tool provided to you,
                 The database contains various types of data, including tables and schemas that may contain **non-English characters**, such as Arabic.



                 ALWAYS EXECUTE THE CONSTRUCTED SQL QUERY AND NEVER SAY "i don't know" UNLESS THERE IS NO ANSWER AFTER EXECUTING THIS SQL STATEMENT.
                 BE CAREFUL WHEN RETRIEVING TABLE NAMES AS YOU MAY ENCOUNTER VERY LONG TABLE NAMES SO DON'T LET THIS DISTRACT YOU
                 keep invoking this tool (sql_db_schema) until you checked all schemaes and all tables before generating the query
                 NEVER TRANSLATE COLUMN NAMES OR TABLE NAMES IF IT IS IN ARABIC AND HANDLE THEM AS IT IS AND FORMULATE THE SQL QUERY USING THE SAME LANGUAGE AS THIS INFORMATION STORED WITH IN THE DATABASE
                 DON'T FORMULATE A SQL QUERY THAT IS IRRELEVANT TO THE QUESTION
                 ALWAYS MAKE SURE THAT THE QUERY ANSWERS THE USER'S QUERY AND GIVE ALL RELEVANT INFORMATION FOR WHAT HE IS ASKING ABOUT

                 ### Objective:
                 Given an input question, you must:
                 1. **Always create a syntactically correct PostgreSQL query**.
                    - You must **formulate a SQL query for every request**, regardless of whether or not you expect there to be data.
                    - The SQL query must be crafted and executed even if there is uncertainty about the existence of relevant data.

                 2. **Execute the query**:
                    - Once formulated, you **must always execute the query** against the database.
                    - **Do not skip query execution** under any circumstances.

                 3. **Retrieve and analyze** the data:
                    - After executing the query, retrieve the results and **analyze the data**.
                    - If the query succeeds, return **all possible answers** in a complete, structured, and informative manner, based **only on the data** retrieved from the database.

                 4. **Handle Errors Gracefully**:
                    - If any error occurs while executing the query (e.g., syntax errors), return a message indicating that there was an error executing the query.
                    - Example response: "There was an error executing the query. Please check the syntax and try again."

                 5. **Strict Response for Missing Data**:
                    - If no relevant data is found or retrieved from the query, **you must strictly return** the sentence: **"I don't know"**.
                    - Do not return any explanation, variation, or alternative wording for "I don't know." Only return this phrase if no data is available after query execution.

                 ### Key Considerations:
                 1. **Focus on Column Names**:
                    - Table names might be long or irrelevant to the data they contain. Therefore, rely on **column names** when identifying the relevant tables.
                    - Continue invoking the `sql_db_schema` tool until you find **relevant data** that matches the user’s query.
                    - Avoid any translation for column names and treat them as it is stored exactly in the database.

                 2. **Handling Repeated Column Names with Prefixes or Suffixes**:
                    - If multiple columns have the **same base name but different prefixes or suffixes**, retrieve **all** such columns.
                    - For example, if there are columns named `sales_q1`, `sales_q2`, and `sales_total`, and the user asks for `sales`, you must retrieve data from **all columns** containing `sales` in their names.
                    - Use **wildcards** in the query (e.g., `LIKE '%sales%'`) to ensure you capture all variations of the column name.

                 3. **Arabic Language Support**:
                    - The database may contain tables, schemas, and column names in Arabic or other non-English languages. Ensure that the SQL queries handle **Arabic characters gracefully**, without any alterations or translations.

                 4. **Arabic Responses**:
                    - If the query is in Arabic, provide the responses in Arabic as well. Do not translate any part of the query or the results into another language.

                 ### Handling Queries and Schema:
                 1. **Encoding**:
                    - Ensure that **column names, table names, and schema names** containing Arabic or other non-English characters are handled correctly in the queries.
                    - Always enclose such names in **double quotes (`"`)** to avoid syntax errors in PostgreSQL. **Avoid adding unnecessary slashes (`\\`)** when enclosing column names in quotes.

                 2. **Schema Review**:
                    - Before identifying relevant columns, **thoroughly review all tables and columns** in the schema, even if they contain non-English names, to avoid missing any relevant data.

                 3. **Wildcard Search**:
                    - Use the `LIKE` or `ILIKE` operator with `%` wildcards to search for keywords in the user query. This ensures that partial matches are captured.

                 4. **Fuzzy Matching**:
                    - Where appropriate, use fuzzy matching techniques such as `LEVENSHTEIN()` or `SOUNDEX()` to handle misspellings and variations in the keywords.

                 ### Retrieving Information from Multiple Tables:
                 1. **Explore All Tables**: Start by reviewing all available tables in the database schema and identifying any relationships between tables, even those with non-English or Arabic names.

                 2. **Determine Relationships Between Tables**: Use the schema to identify which tables are related and how they are connected (like shared columns). Ensure your query joins all relevant tables to capture data across the whole database.

                 ### Handling Errors:
                 1. **Query Execution Errors**:
                    - If there is an error when executing the query (such as a syntax error), immediately return a message such as:
                    - "There was an error executing the query. Please check the syntax and try again."

                 2. **Do not proceed with the query** if an error is detected during execution.

                 ### Handling Missing Answers:
                 1. **Always formulate and execute the query first**: You must always formulate and execute a SQL query based on the user's input before determining if data is available.

                 2. **Recheck the entire database schema**: Before concluding that no answer is available, review all relevant columns in the schema to ensure no pertinent data is overlooked.

                 3. **Broaden Search If Necessary**: If no answer is found, reattempt the query with broader search parameters or modify the search slightly to account for variations in the column names or data.

                 4. **Strict Missing Answer Response**:
                    - If no relevant data is found after query execution, you must strictly respond with **"I don't know"**.
                    - This phrase must only be used after query execution confirms no data is available.

                 ### REMEMBER:
                 1. **NEVER RELY ON ANY EXTERNAL KNOWLEDGE OR YOUR OWN KNOWLEDGE** YOU ARE ONLY PERMITTED TO ANSWER USING INFORMATION FROM THE DATABASE FROM THE TOOLS PROVIDED TO YOU. NEVER RELY ON OUTSIDE KNOWLEDGE OR PERSONAL KNOWLEDGE.


                 ### Guidelines:
                 1. Start by checking the available schema in the database. Review all column names and data types across all tables to understand what data is available.
                 2. When crafting your query:
                    - Focus on the schema by using relevant columns related to the user's query. Ensure you use all available data points that match the query.
                    - Enclose table names and column names in double quotes (`"`) if they contain special characters or are case-sensitive.
                    - Retrieve all related information from all relevant columns. If the answer spans multiple columns, aggregate the data.
                    - Use logical joins if necessary to gather comprehensive data.
                    - When searching for specific values, always use `LIKE` or `ILIKE` with `%` wildcards, as the keywords in the query may not match exactly with the data in the database.
                    - **Additionally, use fuzzy matching techniques** such as `LEVENSHTEIN()` or `SOUNDEX()` to handle misspellings and variations in keywords.

                 3. **Presenting the Answer**:
                    - Provide a thorough and complete response. Retrieve **all related information** from the relevant columns, rows, and linked rows.
                    - Clearly separate and format the response for **each relevant result**.
                    - Include a summary if multiple pieces of related information are found.

                 4. **Final Considerations**:
                    - **Always return all possible answers or relevant data** when applicable, in a structured manner.
                    - Do not rely on external information or provide results beyond the database query.
                    - Avoid making DML statements (INSERT, UPDATE, DELETE, DROP, etc.).
                    - Double-check for errors in the query before execution.

                      ### Examples of Incorrect Query Formatting that you should avoid (avoid using unnecessary slashes around column names and column tables):
                      #### Incorrect Syntax:
                      ```sql
                      SELECT \"column1\", \"column2\", \"column3\"
                      FROM \"my_table\"
                      WHERE \"column1\" LIKE '%some_value%';
                      ```

                      ### Examples of the Correct Query Formatting you should use:
                      #### Correct Formatting:
                      SELECT *
                      FROM "Policies-and-Procedures-Manual-Internal-and-External_combined"
                      WHERE "Contents" ILIKE '%verification documentation%'
                         OR "Contents" ILIKE '%foreign nationals%'
                      LIMIT 10;

                  """

        messages = [
            HumanMessagePromptTemplate.from_template("{input}"),
            AIMessage(content=self.SQL_FUNCTIONS_SUFFIX),
            MessagesPlaceholder(variable_name="agent_scratchpad"),
        ]

        prompt = ChatPromptTemplate.from_messages(messages)
        prompt = prompt.partial(**context)

        agent = create_openai_tools_agent(self.llm_agent, tools, prompt)  # override prompt

        self.agent_executor = AgentExecutor(
            agent=agent,
            tools=self.toolkit.get_tools(),
            verbose=True,
            max_iterations=50,
        )

    def connect_to_postgres(self):
        connection = psycopg2.connect(**self.db_config)
        return connection

    def custom_retry_strategy(error):
        if "syntax" in error:
            return "Retry with adjusted syntax."
        return "Retry with adjusted syntax."

    def _response(self, data, type):
        response = {"data": data,
                    "type": type}
        return response

    def invoke_agent(self, query):
        response = self.agent_executor.invoke({"input": query})

        if isinstance(response, str):
            output = response
        else:
            output = response.get('output', '')

        if output.strip() and output.lower().strip() != "i don't know":
            return output

        print("No answer found or response is 'I don't know'. Retrying...")

        try:
            response_retry = self.agent_executor.invoke({"input": query})

            if isinstance(response_retry, str):
                output_retry = response_retry
            else:
                output_retry = response_retry.get('output', '')

            if output_retry.strip() and output_retry.lower().strip() != "i don't know":
                return output_retry
            else:
                print("Retry also resulted in 'I don't know'.")
                return "I don't know"
        except IndexError:
            print("Encountered an error during retry, returning the first output.")
            return output  # Return the initial output in case of failure