agent version2

app/agents/sql_agent.py

@@ -1,64 +1,384 @@
 
-from langchain_community.utilities import SQLDatabase
-from langchain_groq import ChatGroq
-from langgraph.graph import StateGraph, END, START
-from langchain_core.messages import AIMessage, ToolMessage, AnyMessage, HumanMessage
-from langgraph.graph.message import AnyMessage, add_messages
-from langchain_core.tools import tool
-from typing import Annotated, Literal, TypedDict, Any
-from pydantic import BaseModel, Field
-from langchain_core.runnables import RunnableLambda, RunnableWithFallbacks
-from langgraph.prebuilt import ToolNode
-from langchain_core.prompts import ChatPromptTemplate
-from langchain_community.agent_toolkits import SQLDatabaseToolkit
-from dotenv import load_dotenv
-import os
-from IPython.display import display
-import PIL
-from langgraph.errors import GraphRecursionError
-import os
-import io
-from typing import Annotated, Any, TypedDict
-from langgraph.graph import StateGraph, END, MessagesState
+# from langchain_community.utilities import SQLDatabase
+# from langchain_groq import ChatGroq
+# from langgraph.graph import StateGraph, END, START
+# from langchain_core.messages import AIMessage, ToolMessage, AnyMessage, HumanMessage
+# from langgraph.graph.message import AnyMessage, add_messages
+# from langchain_core.tools import tool
+# from typing import Annotated, Literal, TypedDict, Any
+# from pydantic import BaseModel, Field
+# from langchain_core.runnables import RunnableLambda, RunnableWithFallbacks
+# from langgraph.prebuilt import ToolNode
+# from langchain_core.prompts import ChatPromptTemplate
+# from langchain_community.agent_toolkits import SQLDatabaseToolkit
+# from dotenv import load_dotenv
+# import os
+# from IPython.display import display
+# import PIL
+# from langgraph.errors import GraphRecursionError
+# import os
+# import io
+# from typing import Annotated, Any, TypedDict
+# from langgraph.graph import StateGraph, END, MessagesState
+
+# from IPython.display import Image, display
+# from langchain_core.runnables.graph import MermaidDrawMethod
+# from typing import Optional, Dict
+
+# from langchain_community.utilities import SQLDatabase
+# from langchain_community.agent_toolkits import SQLDatabaseToolkit
+# from langchain_groq import ChatGroq
+# from langchain_core.messages import HumanMessage, AIMessage
+# from langchain_core.prompts import ChatPromptTemplate
+# # from langchain_core.pydantic_v1 import BaseModel, Field
+# from langgraph.graph import StateGraph, END, MessagesState
+# from typing import TypedDict, Annotated, List, Literal, Dict, Any
+# from langchain_google_genai import ChatGoogleGenerativeAI
+# from app.schemas.agent_state import DBQuery, SQLAgentState
+# from app.tools.database_tools import DatabaseTools
+# from app.utils.database_connection import DatabaseConnection
+# from dotenv import load_dotenv
+# load_dotenv()
+# import os
+# os.environ["GROQ_API_KEY"]=os.getenv("GROQ_API_KEY")
+# os.environ["GEMINI_API_KEY"]=os.getenv("GEMINI_API_KEY")
+
+
+# class SQLAgent:
+#     def __init__(self):
+        
+#         # Initialize instance variables
+#         self.db = None
+#         self.toolkit = None
+#         self.tools = None
+#         self.list_tables_tool = None
+#         self.sql_db_query = None
+#         self.get_schema_tool = None
+#         self.app = None
+        
+#         # Setting up LLM
+#         # self.llm = ChatGroq(model=model,api_key = os.getenv("GROQ_API_KEY"))
+#         self.llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash-lite", google_api_key=os.environ["GEMINI_API_KEY"])
+#         # Register the tool method
+#         # self.query_to_database = self._create_query_tool()
+
+    
+
+#     def setup_database_connection(self, connection_string: str):
+#         """Set up database connection and initialize tools"""
+#         try:
+#             # Initialize database connection
+#             # self.db = SQLDatabase.from_uri(connection_string)
+#             # print("Database connection successful!")
+#             self.db = DatabaseConnection(connection_string).db
+#             print("Database connection successful!")
+#             # Initialize toolkit and tools class 
+#             self.db_tools = DatabaseTools(db=self.db, llm=self.llm)
+
+#             try:
+#                 self.initialize_workflow()
+                
+#                 return self.db
+
+#             except Exception as e:
+#                 print(f"Error initializing tools and workflow: {str(e)}")
+#                 raise ValueError(f"Failed to initialize database tools: {str(e)}")
+
+#         except ImportError as e:
+#             print(f"Database driver import error: {str(e)}")
+#             raise ValueError(f"Missing database driver or invalid database type: {str(e)}")
+#         except ValueError as e:
+#             print(f"Invalid connection string or configuration: {str(e)}")
+#             raise
+#         except Exception as e:
+#             print(f"Unexpected error during database connection: {str(e)}")
+#             raise ValueError(f"Failed to establish database connection: {str(e)}")
+
+#     def initialize_workflow(self):
+#         """Initialize the workflow graph"""
+        
+#         print("Intializing Workflow....")
+
+#         def creating_sql_agent_chain():
+#             """Creating a sql agent chain"""
+
+#             # 4. check_query - Check if the query is correct
+#             # - Query checked: {check_query}
+#             # If query generated but not checked, respond with 'check_query'.
+#             #     If query checked but not executed, respond with 'execute_query'.
+#             print("Creating a sql agent chain")
+#             sql_agent_prompt = ChatPromptTemplate.from_messages([
+#                 ("system", """You are a supervisor SQL agent managing tools to get the answer to the user's query.
+                
+#                 Based on the current state, decide which tool should be called next:
+#                 1. list_table_tools - List all tables from the database
+#                 2. get_schema - Get the schema of required tables
+#                 3. generate_query - Generate a SQL query
+                
+#                 4. execute_query - Execute the query
+#                 5. response - Create response for the user
+
+#                 Current state:
+#                 - Tables listed: {tables_list}
+#                 - Schema retrieved: {schema_of_table}
+#                 - Query generated: {query_gen}
+#                 - Query executed: {execute_query}
+#                 - Response created: {response_to_user}
+                
+#                 If no tables are listed, respond with 'list_table_tools'.
+#                 If tables are listed but no schema, respond with 'get_schema'.
+#                 If schema exists but no query generated, respond with 'generate_query'.
+#                 If query generated but not executed, respond with 'execute_query'.
+#                 If query executed but no response, respond with 'response'.
+#                 If everything is complete, respond with 'DONE'.
+                
+#                 Respond with ONLY the tool name or 'DONE'.
+#                 """),
+#                 ("human", "{task}")
+#             ])
+#             return sql_agent_prompt | self.llm
+
+#         def sql_agent(state: SQLAgentState) -> Dict:
+#             """Agent decides which tool to call next"""
+#             messages = state["messages"]
+#             task = messages[-1].content if messages else "No task"
+            
+#             # Store the original query in state if not already stored
+#             if not state.get("query"):
+#                 state["query"] = task
+            
+#             # Check what's been completed (convert to boolean properly)
+#             tables_list = bool(state.get("tables_list", "").strip())
+#             schema_of_table = bool(state.get("schema_of_table", "").strip())
+#             query_gen = bool(state.get("query_gen", "").strip())
+#             # check_query = bool(state.get("check_query", "").strip())
+#             execute_query = bool(state.get("execute_query", "").strip())
+#             response_to_user = bool(state.get("response_to_user", "").strip())
+            
+#             # print(f"State check - Tables: {tables_list}, Schema: {schema_of_table}, Query: {query_gen}, Check: {check_query}, Execute: {execute_query}, Response: {response_to_user}")
+            
+#             chain = creating_sql_agent_chain()
+#             decision = chain.invoke({
+#                 "task": task,
+#                 "tables_list": tables_list,
+#                 "schema_of_table": schema_of_table,
+#                 "query_gen": query_gen,
+#                 # "check_query": check_query,
+#                 "execute_query": execute_query,
+#                 "response_to_user": response_to_user
+#             })
+#             decision_text = decision.content.strip().lower()
+#             print(f"Agent decision: {decision_text}")
+            
+#             if "done" in decision_text:
+#                 next_tool = "end"
+#                 agent_msg = "✅ SQL Agent: All tasks complete!"
+#             elif "list_table_tools" in decision_text:
+#                 next_tool = "list_table_tools"
+#                 agent_msg = "📋 SQL Agent: Listing all tables in database."
+#             elif "get_schema" in decision_text:
+#                 next_tool = "get_schema"
+#                 agent_msg = "📋 SQL Agent: Getting schema of tables."
+#             elif "generate_query" in decision_text:
+#                 next_tool = "generate_query"
+#                 agent_msg = "📋 SQL Agent: Generating SQL query."
+#             # elif "check_query" in decision_text:
+#             #     next_tool = "check_query"
+#             #     agent_msg = "📋 SQL Agent: Checking SQL query."
+#             elif "execute_query" in decision_text:
+#                 next_tool = "execute_query"
+#                 agent_msg = "📋 SQL Agent: Executing query."
+#             elif "response" in decision_text:
+#                 next_tool = "response"
+#                 agent_msg = "📋 SQL Agent: Creating response."
+#             else:
+#                 next_tool = "end"
+#                 agent_msg = "✅ SQL Agent: Task complete."
+            
+#             return {
+#                 "messages": [AIMessage(content=agent_msg)],
+#                 "next_tool": next_tool,
+#                 "current_task": task
+#             }
+        
+#         def router(state: SQLAgentState):
+#             """Route to the next node"""
+#             print("🔁 Entering router...")
+#             next_tool = state.get("next_tool", "")
+#             print(f"➡️ Next tool: {next_tool}")
+            
+#             if next_tool == "end" or state.get("task_complete", False):
+#                 return END
+            
+#             # valid_tools = [
+#             #     "sql_agent", "list_table_tools", "get_schema", "generate_query",
+#             #     "check_query", "execute_query", "response"
+#             # ]
+#             valid_tools = [
+#                 "sql_agent", "list_table_tools", "get_schema", "generate_query",
+#                 "execute_query", "response"
+#             ]
+            
+#             return next_tool if next_tool in valid_tools else "sql_agent"
+        
+#         # Create workflow
+#         workflow = StateGraph(SQLAgentState)
+
+#         # Add nodes
+#         workflow.add_node("sql_agent", sql_agent)
+#         workflow.add_node("list_table_tools", self.db_tools.list_table_tools)
+#         workflow.add_node("get_schema", self.db_tools.get_schema)
+#         workflow.add_node("generate_query", self.db_tools.generate_query)
+#         # workflow.add_node("check_query", self.db_tools.check_query)
+#         workflow.add_node("execute_query", self.db_tools.execute_query)
+#         workflow.add_node("response", self.db_tools.create_response)
+
+#         # Set entry point
+#         workflow.set_entry_point("sql_agent")
+
+#         # Add routing
+#         # for node in ["sql_agent", "list_table_tools", "get_schema", "generate_query", "check_query", "execute_query", "response"]:
+#         for node in ["sql_agent", "list_table_tools", "get_schema", "generate_query", "execute_query", "response"]:
+#             workflow.add_conditional_edges(
+#                 node,
+#                 router,
+#                 {
+#                     "sql_agent": "sql_agent",
+#                     "list_table_tools": "list_table_tools",
+#                     "get_schema": "get_schema",
+#                     "generate_query": "generate_query",
+#                     # "check_query": "check_query",
+#                     "execute_query": "execute_query",
+#                     "response": "response",
+#                     END: END
+#                 }
+#             )
+
+#         # Compile the graph
+#         self.app = workflow.compile()
+#         # self.app.get_graph().draw_mermaid_png(output_file_path="sql_agent_workflow.png", draw_method=MermaidDrawMethod.API)
+                
+                
+  
+#     def is_query_relevant(self, query: str) -> bool:
+#         """Check if the query is relevant to the database using the LLM."""
+        
+#         # Retrieve the schema of the relevant tables
+#         if self.db_tools.list_tables_tool:
+#             relevant_tables = self.db_tools.list_tables_tool.invoke("")
+#             # print(relevant_tables)
+#         table_list= relevant_tables.split(", ")
+#         print(table_list)
+#         # print(agent.get_schema_tool.invoke(table_list[0]))
+#         schema = ""
+#         for table in table_list:
+#             schema+= self.db_tools.get_schema_tool.invoke(table)
+
+#         print(schema)
+            
+#         # if self.get_schema_tool:
+#         #     schema_response = self.get_schema_tool.invoke({})
+#         #     table_schema = schema_response.content  # Assuming this returns the schema as a string
+
+#         relevance_check_prompt = (
+#             """You are an expert SQL agent which takes user query in Natural language and find out it have releavnce with the given schema or not. Please determine if the following query is related to a database.Here is the schema of the tables present in database:\n{schema}\n\n. If the query related to given schema respond with 'yes'. Here is the query: {query}. Answer with only 'yes' or 'no'."""
+#         ).format(schema=relevant_tables, query=query)
+        
+#         response = self.llm.invoke([{"role": "user", "content": relevance_check_prompt}])
+    
+#     # Assuming the LLM returns a simple 'yes' or 'no'
+#         return response.content == "yes"
+
+#     ## called from the fastapi endpoint
+#     def execute_query(self, query: str):
+#         """Execute a query through the workflow"""
+#         if self.db is None:
+#             raise ValueError("Database connection not established. Please set up the connection first.")
+#         if self.app is None:
+#             raise ValueError("Workflow not initialized. Please set up the connection first.")
+#         # First, handle simple queries like "list tables" directly
+#         query_lower = query.lower()
+#         if any(phrase in query_lower for phrase in ["list all the tables", "show tables", "name of tables",
+#                                                     "which tables are present", "how many tables", "list all tables"]):
+#             if self.db_tools.list_tables_tool:
+#                 tables = self.db_tools.list_tables_tool.invoke("")
+#                 return f"The tables in the database are: {tables}"
+#             else:
+#                 return "Error: Unable to list tables. The list_tables_tool is not initialized."
+
+#         # Check if the query is relevant to the database
+#         if not self.is_query_relevant(query):
+#             print("Not relevent to database.")
+#             # If not relevant, let the LLM answer the question directly
+#             non_relevant_prompt = (
+#                 """You are an expert SQL agent created by Kshitij Kumrawat. You can only assist with questions related to databases so repond the user with the following example resonse and Do not answer any questions that are not related to databases.:  
+#                 Please ask a question that pertains to database operations, such as querying tables, retrieving data, or understanding the database schema. """
+#             )
+    
+#     # Invoke the LLM with the non-relevant prompt
+#             response = self.llm.invoke([{"role": "user", "content": non_relevant_prompt}])
+#             # print(response.content)
+#             return response.content
+        
+#         # If relevant, proceed with the SQL workflow
+#         # response = self.app.invoke({"messages": [HumanMessage(content=query, role="user")]})
+#         response = self.app.invoke({
+#         "messages": [HumanMessage(content=query)],
+#         "query": query
+#         })
+
+#         return response["messages"][-1].content
+
+#         # # More robust final answer extraction
+#         # if (
+#         #     response
+#         #     and response["messages"]
+#         #     and response["messages"][-1].tool_calls
+#         #     and len(response["messages"][-1].tool_calls) > 0
+#         #     and "args" in response["messages"][-1].tool_calls[0]
+#         #     and "final_answer" in response["messages"][-1].tool_calls[0]["args"]
+#         # ):
+#         #     return response["messages"][-1].tool_calls[0]["args"]["final_answer"]
+#         # else:
+#         #     return "Error: Could not extract final answer."
+
 
