Update pre-commit configurations and bump to 0.2.0

.github/workflows/tests.yml

@@ -29,6 +29,8 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install -e .[dev]
+      - name: Run pre-commit
+        uses: pre-commit/action@v3.0.1
       - name: Check Python headers
         run: |
           python scripts/check_python_headers.py

.gitignore

@@ -183,9 +183,9 @@ cython_debug/
 .abstra/
 
 # Visual Studio Code
-#  Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore 
+#  Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
 #  that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
-#  and can be added to the global gitignore or merged into this file. However, if you prefer, 
+#  and can be added to the global gitignore or merged into this file. However, if you prefer,
 #  you could uncomment the following to ignore the enitre vscode folder
 .vscode/
 

.pre-commit-config.yaml

@@ -1,4 +1,15 @@
 repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v6.0.0
+    hooks:
+      - id: end-of-file-fixer
+      - id: trailing-whitespace
+      - id: check-yaml
+        exclude: ^mkdocs\.yml$
+      - id: check-toml
+      - id: check-added-large-files
+      - id: check-shebang-scripts-are-executable
+      - id: detect-private-key
   - repo: https://github.com/pycqa/isort
     rev: 6.0.1
     hooks:

LICENSE

@@ -16,4 +16,4 @@ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
-THE SOFTWARE.
+THE SOFTWARE.

RAI_README.md

@@ -19,7 +19,7 @@ Agent Lightning was not designed or evaluated for all possible downstream purpos
 
 Agent Lightning should not be used in highly regulated domains where inaccurate outputs could suggest actions that lead to injury or negatively impact an individual's legal, financial, or life opportunities.
 
-We do not recommend using Agent Lightning in the context of high-risk decision making (e.g. in law enforcement, legal, finance, or healthcare). 
+We do not recommend using Agent Lightning in the context of high-risk decision making (e.g. in law enforcement, legal, finance, or healthcare).
 
 ## HOW TO GET STARTED
 To begin using Agent Lightning, here are some instructions.
@@ -33,7 +33,7 @@ To begin using Agent Lightning, here are some instructions.
 Agent Lightning was evaluated on its ability to correctly complete 3 example tasks: (1) Math. The model needs to answer some math questions, and when answering one question, the model can use the calculator as its tool to help answer. (2) Text2SQL. The model is given a question related to the database, and it is required to generate a SQL which can query the database, find the information to answer the question. (3) Retrieval-Augmented Generation (RAG). The model is given a question which needs some information from Wikipedia to answer. The model is required to generate some queries to find the related information in Wikipedia, and answer the question according to retrieved documents.
 
 ### EVALUATION METHODS AND RESULTS
