chore: import upstream snapshot with attribution

examples/spider/README.md

@@ -0,0 +1,50 @@
+# Spider Example
+
+[![spider CI status](https://github.com/microsoft/agent-lightning/actions/workflows/examples-spider.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-spider.yml)
+
+This example demonstrates how to train a text-to-SQL agent on the Spider dataset using Agent-Lightning with reinforcement learning. It's compatible with Agent-lightning v0.2 or later.
+
+## Requirements
+
+This example depends on LangChain v0.x and several SQL-related libraries. Install the required dependencies with:
+
+```bash
+pip install "langgraph<1.0" "langchain[openai]<1.0" "langchain-community" "langchain-text-splitters<1.0" "sqlparse" "nltk"
+```
+
+Additionally, follow the [installation guide](../../docs/tutorials/installation.md) to install Agent-Lightning and VERL-related dependencies.
+
+## Dataset
+
+Detailed dataset preparation instructions are available in the [How to Train a SQL Agent](../../docs/how-to/train-sql-agent.md) guide.
+
+## Included Files
+
+| File/Directory | Description |
+|----------------|-------------|
+| `train_sql_agent.py` | Training script for SQL agents with support for multiple model configurations (Qwen, LLaMA, fast mode for CI) |
+| `sql_agent.py` | SQL agent implementation using LangGraph and LangChain, with debugging capabilities |
+| `data/` | Directory containing the Spider dataset files |
+| `spider_eval/` | Evaluation utilities for assessing SQL agent performance |
+
+## Running Examples
+
+### Training
+
+Train a SQL agent using the Qwen2.5-Coder-1.5B-Instruct model with the following command. This requires a single node with at least one 40GB GPU:
+
+```bash
+python train_sql_agent.py qwen
+```
+
+If you want to use an NPU for training, please refer to the **Launch Training with NPUS** section in [How to Train a SQL Agent](../../docs/how-to/train-sql-agent.md).
+
+### Debugging
+
+To test and debug the SQL agent interactively:
+
+```bash
+python sql_agent.py
+```
+
+This command requires an OpenAI-compatible API service. Configure your service endpoint and credentials using the `OPENAI_API_BASE` and `OPENAI_API_KEY` environment variables.

examples/spider/spider_eval/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Microsoft. All rights reserved.

examples/spider/spider_eval/async_utils.py

@@ -0,0 +1,38 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import queue
+import threading
+from typing import Any, Coroutine
+
+
+def run_sync_ephemeral(coro: Coroutine[Any, Any, Any]) -> Any:
+    """
+    Run an async coroutine from sync code.
+    - If no loop in this thread: use asyncio.run() directly.
+    - If already in an event loop: spawn a worker thread that calls asyncio.run()
+      (which creates and closes a brand-new event loop per call).
+    """
+    try:
+        asyncio.get_running_loop()
+    except RuntimeError:
+        # No running loop in this thread; safe to use asyncio.run
+        return asyncio.run(coro)
+
+    # Already in a running loop -> execute in a worker thread
+    q = queue.Queue[Any]()
+
+    def worker():
+        try:
+            result = asyncio.run(coro)  # creates & closes its own loop
+            q.put((True, result))
+        except BaseException as e:
+            q.put((False, e))
+
+    t = threading.Thread(target=worker, daemon=True)
+    t.start()
+    ok, payload = q.get()
+    t.join()
+    if ok:
+        return payload
+    raise payload

examples/spider/spider_eval/convert_dataset.py

@@ -0,0 +1,35 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# type: ignore
+
+import os
+
+import pandas as pd
+
+if __name__ == "__main__":
+    data_dir = "data"
+    target_data_dir = "data"
+    columns = ["db_id", "question", "query"]
+
+    dev_path = os.path.join(data_dir, "dev.json")
+    dev_df = pd.read_json(dev_path)
+    print(dev_df)
+    dev_df[columns].to_parquet(os.path.join(target_data_dir, "dev.parquet"), index=False)
+
+    train_path = os.path.join(data_dir, "train_spider.json")
+    train_df = pd.read_json(train_path)
+    print(train_df)
+    train_df[columns].to_parquet(os.path.join(target_data_dir, "train_spider.parquet"), index=False)
+
+    test_path = os.path.join(data_dir, "test.json")
+    test_df = pd.read_json(test_path)
+    print(test_df)
+    test_df[columns].to_parquet(os.path.join(target_data_dir, "test.parquet"), index=False)
+
+    # Select 100 of test df as test_dev
+    test_dev_df = test_df.sample(n=100, random_state=42)
+    test_dev_df[columns].to_parquet(os.path.join(target_data_dir, "test_dev.parquet"), index=False)
+
+    # Select 500 of test df as test_dev
+    test_dev_df = test_df.sample(n=500, random_state=0)
+    test_dev_df[columns].to_parquet(os.path.join(target_data_dir, "test_dev_500.parquet"), index=False)

examples/spider/spider_eval/exec_eval.py

@@ -0,0 +1,250 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# type: ignore
+# The evaluation code is from https://github.com/taoyds/test-suite-sql-eval
+
+import asyncio
+import os
+import pickle as pkl
+import random
+import re
+import sqlite3
+import subprocess
+import threading
+import time
+from collections import defaultdict
+from itertools import chain, product
+from typing import Any, List, Set, Tuple
+
+import tqdm
+
+from .async_utils import run_sync_ephemeral
+from .parse import get_all_preds_for_execution, remove_distinct
+
+threadLock = threading.Lock()
+TIMEOUT = 60
+EXEC_TMP_DIR = "/tmp/"
+
+
+def permute_tuple(element: Tuple, perm: Tuple) -> Tuple:
+    assert len(element) == len(perm)
+    return tuple([element[i] for i in perm])
+
+
+def unorder_row(row: Tuple) -> Tuple:
+    return tuple(sorted(row, key=lambda x: str(x) + str(type(x))))
+
+
+# unorder each row in the table
+# [result_1 and result_2 has the same bag of unordered row]
+# is a necessary condition of
+# [result_1 and result_2 are equivalent in denotation]
+def quick_rej(result1: List[Tuple], result2: List[Tuple], order_matters: bool) -> bool:
+    s1 = [unorder_row(row) for row in result1]
+    s2 = [unorder_row(row) for row in result2]
+    if order_matters:
+        return s1 == s2
+    else:
+        return set(s1) == set(s2)
+
+
+# return whether two bag of relations are equivalent
+def multiset_eq(l1: List, l2: List) -> bool:
+    if len(l1) != len(l2):
+        return False
+    d = defaultdict(int)
+    for e in l1:
+        d[e] = d[e] + 1
+    for e in l2:
+        d[e] = d[e] - 1
+        if d[e] < 0:
+            return False
+    return True
+
+
+def get_constraint_permutation(tab1_sets_by_columns: List[Set], result2: List[Tuple]):
+    num_cols = len(result2[0])
+    perm_constraints = [{i for i in range(num_cols)} for _ in range(num_cols)]
+    if num_cols <= 3:
+        return product(*perm_constraints)
+
+    # we sample 20 rows and constrain the space of permutations
+    for _ in range(20):
+        random_tab2_row = random.choice(result2)
+
+        for tab1_col in range(num_cols):
+            for tab2_col in set(perm_constraints[tab1_col]):
+                if random_tab2_row[tab2_col] not in tab1_sets_by_columns[tab1_col]:
+                    perm_constraints[tab1_col].remove(tab2_col)
+    return product(*perm_constraints)
+
+
+# check whether two denotations are correct
+def result_eq(result1: List[Tuple], result2: List[Tuple], order_matters: bool) -> bool:
+    if len(result1) == 0 and len(result2) == 0:
+        return True
+
+    # if length is not the same, then they are definitely different bag of rows
+    if len(result1) != len(result2):
+        return False
+
+    num_cols = len(result1[0])
+
+    # if the results do not have the same number of columns, they are different
+    if len(result2[0]) != num_cols:
+        return False
+
+    # unorder each row and compare whether the denotation is the same
+    # this can already find most pair of denotations that are different
+    if not quick_rej(result1, result2, order_matters):
+        return False
+
+    # the rest of the problem is in fact more complicated than one might think
+    # we want to find a permutation of column order and a permutation of row order,
+    # s.t. result_1 is the same as result_2
+    # we return true if we can find such column & row permutations
+    # and false if we cannot
+    tab1_sets_by_columns = [{row[i] for row in result1} for i in range(num_cols)]
+
+    # on a high level, we enumerate all possible column permutations that might make result_1 == result_2
+    # we decrease the size of the column permutation space by the function get_constraint_permutation
+    # if one of the permutation make result_1, result_2 equivalent, then they are equivalent
+    for perm in get_constraint_permutation(tab1_sets_by_columns, result2):
+        if len(perm) != len(set(perm)):
+            continue
+        if num_cols == 1:
+            result2_perm = result2
+        else:
+            result2_perm = [permute_tuple(element, perm) for element in result2]
+        if order_matters:
+            if result1 == result2_perm:
+                return True
+        else:
+            # in fact the first condition must hold if the second condition holds
+            # but the first is way more efficient implementation-wise
+            # and we use it to quickly reject impossible candidates
+            if set(result1) == set(result2_perm) and multiset_eq(result1, result2_perm):
+                return True
+    return False
+
+
+def replace_cur_year(query: str) -> str:
+    return re.sub("YEAR\s*\(\s*CURDATE\s*\(\s*\)\s*\)\s*", "2020", query, flags=re.IGNORECASE)
+
+
+# get the database cursor for a sqlite database path
+def get_cursor_from_path(sqlite_path: str):
+    try:
+        if not os.path.exists(sqlite_path):
+            print("Opening a new connection %s" % sqlite_path)
+        connection = sqlite3.connect(sqlite_path)
+    except Exception as e:
+        print(sqlite_path)
+        raise e
+    connection.text_factory = lambda b: b.decode(errors="ignore")
+    cursor = connection.cursor()
+    return cursor
+
+
+async def exec_on_db_(sqlite_path: str, query: str) -> Tuple[str, Any]:
+    query = replace_cur_year(query)
+    cursor = get_cursor_from_path(sqlite_path)
+    try:
+        cursor.execute(query)
+        result = cursor.fetchall()
+        cursor.close()
+        cursor.connection.close()
+        return "result", result
+    except Exception as e:
+        cursor.close()
+        cursor.connection.close()
+        return "exception", e
+
+
+async def exec_on_db(sqlite_path: str, query: str, process_id: str = "", timeout: int = TIMEOUT) -> Tuple[str, Any]:
+    try:
+        return await asyncio.wait_for(exec_on_db_(sqlite_path, query), timeout)
+    except asyncio.TimeoutError:
+        return ("exception", TimeoutError)
+    except Exception as e:
+        return ("exception", e)
+
+
+# postprocess the model predictions to avoid execution errors
+# e.g. removing spaces between ">" and "="
+def postprocess(query: str) -> str:
+    query = query.replace("> =", ">=").replace("< =", "<=").replace("! =", "!=")
+    return query
+
+
+# approximate whether p_str and g_str are semantically equivalent
+# db is the database path
+# we are going to evaluate whether they are equivalent in all the databases
+# that are in the same directory as db
+# 0 if denotationally equivalent
+# 1 otherwise
+# the meaning of each auxiliary argument can be seen in the parser definition in evaluation.py
+def eval_exec_match(
+    db: str, p_str: str, g_str: str, plug_value: bool, keep_distinct: bool, progress_bar_for_each_datapoint: bool
+) -> int:
+    # post-process the prediction.
+    # e.g. removing spaces between ">" and "="
+    p_str, g_str = postprocess(p_str), postprocess(g_str)
+    if not keep_distinct:
+        p_str = remove_distinct(p_str)
+        g_str = remove_distinct(g_str)
+
+    # we decide whether two denotations are equivalent based on "bag semantics"
+    # https://courses.cs.washington.edu/courses/cse444/10sp/lectures/lecture16.pdf
+    # if there is order by in query, then we assume order of the rows matter
+    # order by might also be used to find the max/min instead of sorting,
+    # but in that case the result mostly only contains one row and hence order_matters does not make a difference
+    order_matters = "order by" in g_str.lower()
+
+    # find all databases in the same directory
+    db_dir = os.path.dirname(db)
+    db_paths = [os.path.join(db_dir, basename) for basename in os.listdir(db_dir) if ".sqlite" in basename]
+
+    preds = [p_str]
+    # if plug in value (i.e. we do not consider value prediction correctness)
+    # enumerate all ways to plug in values in the gold query to the model predictions
+    # otherwise, we only evaluate the predicted query with its own value prediction
+    if plug_value:
+        _, preds = get_all_preds_for_execution(g_str, p_str)
+        # we did not add this line in our EMNLP work
+        # this reduces "false negatives" when value is substituted
+        preds = chain([p_str], preds)
+
+    for pred in preds:
+
+        pred_passes = 1
+        # compare the gold and predicted denotations on each database in the directory
+        # wrap with progress bar if required
+        if progress_bar_for_each_datapoint:
+            ranger = tqdm.tqdm(db_paths)
+        else:
+            ranger = db_paths
+
+        for db_path in ranger:
+            g_flag, g_denotation = run_sync_ephemeral(exec_on_db(db_path, g_str))
+            p_flag, p_denotation = run_sync_ephemeral(exec_on_db(db_path, pred))
+
+            # we should expect the gold to be successfully executed on the database
+            assert g_flag != "exception", "gold query %s has error on database file %s" % (g_str, db_path)
+
+            # wrong if execution fails
+            if p_flag == "exception":
+                pred_passes = 0
+
+            # if denotations are not equivalent, the prediction must be wrong
+            elif not result_eq(g_denotation, p_denotation, order_matters=order_matters):
+                pred_passes = 0
+            if pred_passes == 0:
+                break
+
+        # the model prediction has the same denotation as the gold for all databases
+        if pred_passes == 1:
+            return 1
+
+    # none of the predictions passed
+    return 0

examples/spider/spider_eval/parse.py

@@ -0,0 +1,231 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# type: ignore
+# The evaluation code is from https://github.com/taoyds/test-suite-sql-eval
+
+import itertools
+import re
+from collections import namedtuple
+from typing import Any, Dict, Iterator, List, Set, Tuple, Union
+
+import sqlparse
+from sqlparse.sql import Comparison, Identifier
+from sqlparse.tokens import Whitespace
+
+Token = namedtuple("Token", ["ttype", "value"])
+VALUE_NUM_SYMBOL = "VALUERARE"
+QUOTE_CHARS = {"`", "'", '"'}
+
+
+def tokenize(query: str) -> List[Token]:
+    tokens = list([Token(t.ttype, t.value) for t in sqlparse.parse(query)[0].flatten()])
+    return tokens
+
+
+def join_tokens(tokens: List[Token]) -> str:
+    return "".join([x.value for x in tokens]).strip().replace("  ", " ")
+
+
+def round_trip_test(query: str) -> None:
+    tokens = tokenize(query)
+    reconstructed = "".join([token.value for token in tokens])
+    assert query == reconstructed, "Round trip test fails for string %s" % query
+
+
+def postprocess(query: str) -> str:
+    query = query.replace("> =", ">=").replace("< =", "<=").replace("! =", "!=")
+    return query
+
+
+# strip_query, reformat_query and replace values
+# were implemented by Yu Tao for processing CoSQL
+def strip_query(query: str) -> Tuple[List[str], List[str]]:
+    query_keywords, all_values = [], []
+
+    # then replace all stuff enclosed by "" with a numerical value to get it marked as {VALUE}
+
+    # Tao's implementation is commented out here.
+    """
+    str_1 = re.findall("\"[^\"]*\"", query)
+    str_2 = re.findall("\'[^\']*\'", query)
+    values = str_1 + str_2
+        """
+
+    toks = sqlparse.parse(query)[0].flatten()
+    values = [
+        t.value
+        for t in toks
+        if t.ttype == sqlparse.tokens.Literal.String.Single or t.ttype == sqlparse.tokens.Literal.String.Symbol
+    ]
+
+    for val in values:
+        all_values.append(val)
+        query = query.replace(val.strip(), VALUE_NUM_SYMBOL)
+
+    query_tokenized = query.split()
+    float_nums = re.findall("[-+]?\d*\.\d+", query)
+    all_values += [qt for qt in query_tokenized if qt in float_nums]
+    query_tokenized = [VALUE_NUM_SYMBOL if qt in float_nums else qt for qt in query_tokenized]
+
+    query = " ".join(query_tokenized)
+    int_nums = [i.strip() for i in re.findall("[^tT]\d+", query)]
+
+    all_values += [qt for qt in query_tokenized if qt in int_nums]
+    query_tokenized = [VALUE_NUM_SYMBOL if qt in int_nums else qt for qt in query_tokenized]
+    # print int_nums, query, query_tokenized
+
+    for tok in query_tokenized:
+        if "." in tok:
+            table = re.findall("[Tt]\d+\.", tok)
+            if len(table) > 0:
+                to = tok.replace(".", " . ").split()
+                to = [t.lower() for t in to if len(t) > 0]
+                query_keywords.extend(to)
+            else:
+                query_keywords.append(tok.lower())
+
+        elif len(tok) > 0:
+            query_keywords.append(tok.lower())
+    return query_keywords, all_values
+
+
+def reformat_query(query: str) -> str:
+    query = query.strip().replace(";", "").replace("\t", "")
+    query = " ".join([t.value for t in tokenize(query) if t.ttype != sqlparse.tokens.Whitespace])
+    t_stars = ["t1.*", "t2.*", "t3.*", "T1.*", "T2.*", "T3.*"]
+    for ts in t_stars:
+        query = query.replace(ts, "*")
+    return query
+
+
+def replace_values(sql: str) -> Tuple[List[str], Set[str]]:
+    sql = sqlparse.format(sql, reindent=False, keyword_case="upper")
+    # sql = re.sub(r"(<=|>=|!=|=|<|>|,)", r" \1 ", sql)
+    sql = re.sub(r"(T\d+\.)\s", r"\1", sql)
+    query_toks_no_value, values = strip_query(sql)
+    return query_toks_no_value, set(values)
+
+
+# extract the non-value tokens and the set of values
+# from a sql query
+def extract_query_values(sql: str) -> Tuple[List[str], Set[str]]:
+    reformatted = reformat_query(query=sql)
+    query_value_replaced, values = replace_values(reformatted)
+    return query_value_replaced, values
+
+
+# plug in the values into query with value slots
+def plugin(query_value_replaced: List[str], values_in_order: List[str]) -> str:
+    q_length = len(query_value_replaced)
+    query_w_values = query_value_replaced[:]
+    value_idx = [idx for idx in range(q_length) if query_value_replaced[idx] == VALUE_NUM_SYMBOL.lower()]
+    assert len(value_idx) == len(values_in_order)
+
+    for idx, value in zip(value_idx, values_in_order):
+        query_w_values[idx] = value
+    return " ".join(query_w_values)
+
+
+# a generator generating all possible ways of
+# filling values into predicted query
+def plugin_all_permutations(query_value_replaced: List[str], values: Set[str]) -> Iterator[str]:
+    num_slots = len([v for v in query_value_replaced if v == VALUE_NUM_SYMBOL.lower()])
+    for values in itertools.product(*[list(values) for _ in range(num_slots)]):
+        yield plugin(query_value_replaced, list(values))
+
+
+# given the gold query and the model prediction
+# extract values from the gold, extract predicted sql with value slots
+# return 1) number of possible ways to plug in gold values and 2) an iterator of predictions with value plugged in
+def get_all_preds_for_execution(gold: str, pred: str) -> Tuple[int, Iterator[str]]:
+    _, gold_values = extract_query_values(gold)
+    pred_query_value_replaced, _ = extract_query_values(pred)
+    num_slots = len([v for v in pred_query_value_replaced if v == VALUE_NUM_SYMBOL.lower()])
+    num_alternatives = len(gold_values) ** num_slots
+    return num_alternatives, plugin_all_permutations(pred_query_value_replaced, gold_values)
+
+
+def remove_distinct(s):
+    toks = [t.value for t in list(sqlparse.parse(s)[0].flatten())]
+    return "".join([t for t in toks if t.lower() != "distinct"])
+
+
+def extract_all_comparison_from_node(node: Token) -> List[Comparison]:
+    comparison_list = []
+    if hasattr(node, "tokens"):
+        for t in node.tokens:
+            comparison_list.extend(extract_all_comparison_from_node(t))
+    if type(node) == Comparison:
+        comparison_list.append(node)
+    return comparison_list
+
+
+def extract_all_comparison(query: str) -> List[Comparison]:
+    tree = sqlparse.parse(query)[0]
+    comparison_list = extract_all_comparison_from_node(tree)
+    return comparison_list
+
+
+def extract_toks_from_comparison(comparison_node: Comparison) -> List[Token]:
+    tokens = [t for t in comparison_node.tokens if t.ttype != Whitespace]
+    return tokens
+
+
+def extract_info_from_comparison(comparison_node: Comparison) -> Dict[str, Any]:
+    tokens = extract_toks_from_comparison(comparison_node)
+    left, op, right = tokens
+
+    returned_dict = {"left": left, "op": op.value, "right": right}
+
+    if type(left) != Identifier:
+        return returned_dict
+
+    table = None
+    if len(left.tokens) == 3 and re.match("^[tT][0-9]$", left.tokens[0].value) is None:
+        table = left.tokens[0].value.lower()
+    col = left.tokens[-1].value
+
+    if type(right) == Identifier:
+        if len(right.tokens) == 1 and type(right.tokens[0]) == sqlparse.sql.Token:
+            right_val = right.tokens[0].value
+        else:
+            return returned_dict
+    elif type(right) == sqlparse.sql.Token:
+        right_val = right.value
+    else:
+        return returned_dict
+
+    returned_dict["table_col"], returned_dict["val"] = (table, col.upper()), process_str_value(right_val)
+
+    return returned_dict
+
+
+def extract_all_comparison_from_query(query: str) -> List[Dict[str, Any]]:
+    comparison_list = extract_all_comparison(query)
+    return [extract_info_from_comparison(c) for c in comparison_list]
+
+
+def extract_typed_value_in_comparison_from_query(query: str) -> List[Tuple[Tuple[Union[str, None], str], str]]:
+    cmps = extract_all_comparison_from_query(query)
+    typed_values = [(cmp["table_col"], cmp["val"]) for cmp in cmps if "table_col" in cmp]
+    for table, col, val1, val2 in re.findall(
+        "(?:([^\.\s]*)\.)?([^\.\s]+) between ([^\s;]+) and ([^\s;]+)", query, re.IGNORECASE
+    ):
+        if table == "":
+            table = None
+        else:
+            table = table.lower()
+        col = col.upper()
+        for v in [val1, val2]:
+            typed_values.append(((table, col), v))
+    return typed_values
+
+
+def process_str_value(v: str) -> str:
+    if len(v) > 0 and v[0] in QUOTE_CHARS:
+        v = v[1:]
+    if len(v) > 0 and v[-1] in QUOTE_CHARS:
+        v = v[:-1]
+    for c in QUOTE_CHARS:
+        v = v.replace(c + c, c)
+    return v

examples/spider/spider_eval/process_sql.py

@@ -0,0 +1,578 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# type: ignore
+# The evaluation code is from https://github.com/taoyds/test-suite-sql-eval
+
+################################
+# Assumptions:
+#   1. sql is correct
+#   2. only table name has alias
+#   3. only one intersect/union/except
+#
+# val: number(float)/string(str)/sql(dict)
+# col_unit: (agg_id, col_id, isDistinct(bool))
+# val_unit: (unit_op, col_unit1, col_unit2)
+# table_unit: (table_type, col_unit/sql)
+# cond_unit: (not_op, op_id, val_unit, val1, val2)
+# condition: [cond_unit1, 'and'/'or', cond_unit2, ...]
+# sql {
+#   'select': (isDistinct(bool), [(agg_id, val_unit), (agg_id, val_unit), ...])
+#   'from': {'table_units': [table_unit1, table_unit2, ...], 'conds': condition}
+#   'where': condition
+#   'groupBy': [col_unit1, col_unit2, ...]
+#   'orderBy': ('asc'/'desc', [val_unit1, val_unit2, ...])
+#   'having': condition
+#   'limit': None/limit value
+#   'intersect': None/sql
+#   'except': None/sql
+#   'union': None/sql
+# }
+################################
+
+import json
+import sqlite3
+
+from nltk import word_tokenize
+
+CLAUSE_KEYWORDS = ("select", "from", "where", "group", "order", "limit", "intersect", "union", "except")
+JOIN_KEYWORDS = ("join", "on", "as")
+
+WHERE_OPS = ("not", "between", "=", ">", "<", ">=", "<=", "!=", "in", "like", "is", "exists")
+UNIT_OPS = ("none", "-", "+", "*", "/")
+AGG_OPS = ("none", "max", "min", "count", "sum", "avg")
+TABLE_TYPE = {
+    "sql": "sql",
+    "table_unit": "table_unit",
+}
+
+COND_OPS = ("and", "or")
+SQL_OPS = ("intersect", "union", "except")
+ORDER_OPS = ("desc", "asc")
+
+
+class Schema:
+    """
+    Simple schema which maps table&column to a unique identifier
+    """
+
+    def __init__(self, schema):
+        self._schema = schema
+        self._idMap = self._map(self._schema)
+
+    @property
+    def schema(self):
+        return self._schema
+
+    @property
+    def idMap(self):
+        return self._idMap
+
+    def _map(self, schema):
+        idMap = {"*": "__all__"}
+        id = 1
+        for key, vals in schema.items():
+            for val in vals:
+                idMap[key.lower() + "." + val.lower()] = "__" + key.lower() + "." + val.lower() + "__"
+                id += 1
+
+        for key in schema:
+            idMap[key.lower()] = "__" + key.lower() + "__"
+            id += 1
+
+        return idMap
+
+
+def get_schema(db):
+    """
+    Get database's schema, which is a dict with table name as key
+    and list of column names as value
+    :param db: database path
+    :return: schema dict
+    """
+
+    schema = {}
+    conn = sqlite3.connect(db)
+    cursor = conn.cursor()
+
+    # fetch table names
+    cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
+    tables = [str(table[0].lower()) for table in cursor.fetchall()]
+
+    # fetch table info
+    for table in tables:
+        cursor.execute("PRAGMA table_info({})".format(table))
+        schema[table] = [str(col[1].lower()) for col in cursor.fetchall()]
+
+    return schema
+
+
+def get_schema_from_json(fpath):
+    with open(fpath) as f:
+        data = json.load(f)
+
+    schema = {}
+    for entry in data:
+        table = str(entry["table"].lower())
+        cols = [str(col["column_name"].lower()) for col in entry["col_data"]]
+        schema[table] = cols
+
+    return schema
+
+
+def tokenize(string):
+    string = str(string)
+    string = string.replace("'", '"')  # ensures all string values wrapped by "" problem??
+    quote_idxs = [idx for idx, char in enumerate(string) if char == '"']
+    assert len(quote_idxs) % 2 == 0, "Unexpected quote"
+
+    # keep string value as token
+    vals = {}
+    for i in range(len(quote_idxs) - 1, -1, -2):
+        qidx1 = quote_idxs[i - 1]
+        qidx2 = quote_idxs[i]
+        val = string[qidx1 : qidx2 + 1]
+        key = "__val_{}_{}__".format(qidx1, qidx2)
+        string = string[:qidx1] + key + string[qidx2 + 1 :]
+        vals[key] = val
+
+    toks = [word.lower() for word in word_tokenize(string)]
+    # replace with string value token
+    for i in range(len(toks)):
+        if toks[i] in vals:
+            toks[i] = vals[toks[i]]
+
+    # find if there exists !=, >=, <=
+    eq_idxs = [idx for idx, tok in enumerate(toks) if tok == "="]
+    eq_idxs.reverse()
+    prefix = ("!", ">", "<")
+    for eq_idx in eq_idxs:
+        pre_tok = toks[eq_idx - 1]
+        if pre_tok in prefix:
+            toks = toks[: eq_idx - 1] + [pre_tok + "="] + toks[eq_idx + 1 :]
+
+    return toks
+
+
+def scan_alias(toks):
+    """Scan the index of 'as' and build the map for all alias"""
+    as_idxs = [idx for idx, tok in enumerate(toks) if tok == "as"]
+    alias = {}
+    for idx in as_idxs:
+        alias[toks[idx + 1]] = toks[idx - 1]
+    return alias
+
+
+def get_tables_with_alias(schema, toks):
+    tables = scan_alias(toks)
+    for key in schema:
+        assert key not in tables, "Alias {} has the same name in table".format(key)
+        tables[key] = key
+    return tables
+
+
+def parse_col(toks, start_idx, tables_with_alias, schema, default_tables=None):
+    """
+    :returns next idx, column id
+    """
+    tok = toks[start_idx]
+    if tok == "*":
+        return start_idx + 1, schema.idMap[tok]
+
+    if "." in tok:  # if token is a composite
+        alias, col = tok.split(".")
+        key = tables_with_alias[alias] + "." + col
+        return start_idx + 1, schema.idMap[key]
+
+    assert default_tables is not None and len(default_tables) > 0, "Default tables should not be None or empty"
+
+    for alias in default_tables:
+        table = tables_with_alias[alias]
+        if tok in schema.schema[table]:
+            key = table + "." + tok
+            return start_idx + 1, schema.idMap[key]
+
+    assert False, "Error col: {}".format(tok)
+
+
+def parse_col_unit(toks, start_idx, tables_with_alias, schema, default_tables=None):
+    """
+    :returns next idx, (agg_op id, col_id)
+    """
+    idx = start_idx
+    len_ = len(toks)
+    isBlock = False
+    isDistinct = False
+    if toks[idx] == "(":
+        isBlock = True
+        idx += 1
+
+    if toks[idx] in AGG_OPS:
+        agg_id = AGG_OPS.index(toks[idx])
+        idx += 1
+        assert idx < len_ and toks[idx] == "("
+        idx += 1
+        if toks[idx] == "distinct":
+            idx += 1
+            isDistinct = True
+        idx, col_id = parse_col(toks, idx, tables_with_alias, schema, default_tables)
+        assert idx < len_ and toks[idx] == ")"
+        idx += 1
+        return idx, (agg_id, col_id, isDistinct)
+
+    if toks[idx] == "distinct":
+        idx += 1
+        isDistinct = True
+    agg_id = AGG_OPS.index("none")
+    idx, col_id = parse_col(toks, idx, tables_with_alias, schema, default_tables)
+
+    if isBlock:
+        assert toks[idx] == ")"
+        idx += 1  # skip ')'
+
+    return idx, (agg_id, col_id, isDistinct)
+
+
+def parse_val_unit(toks, start_idx, tables_with_alias, schema, default_tables=None):
+    idx = start_idx
+    len_ = len(toks)
+    isBlock = False
+    if toks[idx] == "(":
+        isBlock = True
+        idx += 1
+
+    col_unit1 = None
+    col_unit2 = None
+    unit_op = UNIT_OPS.index("none")
+
+    idx, col_unit1 = parse_col_unit(toks, idx, tables_with_alias, schema, default_tables)
+    if idx < len_ and toks[idx] in UNIT_OPS:
+        unit_op = UNIT_OPS.index(toks[idx])
+        idx += 1
+        idx, col_unit2 = parse_col_unit(toks, idx, tables_with_alias, schema, default_tables)
+
+    if isBlock:
+        assert toks[idx] == ")"
+        idx += 1  # skip ')'
+
+    return idx, (unit_op, col_unit1, col_unit2)
+
+
+def parse_table_unit(toks, start_idx, tables_with_alias, schema):
+    """
+    :returns next idx, table id, table name
+    """
+    idx = start_idx
+    len_ = len(toks)
+    key = tables_with_alias[toks[idx]]
+
+    if idx + 1 < len_ and toks[idx + 1] == "as":
+        idx += 3
+    else:
+        idx += 1
+
+    return idx, schema.idMap[key], key
+
+
+def parse_value(toks, start_idx, tables_with_alias, schema, default_tables=None):
+    idx = start_idx
+    len_ = len(toks)
+
+    isBlock = False
+    if toks[idx] == "(":
+        isBlock = True
+        idx += 1
+
+    if toks[idx] == "select":
+        idx, val = parse_sql(toks, idx, tables_with_alias, schema)
+    elif '"' in toks[idx]:  # token is a string value
+        val = toks[idx]
+        idx += 1
+    else:
+        try:
+            val = float(toks[idx])
+            idx += 1
+        except:
+            end_idx = idx
+            while (
+                end_idx < len_
+                and toks[end_idx] != ","
+                and toks[end_idx] != ")"
+                and toks[end_idx] != "and"
+                and toks[end_idx] not in CLAUSE_KEYWORDS
+                and toks[end_idx] not in JOIN_KEYWORDS
+            ):
+                end_idx += 1
+
+            idx, val = parse_col_unit(toks[start_idx:end_idx], 0, tables_with_alias, schema, default_tables)
+            idx = end_idx
+
+    if isBlock:
+        assert toks[idx] == ")"
+        idx += 1
+
+    return idx, val
+
+
+def parse_condition(toks, start_idx, tables_with_alias, schema, default_tables=None):
+    idx = start_idx
+    len_ = len(toks)
+    conds = []
+
+    while idx < len_:
+        idx, val_unit = parse_val_unit(toks, idx, tables_with_alias, schema, default_tables)
+        not_op = False
+        if toks[idx] == "not":
+            not_op = True
+            idx += 1
+
+        assert idx < len_ and toks[idx] in WHERE_OPS, "Error condition: idx: {}, tok: {}".format(idx, toks[idx])
+        op_id = WHERE_OPS.index(toks[idx])
+        idx += 1
+        val1 = val2 = None
+        if op_id == WHERE_OPS.index("between"):  # between..and... special case: dual values
+            idx, val1 = parse_value(toks, idx, tables_with_alias, schema, default_tables)
+            assert toks[idx] == "and"
+            idx += 1
+            idx, val2 = parse_value(toks, idx, tables_with_alias, schema, default_tables)
+        else:  # normal case: single value
+            idx, val1 = parse_value(toks, idx, tables_with_alias, schema, default_tables)
+            val2 = None
+
+        conds.append((not_op, op_id, val_unit, val1, val2))
+
+        if idx < len_ and (toks[idx] in CLAUSE_KEYWORDS or toks[idx] in (")", ";") or toks[idx] in JOIN_KEYWORDS):
+            break
+
+        if idx < len_ and toks[idx] in COND_OPS:
+            conds.append(toks[idx])
+            idx += 1  # skip and/or
+
+    return idx, conds
+
+
+def parse_select(toks, start_idx, tables_with_alias, schema, default_tables=None):
+    idx = start_idx
+    len_ = len(toks)
+
+    assert toks[idx] == "select", "'select' not found"
+    idx += 1
+    isDistinct = False
+    if idx < len_ and toks[idx] == "distinct":
+        idx += 1
+        isDistinct = True
+    val_units = []
+
+    while idx < len_ and toks[idx] not in CLAUSE_KEYWORDS:
+        agg_id = AGG_OPS.index("none")
+        if toks[idx] in AGG_OPS:
+            agg_id = AGG_OPS.index(toks[idx])
+            idx += 1
+        idx, val_unit = parse_val_unit(toks, idx, tables_with_alias, schema, default_tables)
+        val_units.append((agg_id, val_unit))
+        if idx < len_ and toks[idx] == ",":
+            idx += 1  # skip ','
+
+    return idx, (isDistinct, val_units)
+
+
+def parse_from(toks, start_idx, tables_with_alias, schema):
+    """
+    Assume in the from clause, all table units are combined with join
+    """
+    assert "from" in toks[start_idx:], "'from' not found"
+
+    len_ = len(toks)
+    idx = toks.index("from", start_idx) + 1
+    default_tables = []
+    table_units = []
+    conds = []
+
+    while idx < len_:
+        isBlock = False
+        if toks[idx] == "(":
+            isBlock = True
+            idx += 1
+
+        if toks[idx] == "select":
+            idx, sql = parse_sql(toks, idx, tables_with_alias, schema)
+            table_units.append((TABLE_TYPE["sql"], sql))
+        else:
+            if idx < len_ and toks[idx] == "join":
+                idx += 1  # skip join
+            idx, table_unit, table_name = parse_table_unit(toks, idx, tables_with_alias, schema)
+            table_units.append((TABLE_TYPE["table_unit"], table_unit))
+            default_tables.append(table_name)
+        if idx < len_ and toks[idx] == "on":
+            idx += 1  # skip on
+            idx, this_conds = parse_condition(toks, idx, tables_with_alias, schema, default_tables)
+            if len(conds) > 0:
+                conds.append("and")
+            conds.extend(this_conds)
+
+        if isBlock:
+            assert toks[idx] == ")"
+            idx += 1
+        if idx < len_ and (toks[idx] in CLAUSE_KEYWORDS or toks[idx] in (")", ";")):
+            break
+
+    return idx, table_units, conds, default_tables
+
+
+def parse_where(toks, start_idx, tables_with_alias, schema, default_tables):
+    idx = start_idx
+    len_ = len(toks)
+
+    if idx >= len_ or toks[idx] != "where":
+        return idx, []
+
+    idx += 1
+    idx, conds = parse_condition(toks, idx, tables_with_alias, schema, default_tables)
+    return idx, conds
+
+
+def parse_group_by(toks, start_idx, tables_with_alias, schema, default_tables):
+    idx = start_idx
+    len_ = len(toks)
+    col_units = []
+
+    if idx >= len_ or toks[idx] != "group":
+        return idx, col_units
+
+    idx += 1
+    assert toks[idx] == "by"
+    idx += 1
+
+    while idx < len_ and not (toks[idx] in CLAUSE_KEYWORDS or toks[idx] in (")", ";")):
+        idx, col_unit = parse_col_unit(toks, idx, tables_with_alias, schema, default_tables)
+        col_units.append(col_unit)
+        if idx < len_ and toks[idx] == ",":
+            idx += 1  # skip ','
+        else:
+            break
+
+    return idx, col_units
+
+
+def parse_order_by(toks, start_idx, tables_with_alias, schema, default_tables):
+    idx = start_idx
+    len_ = len(toks)
+    val_units = []
+    order_type = "asc"  # default type is 'asc'
+
+    if idx >= len_ or toks[idx] != "order":
+        return idx, val_units
+
+    idx += 1
+    assert toks[idx] == "by"
+    idx += 1
+
+    while idx < len_ and not (toks[idx] in CLAUSE_KEYWORDS or toks[idx] in (")", ";")):
+        idx, val_unit = parse_val_unit(toks, idx, tables_with_alias, schema, default_tables)
+        val_units.append(val_unit)
+        if idx < len_ and toks[idx] in ORDER_OPS:
+            order_type = toks[idx]
+            idx += 1
+        if idx < len_ and toks[idx] == ",":
+            idx += 1  # skip ','
+        else:
+            break
+
+    return idx, (order_type, val_units)
+
+
+def parse_having(toks, start_idx, tables_with_alias, schema, default_tables):
+    idx = start_idx
+    len_ = len(toks)
+
+    if idx >= len_ or toks[idx] != "having":
+        return idx, []
+
+    idx += 1
+    idx, conds = parse_condition(toks, idx, tables_with_alias, schema, default_tables)
+    return idx, conds
+
+
+def parse_limit(toks, start_idx):
+    idx = start_idx
+    len_ = len(toks)
+
+    if idx < len_ and toks[idx] == "limit":
+        idx += 2
+        # make limit value can work, cannot assume put 1 as a fake limit number
+        if type(toks[idx - 1]) != int:
+            return idx, 1
+
+        return idx, int(toks[idx - 1])
+
+    return idx, None
+
+
+def parse_sql(toks, start_idx, tables_with_alias, schema):
+    isBlock = False  # indicate whether this is a block of sql/sub-sql
+    len_ = len(toks)
+    idx = start_idx
+
+    sql = {}
+    if toks[idx] == "(":
+        isBlock = True
+        idx += 1
+
+    # parse from clause in order to get default tables
+    from_end_idx, table_units, conds, default_tables = parse_from(toks, start_idx, tables_with_alias, schema)
+    sql["from"] = {"table_units": table_units, "conds": conds}
+    # select clause
+    _, select_col_units = parse_select(toks, idx, tables_with_alias, schema, default_tables)
+    idx = from_end_idx
+    sql["select"] = select_col_units
+    # where clause
+    idx, where_conds = parse_where(toks, idx, tables_with_alias, schema, default_tables)
+    sql["where"] = where_conds
+    # group by clause
+    idx, group_col_units = parse_group_by(toks, idx, tables_with_alias, schema, default_tables)
+    sql["groupBy"] = group_col_units
+    # having clause
+    idx, having_conds = parse_having(toks, idx, tables_with_alias, schema, default_tables)
+    sql["having"] = having_conds
+    # order by clause
+    idx, order_col_units = parse_order_by(toks, idx, tables_with_alias, schema, default_tables)
+    sql["orderBy"] = order_col_units
+    # limit clause
+    idx, limit_val = parse_limit(toks, idx)
+    sql["limit"] = limit_val
+
+    idx = skip_semicolon(toks, idx)
+    if isBlock:
+        assert toks[idx] == ")"
+        idx += 1  # skip ')'
+    idx = skip_semicolon(toks, idx)
+
+    # intersect/union/except clause
+    for op in SQL_OPS:  # initialize IUE
+        sql[op] = None
+    if idx < len_ and toks[idx] in SQL_OPS:
+        sql_op = toks[idx]
+        idx += 1
+        idx, IUE_sql = parse_sql(toks, idx, tables_with_alias, schema)
+        sql[sql_op] = IUE_sql
+    return idx, sql
+
+
+def load_data(fpath):
+    with open(fpath) as f:
+        data = json.load(f)
+    return data
+
+
+def get_sql(schema, query):
+    toks = tokenize(query)
+    tables_with_alias = get_tables_with_alias(schema.schema, toks)
+    _, sql = parse_sql(toks, 0, tables_with_alias, schema)
+
+    return sql
+
+
+def skip_semicolon(toks, start_idx):
+    idx = start_idx
+    while idx < len(toks) and toks[idx] == ";":
+        idx += 1
+    return idx

examples/spider/sql_agent.py

@@ -0,0 +1,546 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Sample code that demonstrates an SQL agent using LangGraph and LangChain,
+trainable with Agent-lightning.
+
+Adapted from https://python.langchain.com/docs/tutorials/sql_qa/
+as well as https://langchain-ai.github.io/langgraph/tutorials/sql-agent/
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import re
+import shutil
+import tempfile
+import time
+from typing import Any, Dict, List, Literal, Optional, cast
+
+import pandas as pd
+import termcolor
+from langchain.chat_models import init_chat_model
+from langchain_community.tools.sql_database.tool import QuerySQLDatabaseTool
+from langchain_community.utilities import SQLDatabase
+from langchain_core.messages import AnyMessage, BaseMessage, HumanMessage
+from langchain_core.prompts import ChatPromptTemplate
+from langgraph.graph import END, START, MessagesState, StateGraph
+from langgraph.graph.state import CompiledStateGraph
+from spider_eval.exec_eval import eval_exec_match
+
+import agentlightning as agl
+
+agl.setup_logging(apply_to=[__name__])
+
+logger = logging.getLogger(__name__)
+
+
+WRITE_QUERY_PROMPT = ChatPromptTemplate(
+    [
+        (
+            "system",
+            """
+You are an agent designed to interact with a SQL database.
+     Given an input question, create a syntactically correct {dialect} query to run to help find the answer.
+
+Pay attention to use only the column names that you can see in the schema description.
+Be careful to not query for columns that do not exist.
+Also, pay attention to which column is in which table.
+
+## Table Schema ##
+
+Only use the following tables:
+{table_info}
+
+## Output Format ##
+
+Respond in the following format:
+
+```{dialect}
+GENERATED QUERY
+```
+""".strip(),
+        ),
+        ("user", "Question: {input}"),
+    ]
+)
+
+
+CHECK_QUERY_PROMPT = ChatPromptTemplate(
+    [
+        (
+            "system",
+            """
+You are a SQL expert with a strong attention to detail.
+Double check the {dialect} query for common mistakes, including:
+- Using NOT IN with NULL values
+- Using UNION when UNION ALL should have been used
+- Using BETWEEN for exclusive ranges
+- Data type mismatch in predicates
+- Properly quoting identifiers
+- Using the correct number of arguments for functions
+- Casting to the correct data type
+- Using the proper columns for joins
+- Explicit query execution failures
+- Clearly unreasoable query execution results
+
+## Table Schema ##
+
+{table_info}
+
+## Output Format ##
+
+If any mistakes from the list above are found, list each error clearly.
+After listing mistakes (if any), conclude with **ONE** of the following exact phrases in all caps and without surrounding quotes:
+- If mistakes are found: `THE QUERY IS INCORRECT.`
+- If no mistakes are found: `THE QUERY IS CORRECT.`
+
+DO NOT write the corrected query in the response. You only need to report the mistakes.
+""".strip(),
+        ),
+        (
+            "user",
+            """Question: {input}
+
+Query:
+
+```{dialect}
+{query}
+```
+
+Execution result:
+
+```
+{execution}
+```""",
+        ),
+    ]
+)
+
+
+REWRITE_QUERY_PROMPT = ChatPromptTemplate(
+    [
+        (
+            "system",
+            """
+You are an agent designed to interact with a SQL database.
+Rewrite the previous {dialect} query to fix errors based on the provided feedback.
+The goal is to answer the original question.
+Make sure to address all points in the feedback.
+
+Pay attention to use only the column names that you can see in the schema description.
+Be careful to not query for columns that do not exist.
+Also, pay attention to which column is in which table.
+
+## Table Schema ##
+
+Only use the following tables:
+{table_info}
+
+## Output Format ##
+
+Respond in the following format:
+
+```{dialect}
+REWRITTEN QUERY
+```
+""".strip(),
+        ),
+        (
+            "user",
+            """Question: {input}
+
+## Previous query ##
+
+```{dialect}
+{query}
+```
+
+## Previous execution result ##
+
+```
+{execution}
+```
+
+## Feedback ##
+
+{feedback}
+
+Please rewrite the query to address the feedback.""",
+        ),
+    ]
+)
+
+
+class State(MessagesState):
+    question: str
+    query: str
+    execution: str
+    answer: str
+    feedback: str
+    num_turns: int
+    messages: list[AnyMessage]
+
+
+class SQLAgent:
+
+    def __init__(
+        self,
+        db: str,
+        max_turns: int = 5,
+        debug: bool = False,
+        db_schema: str | None = None,
+        endpoint: str | None = None,
+        verl_replacement: Dict[str, Any] | None = None,
+        table_info_truncate: int = 2048,
+        execution_truncate: int = 2048,
+    ):
+        self.db = SQLDatabase.from_uri(db)  # type: ignore
+        self.db_schema = db_schema
+        self.debug = debug
+        self.max_turns = max_turns
+        self.table_info_truncate = table_info_truncate
+        self.execution_truncate = execution_truncate
+        if verl_replacement is not None:
+            self.model_name: str = verl_replacement["model"]  # type: ignore
+            assert endpoint is not None
+            self.llm = init_chat_model(
+                self.model_name,
+                model_provider="openai",
+                openai_api_base=endpoint,
+                openai_api_key=os.environ.get("OPENAI_API_KEY", "dummy"),
+                temperature=verl_replacement["temperature"],
+                max_retries=0,
+                max_tokens=2048,
+            )
+        else:
+            self.model_name: str = os.environ.get("MODEL", "gpt-4.1-mini")
+            self.llm = init_chat_model(
+                self.model_name,
+                model_provider="openai",
+                openai_api_base=endpoint or os.environ["OPENAI_API_BASE"],
+                openai_api_key=os.environ["OPENAI_API_KEY"],
+                temperature=0,
+                max_retries=1,
+                max_tokens=2048,
+            )
+
+    def get_table_info(self) -> str:
+        """Get the table information in a human-readable format."""
+        try:
+            table_info = self.db.get_table_info()
+            if len(table_info) > self.table_info_truncate:
+                table_info = table_info[: self.table_info_truncate] + "\n... (truncated)"
+            return table_info
+        except Exception as e:
+            logger.error(f"Failed to get table info: {e}")
+            if self.db_schema:
+                if len(self.db_schema) > self.table_info_truncate:
+                    return self.db_schema[: self.table_info_truncate] + "\n... (truncated)"
+                return self.db_schema
+            return "No schema available."
+
+    def invoke_prompt(self, prompt: Any) -> AnyMessage:
+        if self.debug:
+            for message in prompt.messages:
+                termcolor.cprint(message.pretty_repr(), "blue")
+
+        try:
+            result = self.llm.invoke(prompt)
+        except Exception as e:
+            logger.error(f"Failed to invoke prompt: {e}")
+            # FIXME: fallback to create a random trajectory
+            result = self.llm.invoke([HumanMessage(content="Please create a random SQL query as an example.")])
+
+        if self.debug:
+            termcolor.cprint(result.pretty_repr(), "green")
+
+        return result  # type: ignore
+
+    def truncate_execution(self, execution: str) -> str:
+        """Truncate the execution result to a reasonable length."""
+        if len(execution) > self.execution_truncate:
+            return execution[: self.execution_truncate] + "\n... (truncated)"
+        return execution
+
+    def parse_query(self, message: AnyMessage) -> str | None:
+        result: str | None = None
+        for match in re.finditer(r".*```\w*\n(.*?)\n```.*", message.content, re.DOTALL):  # type: ignore
+            result = match.group(1).strip()  # type: ignore
+        return result  # type: ignore
+
+    def write_query(self, state: State) -> State:
+        """Generate SQL query to fetch information."""
+        prompt: Any = WRITE_QUERY_PROMPT.invoke(  # type: ignore
+            {
+                "dialect": self.db.dialect,
+                "input": state["question"],
+                "table_info": self.get_table_info(),
+            }
+        )
+        result = self.invoke_prompt(prompt)  # type: ignore
+
+        query = self.parse_query(result) or result.content  # type: ignore
+
+        return {  # type: ignore
+            **state,
+            "query": query,  # type: ignore
+            "num_turns": 1,
+            "messages": [*prompt.messages, result],
+        }
+
+    def execute_query(self, state: State) -> State:
+        """Execute SQL query."""
+        execute_query_tool = QuerySQLDatabaseTool(db=self.db)
+        execution_result = execute_query_tool.invoke(state["query"])  # type: ignore
+        if not isinstance(execution_result, str):
+            # Convert to string if it's not already
+            execution_result = str(execution_result)
+        if self.debug:
+            termcolor.cprint(execution_result, "yellow")
+        return {**state, "execution": execution_result}
+
+    def check_query(self, state: State) -> State:
+        """Check the SQL query for correctness."""
+        prompt: Any = CHECK_QUERY_PROMPT.invoke(  # type: ignore
+            {
+                "dialect": self.db.dialect,
+                "input": state["question"],
+                "query": state["query"],
+                "execution": self.truncate_execution(state["execution"]),
+                "table_info": self.get_table_info(),
+            }
+        )
+        result = self.invoke_prompt(prompt)  # type: ignore
+
+        res = {  # type: ignore
+            **state,
+            "feedback": result.content,  # type: ignore
+            "messages": [*state.get("messages", []), *prompt.messages, result],
+        }
+        return res  # type: ignore
+
+    def rewrite_query(self, state: State) -> State:
+        """Rewrite SQL query if necessary."""
+        prompt: Any = REWRITE_QUERY_PROMPT.invoke(  # type: ignore
+            {
+                "dialect": self.db.dialect,
+                "input": state["question"],
+                "query": state["query"],
+                "execution": self.truncate_execution(state["execution"]),
+                "feedback": state["feedback"],
+                "table_info": self.get_table_info(),
+            }
+        )
+        result = self.invoke_prompt(prompt)  # type: ignore
+
+        rewritten_query = self.parse_query(result)  # type: ignore
+
+        return {
+            **state,
+            "query": rewritten_query or state["query"],
+            "num_turns": state.get("num_turns", 0) + 1,
+            "messages": [*prompt.messages, result],  # clear previous prompts
+        }
+
+    def should_continue(self, state: State) -> Literal[END, "rewrite_query"]:  # type: ignore
+        """Determine if the agent should continue based on the result."""
+        if state["messages"] and isinstance(state["messages"][-1], BaseMessage):  # type: ignore
+            last_message = state["messages"][-1]
+            if "THE QUERY IS CORRECT" in last_message.content:  # type: ignore
+                if "THE QUERY IS INCORRECT" in last_message.content:  # type: ignore
+                    # Both correct and incorrect messages found
+                    # See which is the last one
+                    correct_index = last_message.content.rfind("THE QUERY IS CORRECT")  # type: ignore
+                    incorrect_index = last_message.content.rfind("THE QUERY IS INCORRECT")  # type: ignore
+                    if correct_index > incorrect_index:
+                        return END
+                else:
+                    return END
+
+        if state.get("num_turns", 0) >= self.max_turns:
+            return END
+
+        return "rewrite_query"
+
+    def graph(self) -> CompiledStateGraph[State]:
+        builder = StateGraph(State)
+        builder.add_node(self.write_query)  # type: ignore
+        builder.add_node(self.execute_query)  # type: ignore
+        builder.add_node(self.check_query)  # type: ignore
+        builder.add_node(self.rewrite_query)  # type: ignore
+
+        builder.add_edge(START, "write_query")
+        builder.add_edge("write_query", "execute_query")
+        builder.add_edge("execute_query", "check_query")
+        builder.add_conditional_edges(
+            "check_query",
+            self.should_continue,  # type: ignore
+        )
+        builder.add_edge("rewrite_query", "execute_query")
+
+        return builder.compile()  # type: ignore
+
+
+def evaluate_query(query: str, ground_truth: str, database: str, raise_on_error: bool = True) -> float:
+    # TODO(yuge): Maybe we can evaluate intermediate queries and assign more precise rewards.
+
+    # included in the original evaluation script
+    # query = query.replace("value", "1")
+
+    try:
+        database = os.path.abspath(database)
+        if not os.path.exists(database):
+            raise FileNotFoundError(f"Database file {database} does not exist.")
+
+        # Parameters following the default setting
+        exec_score = eval_exec_match(
+            db=database,
+            p_str=query,
+            g_str=ground_truth,
+            plug_value=False,
+            keep_distinct=False,
+            progress_bar_for_each_datapoint=False,
+        )
+        if exec_score == 1:
+            return 1.0
+        else:
+            return 0.0
+    except Exception as e:
+        if raise_on_error:
+            raise
+        else:
+            logger.exception(f"Error evaluating query: {e}")
+            return 0.0
+
+
+class LitSQLAgent(agl.LitAgent[Dict[str, Any]]):
+
+    def __init__(
+        self,
+        trained_agents: Optional[str] = r"write",
+        val_temperature: Optional[float] = None,
+        max_turns: int = 3,
+        table_info_truncate: int = 2048,
+        execution_truncate: int = 2048,
+    ) -> None:
+        super().__init__(trained_agents=trained_agents)
+        self.val_temperature = val_temperature
+        self.spider_dir = os.environ.get("VERL_SPIDER_DATA_DIR", "data")
+        self.max_turns = max_turns
+        self.table_info_truncate = table_info_truncate
+        self.execution_truncate = execution_truncate
+
+    def rollout(
+        self,
+        task: Dict[str, Any],
+        resources: agl.NamedResources,
+        rollout: agl.Rollout,
+    ) -> float | None:
+        question = task["question"]
+        start_time = time.time()
+        llm: agl.LLM = cast(agl.LLM, resources["main_llm"])
+
+        if rollout.mode == "train":
+            original_db_path = os.path.join(self.spider_dir, "database", task["db_id"], task["db_id"] + ".sqlite")
+        else:
+            original_db_path = os.path.join(self.spider_dir, "test_database", task["db_id"], task["db_id"] + ".sqlite")
+        ground_truth = task["query"]
+
+        if not os.path.exists(original_db_path):
+            logger.error(f"Database {original_db_path} does not exist. Skipping.")
+            return None
+
+        schema_path = os.path.join(os.path.dirname(original_db_path), "schema.sql")
+        if os.path.exists(schema_path):
+            with open(schema_path, "r") as f:
+                schema = f.read()
+        else:
+            logger.error("Schema file not found: %s", schema_path)
+            schema = "No schema available."
+
+        rollout_id = rollout.rollout_id
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            db_path = os.path.join(temp_dir, os.path.basename(original_db_path))
+            shutil.copyfile(original_db_path, db_path)
+            logger.info(f"[Rollout {rollout_id}] Question: {question}")
+            logger.info(f"[Rollout {rollout_id}] Ground Truth: {ground_truth}")
+
+            # Run the agent
+            agent = SQLAgent(
+                "sqlite:///" + db_path,
+                max_turns=self.max_turns,
+                table_info_truncate=self.table_info_truncate,
+                execution_truncate=self.execution_truncate,
+                debug=False,
+                db_schema=schema,
+                endpoint=llm.get_base_url(rollout.rollout_id, rollout.attempt.attempt_id),  # type: ignore
+                verl_replacement=(
+                    {"model": llm.model, **llm.sampling_parameters}
+                    if rollout.mode == "train"
+                    else {
+                        "model": llm.model,
+                        "temperature": (
+                            self.val_temperature
+                            if self.val_temperature is not None
+                            else llm.sampling_parameters.get("temperature", 0.0)
+                        ),
+                    }
+                ),
+            ).graph()
+            try:
+                # Required to make the langchain tracing work
+                handler = self.tracer.get_langchain_handler()
+                result = agent.invoke(  # type: ignore
+                    {"question": question},  # type: ignore
+                    {"callbacks": [handler] if handler else [], "recursion_limit": 100},
+                )
+            except Exception as e:
+                logger.exception(f"[Rollout {rollout_id}] Error during agent invocation: {e}")
+                return
+
+            logger.info(f"[Rollout {rollout_id}] Generated Query: {result['query']}")
+
+        end_time_rollout = time.time()
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            db_path = os.path.join(temp_dir, os.path.basename(original_db_path))
+            shutil.copyfile(original_db_path, db_path)
+
+            reward = evaluate_query(result["query"], ground_truth, db_path, raise_on_error=False)
+            logger.info("[Rollout %s] Reward: %s", rollout_id, reward)
+
+        end_time_eval = time.time()
+
+        logger.info("[Rollout %s] Time taken for rollout: %.2f seconds", rollout_id, end_time_rollout - start_time)
+        logger.info(
+            "[Rollout %s] Time taken for evaluation: %.2f seconds", rollout_id, end_time_eval - end_time_rollout
+        )
+
+        return reward
+
+
+def debug_sql_agent():
+    spider_dev_data_path = os.path.join(os.environ.get("VERL_SPIDER_DATA_DIR", "data"), "dev.parquet")
+    if not os.path.exists(spider_dev_data_path):
+        raise FileNotFoundError(f"Spider dev data file {spider_dev_data_path} does not exist.")
+    df = pd.read_parquet(spider_dev_data_path).head(10)  # type: ignore
+    df = cast(List[Dict[str, Any]], df.to_dict(orient="records"))  # type: ignore
+    print("Debug data:", df)
+
+    trainer = agl.Trainer(
+        n_workers=1,
+        initial_resources={
+            "main_llm": agl.LLM(
+                endpoint=os.environ["OPENAI_API_BASE"],
+                model="gpt-4.1-nano",
+                sampling_parameters={"temperature": 0.7},
+            )
+        },
+    )
+    trainer.dev(LitSQLAgent(), df)
+
+
+if __name__ == "__main__":
+    debug_sql_agent()

examples/spider/train_sql_agent.py

@@ -0,0 +1,217 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Train an SQL agent on the Spider dataset using Agent-lightning.
+
+This module provides a training script for SQL agents using different model configurations.
+The script supports three different training configurations:
+
+1. 'fast' - A lightweight configuration optimized for CI testing with reduced epochs
+2. 'qwen' - Standard configuration using Qwen-2.5-Coder-1.5B-Instruct model
+3. 'llama' - Configuration using LLaMA-3.2-1B-Instruct model with JSON formatting
+
+Usage:
+    python train_sql_agent.py fast    # Fast training for CI/testing
+    python train_sql_agent.py qwen    # Standard Qwen model training
+    python train_sql_agent.py llama   # LLaMA model training
+
+The script uses reinforcement learning with VERL framework
+to train agents on the Spider dataset for text-to-SQL generation tasks.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+from copy import deepcopy
+from datetime import datetime
+from typing import Any, Dict, Optional
+
+import pandas as pd
+from sql_agent import LitSQLAgent
+
+import agentlightning as agl
+
+RL_TRAINING_CONFIG: Dict[str, Any] = {
+    "algorithm": {
+        "adv_estimator": "grpo",
+        "use_kl_in_reward": False,
+    },
+    "data": {
+        "train_files": "data/train_spider.parquet",
+        "val_files": "data/test_dev_500.parquet",
+        "train_batch_size": 32,
+        "max_prompt_length": 4096,
+        "max_response_length": 2048,
+        "truncation": "error",
+    },
+    "actor_rollout_ref": {
+        "rollout": {
+            "tensor_model_parallel_size": 1,
+            "n": 4,
+            "log_prob_micro_batch_size_per_gpu": 4,
+            "multi_turn": {"format": "hermes"},
+            "name": "vllm",
+            "gpu_memory_utilization": 0.8,
+            "engine_kwargs": {
+                "vllm": {
+                    "enable_auto_tool_choice": True,
+                    "tool_call_parser": "hermes",
+                }
+            },
+        },
+        "actor": {
+            "ppo_mini_batch_size": 32,
+            "ppo_micro_batch_size_per_gpu": 4,
+            "optim": {"lr": 1e-6},
+            "use_kl_loss": False,
+            "kl_loss_coef": 0.0,
+            "entropy_coeff": 0,
+            "clip_ratio_low": 0.2,
+            "clip_ratio_high": 0.3,
+            "fsdp_config": {
+                "param_offload": True,
+                "optimizer_offload": True,
+            },
+        },
+        "ref": {
+            "log_prob_micro_batch_size_per_gpu": 8,
+            "fsdp_config": {"param_offload": True},
+        },
+        "model": {
+            "path": "Qwen/Qwen2.5-Coder-1.5B-Instruct",
+            "use_remove_padding": True,
+            "enable_gradient_checkpointing": True,
+        },
+    },
+    "trainer": {
+        "n_gpus_per_node": 1,
+        "val_before_train": True,
+        "critic_warmup": 0,
+        "logger": ["console", "wandb"],
+        "project_name": "AgentLightning",
+        "experiment_name": "spider",
+        "nnodes": 1,
+        "test_freq": 32,
+        "total_epochs": 2,
+    },
+}
+
+
+def config_train_fast() -> Dict[str, Any]:
+    """A fast training run for CI testing purposes."""
+
+    # `EXPERIMENT_NAME="spider_$(date +%Y%m%d%H%M%S)"`
+    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
+    EXPERIMENT_NAME = f"spider_{timestamp}"
+
+    # `PROJECT_NAME=AgentLightningCI`
+    PROJECT_NAME = "AgentLightningCI"
+
+    # Simulate writing to $GITHUB_OUTPUT if it’s set
+    github_output = os.getenv("GITHUB_OUTPUT")
+    if github_output:
+        with open(github_output, "a") as f:
+            f.write(f"project_name={PROJECT_NAME}\n")
+            f.write(f"run_name={EXPERIMENT_NAME}\n")
+
+    print("Set environment variables:")
+    print(f"PROJECT_NAME={PROJECT_NAME}")
+    print(f"EXPERIMENT_NAME={EXPERIMENT_NAME}")
+
+    config = deepcopy(RL_TRAINING_CONFIG)
+    config["actor_rollout_ref"]["rollout"]["gpu_memory_utilization"] = 0.6
+    config["actor_rollout_ref"]["model"]["path"] = "Qwen/Qwen2.5-Coder-0.5B-Instruct"
+    config["data"]["val_files"] = "data/test_dev.parquet"
+    config["trainer"]["total_epochs"] = 1
+    config["trainer"]["total_training_steps"] = 1
+    config["trainer"]["experiment_name"] = EXPERIMENT_NAME
+    config["trainer"]["project_name"] = PROJECT_NAME
+    config["trainer"]["test_freq"] = 1
+    return config
+
+
+def config_train_qwen() -> Dict[str, Any]:
+    """A configuration for training with Qwen-2.5B."""
+
+    config = deepcopy(RL_TRAINING_CONFIG)
+    return config
+
+
+def config_train_npu() -> Dict[str, Any]:
+    """A configuration for training with NPU."""
+
+    config = deepcopy(RL_TRAINING_CONFIG)
+    del config["actor_rollout_ref"]["rollout"]["engine_kwargs"]["vllm"]["enable_auto_tool_choice"]
+    del config["actor_rollout_ref"]["rollout"]["engine_kwargs"]["vllm"]["tool_call_parser"]
+    del config["trainer"]["logger"][1]
+    config["actor_rollout_ref"]["actor"]["use_torch_compile"] = False
+    config["trainer"]["val_before_train"] = False
+    config["trainer"]["save_freq"] = 256
+    config["trainer"]["device"] = "npu"
+    return config
+
+
+def config_train_llama() -> Dict[str, Any]:
+    """A configuration for training with LLaMA-3.2-1B-Instruct.
+
+    You will need a `HF_TOKEN` set to run with this config.
+    """
+
+    config = deepcopy(RL_TRAINING_CONFIG)
+    config["actor_rollout_ref"]["rollout"]["multi_turn"]["format"] = "llama3_json"
+    config["actor_rollout_ref"]["rollout"]["engine_kwargs"]["vllm"]["tool_call_parser"] = "llama3_json"
+    config["actor_rollout_ref"]["model"]["path"] = "meta-llama/Llama-3.2-1B-Instruct"
+    return config
+
+
+def train(config: Dict[str, Any], active_agent: Optional[str]) -> None:
+    """Train the SQL agent with the given configuration."""
+
+    agent = LitSQLAgent()
+    algorithm = agl.VERL(config)
+    trainer = agl.Trainer(n_runners=10, algorithm=algorithm, adapter={"agent_match": active_agent})
+    print("Adapter agent match acknowledged:", trainer.adapter.agent_match)  # type: ignore
+
+    train_data = pd.read_parquet(config["data"]["train_files"]).to_dict(orient="records")  # type: ignore
+    val_data = pd.read_parquet(config["data"]["val_files"]).to_dict(orient="records")  # type: ignore
+    trainer.fit(agent, train_dataset=train_data, val_dataset=val_data)  # type: ignore
+
+
+def main() -> None:
+    """Main function to parse arguments and run training."""
+    parser = argparse.ArgumentParser(
+        description="Train an SQL agent on the Spider dataset using different model configurations"
+    )
+
+    parser.add_argument(
+        "config",
+        choices=["fast", "qwen", "llama", "npu"],
+        help="Training configuration: 'fast' (CI testing), 'qwen' (Qwen-2.5-Coder-1.5B), 'llama' (LLaMA-3.2-3B),'npu' (Train with NPU)",
+    )
+
+    parser.add_argument(
+        "--active-agent", type=str, help="Override the active agent name (default: auto-generated based on config)"
+    )
+
+    args = parser.parse_args()
+
+    # Get the appropriate configuration
+    config_functions = {
+        "fast": config_train_fast,
+        "qwen": config_train_qwen,
+        "llama": config_train_llama,
+        "npu": config_train_npu,
+    }
+    config = config_functions[args.config]()
+
+    # Set active agent - use provided value or default based on config choice
+    active_agent = args.active_agent
+
+    print(f"Starting training with '{args.config}' configuration...")
+    print(f"Active agent: {active_agent}")
+
+    train(config, active_agent)
+
+
+if __name__ == "__main__":
+    main()