My SQL langchain agent seems to be limiting its answers to 10 rows

Question

My data is in MySQL in the form of a database, containing two tables, 1000 rows and 50 rows. I have created a SQL agent:

agent = create_sql_agent(   
llm=llm,
toolkit=toolkit,
verbose=True,
agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
)

Then, I pass in a query, e.g. List all the trials conducted in Berlin, Germany.

query="List all the trials conducted in Berlin, Germany."
response=agent.run(query)

However, for some reason, my agent only queries the first 10 rows. In this case, it returns all the trials conducted in Berlin within the first ten rows. I looked into the SQL query it is generating and it has a 'limit 10' command included, which I do not want.

How do I get my agent to query through all the rows of my tables, not limiting to just 10?

Answer 1

The create_sql_agent function accepts a top_k argument which has default value of 10. You can change this by setting

agent = create_sql_agent(   
llm=llm,
toolkit=toolkit,
verbose=True,
agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
top_k=1000
)

top_k argument help to limit the token consumption if your LLM model going to analyse the results https://api.python.langchain.com/en/latest/agents/langchain.agents.agent_toolkits.sql.base.create_sql_agent.html https://github.com/langchain-ai/langchain/blob/master/libs/langchain/langchain/agents/agent_toolkits/sql/base.py