-from IPython.display import Image, display
-from langchain_core.runnables.graph import MermaidDrawMethod
-from typing import Optional, Dict
 
-from langchain_community.utilities import SQLDatabase
-from langchain_community.agent_toolkits import SQLDatabaseToolkit
 from langchain_groq import ChatGroq
-from langchain_core.messages import HumanMessage, AIMessage
-from langchain_core.prompts import ChatPromptTemplate
-# from langchain_core.pydantic_v1 import BaseModel, Field
-from langgraph.graph import StateGraph, END, MessagesState
-from typing import TypedDict, Annotated, List, Literal, Dict, Any
+from langgraph.graph import StateGraph, END, START, MessagesState
+from langchain_core.messages import AIMessage, ToolMessage, AnyMessage, HumanMessage, SystemMessage
+from dotenv import load_dotenv
+import os
+from IPython.display import display, Image
 from langchain_google_genai import ChatGoogleGenerativeAI
-from app.schemas.agent_state import DBQuery, SQLAgentState
-from app.tools.database_tools import DatabaseTools
+# from app.tools.database_tools import DatabaseTools
 from app.utils.database_connection import DatabaseConnection
-from dotenv import load_dotenv
+from app.tools.database_tools_v2 import DatabaseTools
+
 load_dotenv()
 import os
 os.environ["GROQ_API_KEY"]=os.getenv("GROQ_API_KEY")
 os.environ["GEMINI_API_KEY"]=os.getenv("GEMINI_API_KEY")
 
 
