chore: import upstream snapshot with attribution

examples/rag/README.md

@@ -0,0 +1,142 @@
+# RAG Agent Example
+
+[![rag workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/examples-rag.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/examples-rag.yml)
+
+This example demonstrates training a Retrieval-Augmented Generation (RAG) agent using Agent-Lightning with retrieval capabilities. The agent answers multi-hop questions from a tiny MuSiQue dataset by retrieving and reasoning over Wikipedia passages.
+
+## Overview
+
+This example can run on a single GPU for demonstration purposes.
+
+**Step 1:** Set up the environment. It is recommended to setup with uv and activate the virtual environment with:
+
+```bash
+uv sync --frozen --extra apo --group agents --group torch-gpu-stable --extra verl --group rag
+source .venv/bin/activate
+```
+
+**Step 2:** Prepare the tiny dataset.
+
+```bash
+pip install gdown
+
+# tiny training dataset
+cd examples/rag
+gdown --fuzzy "https://drive.google.com/file/d/1Pq4Ag8zVoN8gUtLu0LcBfY35Dm5zL0hq/view?usp=drive_link" \
+    -O dataset_tiny.parquet
+
+# chunks_candidate_tiny.pkl
+gdown --fuzzy "https://drive.google.com/file/d/1REXCpRLbeZu1KfWWKhIGEQe_WNHUOBkS/view?usp=drive_link" \
+    -O chunks_candidate_tiny.pkl
+
+# index_hnsw_faiss_n32e40_tiny.index
+gdown --fuzzy "https://drive.google.com/file/d/1f6P-h_8KSRhe5pqDHWbRQWvUhTygfZ-c/view?usp=drive_link" \
+    -O index_hnsw_faiss_n32e40_tiny.index
+```
+
+**Step 3:** Start the MCP server. Open a terminal and run:
+
+```bash
+python wiki_retriever_mcp.py
+```
+
+**Step 4:** Start training. Open another terminal and run:
+
+```bash
+python train_rag.py
+```
+
+## Included Files
+
+| File/Directory | Description |
+|----------------|-------------|
+| `rag_agent.py` | RAG agent example using the OpenAI Agents SDK, with debugging utils |
+| `train_rag.py` | Initiates the GRPO training process |
+| `metric_utils.py` | Scoring utilities for exact match, F1 score, and response parsing |
+| `wiki_retriever_mcp.py` | MCP server for Wikipedia retrieval |
+
+## How to Prepare the Retrieval Corpus Yourself
+
+To enable semantic retrieval with this MCP server, you need two files:
+
+1. **FAISS index file** (`.index`)
+2. **Chunk list file** (`.pkl`)
+
+These two files work together: the FAISS index stores the vector embeddings and their mapping to integer IDs, while the pickle file stores the actual text chunks. The integer IDs in the index correspond exactly to the positions in the chunk list.
+
+### Step 1: Collecting Text Chunks
+
+First, you need a collection of text passages (chunks). For example, you can download a Wikipedia-based dataset such as `wiki18_100w.zip` from the [FlashRAG_dataset](https://huggingface.co/datasets/FlashRAG) or use other pre-split corpora.
+
+### Step 2: Creating the FAISS Index (`nq_hnsw_faiss_n32e40.index`)
+
+- Use a sentence embedding model (e.g., `BAAI/bge-large-en-v1.5`) to encode each chunk into a vector.
+- Build a FAISS index from these vectors.
+- In this example, we use an **HNSW index** (Hierarchical Navigable Small World graph), which supports efficient approximate nearest-neighbor search.
+- The index stores only embeddings and integer IDs (no raw text).
+
+### Step 3: Creating the Chunk List (`nq_list.pkl`)
+
+- Store the raw text chunks in a Python list.
+- Save this list with `pickle`.
+- The index ID returned by FAISS corresponds to the list index in this file. For example, if FAISS search returns `I[0][i] = 12345`, then the corresponding text chunk is `chunks[12345]`.
+
+### Example Schema
+
+- **`nq_hnsw_faiss_n32e40.index`**
+  - Type: FAISS HNSW index
+  - Contains:
+    - Vector embeddings
+    - Graph structure for fast search
+    - Integer IDs mapping to chunk positions
+
+- **`nq_list.pkl`**
+  - Type: Pickled Python list
+  - Element type: string (or dict with text + metadata, depending on preprocessing)
+  - Example:
+    ```python
+    [
+        "The Eiffel Tower is located in Paris, France.",
+        "Albert Einstein developed the theory of relativity.",
+        ...
+    ]
+    ```
+
+### Step 4: Code Example - Building Index and Chunk List
+
+**Warning:** The following example demonstrates a small-scale workflow only. In practice, for large datasets, you should encode the text in batches and incrementally add them to the index.
+
+```python
+import faiss
+import pickle
+from sentence_transformers import SentenceTransformer
+
+# 1. Prepare your text chunks (list of strings)
+chunk_texts = [
+    "The Eiffel Tower is located in Paris, France.",
+    "Albert Einstein developed the theory of relativity.",
+    "Python is a popular programming language.",
+    # ... more chunks
+]
+
+# 2. Load embedding model
+model = SentenceTransformer("BAAI/bge-large-en-v1.5")
+
+# 3. Encode text chunks into embeddings
+embeddings = model.encode(chunk_texts, normalize_embeddings=True)
+
+# 4. Build FAISS HNSW index
+dim = embeddings.shape[1]
+index = faiss.IndexHNSWFlat(dim, 32)   # 32 neighbors by default
+index.hnsw.efConstruction = 40         # efConstruction parameter
+index.add(embeddings)
+
+# 5. Save FAISS index
+faiss.write_index(index, "nq_hnsw_faiss_n32e40.index")
+
+# 6. Save chunk list
+with open("nq_list.pkl", "wb") as f:
+    pickle.dump(chunk_texts, f)
+
+print("Index and chunk list saved successfully.")
+```

examples/rag/metric_utils.py

@@ -0,0 +1,443 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# type: ignore
+
+import re
+import string
+from collections import Counter
+from typing import List, Optional, Set, Tuple
+
+ANS_BEGIN = "<answer>"
+ANS_END = "</answer>"
+GEN_BEGIN = "<|im_start|>assistant\n"
+FORMAT_SCORE = 0.1
+FORMAT_PUNISH = -2
+
+
+def normalize_answer(s: str) -> str:
+    def remove_articles(text: str) -> str:
+        return re.sub(r"\b(a|an|the)\b", " ", text)
+
+    def white_space_fix(text: str) -> str:
+        return " ".join(text.split())
+
+    def remove_punc(text: str) -> str:
+        exclude = set(string.punctuation)
+        return "".join(ch for ch in text if ch not in exclude)
+
+    def lower(text: str) -> str:
+        return text.lower()
+
+    return white_space_fix(remove_articles(remove_punc(lower(s))))
+
+
+def f1_score(prediction: str, ground_truth: str) -> Tuple[float, float, float]:
+    normalized_prediction = normalize_answer(prediction)
+    normalized_ground_truth = normalize_answer(ground_truth)
+
+    ZERO_METRIC = (0, 0, 0)
+
+    if normalized_prediction in ["yes", "no", "noanswer"] and normalized_prediction != normalized_ground_truth:
+        return ZERO_METRIC
+    if normalized_ground_truth in ["yes", "no", "noanswer"] and normalized_prediction != normalized_ground_truth:
+        return ZERO_METRIC
+
+    prediction_tokens = normalized_prediction.split()
+    ground_truth_tokens = normalized_ground_truth.split()
+    common = Counter(prediction_tokens) & Counter(ground_truth_tokens)
+    num_same = sum(common.values())
+    if num_same == 0:
+        return ZERO_METRIC
+    precision = 1.0 * num_same / len(prediction_tokens)
+    recall = 1.0 * num_same / len(ground_truth_tokens)
+    f1 = (2 * precision * recall) / (precision + recall)
+    return f1, precision, recall
+
+
+def lenient_f1_score(prediction: str, ground_truth: str) -> Tuple[float, float, float]:
+    normalized_prediction = normalize_answer(prediction)
+    normalized_ground_truth = normalize_answer(ground_truth)
+
+    ZERO_METRIC = (0, 0, 0)
+
+    if normalized_ground_truth in ["yes", "no", "noanswer"] and normalized_prediction != normalized_ground_truth:
+        if normalized_ground_truth == "yes" and ("no" in normalized_prediction or "noanswer" in normalized_prediction):
+            return ZERO_METRIC
+        if normalized_ground_truth == "no" and ("yes" in normalized_prediction or "noanswer" in normalized_prediction):
+            return ZERO_METRIC
+
+    prediction_tokens = normalized_prediction.split()
+    ground_truth_tokens = normalized_ground_truth.split()
+    common = Counter(prediction_tokens) & Counter(ground_truth_tokens)
+    num_same = sum(common.values())
+    if num_same == 0:
+        return ZERO_METRIC
+    precision = 1.0 * num_same / len(prediction_tokens)
+    recall = 1.0 * num_same / len(ground_truth_tokens)
+    f1 = (2 * precision * recall) / (precision + recall)
+    return f1, precision, recall
+
+
+def exact_match_score(prediction: str, ground_truth: str) -> bool:
+    return normalize_answer(prediction) == normalize_answer(ground_truth)
+
+
+def cover_exact_match_score(prediction: str, ground_truth: str) -> bool:
+    return normalize_answer(ground_truth) in normalize_answer(prediction)
+
+
+def extract_answer(response: str) -> str:
+    if ANS_BEGIN not in response or ANS_END not in response:
+        return ""
+    pos1 = response.rfind(ANS_BEGIN)
+    pos2 = response.rfind(ANS_END)
+    assert pos2 != -1
+    if pos1 != -1:
+        ans = response[pos1 + len(ANS_BEGIN) : pos2]
+    else:
+        ans = response[len(ANS_BEGIN) : pos2]
+    return ans
+
+
+def split_response(text: str) -> Tuple[str, str]:
+    start_response = text.rfind(GEN_BEGIN)
+    response = text[start_response + len(GEN_BEGIN) :]
+    prompt = text[: -len(response)]
+    return prompt, response
+
+
+def extract_recall_chunk(prompt: str, response: str) -> Tuple[Set[str], Set[str]]:
+    import re
+
+    # Regular expression to match content after 1. and 2. within each search_step
+    pattern = r"Retrieved sentences:\s*1\.\s*(.*?)\s*2\.\s*(.*?)(?:\n\s*\d+\.|\n\n|$)"
+
+    # Use re.findall to extract all (s1, s2) pairs
+    origin_recall = re.findall(pattern, prompt, re.DOTALL)
+    sequential_recall = re.findall(pattern, response, re.DOTALL)
+    origin_recall_set = set(s for pair in origin_recall for s in pair)
+    sequential_recall_set = set(s for pair in sequential_recall for s in pair)
+
+    return origin_recall_set, sequential_recall_set
+
+
+def extract_retrieved_paragraphs(log_text: str) -> List[str]:
+    # Regular expression to match content after "Retrieved paragraph:"
+    pattern = re.compile(r"Retrieved paragraph:\s*(.*?)\n", re.DOTALL)
+
+    # Extract matched paragraphs
+    matches = pattern.findall(log_text)
+    matches = list(set(matches))
+    return matches
+
+
+def compute_score(
+    prediction: str, gold: str, gold_sentences: Optional[List[str]] = None, data_source: Optional[str] = None
+) -> float:
+    # format acc
+    format_acc = FORMAT_SCORE
+
+    _, response = split_response(prediction)
+    ans = extract_answer(response)
+    if ans == "":
+        # format score 0.1
+        # if '<query>' not in response or '</query>' not in response:
+        #     return 0.0
+        # return 0.0
+        delimiter = "<|im_start|>assistant"
+        last_time_ans = response.split(delimiter)[-1]
+        if "<query" not in last_time_ans or "</query>" not in last_time_ans:
+            return 0.0
+        return format_acc
+
+    # answer acc
+    em, _ = exact_match_score(ans, gold), cover_exact_match_score(ans, gold)
+    f1, _, _ = f1_score(ans, gold)
+
+    if fact_checking_api(prediction, ans):
+        answer_acc = max(float(em), f1)
+    else:
+        answer_acc = 0
+    # # search acc
+    # if gold_sentences and search_weight:
+    #     origin_recall_set, sequential_recall_set = extract_recall_chunk(prompt, response)
+    #     gold_sentences_set = set(gold_sentences) - origin_recall_set
+    #     matched = gold_sentences_set & sequential_recall_set
+    #     search_acc = len(matched) / len(gold_sentences_set) if len(gold_sentences_set) != 0 else 1.0
+    #     # print(f's_acc {search_acc}|a_acc {answer_acc=}| score {format_acc + (1 - format_acc) * (search_weight + (1 - search_weight) * answer_acc)} |m_len {len(matched)}|g_len {len(gold_sentences_set)}|o_len {len(origin_recall_set)}|s_len {len(sequential_recall_set)}|{gold_sentences_set}|{sequential_recall_set}')
+    #     if search_acc < 1:
+    #         return format_acc + (1 - format_acc) * search_weight * search_acc
+    # # print(f'SCORE: {score} | {ans} | {gold} | {prediction}' )
+
+    return format_acc + (1 - format_acc) * answer_acc
+    # return  answer_acc
+
+
+def compute_reward(
+    solution_str: Optional[str] = None,
+    ground_truth: Optional[str] = None,
+    gold_sentences: Optional[List[str]] = None,
+    data_source: Optional[str] = None,
+    extra_info: Optional[str] = None,
+) -> float:
+    prediction = solution_str
+    gold = ground_truth
+    return compute_score(prediction, gold, gold_sentences=gold_sentences, data_source=data_source)
+
+
+def compute_em(
+    solution_str: Optional[str] = None,
+    ground_truth: Optional[str] = None,
+    gold_sentences: Optional[List[str]] = None,
+    data_source: Optional[str] = None,
+    extra_info: Optional[str] = None,
+) -> float:
+    prediction = solution_str
+    gold = ground_truth
+    _, response = split_response(prediction)
+    ans = extract_answer(response)
+    if ans == "":
+        # format score 0.1
+        # if '<query>' not in response or '</query>' not in response:
+        #     return 0.0
+        return 0.0
+
+    # answer acc
+    em = exact_match_score(ans, gold)
+    return em
+
+
+def compute_cem(
+    solution_str=None,
+    ground_truth=None,
+    gold_sentences=None,
+    data_source=None,
+    extra_info=None,
+):
+    prediction = solution_str
+    gold = ground_truth
+    _, response = split_response(prediction)
+    ans = extract_answer(response)
+    if ans == "":
+        return 0.0
+
+    # answer acc
+    cem = cover_exact_match_score(ans, gold)
+    return cem
+
+
+def compute_response_cem(
+    solution_str=None,
+    ground_truth=None,
+    gold_sentences=None,
+    data_source=None,
+    extra_info=None,
+):
+    prediction = solution_str
+    gold = ground_truth
+    _, response = split_response(prediction)
+    ans = response
+    if ans == "":
+        return 0.0
+
+    # answer acc
+    cem = cover_exact_match_score(ans, gold)
+    return cem
+
+
+def compute_lenient_f1(
+    solution_str=None,
+    ground_truth=None,
+    gold_sentences=None,
+    data_source=None,
+    extra_info=None,
+):
+    prediction = solution_str
+    gold = ground_truth
+    _, response = split_response(prediction)
+    ans = extract_answer(response)
+    if ans == "":
+        return 0.0
+
+    # answer acc
+    f1, prec, recall = lenient_f1_score(ans, gold)
+    return f1
+
+
+def compute_lenient_response_f1(
+    solution_str=None,
+    ground_truth=None,
+    gold_sentences=None,
+    data_source=None,
+    extra_info=None,
+):
+    prediction = solution_str
+    gold = ground_truth
+    _, response = split_response(prediction)
+    ans = response
+    if ans == "":
+        return 0.0
+
+    # answer acc
+    f1, prec, recall = lenient_f1_score(ans, gold)
+    return f1
+
+
+def fact_checking_api(prediction: str, ans: str) -> bool:
+    return True  # Placeholder for actual fact-checking logic
+
+
+def compute_f1(
+    solution_str: Optional[str] = None,
+    ground_truth: Optional[str] = None,
+    gold_sentences: Optional[List[str]] = None,
+    data_source: Optional[str] = None,
+    extra_info: Optional[str] = None,
+) -> float:
+    prediction = solution_str
+    gold = ground_truth
+    _, response = split_response(prediction)
+    ans = extract_answer(response)
+    if ans == "":
+        return 0.0
+
+    # answer acc
+    f1, _, _ = f1_score(ans, gold)
+    return f1
+
+
+def compute_format(
+    solution_str: Optional[str] = None,
+    ground_truth: Optional[str] = None,
+    gold_sentences: Optional[List[str]] = None,
+    data_source: Optional[str] = None,
+    extra_info: Optional[str] = None,
+) -> float:
+    prediction = solution_str
+    gold = ground_truth
+    _, response = split_response(prediction)
+    ans = extract_answer(response)
+    if ans == "":
+        delimiter = "<|im_start|>assistant"
+        last_time_ans = response.split(delimiter)[-1]
+        if "<query" not in last_time_ans or "</query>" not in last_time_ans:
+            return 0
+    return FORMAT_SCORE
+
+
+def split_trace(text: str) -> Tuple[str, str]:
+    start_response = text.find(GEN_BEGIN)
+    response = text[start_response + len(GEN_BEGIN) :]
+    prompt = text[: -len(response)]
+    return prompt, response
+
+
+def compute_action_query(
+    solution_str: Optional[str] = None,
+    ground_truth: Optional[str] = None,
+    gold_sentences: Optional[List[str]] = None,
+    data_source: Optional[str] = None,
+    extra_info: Optional[str] = None,
+) -> int:
+    prediction = solution_str
+    gold = ground_truth
+    prompt, trace = split_trace(prediction)
+    res = min(trace.count("<query>") + trace.count("<query,"), trace.count("</query>"))
+    return res
+
+
+def compute_action_bm25(
+    solution_str: Optional[str] = None,
+    ground_truth: Optional[str] = None,
+    gold_sentences: Optional[List[str]] = None,
+    data_source: Optional[str] = None,
+    extra_info: Optional[str] = None,
+) -> int:
+    prediction = solution_str
+    gold = ground_truth
+    prompt, trace = split_trace(prediction)
+    res = min(trace.count("<query keyword"), trace.count("</query>"))
+    return res
+
+
+def compute_action_read_pre(
+    solution_str: Optional[str] = None,
+    ground_truth: Optional[str] = None,
+    gold_sentences: Optional[List[str]] = None,
+    data_source: Optional[str] = None,
+    extra_info: Optional[str] = None,
+) -> int:
+    prediction = solution_str
+    gold = ground_truth
+    prompt, trace = split_trace(prediction)
+    res = min(trace.count("<query previous"), trace.count("</query>"))
+    return res
+
+
+def compute_action_read_nxt(
+    solution_str: Optional[str] = None,
+    ground_truth: Optional[str] = None,
+    gold_sentences: Optional[List[str]] = None,
+    data_source: Optional[str] = None,
+    extra_info: Optional[str] = None,
+) -> int:
+    prediction = solution_str
+    gold = ground_truth
+    prompt, trace = split_trace(prediction)
+    res = min(trace.count("<query next"), trace.count("</query>"))
+    return res
+
+
+def compute_action_continue(
+    solution_str: Optional[str] = None,
+    ground_truth: Optional[str] = None,
+    gold_sentences: Optional[List[str]] = None,
+    data_source: Optional[str] = None,
+    extra_info: Optional[str] = None,
+) -> int:
+    prediction = solution_str
+    gold = ground_truth
+    prompt, trace = split_trace(prediction)
+    res = min(trace.count(", continue"), trace.count("</query>"))
+    return res
+
+
+def compute_action_match(
+    solution_str: Optional[str] = None,
+    ground_truth: Optional[str] = None,
+    gold_sentences: Optional[List[str]] = None,
+    data_source: Optional[str] = None,
+    extra_info: Optional[str] = None,
+) -> int:
+    prediction = solution_str
+    gold = ground_truth
+    prompt, trace = split_trace(prediction)
+    res = min(trace.count(', match_phrase="'), trace.count("</query>"))
+    return res
+
+
+def compute_total_action_number(
+    solution_str: Optional[str] = None,
+    ground_truth: Optional[str] = None,
+    gold_sentences: Optional[List[str]] = None,
+    data_source: Optional[str] = None,
+    extra_info: Optional[str] = None,
+) -> int:
+    prediction = solution_str
+    gold = ground_truth
+    prompt, trace = split_trace(prediction)
+    res = min(trace.count("<query"), trace.count("</query>"))
+    return res
+
+
+# define reward functions for evaluation
+
+
+def compute_scores(answer: str, ground_truth: str) -> float:
+    parsed_answer = extract_answer(answer)
+    if parsed_answer is None:
+        return -0.1
+    f1, precision, recall = f1_score(parsed_answer, ground_truth)
+    # em = float(exact_match_score(parsed_answer, ground_truth))
+    # cem = float(cover_exact_match_score(answer, ground_truth))
+    return f1

examples/rag/rag_agent.py

@@ -0,0 +1,132 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List, cast
+
+import pandas as pd
+from agents import Agent, Runner
+from agents.extensions.models.litellm_model import LitellmModel
+from agents.mcp import MCPServerSse
+from agents.model_settings import ModelSettings
+from metric_utils import compute_scores
+
+import agentlightning as agl
+
+logger = logging.getLogger("rag_agent")
+
+agent_prompt = """You are an assistant who answers questions using Wikipedia retriever. Answer the question using only the retrieved passages. Verify your answer directly against the text.
+
+After each search:
+- Summarize findings.
+- Decide if info is sufficient.
+  - If sufficient: reply in <answer>...</answer> with your answer. The answer must be extremely concise: a single word or a few words only.
+  - If not: suggest the next search needed to fill info gaps. The system will return top 3 relevant Wikipedia chunks.
+- Explain your reasoning for the chosen action.
+
+Repeat as needed. When done, wrap your final, concise answer in <answer> tags."""
+
+
+class RAGAgent(agl.LitAgent[Dict[str, Any]]):
+    """RAGAgent is an agent that relies on a MCP-based retriever to answer questions."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.mcp_server_url = "http://127.0.0.1:8099/sse"
+
+    async def training_rollout_async(
+        self, task: Dict[str, Any], resources: agl.NamedResources, rollout: agl.Rollout
+    ) -> float | None:
+        # llm resources
+        llm = cast(agl.LLM, resources["main_llm"])
+
+        # The rollout should carry an attempt inside
+        rollout = cast(agl.AttemptedRollout, rollout)
+        base_url = llm.get_base_url(rollout.rollout_id, rollout.attempt.attempt_id)
+
+        logger.info(f"Training with model: {llm.model} on endpoint: {base_url}")
+
+        async with MCPServerSse(
+            name="wiki_retriever_mcp",
+            params={"url": self.mcp_server_url},
+        ) as server:
+            agent = Agent(
+                model=LitellmModel(
+                    model="hosted_vllm/" + llm.model,
+                    base_url=base_url,
+                ),
+                model_settings=ModelSettings(
+                    max_tokens=2048,
+                    temperature=0.7,
+                ),
+                name="Assistant",
+                instructions=agent_prompt,
+                mcp_servers=[server],
+            )
+            result = await Runner.run(agent, task["question"])
+            answer = result.final_output
+
+            # reward
+            reward = compute_scores(answer, str(task["answer"]))
+
+            logger.info(
+                "Question: %s\nAnswer: %s\nGround truth: %s\nReward: %s",
+                task["question"],
+                answer,
+                task["answer"],
+                reward,
+            )
+            return float(reward)  # Convert to float for compatibility with the Runner
+
+    async def validation_rollout_async(
+        self, task: Dict[str, Any], resources: agl.NamedResources, rollout: agl.Rollout
+    ) -> float | None:
+        """Validation rollout will share the same logic as the training rollout."""
+        # Same as training rollout, but with different temperature
+        llm = cast(agl.LLM, resources["main_llm"])
+        rollout = cast(agl.AttemptedRollout, rollout)
+
+        # set temperature
+        val_resources: agl.NamedResources = {
+            "main_llm": agl.LLM(
+                endpoint=llm.get_base_url(rollout.rollout_id, rollout.attempt.attempt_id),
+                model=llm.model,
+                sampling_parameters={"temperature": 0.7},
+            )
+        }
+
+        # reuse training rollout for validation
+        return await self.training_rollout_async(task, val_resources, rollout)
+
+
+def debug():
+    """Debug the RAGAgent."""
+
+    agl.setup_logging("DEBUG", apply_to=[logger.name])
+
+    # 1. loading dataset
+    dataset_path = "data/dataset_tiny.parquet"
+    df: pd.DataFrame = pd.read_parquet(dataset_path)  # type: ignore
+    data: List[Dict[str, Any]] = df.head(5).to_dict(orient="records")  # type: ignore
+    # NOTE: The following dummy data can also be used if you don't have the dataset.
+    # data: List[Dict[str, Any]] = [{"question": "What is the capital of France?", "answer": "Paris"}]
+
+    # 2. configuring resources (LLM)
+    # Note: You need to start a local service compatible with the OpenAI API (such as vLLM)
+    # For example: python -m vllm.entrypoints.openai.api_server --model Qwen/Qwen2.5-1.5B-Instruct --port 8000
+    resources: dict[str, agl.ResourceUnion] = {
+        "main_llm": agl.LLM(
+            endpoint="http://localhost:8000/v1",  # Replace with your actual vLLM address
+            model="Qwen/Qwen2.5-1.5B-Instruct",  # Replace with your actual loaded model name
+            sampling_parameters={"temperature": 0.0},
+        )
+    }
+
+    # 3. run agent
+    trainer = agl.Trainer(initial_resources=resources)
+    trainer.dev(RAGAgent(), train_dataset=data)  # type: ignore
+
+
+if __name__ == "__main__":
+    debug()

examples/rag/train_rag.py

@@ -0,0 +1,200 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Train a RAG agent using Agent-lightning.
+
+Usage:
+    python train_rag.py fast        # Fast training for CI/testing
+    python train_rag.py single_gpu  # Optimized for Single GPU (1.5B/7B models)
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import uuid
+from copy import deepcopy
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+import pandas as pd
+from rag_agent import RAGAgent  # Make sure to import your RAGAgent class
+
+import agentlightning as agl
+
+# Base configuration (default configuration, can be overridden)
+RL_TRAINING_CONFIG: Dict[str, Any] = {
+    "algorithm": {
+        "adv_estimator": "grpo",  # Use GRPO algorithm
+        "use_kl_in_reward": False,
+    },
+    "data": {
+        "train_batch_size": 16,  # Default configuration for multi-GPU
+        "max_prompt_length": 8192,
+        "max_response_length": 2048,
+        "truncation": "error",
+    },
+    "actor_rollout_ref": {
+        "rollout": {
+            "tensor_model_parallel_size": 1,
+            "n": 4,  # Generate 4 responses per sampling
+            "log_prob_micro_batch_size_per_gpu": 4,
+            "multi_turn": {"format": "hermes"},  # Ensure using template format matching the model
+            "name": "vllm",
+            "gpu_memory_utilization": 0.6,  # vLLM GPU memory utilization
+            "engine_kwargs": {
+                "vllm": {
+                    "enable_auto_tool_choice": True,
+                    "tool_call_parser": "hermes",
+                }
+            },
+        },
+        "actor": {
+            "ppo_mini_batch_size": 16,
+            "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,  # Enable parameter offloading to save GPU memory
+                "optimizer_offload": True,
+            },
+        },
+        "ref": {
+            "log_prob_micro_batch_size_per_gpu": 8,
+            "fsdp_config": {"param_offload": True},
+        },
+        "model": {
+            "path": "Qwen/Qwen2.5-1.5B-Instruct",  # Default model
+            "use_remove_padding": True,
+            "enable_gradient_checkpointing": True,
+        },
+    },
+    "trainer": {
+        "n_gpus_per_node": 1,
+        "val_before_train": True,
+        "critic_warmup": 0,
+        "logger": ["console"],  # Disable wandb for easier local debugging, add back when needed
+        "project_name": "AgentLightning",
+        "experiment_name": "rag_agent",
+        "nnodes": 1,
+        "test_freq": 10,
+        "total_epochs": 200,
+    },
+}
+
+
+def config_train_fast() -> Dict[str, Any]:
+    """Fast training configuration for CI/testing"""
+    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
+    random_suffix = uuid.uuid4().hex[:8]
+    EXPERIMENT_NAME = f"rag_fast_{timestamp}_{random_suffix}"
+
+    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)
+
+    # Keep it tiny/light without adding new knobs
+    config["actor_rollout_ref"]["rollout"]["gpu_memory_utilization"] = 0.8
+    config["trainer"]["total_epochs"] = 2
+    config["trainer"]["test_freq"] = 5
+    config["trainer"]["experiment_name"] = EXPERIMENT_NAME
+    config["trainer"]["project_name"] = PROJECT_NAME
+    config["trainer"]["logger"] = ["console", "wandb"]
+    return config
+
+
+def config_train_single_gpu() -> Dict[str, Any]:
+    """Single GPU training optimized configuration (optimized for 24GB GPU memory)"""
+
+    config = deepcopy(RL_TRAINING_CONFIG)
+
+    # 1. Reduce vLLM memory usage to leave space for training
+    config["actor_rollout_ref"]["rollout"]["gpu_memory_utilization"] = 0.4
+
+    # 2. Reduce Batch Size to prevent OOM
+    config["data"]["train_batch_size"] = 4
+    config["actor_rollout_ref"]["actor"]["ppo_mini_batch_size"] = 4
+    config["actor_rollout_ref"]["actor"]["ppo_micro_batch_size_per_gpu"] = 1
+    config["actor_rollout_ref"]["rollout"]["log_prob_micro_batch_size_per_gpu"] = 2
+
+    # 3. Ensure Offload is enabled
+    config["actor_rollout_ref"]["actor"]["fsdp_config"]["param_offload"] = True
+    config["actor_rollout_ref"]["actor"]["fsdp_config"]["optimizer_offload"] = True
+
+    return config
+
+
+def train(config: Dict[str, Any], active_agent: Optional[str]) -> None:
+    """Train the RAG agent with the given configuration."""
+
+    # 1. Instantiate your Agent
+    agent = RAGAgent()
+
+    # 2. Initialize algorithm (VERL)
+    algorithm = agl.VERL(config)
+
+    # 3. Initialize Trainer
+    # n_runners=4 means 4 concurrent rollout runners (can be reduced if insufficient memory, or managed internally by VERL)
+    trainer = agl.Trainer(n_runners=4, algorithm=algorithm, adapter={"agent_match": active_agent})
+
+    # 4. Load data
+    # NOTE: Fill in the path to your previously converted parquet file here
+    # For demo purposes, we use the same dataset for training and validation,
+    # which should be avoided in production.
+    train_df: pd.DataFrame = pd.read_parquet("data/dataset_tiny.parquet")  # type: ignore
+    val_df: pd.DataFrame = pd.read_parquet("data/dataset_tiny.parquet")  # type: ignore
+
+    # Keep the rest of the code unchanged
+    train_data: List[Dict[str, Any]] = train_df.to_dict(orient="records")  # type: ignore
+    val_data: List[Dict[str, Any]] = val_df.to_dict(orient="records")  # type: ignore
+
+    # 5. Start training
+    trainer.fit(agent, train_dataset=train_data, val_dataset=val_data)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Train a RAG agent using different configurations")
+
+    parser.add_argument(
+        "config",
+        choices=["fast", "single_gpu"],
+        default="single_gpu",
+        nargs="?",
+        help="Training configuration name",
+    )
+
+    parser.add_argument("--active-agent", type=str, help="Override the active agent name")
+
+    args = parser.parse_args()
+
+    config_functions = {
+        "fast": config_train_fast,
+        "single_gpu": config_train_single_gpu,
+    }
+    config = config_functions[args.config]()
+
+    # Print key information for confirmation
+    print(f"Starting training with '{args.config}' configuration...")
+    print(f"Model: {config['actor_rollout_ref']['model']['path']}")
+    print(f"Batch Size: {config['data']['train_batch_size']}")
+    print(f"GPU Mem Util: {config['actor_rollout_ref']['rollout']['gpu_memory_utilization']}")
+
+    train(config, args.active_agent)
+
+
+if __name__ == "__main__":
+    main()

examples/rag/wiki_retriever_mcp.py

@@ -0,0 +1,52 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# type: ignore
+
+import pickle
+
+import faiss
+from fastmcp import FastMCP
+from sentence_transformers import SentenceTransformer
+
+index = faiss.read_index("data/index_hnsw_faiss_n32e40_tiny.index")
+print("Index loaded successfully.")
+
+model = SentenceTransformer("BAAI/bge-large-en-v1.5")
+print("Model loaded successfully.")
+
+# with open('/mnt/input/agent_lightning/nq_list.pkl', 'rb') as f:
+with open("data/chunks_candidate_tiny.pkl", "rb") as f:
+    chunks = pickle.load(f)
+print("Chunks loaded successfully.")
+
+mcp = FastMCP(name="wiki retrieval mcp")
+
+
+@mcp.tool(
+    name="retrieve",
+    description="retrieve relevant chunks from the wikipedia",
+)
+def retrieve(query: str) -> list:
+    """
+    Retrieve relevant chunks from the Wikipedia dataset.
+
+    Args:
+        query (str): The query string to search for.
+
+    Returns:
+        list: A list of dictionaries containing the retrieved chunks and their metadata.
+    """
+    top_k = 1  # Number of top results to return
+    embedding = model.encode([query], normalize_embeddings=True)
+    D, I = index.search(embedding, top_k)
+
+    results = []
+    for i in range(top_k):
+        if I[0][i] != -1:
+            chunk = chunks[I[0][i]]
+            results.append({"chunk": chunk, "chunk_id": int(I[0][i]), "distance": float(D[0][i])})
+    return results
+
+
+if __name__ == "__main__":
+    mcp.run(transport="sse", host="127.0.0.1", port=8099)