Sort imports (#100)

.github/workflows/tests.yml

@@ -33,8 +33,9 @@ jobs:
         run: |
           python scripts/check_python_headers.py
       - name: Run Black
-        run: |
-          black --check --diff --line-length=120 .
+        run: black --check .
+      - name: Run isort
+        run: isort --check-only .
 
   docs:
     name: Build documentation

.pre-commit-config.yaml

@@ -1,8 +1,13 @@
 repos:
+  - repo: https://github.com/pycqa/isort
+    rev: 6.0.1
+    hooks:
+      - id: isort
+        args: ["."]
   - repo: https://github.com/psf/black
     rev: 25.1.0
     hooks:
-    - id: black
-      pass_filenames: false
-      always_run: true
-      args: ["--line-length=120", "."]
+      - id: black
+        pass_filenames: false
+        always_run: true
+        args: ["."]

agentlightning/algorithm/base.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import weakref
-from typing import Any, Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, Optional
 
 from agentlightning.client import AgentLightningClient
 from agentlightning.types import Dataset

agentlightning/algorithm/verl/interface.py

@@ -2,14 +2,13 @@
 
 from typing import Any, Optional
 
-
-from hydra import initialize, compose
+from hydra import compose, initialize
 from omegaconf import OmegaConf
 
-from agentlightning.verl.entrypoint import run_ppo
-from agentlightning.client import AgentLightningClient
 from agentlightning.algorithm.base import BaseAlgorithm
+from agentlightning.client import AgentLightningClient
 from agentlightning.types import Dataset
+from agentlightning.verl.entrypoint import run_ppo
 
 
 class VERL(BaseAlgorithm):

agentlightning/cli/agentops_server.py

@@ -1,8 +1,8 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import time
-from agentlightning.instrumentation.agentops import AgentOpsServerManager
 
+from agentlightning.instrumentation.agentops import AgentOpsServerManager
 
 if __name__ == "__main__":
     import argparse

agentlightning/cli/vllm.py

@@ -6,7 +6,6 @@ from vllm.entrypoints.cli.main import main
 
 from agentlightning.instrumentation.vllm import instrument_vllm
 
-
 if __name__ == "__main__":
     instrument_vllm()
     main()

agentlightning/client.py

@@ -4,13 +4,12 @@ import asyncio
 import logging
 import time
 import urllib.parse
-from typing import Any, Dict, Optional, List, Union
+from typing import Any, Dict, List, Optional, Union
 
 import aiohttp
 import requests
 
-from .types import Rollout, Task, TaskInput, TaskIfAny, ResourcesUpdate, NamedResources
-
+from .types import NamedResources, ResourcesUpdate, Rollout, Task, TaskIfAny, TaskInput
 
 logger = logging.getLogger(__name__)
 

agentlightning/config.py

@@ -11,20 +11,20 @@ from __future__ import annotations
 import argparse
 import inspect
 import logging
+from typing import _GenericAlias  # type: ignore
 from typing import (
     Any,
+    Callable,
+    Dict,
     List,
+    Tuple,
     Type,
     TypeVar,
     Union,
-    _GenericAlias,  # type: ignore
-    get_origin,
     get_args,
-    Tuple,
-    Callable,
-    overload,
-    Dict,
+    get_origin,
     get_type_hints,
+    overload,
 )
 
 CliConfigurable = Any

agentlightning/instrumentation/agentops.py

@@ -17,8 +17,8 @@ _original_handle_response = None
 
 
 def _patch_new_agentops():
-    import agentops.instrumentation.providers.openai.wrappers.chat
     import agentops.instrumentation.providers.openai.stream_wrapper
+    import agentops.instrumentation.providers.openai.wrappers.chat
     from agentops.instrumentation.providers.openai.wrappers.chat import handle_chat_attributes
 
     global _original_handle_chat_attributes
@@ -56,8 +56,8 @@ def _patch_new_agentops():
 
 
 def _unpatch_new_agentops():
-    import agentops.instrumentation.providers.openai.wrappers.chat
     import agentops.instrumentation.providers.openai.stream_wrapper
+    import agentops.instrumentation.providers.openai.wrappers.chat
 
     global _original_handle_chat_attributes
     if _original_handle_chat_attributes is not None:

agentlightning/instrumentation/agentops_langchain.py

@@ -1,9 +1,9 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from typing import Dict, Any
-from agentops.integration.callbacks.langchain import LangchainCallbackHandler
-from agentops import instrumentation
+from typing import Any, Dict
 
+from agentops import instrumentation
+from agentops.integration.callbacks.langchain import LangchainCallbackHandler
 
 original_on_chain_start = LangchainCallbackHandler.on_chain_start
 langgraph_entry = None

agentlightning/instrumentation/litellm.py

@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from typing import Optional, Any
+from typing import Any, Optional
 
 from litellm.integrations.opentelemetry import OpenTelemetry
 

agentlightning/instrumentation/verl_chat_scheduler.py

@@ -24,7 +24,6 @@ from typing import Any, Dict, List
 import torch
 from openai.types.chat.chat_completion import ChatCompletion
 from tensordict import TensorDict
-
 from verl.protocol import DataProto
 from verl.workers.rollout.async_server import ChatCompletionScheduler
 

agentlightning/instrumentation/vllm.py

@@ -5,8 +5,8 @@ from __future__ import annotations
 import warnings
 from typing import List
 
-from vllm.entrypoints.openai.protocol import ChatCompletionResponse
 import vllm.entrypoints.openai.protocol
+from vllm.entrypoints.openai.protocol import ChatCompletionResponse
 from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
 
 

agentlightning/litagent.py

@@ -6,14 +6,14 @@ import functools
 import inspect
 import logging
 import weakref
-from typing import Any, Callable, Coroutine, List, Dict, Union, Optional, TYPE_CHECKING, TypeVar, Generic
+from typing import TYPE_CHECKING, Any, Callable, Coroutine, Dict, Generic, List, Optional, TypeVar, Union
 
-from .types import LLM, NamedResources, Rollout, Task, TaskInput, Triplet, RolloutRawResult
+from .types import LLM, NamedResources, Rollout, RolloutRawResult, Task, TaskInput, Triplet
 
 if TYPE_CHECKING:
-    from .trainer import Trainer
     from .runner import AgentRunner
     from .tracer import BaseTracer
+    from .trainer import Trainer
 
 
 logger = logging.getLogger(__name__)

agentlightning/reward.py

@@ -3,7 +3,7 @@
 import asyncio
 import inspect
 import warnings
-from typing import TypedDict, Optional
+from typing import Optional, TypedDict
 
 from agentops.sdk.decorators import operation
 

agentlightning/runner.py

@@ -6,17 +6,16 @@ import logging
 import os
 import time
 from contextlib import nullcontext
-from typing import List, Optional, Union, Dict, Any
+from typing import Any, Dict, List, Optional, Union
 
 import agentops
-
 from opentelemetry.sdk.trace import ReadableSpan
+
 from .client import AgentLightningClient
 from .litagent import LitAgent, is_v0_1_rollout_api
-from .types import Rollout, Task, Triplet, RolloutRawResult
-from .types import ParallelWorkerBase
-from .tracer.base import BaseTracer
 from .tracer import TripletExporter
+from .tracer.base import BaseTracer
+from .types import ParallelWorkerBase, Rollout, RolloutRawResult, Task, Triplet
 
 logger = logging.getLogger(__name__)
 

agentlightning/server.py

@@ -2,23 +2,23 @@
 
 import asyncio
 import logging
+import threading
 import time
 import uuid
-import threading
 from contextlib import asynccontextmanager
-from typing import Any, Dict, List, Optional, Literal
+from typing import Any, Dict, List, Literal, Optional
 
 import uvicorn
 from fastapi import FastAPI, HTTPException, Path
 from pydantic import Field
 
 from .types import (
+    GenericResponse,
+    NamedResources,
+    ResourcesUpdate,
     Rollout,
     Task,
     TaskIfAny,
-    NamedResources,
-    GenericResponse,
-    ResourcesUpdate,
 )
 
 logger = logging.getLogger(__name__)

agentlightning/tracer/__init__.py

@@ -1,5 +1,5 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from .base import BaseTracer
 from .agentops import AgentOpsTracer
+from .base import BaseTracer
 from .triplet import TripletExporter

agentlightning/tracer/agentops.py

@@ -5,18 +5,18 @@ from __future__ import annotations
 import logging
 import os
 from contextlib import contextmanager
-from typing import List, Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING, List, Optional
 
-import agentops.sdk.core
 import agentops
+import agentops.sdk.core
 from agentops.sdk.core import TracingCore
 from agentops.sdk.processors import SpanProcessor
 from opentelemetry.sdk.trace import ReadableSpan
 
-from agentlightning.instrumentation.agentops import AgentOpsServerManager
 from agentlightning.instrumentation import instrument_all, uninstrument_all
-from .base import BaseTracer
+from agentlightning.instrumentation.agentops import AgentOpsServerManager
 
+from .base import BaseTracer
 
 if TYPE_CHECKING:
     from agentops.integration.callbacks.langchain import LangchainCallbackHandler

agentlightning/tracer/base.py

@@ -1,9 +1,10 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 from contextlib import contextmanager
-from typing import Iterator, List, Optional, Callable, Any, Awaitable
+from typing import Any, Awaitable, Callable, Iterator, List, Optional
 
 from opentelemetry.sdk.trace import ReadableSpan
+
 from agentlightning.types import ParallelWorkerBase
 
 

agentlightning/tracer/http.py

@@ -1,27 +1,26 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from contextlib import contextmanager
-from typing import Iterator, List, Optional, Any, Dict, Callable, Awaitable
-import logging
-import uuid
-import pickle
-import multiprocessing
 import asyncio
+import logging
+import multiprocessing
+import pickle
 import queue
+import uuid
+from contextlib import contextmanager
+from typing import Any, Awaitable, Callable, Dict, Iterator, List, Optional
 from urllib.parse import urlparse
 
-from .base import BaseTracer
-
 from httpdbg.hooks.all import httprecord
 from httpdbg.records import HTTPRecords
 from opentelemetry.sdk.trace import ReadableSpan
-from opentelemetry.trace import StatusCode, SpanKind, Status
+from opentelemetry.trace import SpanKind, Status, StatusCode
 from opentelemetry.trace.span import (
     SpanContext,
     TraceFlags,
     TraceState,
 )
 
+from .base import BaseTracer
 
 logger = logging.getLogger(__name__)
 

agentlightning/tracer/triplet.py

@@ -3,11 +3,12 @@
 import json
 import re
 from enum import Enum
-from typing import List, Dict, Tuple, Optional, Any
+from typing import Any, Dict, List, Optional, Tuple
 
-from pydantic import BaseModel
 from opentelemetry import trace as trace_api
 from opentelemetry.sdk.trace import ReadableSpan
+from pydantic import BaseModel
+
 from agentlightning.types import Triplet
 
 

agentlightning/trainer.py

@@ -1,24 +1,23 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import asyncio
+import importlib
 import logging
 import multiprocessing
 import os
 import signal
 import time
-from typing import List, Optional, Union
-import importlib
 import warnings
+from typing import List, Optional, Union
 
+from .algorithm.base import BaseAlgorithm
 from .client import AgentLightningClient
 from .litagent import LitAgent
 from .runner import AgentRunner
-from .types import Dataset, ParallelWorkerBase
-from .algorithm.base import BaseAlgorithm
-from .tracer.base import BaseTracer
 from .tracer.agentops import AgentOpsTracer
+from .tracer.base import BaseTracer
 from .tracer.triplet import TripletExporter
-
+from .types import Dataset, ParallelWorkerBase
 
 logger = logging.getLogger(__name__)
 

agentlightning/types.py

@@ -2,10 +2,10 @@
 
 from __future__ import annotations
 
-from typing import Any, Dict, Generic, List, Optional, Protocol, TypeVar, Union, Literal, Annotated
+from typing import Annotated, Any, Dict, Generic, List, Literal, Optional, Protocol, TypeVar, Union
 
-from pydantic import BaseModel, Field, Discriminator
 from opentelemetry.sdk.trace import ReadableSpan
+from pydantic import BaseModel, Discriminator, Field
 
 __all__ = [
     "Triplet",

agentlightning/verl/__init__.py

@@ -1,5 +1,5 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from .trainer import *
 from .daemon import *
 from .dataset import *
+from .trainer import *

agentlightning/verl/async_server.py

@@ -1,13 +1,14 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-import ray
 from copy import deepcopy
 
-from agentlightning.instrumentation.vllm import instrument_vllm, ChatCompletionResponsePatched
+import ray
 from starlette.requests import Request
 from starlette.responses import JSONResponse, StreamingResponse
-from vllm.entrypoints.openai.protocol import ChatCompletionRequest, ErrorResponse
 from verl.workers.rollout.vllm_rollout.vllm_async_server import AsyncvLLMServer
+from vllm.entrypoints.openai.protocol import ChatCompletionRequest, ErrorResponse
+
+from agentlightning.instrumentation.vllm import ChatCompletionResponsePatched, instrument_vllm
 
 
 def _unwrap_ray_remote(cls):

agentlightning/verl/daemon.py

@@ -13,13 +13,13 @@ from typing import Dict, List, Optional
 import numpy as np
 import requests
 import torch
-from agentlightning import LLM, AgentLightningServer, NamedResources, Rollout, configure_logger
 from flask import Flask, Response, abort, request
 from openai.types.chat.chat_completion import ChatCompletion
 from tensordict import TensorDict
-
 from verl import DataProto
 
+from agentlightning import LLM, AgentLightningServer, NamedResources, Rollout, configure_logger
+
 configure_logger()
 
 

agentlightning/verl/entrypoint.py

@@ -1,15 +1,16 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 from typing import Any
+
 import hydra
 import ray
+from verl.trainer.main_ppo import create_rl_sampler
+from verl.trainer.ppo.reward import load_reward_manager
 
 from agentlightning.types import Dataset
 
 from .dataset import AgentDataset, LoadedDataset
 from .trainer import AgentLightningTrainer
-from verl.trainer.ppo.reward import load_reward_manager
-from verl.trainer.main_ppo import create_rl_sampler
 
 
 @hydra.main(config_path="pkg://agentlightning/verl", config_name="config", version_base=None)
@@ -38,7 +39,6 @@ class TaskRunner:
         from pprint import pprint
 
         from omegaconf import OmegaConf
-
         from verl.utils.fs import copy_to_local
 
         pprint(OmegaConf.to_container(config, resolve=True))  # resolve=True will eval symbol values

agentlightning/verl/trainer.py

@@ -3,30 +3,29 @@
 import random
 from contextlib import contextmanager
 from copy import deepcopy
+from pprint import pprint
 from typing import Dict, Tuple
 
 import numpy as np
 import torch
-from omegaconf import OmegaConf
-from pprint import pprint
-from tqdm import tqdm
-
 from codetiming import Timer
+from omegaconf import OmegaConf
+from tqdm import tqdm
 from verl import DataProto
 from verl.protocol import pad_dataproto_to_divisor, unpad_dataproto
-from verl.trainer.ppo.ray_trainer import (
-    RayPPOTrainer,
-    AdvantageEstimator,
-    apply_kl_penalty,
-    compute_advantage,
-    compute_response_mask,
-)
 from verl.trainer.ppo.core_algos import agg_loss
 from verl.trainer.ppo.metric_utils import (
     compute_data_metrics,
     compute_throughout_metrics,
     compute_timing_metrics,
 )
+from verl.trainer.ppo.ray_trainer import (
+    AdvantageEstimator,
+    RayPPOTrainer,
+    apply_kl_penalty,
+    compute_advantage,
+    compute_response_mask,
+)
 from verl.utils.metric import reduce_metrics
 from verl.utils.tracking import Tracking
 

examples/apo/apo_client.py

@@ -1,9 +1,9 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-import dotenv
 import os
 import random
 
+import dotenv
 from openai import OpenAI
 
 from agentlightning import configure_logger

examples/calc_x/calc_agent.py

@@ -2,8 +2,8 @@
 
 import math
 import os
-import string
 import re
+import string
 from typing import Any
 
 import sympy
@@ -12,7 +12,7 @@ from autogen_core.models import ModelFamily
 from autogen_ext.models.openai import OpenAIChatCompletionClient
 from autogen_ext.tools.mcp import McpWorkbench, StdioServerParams
 
-from agentlightning import Trainer, LitAgent, NamedResources, LLM, reward, configure_logger, DevTaskLoader
+from agentlightning import LLM, DevTaskLoader, LitAgent, NamedResources, Trainer, configure_logger, reward
 
 configure_logger()
 

examples/calc_x/calc_agent_dev.py

@@ -1,9 +1,11 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import os
-from agentlightning import Trainer, DevTaskLoader, LLM
+
 from calc_agent import CalcAgent
 
+from agentlightning import LLM, DevTaskLoader, Trainer
+
 
 def dev_task_loader() -> DevTaskLoader:
     return DevTaskLoader(

examples/calc_x/calc_agent_v0_2.py

@@ -1,15 +1,14 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from typing import Any
 import re
+from typing import Any
 
+from autogen_ext.tools.mcp import McpWorkbench, StdioServerParams
+from calc_agent import eval, get_agent
 from datasets import Dataset
 
-from agentlightning import rollout, Trainer, LLM
+from agentlightning import LLM, Trainer, rollout
 from agentlightning.algorithm.verl import VERL
-from autogen_ext.tools.mcp import McpWorkbench, StdioServerParams
-
-from calc_agent import get_agent, eval
 
 calculator_mcp_server = StdioServerParams(command="uvx", args=["mcp-server-calculator"])
 

examples/calc_x/tests/test_agentops.py

@@ -1,9 +1,10 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import agentops
-from agentlightning.reward import reward
 from agentops.sdk.decorators import operation
 
+from agentlightning.reward import reward
+
 
 @reward
 def process_data(data):

examples/calc_x/tests/test_mcp_calculator.py

@@ -1,11 +1,12 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-import os
 import asyncio
 import json
+import os
+
 import openai
 from mcp import ClientSession
-from mcp.client.stdio import stdio_client, StdioServerParameters
+from mcp.client.stdio import StdioServerParameters, stdio_client
 
 
 async def main():

examples/rag/wiki_retriever_mcp/wiki_retriever_mcp.py

@@ -1,9 +1,10 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-import faiss
-from sentence_transformers import SentenceTransformer
 import pickle
+
+import faiss
 from fastmcp import FastMCP
+from sentence_transformers import SentenceTransformer
 
 # index = faiss.read_index("/mnt/input/agent_lightning/nq_hnsw_faiss_n32e40.index")
 index = faiss.read_index("nq_hnsw_faiss_n32e40.index")

examples/spider/spider_eval/evaluation.py

@@ -24,13 +24,13 @@
 # }
 ################################
 
-import os
-import json
-import sqlite3
 import argparse
+import json
+import os
+import sqlite3
 
-from .process_sql import get_schema, Schema, get_sql
 from .exec_eval import eval_exec_match
+from .process_sql import Schema, get_schema, get_sql
 
 # Flag to disable value evaluation
 DISABLE_VALUE = True

examples/spider/spider_eval/exec_eval.py

@@ -3,22 +3,22 @@
 # type: ignore
 # The evaluation code is from https://github.com/taoyds/test-suite-sql-eval
 
-import os
-import re
 import asyncio
-import sqlite3
-import threading
-from typing import Tuple, Any, List, Set
-from itertools import product
-from collections import defaultdict
-import tqdm
-import random
-from .parse import get_all_preds_for_execution, remove_distinct
-import time
+import os
 import pickle as pkl
+import random
+import re
+import sqlite3
 import subprocess
-from itertools import chain
+import threading
+import time
+from collections import defaultdict
+from itertools import chain, product
+from typing import Any, List, Set, Tuple
 
+import tqdm
+
+from .parse import get_all_preds_for_execution, remove_distinct
 
 threadLock = threading.Lock()
 TIMEOUT = 60

examples/spider/spider_eval/parse.py

@@ -3,13 +3,14 @@
 # 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 typing import List, Tuple, Set, Iterator, Dict, Any, Union
 from sqlparse.sql import Comparison, Identifier
 from sqlparse.tokens import Whitespace
-import itertools
-from collections import namedtuple
 
 Token = namedtuple("Token", ["ttype", "value"])
 VALUE_NUM_SYMBOL = "VALUERARE"

examples/spider/spider_eval/process_sql.py

@@ -31,6 +31,7 @@
 
 import json
 import sqlite3
+
 from nltk import word_tokenize
 
 CLAUSE_KEYWORDS = ("select", "from", "where", "group", "order", "limit", "intersect", "union", "except")

examples/spider/sql_agent.py

@@ -9,21 +9,20 @@ from __future__ import annotations
 
 import os
 import re
+import shutil
 import tempfile
 import time
-import shutil
 from typing import Any, Literal, Optional
 
 import dotenv
 import termcolor
-from langchain_core.messages import BaseMessage, SystemMessage, HumanMessage
-from langchain_core.prompts import ChatPromptTemplate
-from langchain_community.utilities import SQLDatabase
 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 BaseMessage, HumanMessage, SystemMessage
+from langchain_core.prompts import ChatPromptTemplate
 from langgraph.graph import END, START, MessagesState, StateGraph
 from langgraph.graph.state import CompiledStateGraph
-from langchain_community.tools.sql_database.tool import QuerySQLDatabaseTool
-
 from spider_eval.exec_eval import eval_exec_match
 
 import agentlightning

pyproject.toml

@@ -25,6 +25,7 @@ dev = [
   "pre-commit",
   "pytest-rerunfailures",
   "black",
+  "isort",
   "mkdocs",
   "mkdocs-material",
   "mkdocstrings[python]",
@@ -61,3 +62,22 @@ include = ["**/*.yaml", "**/*.yml"]
 
 [tool.pytest.ini_options]
 testpaths = ["tests"]
+
+[tool.black]
+line-length = 120
+target-version = ['py312']
+include = '\.pyi?$'
+extend-exclude = '''
+/(
+  _version\.py
+)
+'''
+
+[tool.isort]
+profile = "black"
+line_length = 120
+known_first_party = ["agentlightning"]
+extend_skip_glob = [
+  "data/**",
+  "_version.py",
+]

scripts/validate_example_wandb.py

@@ -1,8 +1,9 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-import wandb
 import sys
 
+import wandb
+
 if len(sys.argv) != 3:
     print("Usage: python validate_example_wandb.py <project> <run_name>")
 

tests/test_client.py

@@ -2,22 +2,22 @@
 
 import asyncio
 import time
-from typing import Any, Dict, AsyncGenerator
+from typing import Any, AsyncGenerator, Dict
 
 import pytest
 import pytest_asyncio
 from httpx import AsyncClient
 
 from agentlightning import (
+    LLM,
     AgentLightningClient,
     AgentLightningServer,
-    Rollout,
-    Triplet,
     NamedResources,
-    LLM,
     PromptTemplate,
     ResourcesUpdate,
+    Rollout,
     Task,
+    Triplet,
 )
 from agentlightning.client import DevTaskLoader
 

tests/test_config.py

@@ -7,25 +7,26 @@ It can contain logically erroroneous expected values.
 Please do not use the file as a reference for the expected behavior of config.
 """
 
-import pytest
 import argparse
+import inspect
+import sys
+from typing import TypeVar  # Added for completeness if testing TypeVars directly
+from typing import _GenericAlias  # type: ignore
 from typing import (
     Any,
-    List,
-    Type,
-    Union,
-    Optional,
-    Tuple,
     Callable,
     Dict,
-    TypeVar,  # Added for completeness if testing TypeVars directly
-    _GenericAlias,  # type: ignore
-    get_origin,
+    List,
+    Optional,
+    Tuple,
+    Type,
+    Union,
     get_args,
+    get_origin,
 )
-import inspect
 from unittest import mock  # For mock.patch.object, mock.call, MagicMock etc.
-import sys
+
+import pytest
 
 from agentlightning import config
 

tests/test_decorator_executability.py

@@ -3,8 +3,10 @@
 """Test that @llm_rollout and @rollout decorators preserve function executability."""
 
 import inspect
+
 import pytest
-from agentlightning.litagent import llm_rollout, rollout, LitAgentLLM
+
+from agentlightning.litagent import LitAgentLLM, llm_rollout, rollout
 
 
 @llm_rollout

tests/test_litagent_hooks.py

@@ -1,10 +1,10 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-import pytest
 from contextlib import contextmanager
 
+import pytest
 
-from agentlightning import LitAgent, Task, ResourcesUpdate
+from agentlightning import LitAgent, ResourcesUpdate, Task
 from agentlightning.runner import AgentRunner
 from agentlightning.tracer import BaseTracer, TripletExporter
 

tests/test_tracer.py

@@ -16,66 +16,55 @@ against the real API with the ``OPENAI_MODEL`` of your choice (``gpt-4.1-nano``
 by default).
 """
 
-import pytest
+import asyncio
+import difflib
 import inspect
 import json
-import asyncio
-import time
 import os
-import re
-import httpx
-import difflib
 import pprint
-from contextlib import contextmanager
-from unittest.mock import Mock, patch, AsyncMock, MagicMock
-from typing import Dict, Any, List, Optional
+import re
 import threading
-import uvicorn
-from fastapi import FastAPI
-from contextlib import asynccontextmanager
-from pydantic import BaseModel, Field
-
-from agentlightning.tracer.agentops import AgentOpsTracer, LightningSpanProcessor
-from agentlightning.tracer.triplet import TraceTree
-from agentlightning.tracer.http import HttpTracer
-from agentlightning.tracer.triplet import TripletExporter
-from agentlightning.reward import reward
-from agentlightning.types import Triplet
-
-import openai
-from openai import OpenAI, AsyncOpenAI
-
-from langchain import hub
-from langchain.chat_models import init_chat_model
-from langchain_openai import ChatOpenAI
-from langchain_core.prompts import ChatPromptTemplate
-from langchain_core.output_parsers import StrOutputParser
-from langchain_core.messages import HumanMessage, BaseMessage, AIMessage, ToolMessage
-from langchain_community.utilities import SQLDatabase
-from langchain_community.agent_toolkits import SQLDatabaseToolkit
-from langchain.agents import tool, AgentExecutor, create_react_agent
-from langchain_core.messages import AIMessage
-from langgraph.graph import END, START, MessagesState, StateGraph
-from langgraph.prebuilt import ToolNode
-from typing import Literal
-
-from langgraph.graph import StateGraph, END, MessagesState
-from typing_extensions import TypedDict
-
-import autogen_agentchat
-from autogen_agentchat.agents import AssistantAgent, UserProxyAgent
-from autogen_agentchat.teams import RoundRobinGroupChat
-from autogen_agentchat.conditions import ExternalTermination, TextMentionTermination
-from autogen_ext.models.openai import OpenAIChatCompletionClient
-from autogen_ext.tools.mcp import McpWorkbench, StdioServerParams
-
-from agents import Agent, Runner, AgentHooks, InputGuardrail, GuardrailFunctionOutput, function_tool, RunConfig
-from agents.mcp import MCPServerStdio
-from agents.models.openai_provider import OpenAIProvider
-
-import litellm
+import time
+from contextlib import asynccontextmanager, contextmanager
+from typing import Any, Dict, List, Literal, Optional
+from unittest.mock import AsyncMock, MagicMock, Mock, patch
 
 import agentops
+import autogen_agentchat
+import httpx
+import litellm
+import openai
+import pytest
+import uvicorn
+from agents import Agent, AgentHooks, GuardrailFunctionOutput, InputGuardrail, RunConfig, Runner, function_tool
+from agents.mcp import MCPServerStdio
+from agents.models.openai_provider import OpenAIProvider
+from autogen_agentchat.agents import AssistantAgent, UserProxyAgent
+from autogen_agentchat.conditions import ExternalTermination, TextMentionTermination
+from autogen_agentchat.teams import RoundRobinGroupChat
+from autogen_ext.models.openai import OpenAIChatCompletionClient
+from autogen_ext.tools.mcp import McpWorkbench, StdioServerParams
+from fastapi import FastAPI
+from langchain import hub
+from langchain.agents import AgentExecutor, create_react_agent, tool
+from langchain.chat_models import init_chat_model
+from langchain_community.agent_toolkits import SQLDatabaseToolkit
+from langchain_community.utilities import SQLDatabase
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, ToolMessage
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_openai import ChatOpenAI
+from langgraph.graph import END, START, MessagesState, StateGraph
+from langgraph.prebuilt import ToolNode
+from openai import AsyncOpenAI, OpenAI
+from pydantic import BaseModel, Field
+from typing_extensions import TypedDict
+
+from agentlightning.reward import reward
+from agentlightning.tracer.agentops import AgentOpsTracer, LightningSpanProcessor
+from agentlightning.tracer.http import HttpTracer
+from agentlightning.tracer.triplet import TraceTree, TripletExporter
+from agentlightning.types import Triplet
 
 USE_OPENAI = os.environ.get("USE_OPENAI", "false").lower() == "true"
 if USE_OPENAI:

tests/test_tracer_http.py

@@ -9,10 +9,12 @@ that the HttpTracer correctly captures HTTP traffic in both
 normal and subprocess modes.
 """
 
-import pytest
 import asyncio
-import requests
+
 import aiohttp
+import pytest
+import requests
+
 from agentlightning.tracer.http import HttpTracer