+from langgraph.graph import MessagesState
+from langgraph.prebuilt import tools_condition, ToolNode
+from langgraph.checkpoint.memory import MemorySaver
+
+
 class SQLAgent:
+
+
     def __init__(self):
         
         # Initialize instance variables
         self.db = None
-        self.toolkit = None
-        self.tools = None
-        self.list_tables_tool = None
-        self.sql_db_query = None
-        self.get_schema_tool = None
-        self.app = None
-        
+        # self.repl = PythonREPL()
+        # self.code = None 
+
         # Setting up LLM
-        # self.llm = ChatGroq(model=model,api_key = os.getenv("GROQ_API_KEY"))
-        self.llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash-lite", google_api_key=os.environ["GEMINI_API_KEY"])
+        self.llm = ChatGroq(model="openai/gpt-oss-120b",api_key = os.getenv("GROQ_API_KEY"))
+        # self.llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash-lite", google_api_key=os.environ["GEMINI_API_KEY"])
         # Register the tool method
         # self.query_to_database = self._create_query_tool()
 
@@ -67,13 +387,16 @@ class SQLAgent:
     def setup_database_connection(self, connection_string: str):
         """Set up database connection and initialize tools"""
         try:
-            # Initialize database connection
-            # self.db = SQLDatabase.from_uri(connection_string)
-            # print("Database connection successful!")
+
             self.db = DatabaseConnection(connection_string).db
             print("Database connection successful!")
-            # Initialize toolkit and tools class 
-            self.db_tools = DatabaseTools(db=self.db, llm=self.llm)
+            self.db_tools = DatabaseTools(db=self.db, llm=self.llm)    
+            self.list_tables_tool = self.db_tools.list_tables       
+            self.schema_tool = self.db_tools.get_schema 
+            self.execute_query_tools = self.db_tools.execute_query
+            self.tools_list = [self.list_tables_tool, self.schema_tool, self.execute_query_tools]
+
+
 
             try:
                 self.initialize_workflow()
@@ -93,252 +416,72 @@ class SQLAgent:
         except Exception as e:
             print(f"Unexpected error during database connection: {str(e)}")
             raise ValueError(f"Failed to establish database connection: {str(e)}")
