How to log Text2SQL/DBAgent queries? #23700
Unanswered
Ayhm-Tenderd
asked this question in
Q&A
Replies: 1 comment 10 replies
-
|
Hello, @Ayhm-Tenderd! I'm here to help you with any bugs, questions, or contributions you might have. Let's dive into the issue you've posted. To log Text2SQL/DBAgent queries as variables so that you can use your custom logger to log them to PostgreSQL, you can utilize the Here's an example of how you can achieve this:
Below is a sample implementation: from langchain.callbacks.base import BaseCallbackManager
from langchain_openai import ChatOpenAI
from langchain_community.agent_toolkits import create_sql_agent
from langchain_community.utilities import SQLDatabase
class CustomCallbackManager(BaseCallbackManager):
def on_query(self, query: str, **kwargs):
# Log the query to PostgreSQL
log_to_postgresql(query)
def log_to_postgresql(query: str):
import psycopg2
conn = psycopg2.connect(
dbname="your_dbname",
user="your_user",
password="your_password",
host="your_host",
port="your_port"
)
cur = conn.cursor()
cur.execute("INSERT INTO query_logs (query) VALUES (%s)", (query,))
conn.commit()
cur.close()
conn.close()
# Create the SQL agent with the custom callback manager
db = SQLDatabase.from_uri("sqlite:///Chinook.db")
llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
callback_manager = CustomCallbackManager()
agent_executor = create_sql_agent(llm, db=db, agent_type="tool-calling", callback_manager=callback_manager, verbose=True)In this example:
This setup ensures that every query executed by the agent is logged to your PostgreSQL database [1]. |
Beta Was this translation helpful? Give feedback.
10 replies
-
|
To ensure that the callback methods are invoked correctly and the queries are logged, you need to make sure that the from langchain_openai import ChatOpenAI
from langchain_community.agent_toolkits import create_sql_agent
from langchain.callbacks.base import BaseCallbackManager
from app.services.bot.agents.base_agent import BaseAgent
from app.services.bot.prompt_service import get_prompt_from_state
from app.services.logging.logging_service import LoggingService
from langchain_core.messages import AIMessage
from langchain_core.callbacks.manager import CallbackManagerForChainRun
logger_db = LoggingService()
class CustomCallbackManager(BaseCallbackManager):
def __init__(self, handlers=None):
super().__init__(handlers or [])
def on_retriever_start(self, serialized, query, **kwargs):
print('RETRIEVER START\n\n\n\n\n\n\n\n\n\n\n\n')
try:
# Log the query to PostgreSQL
logger_db.log_info(f'Retriever Start Query: {query}')
except Exception as e:
print(f"Error logging query to PostgreSQL: {e}")
def on_retriever_end(self, documents, **kwargs):
try:
# Log the end of the query retrieval
logger_db.log_info(f'Retriever End: Retrieved {len(documents)} documents')
except Exception as e:
print(f"Error logging query to PostgreSQL: {e}")
def on_retriever_error(self, error, **kwargs):
try:
# Log the error
logger_db.log_info(f'Retriever Error: {error}')
except Exception as e:
print(f"Error logging query to PostgreSQL: {e}")
def on_agent_action(self, action, **kwargs):
"""Run on agent action. If the tool being used is sql_db_query,
it means we're submitting the SQL and we can record it as the final SQL."""
if action.tool == "sql_db_query":
try:
logger_db.log_info(f'Query: {str(action)}')
except Exception as e:
print(f"Error logging query to PostgreSQL: {e}")
class DatabaseAgent(BaseAgent):
def __init__(self, llm: ChatOpenAI, db_connection, agent_type, verbose):
super().__init__(llm)
self.db_connection = db_connection
self.agent_type = agent_type
self.verbose = verbose
def create_agent(self):
return create_sql_agent(
self.llm,
db=self.db_connection,
agent_type=self.agent_type,
verbose=self.verbose,
callback_manager=self.get_query_logging_callback(),
agent_executor_kwargs={"return_intermediate_steps": True}
)
def create_node(self, state, agent, name, prompt):
user_input, chat_history = get_prompt_from_state(state)
print('db_agent_student', name, user_input)
print('================================================')
print(prompt.format(chat_history=chat_history, input=user_input, top_k=5))
print('================================================')
formatted_prompt = prompt.format(chat_history=chat_history, input=user_input, top_k=5)
result = agent.invoke(formatted_prompt)
response = {"messages": [AIMessage(content=result["output"], name=name)]}
return response
def get_query_logging_callback(self):
return CustomCallbackManager()
# Example usage
llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
db_connection = SQLDatabase.from_uri("sqlite:///Chinook.db")
agent = DatabaseAgent(llm, db_connection, agent_type="tool-calling", verbose=True)
agent_executor = agent.create_agent()Key Changes:
This setup should ensure that your queries are logged correctly to PostgreSQL. If the methods are still not being called, you may need to verify that the |
Beta Was this translation helpful? Give feedback.
-
|
@dosu Unfortunately that still does not work,can you start over. workflow_service.py: the filename of the db agent code is database_agent.py |
Beta Was this translation helpful? Give feedback.
-
|
To ensure that the queries executed by the SQL agent are logged correctly, you can extend the query chain with a validation step that checks for common mistakes. This approach ensures that the queries are checked for common mistakes before execution, thereby logging the validated queries correctly. Here is an example of how to do this:
|
Beta Was this translation helpful? Give feedback.
-
|
Not working, @dosu from langchain_openai import ChatOpenAI
from langchain_community.agent_toolkits import create_sql_agent
from app.services.bot.agents.base_agent import BaseAgent
from app.services.bot.prompt_service import get_prompt_from_state
from langchain_core.messages import AIMessage
class DatabaseAgent(BaseAgent):
def __init__(self, llm: ChatOpenAI, db_connection, agent_type, verbose):
super().__init__(llm)
self.db_connection = db_connection
self.agent_type = agent_type
self.verbose = verbose
def create_agent(self):
return create_sql_agent(self.llm, db=self.db_connection, agent_type=self.agent_type,
verbose=self.verbosez)
def create_node(self, state, agent, name, prompt):
user_input, chat_history = get_prompt_from_state(state)
print('db_agent_student', name, user_input)
print('================================================')
print(prompt.format(chat_history=chat_history, input=user_input, top_k=5))
print('================================================')
formatted_prompt = prompt.format(chat_history=chat_history,
input=user_input, top_k=5)
result = agent.invoke(formatted_prompt)
response = {"messages": [AIMessage(content=result["output"], name=name)]}
return response |
Beta Was this translation helpful? Give feedback.
-
|
To add a working callback to log the SQL queries executed by the
|
Beta Was this translation helpful? Give feedback.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
-
Checked other resources
Commit to Help
Example Code
NoneDescription
I am aware that you can make the agent log queries to console, but I want to somehow access these queries as a variable separately so that I can use my custom logger to log them to PostgreSQL
System Info
fastapi==0.110.2
pydantic==2.0.2
pydantic-settings==2.0.3
langchain==0.2.0
langchain-core==0.2.0
langchain-openai==0.1.7
langchain-community==0.2.0
faiss-cpu==1.8.0
python-dotenv==1.0.1
uvicorn==0.23.1
psycopg2-binary==2.9.9
docx2txt==0.8
pandas==2.2.2
langgraph==0.0.55
openai==1.33.0
slack-sdk==3.30.0
SQLAlchemy==2.0.29
Beta Was this translation helpful? Give feedback.
All reactions