Training the SQL Agent with Example Queries

[feijoes]

Hi, is there a way to train the SQL agent with specific question-to-query mappings? I tried adding them as text in the agent prefix, but around 60% of the time the agent responds with "I don't know" or "max iterations reached," even when I ask exactly the example question.

My Code

query_1 = """SELECT
    pc.name AS categoria_producto,
    COUNT(DISTINCT am.id) AS cantidad_facturas
FROM
    account_move_line AS aml
....
"""
agent_executor: AgentExecutor = create_sql_agent(llm=llm,toolkit=toolkit,
        prefix = ("Note: Some questions may be in Spanish.\n"
          "Here are examples of common questions in Spanish and their corresponding SQL queries:\n"
          f"1. Categoría más vendida por cantidad de facturas:\n {query_1} \n"
          f"2. Categoría más vendida por monto total:\n {query_2} \n"
          f"3. Cantidad de Clientes:\n {query_3} \n"
          "Please generate the appropriate SQL query for the given question."+ "\n"+
          default_prefix ),
        format_instructions=react_prompt.FORMAT_INSTRUCTIONS,
        agent_type='tool-calling',
        top_k=100,
        suffix=SUFFIX
)

[dosubot[bot]]

Hey @feijoes! I'm here to help you with any bugs, questions, or becoming a contributor while you wait for a human maintainer. Let me know how I can assist you!

I found a similar unsolved discussion on this topic: Training SQL Agent ^[1].

Regarding your question, the create_sql_agent function does not inherently support training with specific question-to-query mappings by adding them as text in the agent prefix. The prefix parameter is primarily used to format the initial part of the prompt, providing instructions for the agent on how to interact with a SQL database ^[2][3][4].

If you want to guide the SQL agent with specific examples, you might consider incorporating these examples into the prompt in a way that the language model can use them as context. However, this approach does not "train" the model in the traditional sense but rather provides it with context to generate queries.

To continue talking to Dosu, mention @dosu.