-
-    def initialize_workflow(self):
-        """Initialize the workflow graph"""
         
-        print("Intializing Workflow....")
-
-        def creating_sql_agent_chain():
+    def sql_agent(self, state: MessagesState):
             """Creating a sql agent chain"""
-
-            # 4. check_query - Check if the query is correct
-            # - Query checked: {check_query}
-            # If query generated but not checked, respond with 'check_query'.
-            #     If query checked but not executed, respond with 'execute_query'.
+            
             print("Creating a sql agent chain")
-            sql_agent_prompt = ChatPromptTemplate.from_messages([
-                ("system", """You are a supervisor SQL agent managing tools to get the answer to the user's query.
+            self.llm_with_tools = self.llm.bind_tools(self.tools_list)
+
+            sys_msg = SystemMessage(content = f"""You are a supervisor SQL agent managing tools to get the answer to the user's query.
                 
-                Based on the current state, decide which tool should be called next:
-                1. list_table_tools - List all tables from the database
+                You posses the following tools :
+                1. list_tables - List all tables from the database
                 2. get_schema - Get the schema of required tables
-                3. generate_query - Generate a SQL query
+                3. execute_query - Execute the SQL query
                 
-                4. execute_query - Execute the query
-                5. response - Create response for the user
-
-                Current state:
-                - Tables listed: {tables_list}
-                - Schema retrieved: {schema_of_table}
-                - Query generated: {query_gen}
-                - Query executed: {execute_query}
-                - Response created: {response_to_user}
+                The following are instructions to help you decide which tool to use next:                    
+                - Always breakdown the user query into smaller sub-tasks and decide which tool should be called next to accomplish each sub-task.
+                - Always list down the tables, never assume any table names or believe on users assuming table names because they can be incorrect.
+                - Dont make any schema assumptions, always get the schema using the get_schema tool before generating any query of the required table.
+                - Use the execute_query tool to run the final query and get results.
+                - If a query execution fails, analyze the error message, adjust the query accordingly, and try executing it again.
+                                   
+                Dont do :
+                - Dont go off topic, always stick to the user query.
+                - Dont answer any unwanted queries of user, stick to the database related queries only.
                 
-                If no tables are listed, respond with 'list_table_tools'.
-                If tables are listed but no schema, respond with 'get_schema'.
-                If schema exists but no query generated, respond with 'generate_query'.
-                If query generated but not executed, respond with 'execute_query'.
-                If query executed but no response, respond with 'response'.
-                If everything is complete, respond with 'DONE'.
-                
-                Respond with ONLY the tool name or 'DONE'.
-                """),
-                ("human", "{task}")
-            ])
-            return sql_agent_prompt | self.llm
-
-        def sql_agent(state: SQLAgentState) -> Dict:
-            """Agent decides which tool to call next"""
-            messages = state["messages"]
-            task = messages[-1].content if messages else "No task"
-            
-            # Store the original query in state if not already stored
-            if not state.get("query"):
-                state["query"] = task
-            
-            # Check what's been completed (convert to boolean properly)
-            tables_list = bool(state.get("tables_list", "").strip())
-            schema_of_table = bool(state.get("schema_of_table", "").strip())
-            query_gen = bool(state.get("query_gen", "").strip())
-            # check_query = bool(state.get("check_query", "").strip())
-            execute_query = bool(state.get("execute_query", "").strip())
-            response_to_user = bool(state.get("response_to_user", "").strip())
-            
-            # print(f"State check - Tables: {tables_list}, Schema: {schema_of_table}, Query: {query_gen}, Check: {check_query}, Execute: {execute_query}, Response: {response_to_user}")
+               """)
             
-            chain = creating_sql_agent_chain()
-            decision = chain.invoke({
-                "task": task,
-                "tables_list": tables_list,
-                "schema_of_table": schema_of_table,
-                "query_gen": query_gen,
-                # "check_query": check_query,
-                "execute_query": execute_query,
-                "response_to_user": response_to_user
-            })
-            decision_text = decision.content.strip().lower()
-            print(f"Agent decision: {decision_text}")
-            
-            if "done" in decision_text:
-                next_tool = "end"
-                agent_msg = "✅ SQL Agent: All tasks complete!"
-            elif "list_table_tools" in decision_text:
-                next_tool = "list_table_tools"
-                agent_msg = "📋 SQL Agent: Listing all tables in database."
-            elif "get_schema" in decision_text:
-                next_tool = "get_schema"
-                agent_msg = "📋 SQL Agent: Getting schema of tables."
-            elif "generate_query" in decision_text:
-                next_tool = "generate_query"
-                agent_msg = "📋 SQL Agent: Generating SQL query."
-            # elif "check_query" in decision_text:
-            #     next_tool = "check_query"
-            #     agent_msg = "📋 SQL Agent: Checking SQL query."
-            elif "execute_query" in decision_text:
-                next_tool = "execute_query"
-                agent_msg = "📋 SQL Agent: Executing query."
-            elif "response" in decision_text:
-                next_tool = "response"
-                agent_msg = "📋 SQL Agent: Creating response."
-            else:
-                next_tool = "end"
-                agent_msg = "✅ SQL Agent: Task complete."
-            
-            return {
-                "messages": [AIMessage(content=agent_msg)],
-                "next_tool": next_tool,
-                "current_task": task
-            }
+            return {"messages": [self.llm_with_tools.invoke([sys_msg] + state["messages"])]}    
+    
+    def initialize_workflow(self):
+        """Initialize the workflow graph"""
         
-        def router(state: SQLAgentState):
-            """Route to the next node"""
-            print("🔁 Entering router...")
-            next_tool = state.get("next_tool", "")
-            print(f"➡️ Next tool: {next_tool}")
-            
-            if next_tool == "end" or state.get("task_complete", False):
-                return END
-            
-            # valid_tools = [
-            #     "sql_agent", "list_table_tools", "get_schema", "generate_query",
-            #     "check_query", "execute_query", "response"
-            # ]
-            valid_tools = [
-                "sql_agent", "list_table_tools", "get_schema", "generate_query",
-                "execute_query", "response"
-            ]
-            
-            return next_tool if next_tool in valid_tools else "sql_agent"
+        memory = MemorySaver()
         
+        print("Intializing Workflow....")
         # Create workflow
-        workflow = StateGraph(SQLAgentState)
+        workflow = StateGraph(MessagesState)
 
         # Add nodes
-        workflow.add_node("sql_agent", sql_agent)
-        workflow.add_node("list_table_tools", self.db_tools.list_table_tools)
-        workflow.add_node("get_schema", self.db_tools.get_schema)
-        workflow.add_node("generate_query", self.db_tools.generate_query)
-        # workflow.add_node("check_query", self.db_tools.check_query)
-        workflow.add_node("execute_query", self.db_tools.execute_query)
-        workflow.add_node("response", self.db_tools.create_response)
+        workflow.add_node("sql_agent", self.sql_agent)
+        workflow.add_node("tools", ToolNode(tools=self.tools_list))
 
         # Set entry point
-        workflow.set_entry_point("sql_agent")
-
-        # Add routing
-        # for node in ["sql_agent", "list_table_tools", "get_schema", "generate_query", "check_query", "execute_query", "response"]:
-        for node in ["sql_agent", "list_table_tools", "get_schema", "generate_query", "execute_query", "response"]:
-            workflow.add_conditional_edges(
-                node,
-                router,
-                {
-                    "sql_agent": "sql_agent",
-                    "list_table_tools": "list_table_tools",
-                    "get_schema": "get_schema",
-                    "generate_query": "generate_query",
-                    # "check_query": "check_query",
-                    "execute_query": "execute_query",
-                    "response": "response",
-                    END: END
-                }
-            )
+        workflow.add_edge(START, "sql_agent")
+        workflow.add_conditional_edges(
+            "sql_agent",
+            # If the latest message (result) from assistant is a tool call -> tools_condition routes to tools
+            # If the latest message (result) from assistant is a not a tool call -> tools_condition routes to END
+            tools_condition,
+        )
 
+        workflow.add_edge("tools", "sql_agent")
         # Compile the graph
-        self.app = workflow.compile()
-        # self.app.get_graph().draw_mermaid_png(output_file_path="sql_agent_workflow.png", draw_method=MermaidDrawMethod.API)
+        self.app = workflow.compile(checkpointer = memory)
+        # display(Image(self.app.get_graph(xray=True).draw_mermaid_png()))                
                 
-                
-  
-    def is_query_relevant(self, query: str) -> bool:
-        """Check if the query is relevant to the database using the LLM."""
-        
-        # Retrieve the schema of the relevant tables
-        if self.db_tools.list_tables_tool:
-            relevant_tables = self.db_tools.list_tables_tool.invoke("")
-            # print(relevant_tables)
-        table_list= relevant_tables.split(", ")
-        print(table_list)
-        # print(agent.get_schema_tool.invoke(table_list[0]))
-        schema = ""
-        for table in table_list:
-            schema+= self.db_tools.get_schema_tool.invoke(table)
-
-        print(schema)
-            
-        # if self.get_schema_tool:
-        #     schema_response = self.get_schema_tool.invoke({})
-        #     table_schema = schema_response.content  # Assuming this returns the schema as a string
-
-        relevance_check_prompt = (
-            """You are an expert SQL agent which takes user query in Natural language and find out it have releavnce with the given schema or not. Please determine if the following query is related to a database.Here is the schema of the tables present in database:\n{schema}\n\n. If the query related to given schema respond with 'yes'. Here is the query: {query}. Answer with only 'yes' or 'no'."""
-        ).format(schema=relevant_tables, query=query)
-        
-        response = self.llm.invoke([{"role": "user", "content": relevance_check_prompt}])
-    
-    # Assuming the LLM returns a simple 'yes' or 'no'
-        return response.content == "yes"
-
     ## called from the fastapi endpoint
-    def execute_query(self, query: str):
+    def execute_query(self, query: str, config: dict):
         """Execute a query through the workflow"""
         if self.db is None:
             raise ValueError("Database connection not established. Please set up the connection first.")
         if self.app is None:
             raise ValueError("Workflow not initialized. Please set up the connection first.")
-        # First, handle simple queries like "list tables" directly
-        query_lower = query.lower()
-        if any(phrase in query_lower for phrase in ["list all the tables", "show tables", "name of tables",
-                                                    "which tables are present", "how many tables", "list all tables"]):
-            if self.db_tools.list_tables_tool:
-                tables = self.db_tools.list_tables_tool.invoke("")
-                return f"The tables in the database are: {tables}"
-            else:
-                return "Error: Unable to list tables. The list_tables_tool is not initialized."
-
-        # Check if the query is relevant to the database
-        if not self.is_query_relevant(query):
-            print("Not relevent to database.")
-            # If not relevant, let the LLM answer the question directly
-            non_relevant_prompt = (
-                """You are an expert SQL agent created by Kshitij Kumrawat. You can only assist with questions related to databases so repond the user with the following example resonse and Do not answer any questions that are not related to databases.:  
-                Please ask a question that pertains to database operations, such as querying tables, retrieving data, or understanding the database schema. """
-            )
-    
-    # Invoke the LLM with the non-relevant prompt
-            response = self.llm.invoke([{"role": "user", "content": non_relevant_prompt}])
-            # print(response.content)
-            return response.content
-        
-        # If relevant, proceed with the SQL workflow
-        # response = self.app.invoke({"messages": [HumanMessage(content=query, role="user")]})
+
         response = self.app.invoke({
-        "messages": [HumanMessage(content=query)],
-        "query": query
-        })
+        "messages": [HumanMessage(content=query)]
+        }, config=config)
 
         return response["messages"][-1].content
