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"