agent.py

import json
from typing import Any, Literal

from langgraph.graph import END, StateGraph
from langchain_community.document_loaders import JSONLoader
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser

from langgraph.store.memory import InMemoryStore
from langchain.chat_models.base import BaseChatModel
from langgraph.checkpoint.memory import MemorySaver
from langchain_community.utilities import SQLDatabase

from constant import DB_PATH
from state import AgentState
import logging

# Initialize logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class SQLAgentRAG:
    def __init__(
        self,
        llm: BaseChatModel,
        tools: Any,
        db_uri: str = DB_PATH,
        table_json_path: str = "table.jsonl",
        column_json_path: str = "column.jsonl"
    ):
        self.llm = llm
        self.table_json_path = table_json_path
        self.column_json_path = column_json_path
        self.db = SQLDatabase.from_uri(db_uri)
        self.schema = None
        self.retriever = tools

         # add nodes
        graph = StateGraph(AgentState)
        graph.add_node("router", self.router_node)
        graph.add_node("general_asistant", self._general_asistant)
        graph.add_node("sql_gen", self._sql_gen)
        graph.add_node("validate_sql", self._validate_sql)
        graph.add_node("solve_error", self._solve_error)
        graph.add_node("response", self._query_gen_node)

        # add edges
        graph.set_entry_point("router")
        graph.add_edge("sql_gen", "validate_sql")
        graph.add_conditional_edges(
            "router",
            self.router,
            {
                "SQL": "sql_gen",
                "GENERAL": "general_asistant"
            }
        )
        graph.add_conditional_edges(
            "validate_sql",
            self._should_continue
        )
        graph.add_edge("solve_error", "validate_sql")
        graph.add_edge("response", END)
        graph.add_edge("general_asistant", END)

        # compile
        store = InMemoryStore()
        checkpointer = MemorySaver()
        self.graph = graph.compile(checkpointer=checkpointer, store=store)


    def _indexing_table(self, query: str):
        """
        Index and retrieve relevant tables based on the input query.
        """
        logger.info("Indexing Table...")
        docs_table = JSONLoader(
            file_path=self.table_json_path,
            jq_schema='.',
            text_content=False,
            json_lines=True
        ).load()

        retriever = self.retriever(docs_table, k=5, search_type='mmr', lambda_mult=1)

        matched_documents_table = retriever.invoke(query)
        matched_tables = [
            json.loads(doc.page_content)["table"] for doc in matched_documents_table
        ]

        return matched_tables
    
    def _indexing_column(self, matched_tables, query: str):
        """
        Index and retrieve relevant columns based on the matched tables and query.
        """
        logger.info("Get matched schema...")
        docs_column = JSONLoader(
            file_path=self.column_json_path,
            jq_schema='.',
            text_content=False,
            json_lines=True
        ).load()

        retriever = self.retriever(docs_column, k=20, search_type='similarity')

        matched_columns = retriever.invoke(query)
        matched_columns_filtered = [
            json.loads(doc.page_content) for doc in matched_columns
            if json.loads(doc.page_content)["table_name"] in matched_tables
        ]

        matched_columns_cleaned = [
            f'table_name={doc["table_name"]}|column_name={doc["column_name"]}|data_type={doc["data_type"]}'
            for doc in matched_columns_filtered
        ]

        return matched_columns_cleaned
    
    def _sql_gen(self, state: AgentState):
        """
        Generates a SQL query based on the input provided by the user.
        This function uses the LLM to construct the query from matched tables and columns.
        """
        logger.info("Generate SQL Query...")
        messages = state["messages"][-1].content
        matched_table = self._indexing_table(messages)
        self.schema = self._indexing_column(matched_table, messages)

        prompt = PromptTemplate(
            template="""
              You are a SQL master expert specializing in writing complex SQL queries for SQLite. Your task is to construct a SQL query based on the provided information. Follow these strict rules:

              QUERY: {query}
              -------
              MATCHED_SCHEMA: {matched_schema}
              -------

              Please construct a SQL query using the MATCHED_SCHEMA and the QUERY provided above.
              IMPORTANT: Use ONLY the column names (column_name) mentioned in MATCHED_SCHEMA. DO NOT USE any other column names outside of this.
              IMPORTANT: Associate column_name mentioned in MATCHED_SCHEMA only to the table_name specified under MATCHED_SCHEMA.
              NOTE: Use SQL 'AS' statement to assign a new name temporarily to a table column or even a table wherever needed.
              Generate ONLY the SQL query. Do not provide any explanations, comments, or additional text.

            """,
            input_variables=["query", "matched_schema"]
        )

        sql_gen = prompt | self.llm | StrOutputParser()
        result_sql = sql_gen.invoke({"query": messages, "matched_schema": self.schema})
        return {"sql_query": [result_sql]}
    
    def _validate_sql(self, state: AgentState):
        """
        Validates the generated SQL query by attempting to execute it.
        Returns "" if no errors, otherwise returns the error message.
        """
        logger.info("Validate SQL...")
        query = state["sql_query"][-1]
        try:
            logger.info(f"Query:\n{query}")
            self.db.run(query)
            return {"error_str": ""}

        except Exception as e:
            return {"error_str": [f"Unexpected Error: {str(e)}"]}
        
    def _solve_error(self, state: AgentState):
        """
        Called with the error code and error description as the argument to get guidance on how to solve the error
        """
        error_string = state.get('error_str', "")
        sql_query = state.get('sql_query', [None])[-1]

        logger.info(f"Error{error_string}")
        logger.info(f"SQL query {sql_query}")

        prompt = PromptTemplate(
            template="""
              First, identify the main issues with the given SQL query based on the error message.
              {error_string}

              Next, examine the schema and current SQL query to locate potential sources of the error.
              {schema}

              Then, modify the current SQL query to fix the error and avoid similar issues in the future.
              {sql_query}

              Finally, ensure the revised SQL query conforms to the requirements outlined in the original task and provide the corrected SQL query.
              Generate ONLY the SQL query. Do not provide any explanations, comments, or additional text.

            """,
            input_variables=["error_string", "schema", "sql_query"]
        )
        resolver = prompt | self.llm | StrOutputParser()
        corrected_query = resolver.invoke({
            "error_string": error_string, 
            "schema": self.schema, 
            "sql_query": sql_query
        })
        return {"sql_query": [corrected_query]}
    
    def _query_gen_node(self, state: AgentState):
        """
        Generates a final response after executing the SQL query and getting the result.
        """
        logger.info("Generate Response...")
        query = state.get('sql_query', [None])[-1]
        messages = state["messages"][-1]  
        prompt = PromptTemplate(
            template="""Based on the following SQL result, generate a natural language response:
            Query SQL: {user_query}
            SQL Result: {sql_response}
            """,
            input_variables=["user_query", "sql_response"]
        )

        sql_response = self.db.run(query)
        gen_llm = prompt | self.llm | StrOutputParser()
        response = gen_llm.invoke({"user_query": messages, "sql_response": sql_response})
        return {"messages": [response]}
    
    def _general_asistant(self, state: AgentState):
        """Assistan handle question out of context"""
        logger.info("Assistant...")
        messages = state["messages"]
        response = self.llm.invoke(messages)   

        return {"messages": [response], "sql_query": [""]}

    def router(self, state: AgentState):
        """Router"""
        logger.info("Router...")
        return state["question_type"]

    def router_node(self, state: AgentState):
        """Router Node"""
        question = state["messages"][-1].content
        prompt = PromptTemplate(
            template= """
            You are a senior specialist of analytical support. Classify incoming questions into one of two types:
              - SQL: Related to flight information, schedules, hotels, rentals, recommendations, and anything about vacations
              - GENERAL: General questions
            Return only one word: SQL, or GENERAL.

            {question}
        """,
            input_variables=["question"]
        )
        router = prompt | self.llm | StrOutputParser()
        question_type = router.invoke({"question": question})
        return {"question_type": question_type}

    def _should_continue(self, state: AgentState) -> Literal["response", "solve_error"]:
        """
        Decides whether to proceed based on SQL validation results.
        If the last message contains "error_str", 
        go to error-solving or retry. Otherwise,continue to response. 
        """
        if state.get("error_str"):
            # If error exists, go to solve_error
            return "solve_error"
        return "response"