-
-        # # More robust final answer extraction
-        # if (
-        #     response
-        #     and response["messages"]
-        #     and response["messages"][-1].tool_calls
-        #     and len(response["messages"][-1].tool_calls) > 0
-        #     and "args" in response["messages"][-1].tool_calls[0]
-        #     and "final_answer" in response["messages"][-1].tool_calls[0]["args"]
-        # ):
-        #     return response["messages"][-1].tool_calls[0]["args"]["final_answer"]
-        # else:
-        #     return "Error: Could not extract final answer."
-

app/api/v1/endpoints/sql_query.py

@@ -1,37 +1,67 @@
+# # from fastapi import APIRouter, HTTPException
+# # from app.models import SQLQueryRequest, SQLQueryResponse
+# # from app.services.sql_agent import execute_query
+
+# # router = APIRouter()
+
+# # @router.post("/query", response_model=SQLQueryResponse)
+# # async def query_database(request: SQLQueryRequest):
+# #     try:
+# #         result = execute_query(request.query)
+# #         return SQLQueryResponse(result=result)
+# #     except ValueError as e:
+# #         raise HTTPException(status_code=400, detail=str(e))
+# #     except Exception as e:
+# #         raise HTTPException(status_code=500, detail=str(e))
+
+# # app/api/v1/endpoints/sql_query.py
 # from fastapi import APIRouter, HTTPException