-For detailed evaluation methods and results, please refer to the latest version of our  [technical report](https://arxiv.org/abs/2508.03680). 
+For detailed evaluation methods and results, please refer to the latest version of our  [technical report](https://arxiv.org/abs/2508.03680).
 
 
 ## LIMITATIONS
@@ -43,7 +43,7 @@ Agent Lightning was designed and tested using the English language. Performance
 
 Outputs generated by AI may include factual errors, fabrication, or speculation. Users are responsible for assessing the accuracy of generated content. All decisions leveraging outputs of the system should be made with human oversight and not be based solely on system outputs.
 Agent Lightning inherits any biases, errors, or omissions produced by its base model. Developers are advised to choose an appropriate base LLM/MLLM carefully, depending on the intended use case.
-We use some demo cases to show the effectiveness of our training framework. See their links to understand the capabilities and limitations of this model. 
+We use some demo cases to show the effectiveness of our training framework. See their links to understand the capabilities and limitations of this model.
 
 ## BEST PRACTICES
 Better performance can be achieved by following the instructions in how to get started section.
@@ -54,10 +54,10 @@ We strongly encourage users to use LLMs/MLLMs that support robust Responsible AI
 - [Azure OpenAI Transparency Note](https://learn.microsoft.com/en-us/legal/cognitive-services/openai/transparency-note)
 - [OpenAI’s Usage policies](https://openai.com/policies/usage-policies)
 - [Azure OpenAI’s Code of Conduct](https://learn.microsoft.com/en-us/legal/cognitive-services/openai/code-of-conduct)
-
-Users are responsible for sourcing their datasets legally and ethically. This could include securing appropriate rights, ensuring consent for use of audio/images, and/or the anonymization of data prior to use in research.   
 
-Users are reminded to be mindful of data privacy concerns and are encouraged to review the privacy policies associated with any models and data storage solutions interfacing with Agent Lightning. 
+Users are responsible for sourcing their datasets legally and ethically. This could include securing appropriate rights, ensuring consent for use of audio/images, and/or the anonymization of data prior to use in research.
+
+Users are reminded to be mindful of data privacy concerns and are encouraged to review the privacy policies associated with any models and data storage solutions interfacing with Agent Lightning.
 
 It is the user’s responsibility to ensure that the use of Agent Lightning complies with relevant data protection regulations and organizational guidelines.
 

README.md

@@ -143,13 +143,13 @@ If you find Agent Lightning useful in your research or projects, please cite our
 
 ```bibtex
 @misc{luo2025agentlightningtrainai,
-      title={Agent Lightning: Train ANY AI Agents with Reinforcement Learning}, 
+      title={Agent Lightning: Train ANY AI Agents with Reinforcement Learning},
       author={Xufang Luo and Yuge Zhang and Zhiyuan He and Zilong Wang and Siyun Zhao and Dongsheng Li and Luna K. Qiu and Yuqing Yang},
       year={2025},
       eprint={2508.03680},
       archivePrefix={arXiv},
       primaryClass={cs.AI},
-      url={https://arxiv.org/abs/2508.03680}, 
+      url={https://arxiv.org/abs/2508.03680},
 }
 ```
 

SECURITY.md

@@ -11,4 +11,4 @@ For security reporting information, locations, contact information, and policies
 please review the latest guidance for Microsoft repositories at
 [https://aka.ms/SECURITY.md](https://aka.ms/SECURITY.md).
 
-<!-- END MICROSOFT SECURITY.MD BLOCK -->
+<!-- END MICROSOFT SECURITY.MD BLOCK -->

agentlightning/__init__.py

@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-__version__ = "0.1.2"
+__version__ = "0.2.0"
 
 from .client import AgentLightningClient, DevTaskLoader
 from .config import lightning_cli

docs/deep-dive/server-client-architecture.md

@@ -28,4 +28,4 @@ sequenceDiagram
         TS->>RL: Send Batch of Traces (11)
         RL->>TS: Return Updated Model (12)
     end
-```
+```

docs/index.md

@@ -31,13 +31,13 @@ If you find Agent Lightning useful in your research or projects, please cite our
 
 ```bibtex
 @misc{luo2025agentlightningtrainai,
-      title={Agent Lightning: Train ANY AI Agents with Reinforcement Learning}, 
+      title={Agent Lightning: Train ANY AI Agents with Reinforcement Learning},
       author={Xufang Luo and Yuge Zhang and Zhiyuan He and Zilong Wang and Siyun Zhao and Dongsheng Li and Luna K. Qiu and Yuqing Yang},
       year={2025},
       eprint={2508.03680},
       archivePrefix={arXiv},
       primaryClass={cs.AI},
-      url={https://arxiv.org/abs/2508.03680}, 
+      url={https://arxiv.org/abs/2508.03680},
 }
 ```
 

docs/js/favicon-theme.js

@@ -4,35 +4,35 @@
         // Remove existing favicon links
         const existingFavicons = document.querySelectorAll('link[rel*="icon"]');
         existingFavicons.forEach(link => link.remove());
-        
+
         // Create new favicon link
         const favicon = document.createElement('link');
         favicon.rel = 'icon';
         favicon.type = 'image/png';
-        
+
         // Get the site root by finding how many levels deep we are
         const pathSegments = window.location.pathname.split('/').filter(s => s);
         const siteRoot = window.location.origin + '/' + pathSegments[0] + '/';
-        
+
         // Choose favicon based on theme
         if (isDark) {
             favicon.href = siteRoot + 'assets/logo-dark.png';
         } else {
             favicon.href = siteRoot + 'assets/logo-light.png';
         }
-        
+
         // Add to document head
         document.head.appendChild(favicon);
     }
-    
+
     function updateFavicon() {
         // Check system preference
         const prefersDark = window.matchMedia('(prefers-color-scheme: dark)').matches;
-        
+
         // Check if user has manually selected a theme
         const palette = document.querySelector('[data-md-color-scheme]');
         const scheme = palette ? palette.getAttribute('data-md-color-scheme') : null;
-        
+
         let isDark = false;
         if (scheme === 'slate') {
             isDark = true;
@@ -42,33 +42,33 @@
             // Fall back to system preference
             isDark = prefersDark;
         }
-        
+
         setFavicon(isDark);
     }
-    
+
     // Initial favicon set
     updateFavicon();
-    
+
     // Listen for system theme changes
     window.matchMedia('(prefers-color-scheme: dark)').addEventListener('change', updateFavicon);
-    
+
     // Listen for manual theme changes in MkDocs Material
     const observer = new MutationObserver(function(mutations) {
         mutations.forEach(function(mutation) {
-            if (mutation.type === 'attributes' && 
-                (mutation.attributeName === 'data-md-color-scheme' || 
+            if (mutation.type === 'attributes' &&
+                (mutation.attributeName === 'data-md-color-scheme' ||
                  mutation.attributeName === 'data-md-color-primary')) {
                 updateFavicon();
             }
         });
     });
-    
+
     // Observe the document body for theme changes
     observer.observe(document.body, {
         attributes: true,
         attributeFilter: ['data-md-color-scheme', 'data-md-color-primary']
     });
-    
+
     // Also listen for palette toggle clicks
     document.addEventListener('DOMContentLoaded', function() {
         const toggles = document.querySelectorAll('[data-md-color-scheme]');
@@ -79,4 +79,4 @@
             });
         });
     });
-})();
+})();