-# from app.models import SQLQueryRequest, SQLQueryResponse
-# from app.services.sql_agent import execute_query
+# from pydantic import BaseModel
+# from app.services.sql_agent_instance import sql_agent
 
 # router = APIRouter()
 
+# class SQLQueryRequest(BaseModel):
+#     query: str
+
+# class SQLQueryResponse(BaseModel):
+#     result: str
+
 # @router.post("/query", response_model=SQLQueryResponse)
 # async def query_database(request: SQLQueryRequest):
 #     try:
-#         result = execute_query(request.query)
+#         result = sql_agent.execute_query(request.query)
 #         return SQLQueryResponse(result=result)
 #     except ValueError as e:
 #         raise HTTPException(status_code=400, detail=str(e))
 #     except Exception as e:
 #         raise HTTPException(status_code=500, detail=str(e))
 
-# app/api/v1/endpoints/sql_query.py
 from fastapi import APIRouter, HTTPException
 from pydantic import BaseModel
 from app.services.sql_agent_instance import sql_agent
-
+from typing import Optional
+import uuid
 router = APIRouter()
 
 class SQLQueryRequest(BaseModel):
     query: str
+    thread_id: Optional[str] = None
 
 class SQLQueryResponse(BaseModel):
     result: str
+    thread_id: str ## client can use this to continue the conversation
 
 @router.post("/query", response_model=SQLQueryResponse)
 async def query_database(request: SQLQueryRequest):
     try:
-        result = sql_agent.execute_query(request.query)
-        return SQLQueryResponse(result=result)
+        ## generate if not provided thread id 
+        thread_id = request.thread_id or str(uuid.uuid4())
+        ## add debug 
+        print(f"Thread ID: {thread_id}, Query: {request.query}")
+        result = sql_agent.execute_query(request.query, config={"configurable": {"thread_id": thread_id}})
+        print(f"Result: {result}")
+        return SQLQueryResponse(result=result, thread_id=thread_id)
     except ValueError as e:
         raise HTTPException(status_code=400, detail=str(e))
     except Exception as e:

app/tools/database_tools.py

@@ -1,9 +1,245 @@
+# ## creating database tools 
+# from langchain_core.tools import tool
+# from app.schemas.agent_state import SQLAgentState
+# from typing import Dict
+# from langchain_core.messages import AIMessage, HumanMessage
+# from app.utils.database_connection import DatabaseConnection
+# from langchain_community.agent_toolkits import SQLDatabaseToolkit
+# from app.schemas.agent_state import DBQuery
+# from langchain_core.prompts import ChatPromptTemplate
+
+# class DatabaseTools:
+#     def __init__(self,db = None, llm = None):
+#         self.db = db 
+#         self.llm = llm
+#         self._create_query_tool = self._create_query_tool()
+#         try:
+#                 # Initialize toolkit and tools
+#                 self.toolkit = SQLDatabaseToolkit(db=self.db, llm=self.llm)
+#                 self.tools = self.toolkit.get_tools()
+#                 for tool in self.tools:
+#                     print(f"Initialized tool: {tool.name}")
+
+#                 # Create instances of the tools
+#                 self.list_tables_tool = next((tool for tool in self.tools if tool.name == "sql_db_list_tables"), None)
+#                 self.query_tool = next((tool for tool in self.tools if tool.name == "sql_db_query"), None)
+#                 self.get_schema_tool = next((tool for tool in self.tools if tool.name == "sql_db_schema"), None)
+#                 self.query_checker_tool = next((tool for tool in self.tools if tool.name == "sql_db_query_checker"), None)
+#                 if not all([self.list_tables_tool, self.query_tool, self.get_schema_tool, self.query_checker_tool]):
+#                     raise ValueError("Failed to initialize one or more required database tools")
+
+#                 # # Initialize workflow and compile it into an app
+#                 # self.initialize_workflow()
+                
+#         except Exception as e:
+#             print(f"Error initializing tools and workflow: {str(e)}")
+#             raise ValueError(f"Failed to initialize database tools: {str(e)}")
+    
+#     def _create_query_tool(self):
+#             """Create the query tool bound to this instance"""
+#             print("creating _create_query_tool")
+#             @tool
+#             def query_to_database(query: str) -> str:
+#                 """
+#                 Execute a SQL query against the database and return the result.
+#                 If the query is invalid or returns no result, an error message will be returned.
+#                 In case of an error, the user is advised to rewrite the query and try again.
+#                 """
+#                 if self.db is None:
+#                     return "Error: Database connection not established. Please set up the connection first."
+#                 result = self.db.run_no_throw(query)
+#                 if not result:
+#                     return "Error: Query failed. Please rewrite your query and try again."
+#                 return result
+            
+#             return query_to_database
+
+#     def list_table_tools(self, state: SQLAgentState = None) -> Dict:
+#             """List all the tables"""
+#             tables_list = self.list_tables_tool.invoke("")
+#             print(f"Tables found: {tables_list}")
+#             return {
+#                 "messages": [AIMessage(content=f"Tables found: {tables_list}")],
+#                 "tables_list": tables_list,
+#                 "next_tool": "sql_agent"
+#             }
+    
+#     def get_schema(self,state: SQLAgentState) -> Dict:
+#             """Get the schema of required tables"""
+#             print("📘 Getting schema...")
+#             tables_list = state.get("tables_list", "")
+#             if not tables_list:
+#                 tables_list = self.list_tables_tool.invoke("")
+            
+#             tables = [table.strip() for table in tables_list.split(",")]
+#             full_schema = ""
+            
+#             for table in tables:
+#                 try:
+#                     schema = self.get_schema_tool.invoke(table)
+#                     full_schema += f"\nTable: {table}\n{schema}\n"
+#                 except Exception as e:
+#                     print(f"Error getting schema for {table}: {e}")
+            
+#             print(f"📘 Schema collected for tables: {tables}")
+#             return {
+#                 "messages": [AIMessage(content=f"Schema retrieved: {full_schema}")],
+#                 "schema_of_table": full_schema,
+#                 "tables_list": tables_list,
+#                 "next_tool": "sql_agent"
+#             }
+#     def generate_query(self, state: SQLAgentState) -> Dict:
+#             """Generate a SQL Query according to the user query"""
+#             schema = state.get("schema_of_table", "")
+#             human_query = state.get("query", "")
+#             tables = state.get("tables_list", "")
+            
+#             print(f"Generating query for: {human_query}")
+            
+#             generate_query_system_prompt = """You are a SQL expert that generates precise SQL queries based on user questions.
+            
+#             You will be provided with:
+#             - User's question
+#             - Available tables
+#             - Complete schema information
+            
+#             Generate a SQL query that:
+#             - Uses correct column names from schema
+#             - Properly joins tables if needed
+#             - Includes appropriate WHERE clauses
+#             - Uses proper aggregation functions when needed
+            
+#             Respond ONLY with the SQL query. Do not explain."""
+            
+#             combined_input = f"""
+#             User Question: {human_query}
+#             Tables: {tables}
+#             Schema: {schema}
+#             """
+            
+#             generate_query_prompt = ChatPromptTemplate.from_messages([
+#                 ("system", generate_query_system_prompt),
+#                 ("human", "{input}")
+#             ])
+            
+#             try:
+#                 formatted_prompt = generate_query_prompt.invoke({"input": combined_input})
+#                 generate_query_llm = self.llm.with_structured_output(DBQuery)
+#                 result = generate_query_llm.invoke(formatted_prompt)
+                
+#                 print(f"✅ Query generated: {result.query}")
+#                 return {
+#                     "messages": [AIMessage(content=f"Query generated: {result.query}")],
+#                     "query_gen": result.query,
+#                     "next_tool": "sql_agent"
+#                 }
+#             except Exception as e:
+#                 print(f"❌ Failed to generate query: {e}")
+#                 return {
+#                     "messages": [AIMessage(content="⚠️ Failed to generate SQL query.")],
+#                     "query_gen": "",
+#                     "next_tool": "sql_agent"
+#                 }
+            
+#     def check_query(self,state: SQLAgentState) -> Dict:
+#             """Check if the query is correct"""
+#             query = state.get("query_gen", "")
+#             print(f"Checking query: {query}")
+            
+#             if not query:
+#                 return {
+#                     "messages": [AIMessage(content="No query to check")],
+#                     "check_query": "",
+#                     "next_tool": "sql_agent"
+#                 }
+            
+#             try:
+#                 checked_query = self.query_checker_tool.invoke(query)
+#                 ## if checked query contains ``` anywhere remove it 
+#                 if "```" in checked_query:
+#                     checked_query = checked_query.replace("```", "")
+#                 print(f"Query checked: {checked_query}")
+#                 return {
+#                     "messages": [AIMessage(content=f"Query checked: {checked_query}")],
+#                     "check_query": checked_query if checked_query else query,
+#                     "next_tool": "sql_agent"
+#                 }
+#             except Exception as e:
+#                 print(f"Error checking query: {e}")
+#                 return {
+#                     "messages": [AIMessage(content="Query check failed, using original query")],
+#                     "check_query": query,
+#                     "next_tool": "sql_agent"
+#                 }
+            
+#     def execute_query(self,state: SQLAgentState) -> Dict:
+#             """Execute the SQL query"""
+#             query = state.get("check_query", "") or state.get("query_gen", "")
+#             print(f"Executing query: {query}")
+            
+#             if not query:
+#                 return {
+#                     "messages": [AIMessage(content="No query to execute")],
+#                     "execute_query": "",
+#                     "next_tool": "sql_agent"
+#                 }
+            
+#             try:
+#                 results = self.query_tool.invoke(query)
+#                 print(f"Query results: {results}")
+#                 return {
+#                     "messages": [AIMessage(content=f"Query executed successfully: {results}")],
+#                     "execute_query": results,
+#                     "next_tool": "sql_agent"
+#                 }
+#             except Exception as e:
+#                 print(f"Error executing query: {e}")
+#                 return {
+#                     "messages": [AIMessage(content=f"Query execution failed: {e}")],
+#                     "execute_query": "",
+#                     "next_tool": "sql_agent"
+#                 }
+#     def create_response(self,state: SQLAgentState) -> Dict:
+#             """Create a final response for the user"""
+#             print("Creating final response...")
+            
+#             query = state.get("check_query", "") or state.get("query_gen", "")
+#             result = state.get("execute_query", "")
+#             human_query = state.get("query", "")
+            
+#             response_prompt = f"""Create a clear, concise response for the user based on:
+            
+#             User Question: {human_query}
+#             SQL Query: {query}
+#             Query Result: {result}
+            
+#             Provide a natural language answer that directly addresses the user's question. Make sure to provide only answer to human question, no any internal process results and explaination, just answer related to the human query."""
+            
+#             try:
+#                 response = self.llm.invoke([HumanMessage(content=response_prompt)])
+#                 print(f"Response created: {response.content}")
+                
+#                 return {
+#                     "messages": [response],
+#                     "response_to_user": response.content,
+#                     "next_tool": "sql_agent",
+#                     "task_complete": True
+#                 }
+#             except Exception as e:
+#                 print(f"Error creating response: {e}")
+#                 return {
+#                     "messages": [AIMessage(content="Failed to create response")],
+#                     "response_to_user": "",
+#                     "next_tool": "sql_agent",
+#                     "task_complete": True
+#                 }
+
+
 ## creating database tools 
 from langchain_core.tools import tool
 from app.schemas.agent_state import SQLAgentState
 from typing import Dict
-from langchain_core.messages import AIMessage, HumanMessage
-from app.utils.database_connection import DatabaseConnection
+from langchain_core.messages import AIMessage
 from langchain_community.agent_toolkits import SQLDatabaseToolkit
 from app.schemas.agent_state import DBQuery
 from langchain_core.prompts import ChatPromptTemplate
@@ -12,7 +248,8 @@ class DatabaseTools:
     def __init__(self,db = None, llm = None):
         self.db = db 
         self.llm = llm
-        self._create_query_tool = self._create_query_tool()
+        # self._create_query_tool = self._create_query_tool()
+        self.tools = self.get_all_tools()
         try:
                 # Initialize toolkit and tools
                 self.toolkit = SQLDatabaseToolkit(db=self.db, llm=self.llm)
@@ -34,60 +271,51 @@ class DatabaseTools:
         except Exception as e:
             print(f"Error initializing tools and workflow: {str(e)}")
             raise ValueError(f"Failed to initialize database tools: {str(e)}")