docs/quickstart/getting-started.md

@@ -119,7 +119,7 @@ For each prompt, queue a task and wait for results. The `{"prompt": ...}` format
 ```python
 # Queue a task for clients to process
 task_id = await server.queue_task(
-    sample={"prompt": "What is the capital of France?"}, 
+    sample={"prompt": "What is the capital of France?"},
     mode="train"
 )
 

examples/rag/README.md

@@ -17,49 +17,49 @@ This example originally runs on a single node with four GPUs, each requiring at
 
 To enable semantic retrieval with this mcp server, we need two files:
 
-1. **FAISS index file** (`.index`)  
-2. **Chunk list file** (`.pkl`)  
+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
 
-You first need a collection of text passages (chunks). For example, you can download a Wikipedia-based dataset such as `wiki18_100w.zip` in the [FlashRAG_dataset](https://huggingface.co/datasets/FlashRAG) or use other pre-split corpora.  
+You first need a collection of text passages (chunks). For example, you can download a Wikipedia-based dataset such as `wiki18_100w.zip` in 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.  
+- 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 only stores 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`.  
+- 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:  
+- **`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.",

examples/rag/wiki_retriever_mcp/wiki_retriever_install.sh

@@ -1,3 +1,3 @@
 conda create -n mcp_server python=3.12 -y
 conda activate mcp_server
-pip install faiss-cpu==1.11.0 fastmcp==2.5.1 sentence-transformers==4.1.0
+pip install faiss-cpu==1.11.0 fastmcp==2.5.1 sentence-transformers==4.1.0

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "agentlightning"
-version = "0.1.2"
+version = "0.2.0"
 description = "Agent Lightning is the absolute trainer to light up AI agents."
 readme = "README.md"
 requires-python = ">=3.10"

tests/test_tracer.py

@@ -316,7 +316,7 @@ def call_get_schema(state: MessagesState):
     # Generate SQL Query
     def generate_query(state: MessagesState):
         prompt = f"""
-    You are an agent for SQL ({db.dialect}). 
+    You are an agent for SQL ({db.dialect}).
     Write a query to answer the user. Limit results to 5. Do not modify data.
     """
         msg = {"role": "system", "content": prompt}