-    
-    def _create_query_tool(self):
-            """Create the query tool bound to this instance"""
-            print("creating _create_query_tool")
-            @tool
-            def query_to_database(query: str) -> str:
-                """
-                Execute a SQL query against the database and return the result.
-                If the query is invalid or returns no result, an error message will be returned.
-                In case of an error, the user is advised to rewrite the query and try again.
-                """
-                if self.db is None:
-                    return "Error: Database connection not established. Please set up the connection first."
-                result = self.db.run_no_throw(query)
-                if not result:
-                    return "Error: Query failed. Please rewrite your query and try again."
-                return result
-            
-            return query_to_database
-
-    def list_table_tools(self, state: SQLAgentState = None) -> Dict:
+    # @tool
+    # def _create_query_tool(self):
+    #         """Create the query tool bound to this instance"""
+    #         print("creating _create_query_tool")
+    #         @tool
+    #         def query_to_database(query: str) -> str:
+    #             """
+    #             Execute a SQL query against the database and return the result.
+    #             If the query is invalid or returns no result, an error message will be returned.
+    #             In case of an error, the user is advised to rewrite the query and try again.
+    #             """
+    #             if self.db is None:
+    #                 return "Error: Database connection not established. Please set up the connection first."
+    #             result = self.db.run_no_throw(query)
+    #             if not result:
+    #                 return "Error: Query failed. Please rewrite your query and try again."
+    #             return result
+            
+    #         return query_to_database
+    def list_tables(self) -> Dict:
             """List all the tables"""
             tables_list = self.list_tables_tool.invoke("")
             print(f"Tables found: {tables_list}")
-            return {
-                "messages": [AIMessage(content=f"Tables found: {tables_list}")],
-                "tables_list": tables_list,
-                "next_tool": "sql_agent"
-            }
+            return tables_list
     
-    def get_schema(self,state: SQLAgentState) -> Dict:
+    def get_schema(self, table_name: list[str]) -> Dict:
             """Get the schema of required tables"""
             print("📘 Getting schema...")
-            tables_list = state.get("tables_list", "")
-            if not tables_list:
-                tables_list = self.list_tables_tool.invoke("")
+            tables_list = self.list_tables_tool.invoke("")
+            if any(table not in tables_list for table in table_name):
+                 return "Table not exits in database"
             
             tables = [table.strip() for table in tables_list.split(",")]
-            full_schema = ""
+            required_schema = ""
             
             for table in tables:
                 try:
                     schema = self.get_schema_tool.invoke(table)
-                    full_schema += f"\nTable: {table}\n{schema}\n"
+                    required_schema += f"\nTable: {table}\n{schema}\n"
                 except Exception as e:
                     print(f"Error getting schema for {table}: {e}")
             
-            print(f"📘 Schema collected for tables: {tables}")
-            return {
-                "messages": [AIMessage(content=f"Schema retrieved: {full_schema}")],
-                "schema_of_table": full_schema,
-                "tables_list": tables_list,
-                "next_tool": "sql_agent"
-            }
+            return required_schema
+    
+
     def generate_query(self, state: SQLAgentState) -> Dict:
             """Generate a SQL Query according to the user query"""
             schema = state.get("schema_of_table", "")
@@ -141,95 +369,25 @@ class DatabaseTools:
                     "next_tool": "sql_agent"
                 }
             
-    def check_query(self,state: SQLAgentState) -> Dict:
-            """Check if the query is correct"""
-            query = state.get("query_gen", "")
-            print(f"Checking query: {query}")
-            
-            if not query:
-                return {
-                    "messages": [AIMessage(content="No query to check")],
-                    "check_query": "",
-                    "next_tool": "sql_agent"
-                }
-            
-            try:
-                checked_query = self.query_checker_tool.invoke(query)
-                ## if checked query contains ``` anywhere remove it 
-                if "```" in checked_query:
-                    checked_query = checked_query.replace("```", "")
-                print(f"Query checked: {checked_query}")
-                return {
-                    "messages": [AIMessage(content=f"Query checked: {checked_query}")],
-                    "check_query": checked_query if checked_query else query,
-                    "next_tool": "sql_agent"
-                }
-            except Exception as e:
-                print(f"Error checking query: {e}")
-                return {
-                    "messages": [AIMessage(content="Query check failed, using original query")],
-                    "check_query": query,
-                    "next_tool": "sql_agent"
-                }
-            
-    def execute_query(self,state: SQLAgentState) -> Dict:
-            """Execute the SQL query"""
-            query = state.get("check_query", "") or state.get("query_gen", "")
-            print(f"Executing query: {query}")
+    
+    def execute_query(self,query: str) -> Dict:
+            """Execute the SQL query
             
-            if not query:
-                return {
-                    "messages": [AIMessage(content="No query to execute")],
-                    "execute_query": "",
-                    "next_tool": "sql_agent"
-                }
+            Arguments:
+            query -- The SQL query to execute
+
+            returns:
+            execution results
+            """
             
             try:
                 results = self.query_tool.invoke(query)
                 print(f"Query results: {results}")
-                return {
-                    "messages": [AIMessage(content=f"Query executed successfully: {results}")],
-                    "execute_query": results,
-                    "next_tool": "sql_agent"
-                }
+                return results
             except Exception as e:
                 print(f"Error executing query: {e}")
-                return {
-                    "messages": [AIMessage(content=f"Query execution failed: {e}")],
-                    "execute_query": "",
-                    "next_tool": "sql_agent"
-                }
-    def create_response(self,state: SQLAgentState) -> Dict:
-            """Create a final response for the user"""
-            print("Creating final response...")
-            
-            query = state.get("check_query", "") or state.get("query_gen", "")
-            result = state.get("execute_query", "")
-            human_query = state.get("query", "")
-            
-            response_prompt = f"""Create a clear, concise response for the user based on:
+                return "Query execution failed."
             
-            User Question: {human_query}
-            SQL Query: {query}
-            Query Result: {result}
-            
-            Provide a natural language answer that directly addresses the user's question. Make sure to provide only answer to human question, no any internal process results and explaination, just answer related to the human query."""
-            
-            try:
-                response = self.llm.invoke([HumanMessage(content=response_prompt)])
-                print(f"Response created: {response.content}")
-                
-                return {
-                    "messages": [response],
-                    "response_to_user": response.content,
-                    "next_tool": "sql_agent",
-                    "task_complete": True
-                }
-            except Exception as e:
-                print(f"Error creating response: {e}")
-                return {
-                    "messages": [AIMessage(content="Failed to create response")],
-                    "response_to_user": "",
-                    "next_tool": "sql_agent",
-                    "task_complete": True
-                }
+    def get_all_tools(self):
+         return [self.list_tables, self.get_schema, self.execute_query]
+