Fix pyright issues (#101)

.github/workflows/tests.yml

@@ -16,8 +16,8 @@ on:
 
 jobs:
 
-  lint:
-    name: Lint with Black
+  lint-fast:
+    name: Lint - Fast
     runs-on: ubuntu-latest
     timeout-minutes: 10
     steps:
@@ -36,6 +36,27 @@ jobs:
         run: black --check .
       - name: Run isort
         run: isort --check-only .
+      - name: Run pyright
+        run: pyright -p pyrightconfig.fast.json
+
+  lint-slow:
+    name: Lint - Slow
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+    steps:
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.12'
+      - name: Install dependencies
+        run: |
+          ./scripts/setup_stable_full.sh
+      - name: Run Black
+        run: black --check .
+      - name: Run isort
+        run: isort --check-only .
+      - name: Run pyright
+        run: pyright -p pyrightconfig.json
 
   docs:
     name: Build documentation

agentlightning/__init__.py

@@ -10,3 +10,14 @@ from .reward import reward
 from .server import AgentLightningServer
 from .trainer import Trainer
 from .types import *
+
+__all__ = [
+    "AgentLightningClient",
+    "DevTaskLoader",
+    "lightning_cli",
+    "configure_logger",
+    "reward",
+    "AgentLightningServer",
+    "Trainer",
+    "__version__",
+]

agentlightning/algorithm/verl/interface.py

@@ -12,7 +12,7 @@ from agentlightning.verl.entrypoint import run_ppo
 
 
 class VERL(BaseAlgorithm):
-    def __init__(self, config: dict):
+    def __init__(self, config: dict[str, Any]):
         super().__init__()
 
         # Compose the base config exactly like your decorator:

agentlightning/cli/vllm.py

@@ -1,7 +1,5 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from typing import List
-
 from vllm.entrypoints.cli.main import main
 
 from agentlightning.instrumentation.vllm import instrument_vllm

agentlightning/client.py

@@ -83,7 +83,7 @@ class AgentLightningClient:
                 logger.debug(f"Async POST request failed for {url}: {e}")
                 return None
 
-    async def poll_next_task_async(self) -> Task:
+    async def poll_next_task_async(self) -> Optional[Task]:
         """Polls the server asynchronously for the next task until one is available.
 
         Returns:
@@ -186,7 +186,7 @@ class AgentLightningClient:
             logger.debug(f"Sync POST request failed for {url}: {e}")
             return None
 
-    def poll_next_task(self) -> Task:
+    def poll_next_task(self) -> Optional[Task]:
         """Polls the server synchronously for the next task until one is available.
 
         Returns:
@@ -303,7 +303,7 @@ class DevTaskLoader(AgentLightningClient):
         """Return rollouts that have been posted back to the loader."""
         return self._rollouts
 
-    def poll_next_task(self) -> Task:
+    def poll_next_task(self) -> Optional[Task]:
         """Returns the next task from the local queue.
 
         If tasks are TaskInput objects, assembles them into Task objects.
@@ -350,7 +350,7 @@ class DevTaskLoader(AgentLightningClient):
         self._rollouts.append(rollout)
         return {"status": "received", "rollout_id": rollout.rollout_id}
 
-    async def poll_next_task_async(self) -> Task:
+    async def poll_next_task_async(self) -> Optional[Task]:
         return self.poll_next_task()
 
     async def get_resources_by_id_async(self, resource_id: str) -> Optional[ResourcesUpdate]:

agentlightning/config.py

@@ -69,8 +69,8 @@ def nullable_float(value: str) -> float | None:
 
 def _str_to_bool(v: str) -> bool:
     """Converts common string representations of bool to Python bool (case-insensitive)."""
-    if isinstance(v, bool):  # Allow passing bools directly if used programmatically
-        return v
+    if isinstance(v, bool):  # type: ignore
+        return v  # Allow passing bools directly if used programmatically
     lowered_v = v.lower()
     if lowered_v in ("yes", "true", "t", "y", "1"):
         return True
@@ -307,7 +307,7 @@ def lightning_cli(cls1: Type[_C1], cls2: Type[_C2], cls3: Type[_C3], cls4: Type[
 def lightning_cli(*classes: Type[CliConfigurable]) -> Tuple[CliConfigurable, ...]: ...
 
 
-def lightning_cli(*classes: Type[CliConfigurable]) -> CliConfigurable | Tuple[CliConfigurable, ...]:
+def lightning_cli(*classes: Type[CliConfigurable]) -> CliConfigurable | Tuple[CliConfigurable, ...]:  # type: ignore
     """
     Parses command-line arguments to configure and instantiate provided CliConfigurable classes.
 

agentlightning/instrumentation/__init__.py

@@ -2,22 +2,22 @@
 
 import warnings
 
-AGENTOPS_INSTALLED = False
-AGENTOPS_LANGCHAIN_INSTALLED = False
-LITELLM_INSTALLED = False
-VLLM_INSTALLED = False
+AGENTOPS_INSTALLED: bool = False
+AGENTOPS_LANGCHAIN_INSTALLED: bool = False
+LITELLM_INSTALLED: bool = False
+VLLM_INSTALLED: bool = False
 
 try:
-    from . import agentops
+    from . import agentops  # type: ignore
 
-    AGENTOPS_INSTALLED = True
+    AGENTOPS_INSTALLED = True  # type: ignore
 except ImportError:
     pass
 
 try:
-    from . import litellm
+    from . import litellm  # type: ignore
 
-    LITELLM_INSTALLED = True
+    LITELLM_INSTALLED = True  # type: ignore
 except ImportError:
     pass
 
@@ -32,9 +32,9 @@ except ImportError:
 
 
 try:
-    from . import agentops_langchain
+    from . import agentops_langchain  # type: ignore
 
-    AGENTOPS_LANGCHAIN_INSTALLED = True
+    AGENTOPS_LANGCHAIN_INSTALLED = True  # type: ignore
 except ImportError:
     pass
 

agentlightning/instrumentation/agentops.py

@@ -5,6 +5,7 @@ import multiprocessing
 import signal
 import socket
 import time
+from typing import Any, Callable
 
 import flask
 import setproctitle
@@ -12,14 +13,14 @@ import setproctitle
 logger = logging.getLogger(__name__)
 
 # Module-level storage for originals
-_original_handle_chat_attributes = None
-_original_handle_response = None
+_original_handle_chat_attributes: Callable[..., Any] | None = None
+_original_handle_response: Callable[..., Any] | None = None
 
 
 def _patch_new_agentops():
     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
+    from agentops.instrumentation.providers.openai.wrappers.chat import handle_chat_attributes  # type: ignore
 
     global _original_handle_chat_attributes
 
@@ -27,23 +28,28 @@ def _patch_new_agentops():
         logger.warning("AgentOps already patched. Skipping.")
         return True
 
-    _original_handle_chat_attributes = handle_chat_attributes
+    _original_handle_chat_attributes = handle_chat_attributes  # type: ignore
 
-    def _handle_chat_attributes_with_tokens(args=None, kwargs=None, return_value=None, **kws):
-        attributes = _original_handle_chat_attributes(args=args, kwargs=kwargs, return_value=return_value, **kws)
-        if hasattr(return_value, "prompt_token_ids"):
-            attributes["prompt_token_ids"] = list(return_value.prompt_token_ids)
-        if hasattr(return_value, "response_token_ids"):
-            attributes["response_token_ids"] = list(return_value.response_token_ids[0])
+    def _handle_chat_attributes_with_tokens(args=None, kwargs=None, return_value=None, **kws):  # type: ignore
+        attributes = _original_handle_chat_attributes(args=args, kwargs=kwargs, return_value=return_value, **kws)  # type: ignore
+        if return_value is not None and hasattr(return_value, "prompt_token_ids"):  # type: ignore
+            attributes["prompt_token_ids"] = list(return_value.prompt_token_ids)  # type: ignore
+        if return_value is not None and hasattr(return_value, "response_token_ids"):  # type: ignore
+            attributes["response_token_ids"] = list(return_value.response_token_ids[0])  # type: ignore
 
         # For LiteLLM, response is a openai._legacy_response.LegacyAPIResponse
-        if hasattr(return_value, "http_response") and hasattr(return_value.http_response, "json"):
-            json_data = return_value.http_response.json()
+        if (
+            return_value is not None
+            and hasattr(return_value, "http_response")  # type: ignore
+            and return_value.http_response is not None  # type: ignore
+            and hasattr(return_value.http_response, "json")  # type: ignore
+        ):
+            json_data = return_value.http_response.json()  # type: ignore
             if isinstance(json_data, dict):
                 if "prompt_token_ids" in json_data:
-                    attributes["prompt_token_ids"] = list(json_data["prompt_token_ids"])
+                    attributes["prompt_token_ids"] = list(json_data["prompt_token_ids"])  # type: ignore
                 if "response_token_ids" in json_data:
-                    attributes["response_token_ids"] = list(json_data["response_token_ids"][0])
+                    attributes["response_token_ids"] = list(json_data["response_token_ids"][0])  # type: ignore
 
         return attributes
 
@@ -72,40 +78,40 @@ def _unpatch_new_agentops():
 
 
 def _patch_old_agentops():
-    import opentelemetry.instrumentation.openai.shared.chat_wrappers
-    from opentelemetry.instrumentation.openai.shared.chat_wrappers import _handle_response, dont_throw
+    import opentelemetry.instrumentation.openai.shared.chat_wrappers  # type: ignore
+    from opentelemetry.instrumentation.openai.shared.chat_wrappers import _handle_response, dont_throw  # type: ignore
 
     global _original_handle_response
-    _original_handle_response = _handle_response
+    _original_handle_response = _handle_response  # type: ignore
 
-    @dont_throw
-    def _handle_response_with_tokens(response, span, *args, **kwargs):
-        _original_handle_response(response, span, *args, **kwargs)
-        if hasattr(response, "prompt_token_ids"):
-            span.set_attribute("prompt_token_ids", list(response.prompt_token_ids))
-        if hasattr(response, "response_token_ids"):
-            span.set_attribute("response_token_ids", list(response.response_token_ids[0]))
+    @dont_throw  # type: ignore
+    def _handle_response_with_tokens(response, span, *args, **kwargs):  # type: ignore
+        _original_handle_response(response, span, *args, **kwargs)  # type: ignore
+        if hasattr(response, "prompt_token_ids"):  # type: ignore
+            span.set_attribute("prompt_token_ids", list(response.prompt_token_ids))  # type: ignore
+        if hasattr(response, "response_token_ids"):  # type: ignore
+            span.set_attribute("response_token_ids", list(response.response_token_ids[0]))  # type: ignore
 
         # For LiteLLM, response is a openai._legacy_response.LegacyAPIResponse
-        if hasattr(response, "http_response") and hasattr(response.http_response, "json"):
-            json_data = response.http_response.json()
+        if hasattr(response, "http_response") and hasattr(response.http_response, "json"):  # type: ignore
+            json_data = response.http_response.json()  # type: ignore
             if isinstance(json_data, dict):
                 if "prompt_token_ids" in json_data:
-                    span.set_attribute("prompt_token_ids", list(json_data["prompt_token_ids"]))
+                    span.set_attribute("prompt_token_ids", list(json_data["prompt_token_ids"]))  # type: ignore
                 if "response_token_ids" in json_data:
-                    span.set_attribute("response_token_ids", list(json_data["response_token_ids"][0]))
+                    span.set_attribute("response_token_ids", list(json_data["response_token_ids"][0]))  # type: ignore
 
-    opentelemetry.instrumentation.openai.shared.chat_wrappers._handle_response = _handle_response_with_tokens
+    opentelemetry.instrumentation.openai.shared.chat_wrappers._handle_response = _handle_response_with_tokens  # type: ignore
     logger.info("Patched earlier version of agentops using _handle_response")
     return True
 
 
 def _unpatch_old_agentops():
-    import opentelemetry.instrumentation.openai.shared.chat_wrappers
+    import opentelemetry.instrumentation.openai.shared.chat_wrappers  # type: ignore
 
     global _original_handle_response
     if _original_handle_response is not None:
-        opentelemetry.instrumentation.openai.shared.chat_wrappers._handle_response = _original_handle_response
+        opentelemetry.instrumentation.openai.shared.chat_wrappers._handle_response = _original_handle_response  # type: ignore
         _original_handle_response = None
         logger.info("Unpatched earlier version of agentops using _handle_response")
 
@@ -151,18 +157,18 @@ def agentops_local_server():
     app = flask.Flask(__name__)
 
     @app.route("/v3/auth/token", methods=["POST"])
-    def fetch_token():
+    def fetch_token():  # type: ignore
         return {"token": "dummy", "project_id": "dummy"}
 
     @app.route("/", defaults={"path": ""}, methods=["GET", "POST"])
     @app.route("/<path:path>", methods=["GET", "POST"])
-    def catch_all(path):
+    def catch_all(path: str):  # type: ignore
         return {"path": path}
 
     return app
 
 
-def _run_server(**kwargs):
+def _run_server(**kwargs: Any):  # type: ignore
     """
     Internal function to run the Flask server.
     This is used to avoid issues with multiprocessing and Flask's reloader.
@@ -215,7 +221,7 @@ class AgentOpsServerManager:
         return False
 
     def stop(self):
-        if self.is_alive():
+        if self.server_process is not None and self.server_process.is_alive():
             logger.info(f"Stopping AgentOps local server (PID: {self.server_process.pid})...")
             self.server_process.terminate()  # Send SIGTERM
             self.server_process.join(timeout=5)  # Wait for clean exit

agentlightning/instrumentation/agentops_langchain.py

@@ -9,14 +9,14 @@ original_on_chain_start = LangchainCallbackHandler.on_chain_start
 langgraph_entry = None
 
 
-def on_chain_start(self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any) -> None:
+def on_chain_start(self: Any, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any) -> None:
     if "name" in kwargs:
-        if serialized is None:
+        if serialized is None:  # type: ignore
             serialized = {}
         serialized = serialized.copy()
         serialized["name"] = kwargs["name"]
     if "run_id" in kwargs:
-        if serialized is None:
+        if serialized is None:  # type: ignore
             serialized = {}
         serialized = serialized.copy()
         if "id" not in serialized:

agentlightning/instrumentation/litellm.py

@@ -8,15 +8,15 @@ from litellm.integrations.opentelemetry import OpenTelemetry
 # It seems that LiteLLM owns its own telemetry from their own entrance
 # https://docs.litellm.ai/docs/observability/agentops_integration
 
-original_set_attributes = OpenTelemetry.set_attributes
+original_set_attributes = OpenTelemetry.set_attributes  # type: ignore
 
 
-def patched_set_attributes(self, span: Any, kwargs, response_obj: Optional[Any]):
+def patched_set_attributes(self: Any, span: Any, kwargs: Any, response_obj: Optional[Any]):
     original_set_attributes(self, span, kwargs, response_obj)
     # Add custom attributes
-    if response_obj.get("prompt_token_ids"):
+    if response_obj is not None and response_obj.get("prompt_token_ids"):
         span.set_attribute("prompt_token_ids", list(response_obj.get("prompt_token_ids")))
-    if response_obj.get("response_token_ids"):
+    if response_obj is not None and response_obj.get("response_token_ids"):
         span.set_attribute("response_token_ids", list(response_obj.get("response_token_ids")[0]))
 
 

agentlightning/instrumentation/verl_chat_scheduler.py

@@ -1,149 +0,0 @@
-# Copyright (c) Microsoft. All rights reserved.
-
-# type: ignore
-
-# https://github.com/volcengine/verl/blob/bd94bd61fe4193e56f2845dc794004afbef7f818/examples/ppo_trainer/naive_chat_scheduler.py
-# This file is part of VERL example. It should be included in the VERL package but it's not currently.
-
-# Copyright 2024 Bytedance Ltd. and/or its affiliates
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import asyncio
-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
-
-
-class NaiveChatCompletionScheduler(ChatCompletionScheduler):
-    """
-    A very naive implementation of ChatCompletionScheduler for demo purpose,
-    only do single-turn chat completion.
-    """
-
-    async def generate_sequences(self, batch: DataProto, **sampling_params) -> DataProto:
-        kwargs = dict(
-            n=self.config.n,
-            max_completion_tokens=self.config.response_length,
-            temperature=self.config.temperature,
-            top_p=self.config.top_p,
-        )
-
-        do_sample = batch.meta_info.get("do_sample", True)
-        is_validate = batch.meta_info.get("validate", False)
-        if not do_sample or is_validate:
-            kwargs["n"] = 1
-            kwargs["temperature"] = 0
-
-        kwargs.update(sampling_params)
-        print(f"[NaiveChatCompletionScheduler] generate_sequences sampling params: {kwargs}")
-
-        async def callback(completions: ChatCompletion, info: Dict[str, Any], exception: Exception):
-            assert exception is None, f"exception: {exception}"
-            conversation, batch_conversations, batch_index = (
-                info["conversation"],
-                info["batch_conversations"],
-                info["batch_index"],
-            )
-
-            conversations = []
-            for choice in completions.choices:
-                chat = conversation.copy()
-                chat.append({"role": choice.message.role, "content": choice.message.content})
-                conversations.append(chat)
-            batch_conversations[batch_index] = conversations
-
-            # NOTE: we can call tools and resubmit chat completions here.
-            # call_tools(completions, info)
-            # await self.submit_chat_completions(callback2, ...)
-
-        # TODO: we may need to control max concurrent requests here, or it will harm prefix cache hit rate.
-        tasks, batch_conversations = [], [None] * len(batch)
-        for batch_index, conversation in enumerate(batch.non_tensor_batch["raw_prompt"]):
-            # raw_prompt: [{"role": "user", "content": ""}, ["role": "assistant", "content"], ...]
-            tasks.append(
-                asyncio.create_task(
-                    self.submit_chat_completions(
-                        callback=callback,
-                        callback_additional_info={
-                            "batch_conversations": batch_conversations,
-                            "batch_index": batch_index,
-                            "conversation": list(conversation),
-                        },
-                        model=self.model_name,
-                        messages=conversation.tolist(),
-                        **kwargs,
-                    )
-                )
-            )
-        await asyncio.gather(*tasks)
-        print("[NaiveChatCompletionScheduler] generate_sequences done")
-
-        return self._postprocess(batch, batch_conversations, kwargs["n"])
-
-    def _postprocess(
-        self, batch: DataProto, batch_conversations: List[List[List[Dict[str, str]]]], n: int
-    ) -> DataProto:
-        # NOTE: consistent with batch version of generate_sequences in vllm_rollout_spmd.py
-        # prompts: left pad
-        # responses: right pad
-        # input_ids: prompt + response
-        # attention_mask: [0,0,0,0,1,1,1,1, | 1,1,1,0,0,0,0,0]
-        # position_ids:   [0,0,0,0,0,1,2,3, | 4,5,6,7,8,9,10,11]
-
-        # prompts: [prompt] from input dataset
-        prompts = [
-            self.tokenizer.apply_chat_template(prompt, add_generation_prompt=True, tokenize=False)
-            for prompt in batch.non_tensor_batch["raw_prompt"]
-        ]
-
-        # flatten batch_conversations if n > 1
-        assert len(batch_conversations) == len(prompts)
-        batch_conversations = [conversation for conversations in batch_conversations for conversation in conversations]
-        assert len(batch_conversations) == len(prompts) * n
-
-        # sequences: [prompt + response]
-        sequences = [
-            self.tokenizer.apply_chat_template(conversation, add_generation_prompt=False, tokenize=False)
-            for conversation in batch_conversations
-        ]
-
-        # responses: [response]
-        # TODO: mask out tools calling tokens?
-        responses = [sequence[len(prompts[i // n]) :] for i, sequence in enumerate(sequences)]
-
-        prompts = self.tokenizer(prompts, return_tensors="pt", padding="longest", padding_side="left")
-        responses = self.tokenizer(responses, return_tensors="pt", padding="longest", padding_side="right")
-        if n > 1:
-            prompts["input_ids"] = prompts["input_ids"].repeat_interleave(n, dim=0)
-            prompts["attention_mask"] = prompts["attention_mask"].repeat_interleave(n, dim=0)
-
-        input_ids = torch.cat([prompts["input_ids"], responses["input_ids"]], dim=1)
-        attention_mask = torch.cat([prompts["attention_mask"], responses["attention_mask"]], dim=1)
-        position_ids = (attention_mask.cumsum(dim=1) - 1) * attention_mask
-
-        batch = TensorDict(
-            {
-                "prompts": prompts["input_ids"],
-                "responses": responses["input_ids"],
-                "input_ids": input_ids,
-                "attention_mask": attention_mask,
-                "position_ids": position_ids,
-            },
-            batch_size=len(input_ids),
-        )
-
-        return DataProto(batch=batch)

agentlightning/instrumentation/vllm.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 import warnings
-from typing import List
+from typing import Any, List
 
 import vllm.entrypoints.openai.protocol
 from vllm.entrypoints.openai.protocol import ChatCompletionResponse
@@ -19,15 +19,15 @@ original_chat_completion_full_generator = OpenAIServingChat.chat_completion_full
 
 
 async def chat_completion_full_generator(
-    self,
-    request,
-    result_generator,
+    self: Any,
+    request: Any,
+    result_generator: Any,
     request_id: str,
     model_name: str,
-    conversation,
-    tokenizer,
-    request_metadata,
-):
+    conversation: Any,
+    tokenizer: Any,
+    request_metadata: Any,
+) -> Any:
     prompt_token_ids: List[int] | None = None
     response_token_ids: List[List[int]] | None = None
 

agentlightning/litagent.py

@@ -6,9 +6,9 @@ import functools
 import inspect
 import logging
 import weakref
-from typing import TYPE_CHECKING, Any, Callable, Coroutine, Dict, Generic, List, Optional, TypeVar, Union
+from typing import TYPE_CHECKING, Any, Callable, Coroutine, Generic, Optional, TypeVar, Union
 
-from .types import LLM, NamedResources, Rollout, RolloutRawResult, Task, TaskInput, Triplet
+from .types import LLM, NamedResources, Rollout, RolloutRawResult, Task
 
 if TYPE_CHECKING:
     from .runner import AgentRunner
@@ -28,7 +28,7 @@ __all__ = [
 ]
 
 
-def is_v0_1_rollout_api(func: Callable) -> bool:
+def is_v0_1_rollout_api(func: Callable[..., Any]) -> bool:
     """Check if the rollout API is v0.1.
     Inspect the function signature to see if it has a rollout_id parameter.
 
@@ -71,13 +71,13 @@ class LitAgent(Generic[T]):
         return (
             (
                 hasattr(self, "training_rollout_async")
-                and self.__class__.training_rollout_async is not LitAgent.training_rollout_async
+                and self.__class__.training_rollout_async is not LitAgent.training_rollout_async  # type: ignore
             )
             or (
                 hasattr(self, "validation_rollout_async")
-                and self.__class__.validation_rollout_async is not LitAgent.validation_rollout_async
+                and self.__class__.validation_rollout_async is not LitAgent.validation_rollout_async  # type: ignore
             )
-            or (hasattr(self, "rollout_async") and self.__class__.rollout_async is not LitAgent.rollout_async)
+            or (hasattr(self, "rollout_async") and self.__class__.rollout_async is not LitAgent.rollout_async)  # type: ignore
         )
 
     def set_trainer(self, trainer: Trainer) -> None:
@@ -321,9 +321,9 @@ class LitAgentLLM(LitAgent[T]):
         self._accepts_rollout = "rollout" in inspect.signature(llm_rollout_func).parameters
 
         # Copy function metadata to preserve type hints and other attributes
-        functools.update_wrapper(self, llm_rollout_func)
+        functools.update_wrapper(self, llm_rollout_func)  # type: ignore
 
-    def __call__(self, *args, **kwargs):
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
         """Make the agent instance callable, preserving the original function behavior."""
         return self.llm_rollout_func(*args, **kwargs)
 
@@ -434,7 +434,7 @@ def llm_rollout(func: LlmRolloutFunc[T], *, trained_agents: Optional[str] = None
     return LitAgentLLM(func, trained_agents=trained_agents)
 
 
-def rollout(func: Union[LlmRolloutFunc[T], Callable], *, trained_agents: Optional[str] = None) -> LitAgent[T]:
+def rollout(func: Union[LlmRolloutFunc[T], Callable[..., Any]], *, trained_agents: Optional[str] = None) -> LitAgent[T]:
     """Create a LitAgent from a function, automatically detecting the appropriate type.
 
     This function inspects the provided callable and creates the appropriate

agentlightning/reward.py

@@ -3,17 +3,20 @@
 import asyncio
 import inspect
 import warnings
-from typing import Optional, TypedDict
+from typing import Any, Callable, Literal, Optional, TypedDict, TypeVar
 
 from agentops.sdk.decorators import operation
 
 
 class RewardSpanData(TypedDict):
-    type: "reward"
+    type: Literal["reward"]
     value: Optional[float]
 
 
-def reward(fn: callable) -> callable:
+FnType = TypeVar("FnType", bound=Callable[..., Any])
+
+
+def reward(fn: FnType) -> FnType:
     """
     A decorator to wrap a function that computes rewards.
     It will automatically handle the input and output of the function.
@@ -25,7 +28,7 @@ def reward(fn: callable) -> callable:
         """
         if result is None:
             return {"type": "reward", "value": None}
-        if not isinstance(result, (float, int)):
+        if not isinstance(result, (float, int)):  # type: ignore
             warnings.warn(f"Reward is ignored because it is not a number: {result}")
             return {"type": "reward", "value": None}
         return {"type": "reward", "value": float(result)}
@@ -35,7 +38,7 @@ def reward(fn: callable) -> callable:
 
     if is_async:
 
-        async def wrapper_async(*args, **kwargs):
+        async def wrapper_async(*args: Any, **kwargs: Any) -> Any:
             result: Optional[float] = None
 
             @operation
@@ -49,11 +52,11 @@ def reward(fn: callable) -> callable:
             await agentops_reward_operation()
             return result
 
-        return wrapper_async
+        return wrapper_async  # type: ignore
 
     else:
 
-        def wrapper(*args, **kwargs):
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
             result: Optional[float] = None
 
             @operation
@@ -65,4 +68,4 @@ def reward(fn: callable) -> callable:
             agentops_reward_operation()
             return result
 
-        return wrapper
+        return wrapper  # type: ignore

agentlightning/runner.py

@@ -1,21 +1,17 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-import asyncio
 import json
 import logging
-import os
 import time
-from contextlib import nullcontext
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, cast
 
-import agentops
 from opentelemetry.sdk.trace import ReadableSpan
 
 from .client import AgentLightningClient
 from .litagent import LitAgent, is_v0_1_rollout_api
 from .tracer import TripletExporter
 from .tracer.base import BaseTracer
-from .types import ParallelWorkerBase, Rollout, RolloutRawResult, Task, Triplet
+from .types import ParallelWorkerBase, Rollout, RolloutRawResult, Triplet
 
 logger = logging.getLogger(__name__)
 
@@ -38,7 +34,7 @@ class AgentRunner(ParallelWorkerBase):
 
     def __init__(
         self,
-        agent: LitAgent,
+        agent: LitAgent[Any],
         client: AgentLightningClient,
         tracer: BaseTracer,
         triplet_exporter: TripletExporter,
@@ -169,8 +165,11 @@ class AgentRunner(ParallelWorkerBase):
                 rollout_method = self.agent.training_rollout if task.mode == "train" else self.agent.validation_rollout
                 # Pass the task input, not the whole task object
                 if is_v0_1_rollout_api(rollout_method):
-                    result = rollout_method(
-                        task.input, rollout_id=rollout_obj.rollout_id, resources=resources_update.resources  # type: ignore
+                    result = cast(
+                        RolloutRawResult,
+                        rollout_method(
+                            task.input, rollout_id=rollout_obj.rollout_id, resources=resources_update.resources  # type: ignore
+                        ),
                     )  # type: ignore
                 else:
                     result = rollout_method(task.input, resources=resources_update.resources, rollout=rollout_obj)
@@ -245,8 +244,11 @@ class AgentRunner(ParallelWorkerBase):
                 )
                 # Pass the task input, not the whole task object
                 if is_v0_1_rollout_api(rollout_method):
-                    result = await rollout_method(
-                        task.input, rollout_id=rollout_obj.rollout_id, resources=resources_update.resources  # type: ignore
+                    result = cast(
+                        RolloutRawResult,
+                        await rollout_method(
+                            task.input, rollout_id=rollout_obj.rollout_id, resources=resources_update.resources  # type: ignore
+                        ),
                     )  # type: ignore
                 else:
                     result = await rollout_method(task.input, resources=resources_update.resources, rollout=rollout_obj)

agentlightning/server.py

@@ -10,7 +10,6 @@ from typing import Any, Dict, List, Literal, Optional
 
 import uvicorn
 from fastapi import FastAPI, HTTPException, Path
-from pydantic import Field
 
 from .types import (
     GenericResponse,
@@ -218,7 +217,7 @@ class AgentLightningServer:
             return
         processing_tasks = self._store.get_processing_tasks()
 
-        for rollout_id, task in processing_tasks.items():
+        for _, task in processing_tasks.items():
             if task.last_claim_time and current_time - task.last_claim_time > self._task_timeout_seconds:
                 await self._store.requeue_task(task)
                 logger.warning(
@@ -229,7 +228,7 @@ class AgentLightningServer:
         """Setup FastAPI routes."""
 
         @self._app.get("/task", response_model=TaskIfAny)
-        async def next_task() -> TaskIfAny:
+        async def next_task() -> TaskIfAny:  # type: ignore
             """Endpoint for clients to poll for the next available task."""
             await self._check_and_requeue_stale_tasks()
 
@@ -245,7 +244,7 @@ class AgentLightningServer:
                 return TaskIfAny(is_available=False)
 
         @self._app.get("/resources/latest", response_model=ResourcesUpdate)
-        async def fetch_latest_resources() -> ResourcesUpdate:
+        async def fetch_latest_resources() -> ResourcesUpdate:  # type: ignore
             """Endpoint for clients to poll for the latest available resources."""
             if not self._store:
                 raise HTTPException(status_code=503, detail="Server not fully initialized.")
@@ -256,7 +255,7 @@ class AgentLightningServer:
             return resources_update
 
         @self._app.get("/resources/{resource_id}", response_model=ResourcesUpdate)
-        async def fetch_resources_by_id(
+        async def fetch_resources_by_id(  # type: ignore
             resource_id: str = Path(..., description="The unique identifier for the resource version.")
         ) -> ResourcesUpdate:
             """Endpoint for clients to fetch a specific version of resources."""
@@ -269,7 +268,7 @@ class AgentLightningServer:
             return resources_update
 
         @self._app.post("/rollout", response_model=GenericResponse)
-        async def post_rollout(payload: Rollout) -> GenericResponse:
+        async def post_rollout(payload: Rollout) -> GenericResponse:  # type: ignore
             """Endpoint for clients to report a completed rollout."""
             if not self._store:
                 raise HTTPException(status_code=503, detail="Server not fully initialized.")

agentlightning/tracer/__init__.py

@@ -3,3 +3,5 @@
 from .agentops import AgentOpsTracer
 from .base import BaseTracer
 from .triplet import TripletExporter
+
+__all__ = ["AgentOpsTracer", "BaseTracer", "TripletExporter"]

agentlightning/tracer/agentops.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 import logging
 import os
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, List, Optional
+from typing import TYPE_CHECKING, Any, Iterator, List, Optional
 
 import agentops
 import agentops.sdk.core
@@ -67,12 +67,12 @@ class AgentOpsTracer(BaseTracer):
         logger.debug(f"Getting state for pickling Trainer (PID {os.getpid()}). _agentops_server_manager excluded.")
         return state
 
-    def __setstate__(self, state):
+    def __setstate__(self, state: Any):
         self.__dict__.update(state)
         # In child process, self._agentops_server_manager will be None.
         logger.debug(f"Setting state for unpickled Trainer (PID {os.getpid()}). _agentops_server_manager is None.")
 
-    def init(self, *args, **kwargs):
+    def init(self, *args: Any, **kwargs: Any):
         if self.agentops_managed and self._agentops_server_manager:
             self._agentops_server_manager.start()
             self._agentops_server_port_val = self._agentops_server_manager.get_port()
@@ -124,7 +124,7 @@ class AgentOpsTracer(BaseTracer):
                 )
 
             if not agentops.get_client().initialized:
-                agentops.init()
+                agentops.init()  # type: ignore
                 logger.info(f"[Worker {worker_id}] AgentOps client initialized.")
             else:
                 logger.warning(f"[Worker {worker_id}] AgentOps client was already initialized.")
@@ -134,11 +134,11 @@ class AgentOpsTracer(BaseTracer):
         try:
             # new versions
             instance = agentops.sdk.core.tracer
-            instance.provider.add_span_processor(self._lightning_span_processor)
+            instance.provider.add_span_processor(self._lightning_span_processor)  # type: ignore
         except AttributeError:
             # old versions
-            instance = TracingCore.get_instance()
-            instance._provider.add_span_processor(self._lightning_span_processor)
+            instance = TracingCore.get_instance()  # type: ignore
+            instance._provider.add_span_processor(self._lightning_span_processor)  # type: ignore
 
     def teardown_worker(self, worker_id: int) -> None:
         super().teardown_worker(worker_id)
@@ -148,7 +148,7 @@ class AgentOpsTracer(BaseTracer):
             logger.info(f"[Worker {worker_id}] Instrumentation removed.")
 
     @contextmanager
-    def trace_context(self, name: Optional[str] = None):
+    def trace_context(self, name: Optional[str] = None) -> Iterator[LightningSpanProcessor]:
         """
         Starts a new tracing context. This should be used as a context manager.
 
@@ -209,7 +209,7 @@ class LightningSpanProcessor(SpanProcessor):
         self._spans = []
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any):
         pass
 
     def spans(self) -> List[ReadableSpan]:

agentlightning/tracer/base.py

@@ -67,7 +67,7 @@ class BaseTracer(ParallelWorkerBase):
         """
         raise NotImplementedError()
 
-    def trace_run(self, func: Callable, *args, **kwargs) -> Any:
+    def trace_run(self, func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
         """
         A convenience wrapper to trace the execution of a single synchronous function.
 
@@ -82,7 +82,7 @@ class BaseTracer(ParallelWorkerBase):
         with self.trace_context(name=func.__name__):
             return func(*args, **kwargs)
 
-    async def trace_run_async(self, func: Callable[..., Awaitable], *args, **kwargs) -> Any:
+    async def trace_run_async(self, func: Callable[..., Awaitable[Any]], *args: Any, **kwargs: Any) -> Any:
         """
         A convenience wrapper to trace the execution of a single asynchronous function.
 

agentlightning/tracer/http.py

@@ -3,11 +3,10 @@
 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 typing import Any, Awaitable, Callable, Dict, Iterator, List, Optional, Tuple
 from urllib.parse import urlparse
 
 from httpdbg.hooks.all import httprecord
@@ -59,14 +58,14 @@ class HttpTracer(BaseTracer):
         subprocess_timeout: float = 3600.0,
     ):
         super().__init__()
-        self._last_records = None
+        self._last_records: Optional[HTTPRecords] = None
         self.include_headers = include_headers
         self.include_body = include_body
         self.include_agentlightning_requests = include_agentlightning_requests
         self.subprocess_mode = subprocess_mode
         self.subprocess_timeout = subprocess_timeout
 
-    def init_worker(self, worker_id: int):
+    def init_worker(self, worker_id: int) -> None:
         """
         Initialize the tracer in a worker process.
 
@@ -114,7 +113,7 @@ class HttpTracer(BaseTracer):
         Returns:
             A list of ReadableSpan objects representing the HTTP activities.
         """
-        spans = []
+        spans: List[ReadableSpan] = []
 
         # Create a trace ID that will be shared by all spans in this trace
         trace_id = int(uuid.uuid4().hex[:16], 16)
@@ -157,7 +156,7 @@ class HttpTracer(BaseTracer):
                 "http.host": parsed_url.netloc,
             }
 
-            if status_code is not None and status_code > 0:
+            if status_code is not None and status_code > 0:  # type: ignore
                 attributes["http.status_code"] = status_code
 
             # Calculate duration - from begin time to last update
@@ -221,7 +220,7 @@ class HttpTracer(BaseTracer):
 
         return spans
 
-    def trace_run(self, func: Callable, *args, **kwargs) -> Any:
+    def trace_run(self, func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
         """
         A convenience wrapper to trace the execution of a single synchronous function.
 
@@ -241,7 +240,7 @@ class HttpTracer(BaseTracer):
         else:
             return super().trace_run(func, *args, **kwargs)
 
-    async def trace_run_async(self, func: Callable[..., Awaitable], *args, **kwargs) -> Any:
+    async def trace_run_async(self, func: Callable[..., Awaitable[Any]], *args: Any, **kwargs: Any) -> Any:
         """
         A convenience wrapper to trace the execution of a single asynchronous function.
 
@@ -264,7 +263,13 @@ class HttpTracer(BaseTracer):
         else:
             return await super().trace_run_async(func, *args, **kwargs)
 
-    def _trace_run_subprocess(self, func: Callable, args=None, kwargs=None, is_async: bool = False) -> Any:
+    def _trace_run_subprocess(
+        self,
+        func: Callable[..., Any],
+        args: Optional[Tuple[Any, ...]] = None,
+        kwargs: Optional[Dict[str, Any]] = None,
+        is_async: bool = False,
+    ) -> Any:
         """
         Execute a function in a subprocess with HTTP tracing.
 
@@ -283,23 +288,23 @@ class HttpTracer(BaseTracer):
             kwargs = {}
 
         # Create a queue to receive results from the subprocess
-        result_queue = multiprocessing.Queue()
+        result_queue = multiprocessing.Queue()  # type: ignore
 
         # Create and start the subprocess
         process = multiprocessing.Process(
-            target=self._subprocess_worker, args=(func, args, kwargs, result_queue, is_async)
+            target=self._subprocess_worker, args=(func, args, kwargs, result_queue, is_async)  # type: ignore
         )
         process.start()
 
         try:
             # Wait for the process to complete and get the result
             process.join(timeout=self.subprocess_timeout)
-            result = result_queue.get_nowait()
+            result = result_queue.get_nowait()  # type: ignore
 
             if result["success"]:
                 # Store the captured records for get_last_trace()
                 self._last_records = result["records"]
-                return result["return_value"]
+                return result["return_value"]  # type: ignore
             else:
                 if "records" in result:
                     self._last_records = result["records"]
@@ -317,7 +322,14 @@ class HttpTracer(BaseTracer):
                 process.terminate()
                 process.join()
 
-    def _subprocess_worker(self, func: Callable, args, kwargs, result_queue: multiprocessing.Queue, is_async: bool):
+    def _subprocess_worker(
+        self,
+        func: Callable[..., Any],
+        args: Tuple[Any, ...],
+        kwargs: Dict[str, Any],
+        result_queue: multiprocessing.Queue,  # type: ignore
+        is_async: bool,
+    ) -> None:
         """
         Worker function that runs in the subprocess to execute the traced function.
 
@@ -355,7 +367,7 @@ class HttpTracer(BaseTracer):
             records = subprocess_tracer._last_records
 
             # Send success result back to parent
-            result_queue.put({"success": True, "return_value": return_value, "records": records})
+            result_queue.put({"success": True, "return_value": return_value, "records": records})  # type: ignore
 
         except Exception as e:
             # Log the exception
@@ -364,4 +376,4 @@ class HttpTracer(BaseTracer):
             # Get the captured records even when there's an exception
             records = subprocess_tracer._last_records
             # Send error result back to parent
-            result_queue.put({"success": False, "exception": e, "records": records})
+            result_queue.put({"success": False, "exception": e, "records": records})  # type: ignore

agentlightning/tracer/triplet.py

@@ -3,7 +3,7 @@
 import json
 import re
 from enum import Enum
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple, cast
 
 from opentelemetry import trace as trace_api
 from opentelemetry.sdk.trace import ReadableSpan
@@ -85,7 +85,7 @@ class TraceTree:
 
         dot = graphviz.Digraph(comment="Trace Tree")
 
-        should_visit_cache = {}
+        should_visit_cache: Dict[str, bool] = {}
 
         def should_visit(node: "TraceTree") -> bool:
             if node.id in should_visit_cache:
@@ -111,14 +111,14 @@ class TraceTree:
             vis_name = node.id[:8] + " (" + node.span.name + ")"
             if agent_name is not None:
                 vis_name += " [" + agent_name + "]"
-            dot.node(node.id, vis_name)
+            dot.node(node.id, vis_name)  # type: ignore
             for child in node.children:
                 if visit(child):
-                    dot.edge(node.id, child.id)
+                    dot.edge(node.id, child.id)  # type: ignore
             return True
 
         visit(self)
-        dot.render(filename, format="png", cleanup=True)
+        dot.render(filename, format="png", cleanup=True)  # type: ignore
 
     def names_tuple(self) -> Tuple[str, List[Any]]:
         """Return the span name, and a list of children.
@@ -129,7 +129,7 @@ class TraceTree:
         agent_name = self.agent_name()
         if agent_name is not None:
             name += " [" + agent_name + "]"
-        children_names = []
+        children_names: List[Tuple[str, List[Any]]] = []
         for child in self.children:
             child_name, child_children = child.names_tuple()
             children_names.append((child_name, child_children))
@@ -163,17 +163,18 @@ class TraceTree:
             raise ValueError("No spans provided to create TraceTree.")
 
         # Process trace items in topological order
-        id_to_span = {span.get_span_context().span_id: span for span in spans}
+        id_to_span = {span.get_span_context().span_id: span for span in spans}  # type: ignore
 
         forward_graph: dict[int, list[int]] = {}
         root_ids: list[int] = []
         for span in spans:
+            span_id: int = span.get_span_context().span_id  # type: ignore
             if span.parent is None:
-                root_ids.append(span.get_span_context().span_id)
+                root_ids.append(span_id)
             else:
                 if span.parent.span_id not in forward_graph:
                     forward_graph[span.parent.span_id] = []
-                forward_graph[span.parent.span_id].append(span.get_span_context().span_id)
+                forward_graph[span.parent.span_id].append(span_id)
 
         # Diff between span with data and forward_graph keys
         # Sometimes the top-level session span is lost.
@@ -181,7 +182,7 @@ class TraceTree:
         for unfound_root in unfound_roots:
             root_ids.append(unfound_root)
 
-        def visit(node_id):
+        def visit(node_id: int) -> "TraceTree":
             children: list[TraceTree] = []
             if node_id in forward_graph:
                 for child_id in forward_graph[node_id]:
@@ -191,15 +192,15 @@ class TraceTree:
                 assert len(children) > 0
                 virtual_span = ReadableSpan(
                     context=trace_api.SpanContext(
-                        trace_id=children[0].span.get_span_context().trace_id,
+                        trace_id=children[0].span.get_span_context().trace_id,  # type: ignore
                         span_id=node_id,
                         is_remote=False,
                     ),
                     name="virtual-node",
                     kind=trace_api.SpanKind.INTERNAL,
                     attributes={},
-                    start_time=min(child.start_time for child in children),
-                    end_time=max(child.end_time for child in children),
+                    start_time=min(child.start_time for child in children),  # type: ignore
+                    end_time=max(child.end_time for child in children),  # type: ignore
                 )
                 return cls(trace_api.format_span_id(node_id), virtual_span, children=children)
             else:
@@ -216,7 +217,7 @@ class TraceTree:
                 id="virtual-root",
                 span=ReadableSpan(
                     context=trace_api.SpanContext(
-                        trace_id=root_spans[0].span.get_span_context().trace_id,
+                        trace_id=root_spans[0].span.get_span_context().trace_id,  # type: ignore
                         span_id=0,
                         is_remote=False,
                     ),
@@ -239,31 +240,34 @@ class TraceTree:
     def agent_name(self) -> Optional[str]:
         """Return the name of agent span. Return the agent or None (not an agent at all).
         Extend this function to support more agent frameworks."""
+        attributes = self.span.attributes
+        if attributes is None:
+            return None
 
         # Case 1: OpenAI Agent SDK
-        agent_name = self.span.attributes.get("agent.name")
+        agent_name = cast(Optional[str], attributes.get("agent.name"))
         if agent_name is not None:
             return agent_name
 
         # Case 2: Agentops decorator @agent
-        is_agent = self.span.attributes.get("agentops.span.kind") == "agent"
+        is_agent = attributes.get("agentops.span.kind") == "agent"
         if is_agent:
-            agent_name = self.span.attributes.get("operation.name")
+            agent_name = cast(Optional[str], attributes.get("operation.name"))
             if agent_name is not None:
                 return agent_name
 
         # Case 3: Autogen team
-        agent_name = self.span.attributes.get("recipient_agent_type")
+        agent_name = cast(Optional[str], attributes.get("recipient_agent_type"))
         if agent_name is not None:
             return agent_name
 
         # Case 4: LangGraph
-        agent_name = self.span.attributes.get("langchain.chain.type")
+        agent_name = cast(Optional[str], attributes.get("langchain.chain.type"))
         if agent_name is not None:
             return agent_name
 
         # Case 5: agent-framework
-        agent_name = self.span.attributes.get("executor.id")
+        agent_name = cast(Optional[str], attributes.get("executor.id"))
         if agent_name is not None:
             return agent_name
 
@@ -272,7 +276,7 @@ class TraceTree:
             "agentops.task.output",  # newer versions of agentops
             "agentops.entity.output",
         ]:
-            output = self.span.attributes.get(key)
+            output = self.span.attributes.get(key)  # type: ignore
             if output:
                 if isinstance(output, dict):
                     return output
@@ -285,7 +289,7 @@ class TraceTree:
 
     def is_reward_span(self) -> bool:
         maybe_reward = self.maybe_reward_dict()
-        return maybe_reward and maybe_reward.get("type") == "reward"
+        return maybe_reward and maybe_reward.get("type") == "reward"  # type: ignore
 
     def find_llm_calls(
         self,
@@ -315,7 +319,7 @@ class TraceTree:
             is_llm_call = False
         if is_llm_call:
             # Check the response id
-            response_id = self.span.attributes.get("gen_ai.response.id")
+            response_id: Optional[str] = self.span.attributes.get("gen_ai.response.id")  # type: ignore
             if response_id is None and within_llm_call is True:
                 is_llm_call = False
             if (
@@ -326,7 +330,7 @@ class TraceTree:
                 is_llm_call = False
 
             if is_llm_call:
-                llm_calls.append((self, within_matching_subtree))
+                llm_calls.append((self, within_matching_subtree))  # type: ignore
                 existing_llm_call_response_ids = existing_llm_call_response_ids or set()
                 if response_id is not None:
                     existing_llm_call_response_ids.add(response_id)
@@ -383,10 +387,10 @@ class TraceTree:
                     continue
                 if node is self:
                     continue
-                if node.start_time <= repair_node.start_time and node.end_time >= repair_node.end_time:
-                    duration_delta = node.end_time - repair_node.end_time + repair_node.start_time - node.start_time
+                if node.start_time <= repair_node.start_time and node.end_time >= repair_node.end_time:  # type: ignore
+                    duration_delta = node.end_time - repair_node.end_time + repair_node.start_time - node.start_time  # type: ignore
                     if duration_delta > 0 and duration_delta < closest_duration:
-                        closest_duration = duration_delta
+                        closest_duration = duration_delta  # type: ignore
                         closest_parent = node
 
             # Repair the hierarchy
@@ -400,18 +404,18 @@ class TraceTree:
         rewards: dict[str, Optional[float]] = {}
 
         if reward_match == RewardMatchPolicy.FIRST_OCCURRENCE:
-            time_sorted: List[TraceTree] = sorted(self.traverse(), key=lambda x: x.start_time)
-            assign_to: List[Tuple[str, int]] = []
+            time_sorted: List[TraceTree] = cast(List[TraceTree], sorted(self.traverse(), key=lambda x: x.start_time))  # type: ignore
+            assign_to: List[Tuple[str, int]] = []  # type: ignore
             for item in time_sorted:
                 if item.id in llm_call_ids:
-                    assign_to.append((item.id, item.end_time))
+                    assign_to.append((item.id, item.end_time))  # type: ignore
 
                 # get reward
                 agentops_output = item.maybe_reward_dict()
                 if agentops_output and agentops_output.get("type") == "reward":
                     for assign_to_id, assign_to_end_time in reversed(assign_to):
                         # This reward happens before the end of the LLM call.
-                        if assign_to_end_time > item.start_time:
+                        if assign_to_end_time > item.start_time:  # type: ignore
                             continue
                         # Ok, we found someone to assign to
                         if assign_to_id in rewards:
@@ -425,12 +429,12 @@ class TraceTree:
                 assign_to: List[Tuple[str, int]] = []
                 for child in item.children:
                     if child.id in llm_call_ids:
-                        assign_to.append(child.id)
+                        assign_to.append(child.id)  # type: ignore
 
                     agentops_output = item.maybe_reward_dict()
                     if agentops_output and agentops_output.get("type") == "reward":
                         for assign_to_id, assign_to_end_time in reversed(assign_to):
-                            if assign_to_end_time > item.start_time:
+                            if assign_to_end_time > item.start_time:  # type: ignore
                                 # This reward happens before the end of the LLM call.
                                 continue
                             if assign_to_id in rewards:
@@ -476,11 +480,11 @@ class TraceTree:
             (
                 llm_call.id,
                 Triplet(
-                    prompt={"token_ids": llm_call.span.attributes.get("prompt_token_ids", [])},
-                    response={"token_ids": llm_call.span.attributes.get("response_token_ids", [])},
+                    prompt={"token_ids": llm_call.span.attributes.get("prompt_token_ids", [])},  # type: ignore
+                    response={"token_ids": llm_call.span.attributes.get("response_token_ids", [])},  # type: ignore
                     reward=None,
                     metadata=dict(
-                        response_id=llm_call.span.attributes.get(
+                        response_id=llm_call.span.attributes.get(  # type: ignore
                             "gen_ai.response.id", None
                         ),  # it works at least for OpenAI
                         agent_name=agent_name,

agentlightning/trainer.py

@@ -4,11 +4,10 @@ import asyncio
 import importlib
 import logging
 import multiprocessing
-import os
 import signal
 import time
 import warnings
-from typing import List, Optional, Union
+from typing import Any, Dict, List, Optional, TypeVar, Union
 
 from .algorithm.base import BaseAlgorithm
 from .client import AgentLightningClient
@@ -21,6 +20,8 @@ from .types import Dataset, ParallelWorkerBase
 
 logger = logging.getLogger(__name__)
 
+T_co = TypeVar("T_co", covariant=True)
+
 
 class Trainer(ParallelWorkerBase):
     """Orchestrates the distributed execution of agent rollouts.
@@ -52,9 +53,9 @@ class Trainer(ParallelWorkerBase):
         n_workers: int = 1,
         max_tasks: Optional[int] = None,
         daemon: bool = True,
-        tracer: Union[BaseTracer, str, dict, None] = None,
-        triplet_exporter: Union[TripletExporter, dict, None] = None,
-        algorithm: Union[BaseAlgorithm, str, dict, None] = None,
+        tracer: Union[BaseTracer, str, Dict[str, Any], None] = None,
+        triplet_exporter: Union[TripletExporter, Dict[str, Any], None] = None,
+        algorithm: Union[BaseAlgorithm, str, Dict[str, Any], None] = None,
     ):
         super().__init__()
         self.n_workers = n_workers
@@ -84,7 +85,7 @@ class Trainer(ParallelWorkerBase):
                 "The cleanup must be handled manually."
             )
 
-    def _make_tracer(self, tracer: Union[BaseTracer, str, dict, None]) -> BaseTracer:
+    def _make_tracer(self, tracer: Union[BaseTracer, str, Dict[str, Any], None]) -> BaseTracer:
         """Creates a tracer instance based on the provided configuration."""
         if isinstance(tracer, BaseTracer):
             return tracer
@@ -107,7 +108,7 @@ class Trainer(ParallelWorkerBase):
             return AgentOpsTracer(agentops_managed=True, instrument_managed=True, daemon=self.daemon)
         raise ValueError(f"Invalid tracer type: {type(tracer)}. Expected BaseTracer, str, dict, or None.")
 
-    def _make_algorithm(self, algorithm: Union[BaseAlgorithm, str, dict, None]) -> Optional[BaseAlgorithm]:
+    def _make_algorithm(self, algorithm: Union[BaseAlgorithm, str, Dict[str, Any], None]) -> Optional[BaseAlgorithm]:
         """Creates an algorithm instance based on the provided configuration."""
         if isinstance(algorithm, BaseAlgorithm):
             return algorithm
@@ -131,7 +132,7 @@ class Trainer(ParallelWorkerBase):
         raise ValueError(f"Invalid algorithm type: {type(algorithm)}. Expected BaseAlgorithm, str, dict, or None.")
 
     def _extract_client_from_data(
-        self, data: Union[str, AgentLightningClient, Dataset]
+        self, data: Union[str, AgentLightningClient, Dataset[Any]]
     ) -> Optional[AgentLightningClient]:
         """Extract client from data if it's a string URL or AgentLightningClient."""
         if isinstance(data, str):
@@ -142,7 +143,9 @@ class Trainer(ParallelWorkerBase):
             return data
         return None
 
-    def _extract_dataset_from_data(self, data: Union[str, AgentLightningClient, Dataset]) -> Optional[Dataset]:
+    def _extract_dataset_from_data(
+        self, data: Union[str, AgentLightningClient, Dataset[Any]]
+    ) -> Optional[Dataset[Any]]:
         """Extract dataset from data if it's a Dataset."""
         if isinstance(data, str) or isinstance(data, AgentLightningClient):
             return None
@@ -150,8 +153,8 @@ class Trainer(ParallelWorkerBase):
 
     def _determine_backend(
         self,
-        train_data: Union[str, AgentLightningClient, Dataset],
-        dev_data: Union[str, AgentLightningClient, Dataset, None] = None,
+        train_data: Union[str, AgentLightningClient, Dataset[Any]],
+        dev_data: Union[str, AgentLightningClient, Dataset[Any], None] = None,
     ) -> Union[str, AgentLightningClient]:
         """Determine which backend to use for initialization."""
         if self.dev:
@@ -172,6 +175,10 @@ class Trainer(ParallelWorkerBase):
                 client = self.algorithm.get_client()
                 logger.info(f"Algorithm created client: {client}")
                 return client
+            if client is None:
+                raise ValueError(
+                    "train_data must be a string URL or AgentLightningClient when no algorithm is provided."
+                )
             return client
 
     def init(self, backend: Union[str, AgentLightningClient]) -> None:
@@ -203,7 +210,7 @@ class Trainer(ParallelWorkerBase):
                 self._client = backend
             else:
                 logger.info(f"Initializing AgentLightningClient with endpoint: {backend}")
-                if not isinstance(backend, str):
+                if not isinstance(backend, str):  # type: ignore
                     raise ValueError("backend must be a string URL or an AgentLightningClient instance.")
                 if not backend.startswith("http://") and not backend.startswith("https://"):
                     raise ValueError("backend must be a valid URL starting with http:// or https://")
@@ -213,7 +220,7 @@ class Trainer(ParallelWorkerBase):
             logger.warning("AgentLightningClient already initialized. Returning existing instance.")
         return self._client
 
-    def _worker_main_loop(self, agent: LitAgent, worker_id: int, is_async: bool):
+    def _worker_main_loop(self, agent: LitAgent[Any], worker_id: int, is_async: bool):
         """The main function for each worker process.
 
         This function initializes the client and the loop, then starts the
@@ -283,7 +290,7 @@ class Trainer(ParallelWorkerBase):
         """
         import psutil
 
-        for proc in psutil.process_iter():
+        for proc in psutil.process_iter():  # type: ignore
             # check whether the process name matches
             if proc.name().startswith("AgentLightning-"):
                 proc.kill()
@@ -308,11 +315,11 @@ class Trainer(ParallelWorkerBase):
 
     def fit(
         self,
-        agent: LitAgent,
-        train_data: Union[str, AgentLightningClient, Dataset],
+        agent: LitAgent[T_co],
+        train_data: Union[str, AgentLightningClient, Dataset[T_co]],
         *,
-        val_data: Union[str, AgentLightningClient, Dataset, None] = None,
-        dev_data: Union[str, AgentLightningClient, Dataset, None] = None,
+        val_data: Union[str, AgentLightningClient, Dataset[T_co], None] = None,
+        dev_data: Union[str, AgentLightningClient, Dataset[T_co], None] = None,
         dev_backend: Union[str, AgentLightningClient, None] = None,
     ):
         """Train the agent using the provided data.
@@ -443,7 +450,7 @@ class Trainer(ParallelWorkerBase):
             self._terminate_processes(processes)
             logger.info(f"Workers terminated or single worker interrupted.")
             raise
-        except Exception as e:
+        except Exception:
             logger.exception(f"Unhandled exception in fit method.")
             self._terminate_processes(processes)
             logger.info(f"Workers terminated or single worker interrupted.")

agentlightning/types.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 from typing import Annotated, Any, Dict, Generic, List, Literal, Optional, Protocol, TypeVar, Union
 
 from opentelemetry.sdk.trace import ReadableSpan
-from pydantic import BaseModel, Discriminator, Field
+from pydantic import BaseModel, Field
 
 __all__ = [
     "Triplet",
@@ -25,7 +25,7 @@ __all__ = [
     "Dataset",
 ]
 
-T = TypeVar("T")
+T_co = TypeVar("T_co", covariant=True)
 
 
 class Triplet(BaseModel):
@@ -214,13 +214,13 @@ class ParallelWorkerBase:
         pass
 
 
-class Dataset(Protocol, Generic[T]):
+class Dataset(Protocol, Generic[T_co]):
     """The general interface for a dataset.
 
     It's currently implemented as a protocol, having a similar interface to torch.utils.data.Dataset.
     You don't have to inherit from this class; you can use a simple list if you want to.
     """
 
-    def __getitem__(self, index: int) -> T: ...
+    def __getitem__(self, index: int) -> T_co: ...
 
     def __len__(self) -> int: ...

agentlightning/verl/async_server.py

@@ -1,5 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+# type: ignore
+
 from copy import deepcopy
 
 import ray

agentlightning/verl/daemon.py

@@ -7,14 +7,13 @@ import socket
 import threading
 import time
 import uuid
-from collections.abc import Mapping, Sequence
-from typing import Dict, List, Optional
+from collections.abc import Mapping
+from typing import Any, Dict, List, Optional, Tuple
 
 import numpy as np
 import requests
 import torch
 from flask import Flask, Response, abort, request
-from openai.types.chat.chat_completion import ChatCompletion
 from tensordict import TensorDict
 from verl import DataProto
 
@@ -23,7 +22,9 @@ from agentlightning import LLM, AgentLightningServer, NamedResources, Rollout, c
 configure_logger()
 
 
-def get_left_padded_ids_and_attention_mask(ids: List[int], max_length: int, pad_token_id: int):
+def get_left_padded_ids_and_attention_mask(
+    ids: List[int], max_length: int, pad_token_id: int
+) -> Tuple[List[int], List[int]]:
     """
     Left-pad (or truncate) a sequence of token IDs to a fixed length,
     and build the corresponding attention mask.
@@ -52,7 +53,9 @@ def get_left_padded_ids_and_attention_mask(ids: List[int], max_length: int, pad_
     return padded_ids, attention_mask
 
 
-def get_right_padded_ids_and_attention_mask(ids: List[int], max_length: int, pad_token_id: int):
+def get_right_padded_ids_and_attention_mask(
+    ids: List[int], max_length: int, pad_token_id: int
+) -> Tuple[List[int], List[int]]:
     """
     Right-pad (or truncate) a sequence of token IDs to a fixed length,
     and build the corresponding attention mask.
@@ -87,7 +90,7 @@ def _find_available_port() -> int:
         return s.getsockname()[1]
 
 
-def _to_native(obj):
+def _to_native(obj: Any) -> Any:
     """Convert data retrieved from Parquet to data usable in AGL server."""
     # 1) Arrays -> list (then recurse)
     if isinstance(obj, np.ndarray):
@@ -99,11 +102,11 @@ def _to_native(obj):
 
     # 3) Dict-like -> dict
     if isinstance(obj, Mapping):
-        return {_to_native(k): _to_native(v) for k, v in obj.items()}
+        return {_to_native(k): _to_native(v) for k, v in obj.items()}  # type: ignore
 
     # 4) Lists/Tuples/Sets -> list
     if isinstance(obj, (list, tuple, set)):
-        return [_to_native(x) for x in obj]
+        return [_to_native(x) for x in obj]  # type: ignore
 
     # 5) Anything else: leave as-is
     return obj
@@ -120,14 +123,14 @@ class AgentModeDaemon:
 
     def __init__(
         self,
-        port,
-        train_rollout_n,
-        train_information,
-        tokenizer,
-        mini_batch_size,
-        pad_token_id,
-        reward_fillna_value=0.0,
-        llm_timeout_seconds=1200.0,
+        port: int,
+        train_rollout_n: int,
+        train_information: Dict[str, Any],
+        tokenizer: Any,
+        mini_batch_size: int,
+        pad_token_id: int,
+        reward_fillna_value: float = 0.0,
+        llm_timeout_seconds: float = 1200.0,
     ):
         # Server and Task Configuration
         self.server_port = port
@@ -149,7 +152,7 @@ class AgentModeDaemon:
         self.backend_llm_server_addresses: List[str] = []
         self._total_tasks_queued = 0
         self._completed_rollouts: Dict[str, Rollout] = {}
-        self._task_id_to_original_sample: Dict[str, Dict] = {}
+        self._task_id_to_original_sample: Dict[str, Dict[str, Any]] = {}
         self._server_thread: Optional[threading.Thread] = None
         self._proxy_thread: Optional[threading.Thread] = None
         self.is_train = True
@@ -165,7 +168,7 @@ class AgentModeDaemon:
         last_request_time = 0
 
         @app.route("/v1/<path:path>", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"])
-        def proxy(path):
+        def proxy(path: str):  # type: ignore
             if not self.backend_llm_server_addresses:
                 abort(503, description="No backend LLM servers available.")
 
@@ -190,7 +193,7 @@ class AgentModeDaemon:
                     method=request.method,
                     url=target_url,
                     headers=headers,
-                    params=request.args,
+                    params=request.args,  # type: ignore
                     data=request.get_data(),
                     cookies=request.cookies,
                     allow_redirects=False,
@@ -264,7 +267,7 @@ class AgentModeDaemon:
 
         self._start_proxy_server()
 
-    async def _async_set_up(self, data, server_addresses, is_train=True):
+    async def _async_set_up(self, data: Dict[str, Any], server_addresses: List[str], is_train: bool = True):
         """Async helper to set up data and resources on the server."""
         self.clear_data_and_server()
         self.backend_llm_server_addresses = server_addresses
@@ -290,7 +293,7 @@ class AgentModeDaemon:
             original_sample["data_id"] = data_id
 
             # For training, each sample is rolled out multiple times
-            for j in range(rollouts_per_sample):
+            for _ in range(rollouts_per_sample):
                 task_metadata = {"data_id": data_id, "is_train": is_train}
 
                 # Data ID is different from Rollout ID, as one data can have multiple rollouts.
@@ -304,7 +307,7 @@ class AgentModeDaemon:
                 self._task_id_to_original_sample[rollout_id] = original_sample
                 self._total_tasks_queued += 1
 
-    def set_up_data_and_server(self, data, server_addresses, is_train=True):
+    def set_up_data_and_server(self, data: Dict[str, Any], server_addresses: List[str], is_train: bool = True):
         """Synchronous wrapper for setting up data and server resources."""
         if not self.server.loop or not self.server.startup_event.is_set():
             raise RuntimeError("Server is not running or ready.")
@@ -331,7 +334,7 @@ class AgentModeDaemon:
         elif any(not r.prompt.get("token_ids", []) for r in rollout.triplets):
             print(f"Warning: Rollout {rollout.rollout_id} contains empty prompt: {rollout.triplets}")
 
-    async def _async_run_until_finished(self, verbose=True):
+    async def _async_run_until_finished(self, verbose: bool = True):
         """Async helper to wait for all tasks to complete."""
         while len(self._completed_rollouts) < self._total_tasks_queued:
             completed_batch = await self.server.retrieve_completed_rollouts()
@@ -346,7 +349,7 @@ class AgentModeDaemon:
             await asyncio.sleep(5)
         print("All tasks finished.")
 
-    def run_until_all_finished(self, verbose=True):
+    def run_until_all_finished(self, verbose: bool = True):
         """Synchronously waits for all queued tasks to be completed and reported."""
         if self._total_tasks_queued == 0:
             print("Warning: No tasks were queued.")
@@ -368,8 +371,8 @@ class AgentModeDaemon:
         assert not self.is_train, "This method should only be called during validation."
         assert len(self._completed_rollouts) == self._total_tasks_queued
 
-        sample_stat_list = []
-        for rollout_id, rollout in self._completed_rollouts.items():
+        sample_stat_list: List[Dict[str, Any]] = []
+        for _, rollout in self._completed_rollouts.items():
             final_reward = self._fillna_reward(rollout)
             if not rollout.triplets:
                 print(f"Warning: No triplets found for test rollout {rollout.rollout_id}.")
@@ -397,7 +400,7 @@ class AgentModeDaemon:
             "val/turn_count": np.mean([stat["turn_count"] for stat in stats_w_trace]),
         }
 
-    def get_train_data_batch(self, max_prompt_length, max_response_length, device):
+    def get_train_data_batch(self, max_prompt_length: int, max_response_length: int, device: torch.device):
         """
         Processes completed rollouts to generate a training data batch.
 
@@ -409,8 +412,8 @@ class AgentModeDaemon:
         assert len(self._completed_rollouts) == self._total_tasks_queued
 
         # 1. Reconstruct the `finished_id_to_sample_info` structure from completed rollouts
-        finished_id_to_sample_info = {}
-        finished_id_to_final_reward = {}
+        finished_id_to_sample_info: Dict[str, Dict[str, Any]] = {}
+        finished_id_to_final_reward: Dict[str, float] = {}
         for rollout_id, rollout in self._completed_rollouts.items():
             original_sample = self._task_id_to_original_sample[rollout_id]
 
@@ -446,9 +449,15 @@ class AgentModeDaemon:
         #     discarded here. They are only truncated and marked, to be discarded later.
         #     This is for the correctness of the advantage calculation.
         #   - The discard for the PPO mini-batch should also be handled this way.
-        input_ids_list, input_attention_mask_list = [], []
-        response_ids_list, response_attention_mask_list = [], []
-        reward_list, data_id_list, rollout_id_list, turn_index_list, is_drop_list = [], [], [], [], []
+        input_ids_list: List[List[int]] = []
+        input_attention_mask_list: List[List[int]] = []
+        response_ids_list: List[List[int]] = []
+        response_attention_mask_list: List[List[int]] = []
+        reward_list: List[float] = []
+        data_id_list: List[str] = []
+        rollout_id_list: List[str] = []
+        turn_index_list: List[int] = []
+        is_drop_list: List[bool] = []
         n_trunc_sample_because_of_response = 0
 
         for rollout_id, sample_info in finished_id_to_sample_info.items():
@@ -531,9 +540,9 @@ class AgentModeDaemon:
         }
 
         # Add non-tensor data for advantage calculation and logging
-        data_proto.non_tensor_batch["data_id_list"] = np.array(data_id_list)
-        data_proto.non_tensor_batch["rollout_id_list"] = np.array(rollout_id_list)
-        data_proto.non_tensor_batch["turn_index_list"] = np.array(turn_index_list)
+        data_proto.non_tensor_batch["data_id_list"] = np.array(data_id_list)  # type: ignore
+        data_proto.non_tensor_batch["rollout_id_list"] = np.array(rollout_id_list)  # type: ignore
+        data_proto.non_tensor_batch["turn_index_list"] = np.array(turn_index_list)  # type: ignore
 
         return data_proto, data_metrics
 
@@ -547,9 +556,9 @@ class AgentModeDaemon:
         # This implementation assumes that `set_up_data_and_server` is called
         # for each new run, effectively starting a fresh batch.
 
-    def _fillna_reward(self, rollout):
+    def _fillna_reward(self, rollout: Rollout):
         if rollout.final_reward is None:
-            if self.reward_fillna_value is not None:
+            if self.reward_fillna_value is not None:  # type: ignore
                 final_reward = self.reward_fillna_value
             else:
                 raise ValueError(f"Reward is None for rollout {rollout.rollout_id}, please check the reward function.")

agentlightning/verl/dataset.py

@@ -1,5 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+# type: ignore
+
 import torch
 from datasets import Dataset as HuggingFaceDataset
 from omegaconf import DictConfig

agentlightning/verl/entrypoint.py

@@ -1,5 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+# type: ignore
+
 from typing import Any
 
 import hydra
@@ -18,7 +20,7 @@ def main(config):
     run_ppo(config, None, None)
 
 
-def run_ppo(config: Any, train_dataset: Dataset | None, val_dataset: Dataset | None) -> None:
+def run_ppo(config: Any, train_dataset: Dataset[Any] | None, val_dataset: Dataset[Any] | None) -> None:
     if not ray.is_initialized():
         # this is for local ray cluster
         ray.init(

agentlightning/verl/trainer.py

@@ -1,5 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+# type: ignore
+
 import random
 from contextlib import contextmanager
 from copy import deepcopy

examples/apo/apo.py

@@ -1,6 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import asyncio
+from typing import cast
 
 from agentlightning.server import AgentLightningServer
 from agentlightning.types import NamedResources, PromptTemplate
@@ -19,7 +20,7 @@ async def example_apo():
         "You are a friendly chatbot.",
     ]
 
-    prompt_and_rewards = []
+    prompt_and_rewards: list[tuple[str, float]] = []
 
     for prompt in prompt_candidates:
         # 1. The optimization algorithm updates the prompt template
@@ -38,7 +39,7 @@ async def example_apo():
         assert rollout, "Expected a completed rollout from the client."
         print(f"[Algo] Received Result: {rollout}")
         reward = rollout.final_reward
-        prompt_and_rewards.append((prompt, reward))
+        prompt_and_rewards.append((prompt, cast(float, reward)))
 
     print(f"\n[Algo] All prompts and their rewards: {prompt_and_rewards}")
     best_prompt = max(prompt_and_rewards, key=lambda x: x[1])

examples/apo/apo_client.py

@@ -2,6 +2,7 @@
 
 import os
 import random
+from typing import Any
 
 import dotenv
 from openai import OpenAI
@@ -11,10 +12,10 @@ from agentlightning.litagent import LitAgent
 from agentlightning.trainer import Trainer
 
 
-class SimpleAgent(LitAgent):
+class SimpleAgent(LitAgent[Any]):
 
-    def training_rollout(self, task, rollout_id, resources):
-        print("Resources:", resources)
+    def training_rollout(self, task, rollout_id, resources):  # type: ignore
+        print("Resources:", resources)  # type: ignore
 
         openai = OpenAI(
             api_key=os.environ["OPENAI_API_KEY"],
@@ -24,7 +25,7 @@ class SimpleAgent(LitAgent):
         result = openai.chat.completions.create(
             model="gpt-4.1-mini",
             messages=[
-                {"role": "system", "content": resources["system_prompt"].template},
+                {"role": "system", "content": resources["system_prompt"].template},  # type: ignore
                 {"role": "user", "content": task["prompt"]},
             ],
         )

examples/calc_x/calc_agent.py

@@ -4,7 +4,7 @@ import math
 import os
 import re
 import string
-from typing import Any
+from typing import Any, cast
 
 import sympy
 from autogen_agentchat.agents import AssistantAgent
@@ -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 LLM, DevTaskLoader, LitAgent, NamedResources, Trainer, configure_logger, reward
+from agentlightning import LLM, LitAgent, NamedResources, Trainer, configure_logger, reward
 
 configure_logger()
 
@@ -48,8 +48,8 @@ def float_eval(input_str: str) -> float:
 
 
 def scalar_are_results_same(pred_result: str, true_result: str, rel_tol: float) -> bool:
-    pred_result = str(pred_result) if pred_result is not None else ""
-    true_result = str(true_result) if true_result is not None else ""
+    pred_result = str(pred_result) if pred_result is not None else ""  # type: ignore
+    true_result = str(true_result) if true_result is not None else ""  # type: ignore
 
     if pred_result.strip() == true_result.strip():
         return True
@@ -76,7 +76,7 @@ async def eval(prediction: str, ground_truth: str) -> float:
     return float(scalar_are_results_same(prediction, ground_truth, 1e-2))
 
 
-def get_agent(model, openai_base_url, temperature, workbench):
+def get_agent(model: str, openai_base_url: str, temperature: float, workbench: McpWorkbench) -> AssistantAgent:
     model_client = OpenAIChatCompletionClient(
         model=model,
         base_url=openai_base_url,
@@ -100,10 +100,10 @@ def get_agent(model, openai_base_url, temperature, workbench):
     return calc_agent
 
 
-class CalcAgent(LitAgent):
+class CalcAgent(LitAgent[Any]):
 
-    async def training_rollout_async(self, task: Any, rollout_id: str, resources: NamedResources) -> Any:
-        llm: LLM = resources.get("main_llm")
+    async def training_rollout_async(self, task: Any, rollout_id: str, resources: NamedResources) -> Any:  # type: ignore
+        llm: LLM = cast(LLM, resources.get("main_llm"))
         async with McpWorkbench(calculator_mcp_server) as workbench:
             calc_agent = get_agent(
                 llm.model,
@@ -116,19 +116,19 @@ class CalcAgent(LitAgent):
                 prompt = task["question"] + " " + output_format
                 result = await calc_agent.run(task=prompt)
                 # evaluate
-                answer = re.search(r"###\s*ANSWER:\s*(.+?)(\s*###|$)", result.messages[-1].content)
+                answer = re.search(r"###\s*ANSWER:\s*(.+?)(\s*###|$)", result.messages[-1].content)  # type: ignore
                 if answer:
                     answer = answer.group(1)
                 else:
-                    answer = result.messages[-1].content
+                    answer = result.messages[-1].content  # type: ignore
             except Exception as e:
                 print("Failure:", str(e))
                 answer = "None"
-            reward = await eval(answer, str(task["result"]))  # reward is tracked with the decorator
-            print("answer: {} ground_truth: {} reward: {}".format(answer, task["result"], reward))
+            reward = await eval(answer, str(task["result"]))  # reward is tracked with the decorator  # type: ignore
+            print("answer: {} ground_truth: {} reward: {}".format(answer, task["result"], reward))  # type: ignore
 
-    async def validation_rollout_async(self, task: Any, rollout_id: str, resources: NamedResources) -> Any:
-        llm: LLM = resources.get("main_llm")
+    async def validation_rollout_async(self, task: Any, rollout_id: str, resources: NamedResources) -> Any:  # type: ignore
+        llm: LLM = cast(LLM, resources.get("main_llm"))
         resources = {
             "main_llm": LLM(
                 endpoint=llm.endpoint,

examples/calc_x/calc_agent_v0_2.py

@@ -27,16 +27,16 @@ async def calc_agent(task: Any, llm: LLM) -> Any:
             prompt = task["question"] + " " + output_format
             result = await calc_agent.run(task=prompt)
             # evaluate
-            answer = re.search(r"###\s*ANSWER:\s*(.+?)(\s*###|$)", result.messages[-1].content)
+            answer = re.search(r"###\s*ANSWER:\s*(.+?)(\s*###|$)", result.messages[-1].content)  # type: ignore
             if answer:
                 answer = answer.group(1)
             else:
-                answer = result.messages[-1].content
+                answer = result.messages[-1].content  # type: ignore
         except Exception as e:
             print("Failure:", str(e))
             answer = "None"
-        reward = await eval(answer, str(task["result"]))  # reward is tracked with the decorator
-        print("answer: {} ground_truth: {} reward: {}".format(answer, task["result"], reward))
+        reward = await eval(answer, str(task["result"]))  # reward is tracked with the decorator  # type: ignore
+        print("answer: {} ground_truth: {} reward: {}".format(answer, task["result"], reward))  # type: ignore
 
 
 def main():
@@ -104,16 +104,16 @@ def main():
         },
     }
 
-    train_dataset = Dataset.from_parquet("data/train.parquet").to_list()
-    val_dataset = Dataset.from_parquet("data/test_mini.parquet").to_list()
+    train_dataset = Dataset.from_parquet("data/train.parquet").to_list()  # type: ignore
+    val_dataset = Dataset.from_parquet("data/test_mini.parquet").to_list()  # type: ignore
 
     print("First 5 rows of train dataset:")
-    print(train_dataset[:5])
+    print(train_dataset[:5])  # type: ignore
     print("First 5 rows of val dataset:")
-    print(val_dataset[:5])
+    print(val_dataset[:5])  # type: ignore
 
     trainer = Trainer(algorithm=VERL(rl_training_config), n_workers=2)
-    trainer.fit(calc_agent, train_dataset, val_data=val_dataset)
+    trainer.fit(calc_agent, train_dataset, val_data=val_dataset)  # type: ignore
 
 
 if __name__ == "__main__":

examples/calc_x/tests/test_agentops.py

@@ -7,21 +7,21 @@ from agentlightning.reward import reward
 
 
 @reward
-def process_data(data):
+def process_data(data: str) -> float:
     # Your function logic here
-    processed_result = data.upper()
+    processed_result = data.upper()  # type: ignore
     # agentops.record(Events("Processed Data", result=processed_result)) # Optional: record specific events
     return 1.0
 
 
 @operation
-def process_data2(data):
+def process_data2(data: str) -> str:
     # Your function logic here
-    processed_result = data.upper()
+    processed_result = data.upper()  # type: ignore
     # agentops.record(Events("Processed Data", result=processed_result)) # Optional: record specific events
     return processed_result
 
 
-agentops.init()
+agentops.init()  # type: ignore
 process_data("hello")
 process_data2("hello2")

examples/calc_x/tests/test_mcp_calculator.py

@@ -44,8 +44,8 @@ async def main():
     print(chat_resp)
 
     # 4. Extract the expression argument
-    func_call = chat_resp.choices[0].message.tool_calls[0]
-    expr = json.loads(func_call.function.arguments)["expression"]
+    func_call = chat_resp.choices[0].message.tool_calls[0]  # type: ignore
+    expr = json.loads(func_call.function.arguments)["expression"]  # type: ignore
 
     # 5. Connect to the MCP server and invoke the 'calculate' tool
     async with stdio_client(server_params) as (read, write):
@@ -54,7 +54,7 @@ async def main():
             print("Session initialized.")
             result = await session.call_tool("calculate", arguments={"expression": expr})
             # The structured result is under `.structuredContent`
-            value = result.structuredContent["result"]
+            value = result.structuredContent["result"]  # type: ignore
 
     # 6. Print out the result
     print(f"{expr} = {value}")

examples/rag/rag_agent.py

@@ -2,39 +2,20 @@
 
 from __future__ import annotations
 
-import os
-import re
-import shutil
-import sys
-import tempfile
-import time
-from typing import Any, Literal
+from typing import Any, cast
 
-import dotenv
-import termcolor
-from agents import (
-    Agent,
-    Runner,
-    function_tool,
-    gen_trace_id,
-    set_trace_processors,
-    set_tracing_disabled,
-    trace,
-)
+from agents import Agent, Runner
 from agents.extensions.models.litellm_model import LitellmModel
-from agents.mcp import MCPServer, MCPServerSse
+from agents.mcp import MCPServerSse
 from agents.model_settings import ModelSettings
-from agents.tracing.processors import BatchTraceProcessor, ConsoleSpanExporter
 from utils import compute_scores
 
-import agentlightning
 from agentlightning import (
     LLM,
     LitAgent,
     NamedResources,
     Trainer,
     configure_logger,
-    reward,
 )
 
 configure_logger()
@@ -51,13 +32,13 @@ After each search:
 Repeat as needed. When done, wrap your final, concise answer in <answer> tags."""
 
 
-class RAGAgent(LitAgent):
+class RAGAgent(LitAgent[Any]):
     def __init__(self, trained_agents: str | None = None) -> None:
         super().__init__(trained_agents=trained_agents)
         self.mcp_server_url = "http://127.0.0.1:8099/sse"
 
-    async def training_rollout_async(self, task: Any, rollout_id: str, resources: NamedResources) -> Any:
-        llm: LLM = resources.get("main_llm")
+    async def training_rollout_async(self, task: Any, rollout_id: str, resources: NamedResources) -> Any:  # type: ignore
+        llm: LLM = cast(LLM, resources.get("main_llm"))
         print("Training with model:", llm.model, "on endpoint:", llm.endpoint)
         async with MCPServerSse(
             name="wiki_retriever_mcp",
@@ -74,7 +55,7 @@ class RAGAgent(LitAgent):
                 mcp_servers=[server],
             )
             result = await Runner.run(agent, task["question"])
-            answer = result.final_output
+            answer = result.final_output  # type: ignore
             reward = compute_scores(answer, str(task["answer"]))
             print(
                 "question:{} answer: {} ground_truth: {} reward: {}".format(
@@ -83,8 +64,8 @@ class RAGAgent(LitAgent):
             )
             return reward
 
-    async def validation_rollout_async(self, task: Any, rollout_id: str, resources: NamedResources) -> Any:
-        llm: LLM = resources.get("main_llm")
+    async def validation_rollout_async(self, task: Any, rollout_id: str, resources: NamedResources) -> Any:  # type: ignore
+        llm: LLM = cast(LLM, resources.get("main_llm"))
         resources = {
             "main_llm": LLM(
                 endpoint=llm.endpoint,

examples/rag/utils.py

@@ -1,11 +1,11 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-import json
-import pickle
+# type: ignore
+
 import re
 import string
-import sys
 from collections import Counter
+from typing import List, Optional, Set, Tuple
 
 ANS_BEGIN = "<answer>"
 ANS_END = "</answer>"
@@ -14,24 +14,24 @@ FORMAT_SCORE = 0.1
 FORMAT_PUNISH = -2
 
 
-def normalize_answer(s):
-    def remove_articles(text):
+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):
+    def white_space_fix(text: str) -> str:
         return " ".join(text.split())
 
-    def remove_punc(text):
+    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):
+    def lower(text: str) -> str:
         return text.lower()
 
     return white_space_fix(remove_articles(remove_punc(lower(s))))
 
 
-def f1_score(prediction, ground_truth):
+def f1_score(prediction: str, ground_truth: str) -> Tuple[float, float, float]:
     normalized_prediction = normalize_answer(prediction)
     normalized_ground_truth = normalize_answer(ground_truth)
 
@@ -54,7 +54,7 @@ def f1_score(prediction, ground_truth):
     return f1, precision, recall
 
 
-def lenient_f1_score(prediction, ground_truth):
+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)
 
@@ -78,15 +78,15 @@ def lenient_f1_score(prediction, ground_truth):
     return f1, precision, recall
 
 
-def exact_match_score(prediction, ground_truth):
+def exact_match_score(prediction: str, ground_truth: str) -> bool:
     return normalize_answer(prediction) == normalize_answer(ground_truth)
 
 
-def cover_exact_match_score(prediction, 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):
+def extract_answer(response: str) -> str:
     if ANS_BEGIN not in response or ANS_END not in response:
         return ""
     pos1 = response.rfind(ANS_BEGIN)
@@ -99,14 +99,14 @@ def extract_answer(response):
     return ans
 
 
-def split_response(text):
+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, response):
+def extract_recall_chunk(prompt: str, response: str) -> Tuple[Set[str], Set[str]]:
     import re
 
     # 正则表达式，匹配每个search_step内1.和2.后面的内容
@@ -124,7 +124,7 @@ def extract_recall_chunk(prompt, response):
 import re
 
 
-def extract_retrieved_paragraphs(log_text):
+def extract_retrieved_paragraphs(log_text: str) -> List[str]:
     # 正则表达式匹配 "Retrieved paragraph:" 后的内容
     pattern = re.compile(r"Retrieved paragraph:\s*(.*?)\n", re.DOTALL)
 
@@ -134,11 +134,13 @@ def extract_retrieved_paragraphs(log_text):
     return matches
 
 
-def compute_score(prediction, gold, gold_sentences=None, data_source=None):
+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
 
-    prompt, response = split_response(prediction)
+    _, response = split_response(prediction)
     ans = extract_answer(response)
     if ans == "":
         # format score 0.1
@@ -152,8 +154,8 @@ def compute_score(prediction, gold, gold_sentences=None, data_source=None):
         return format_acc
 
     # answer acc
-    em, cem = exact_match_score(ans, gold), cover_exact_match_score(ans, gold)
-    f1, prec, recall = f1_score(ans, gold)
+    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)
@@ -175,27 +177,27 @@ def compute_score(prediction, gold, gold_sentences=None, data_source=None):
 
 
 def compute_reward(
-    solution_str=None,
-    ground_truth=None,
-    gold_sentences=None,
-    data_source=None,
-    extra_info=None,
-):
+    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=None,
-    ground_truth=None,
-    gold_sentences=None,
-    data_source=None,
-    extra_info=None,
-):
+    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
-    prompt, response = split_response(prediction)
+    _, response = split_response(prediction)
     ans = extract_answer(response)
     if ans == "":
         # format score 0.1
@@ -217,7 +219,7 @@ def compute_cem(
 ):
     prediction = solution_str
     gold = ground_truth
-    prompt, response = split_response(prediction)
+    _, response = split_response(prediction)
     ans = extract_answer(response)
     if ans == "":
         return 0.0
@@ -236,7 +238,7 @@ def compute_response_cem(
 ):
     prediction = solution_str
     gold = ground_truth
-    prompt, response = split_response(prediction)
+    _, response = split_response(prediction)
     ans = response
     if ans == "":
         return 0.0
@@ -255,7 +257,7 @@ def compute_lenient_f1(
 ):
     prediction = solution_str
     gold = ground_truth
-    prompt, response = split_response(prediction)
+    _, response = split_response(prediction)
     ans = extract_answer(response)
     if ans == "":
         return 0.0
@@ -274,7 +276,7 @@ def compute_lenient_response_f1(
 ):
     prediction = solution_str
     gold = ground_truth
-    prompt, response = split_response(prediction)
+    _, response = split_response(prediction)
     ans = response
     if ans == "":
         return 0.0
@@ -284,39 +286,39 @@ def compute_lenient_response_f1(
     return f1
 
 
-def fact_checking_api(prediction, ans):
+def fact_checking_api(prediction: str, ans: str) -> bool:
     return True  # Placeholder for actual fact-checking logic
 
 
 def compute_f1(
-    solution_str=None,
-    ground_truth=None,
-    gold_sentences=None,
-    data_source=None,
-    extra_info=None,
-):
+    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
-    prompt, response = split_response(prediction)
+    _, response = split_response(prediction)
     ans = extract_answer(response)
     if ans == "":
         return 0.0
 
     # answer acc
-    f1, prec, recall = f1_score(ans, gold)
+    f1, _, _ = f1_score(ans, gold)
     return f1
 
 
 def compute_format(
-    solution_str=None,
-    ground_truth=None,
-    gold_sentences=None,
-    data_source=None,
-    extra_info=None,
-):
+    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
-    prompt, response = split_response(prediction)
+    _, response = split_response(prediction)
     ans = extract_answer(response)
     if ans == "":
         delimiter = "<|im_start|>assistant"
@@ -326,7 +328,7 @@ def compute_format(
     return FORMAT_SCORE
 
 
-def split_trace(text):
+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)]
@@ -334,12 +336,12 @@ def split_trace(text):
 
 
 def compute_action_query(
-    solution_str=None,
-    ground_truth=None,
-    gold_sentences=None,
-    data_source=None,
-    extra_info=None,
-):
+    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)
@@ -348,12 +350,12 @@ def compute_action_query(
 
 
 def compute_action_bm25(
-    solution_str=None,
-    ground_truth=None,
-    gold_sentences=None,
-    data_source=None,
-    extra_info=None,
-):
+    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)
@@ -362,12 +364,12 @@ def compute_action_bm25(
 
 
 def compute_action_read_pre(
-    solution_str=None,
-    ground_truth=None,
-    gold_sentences=None,
-    data_source=None,
-    extra_info=None,
-):
+    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)
@@ -376,12 +378,12 @@ def compute_action_read_pre(
 
 
 def compute_action_read_nxt(
-    solution_str=None,
-    ground_truth=None,
-    gold_sentences=None,
-    data_source=None,
-    extra_info=None,
-):
+    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)
@@ -390,12 +392,12 @@ def compute_action_read_nxt(
 
 
 def compute_action_continue(
-    solution_str=None,
-    ground_truth=None,
-    gold_sentences=None,
-    data_source=None,
-    extra_info=None,
-):
+    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)
@@ -404,12 +406,12 @@ def compute_action_continue(
 
 
 def compute_action_match(
-    solution_str=None,
-    ground_truth=None,
-    gold_sentences=None,
-    data_source=None,
-    extra_info=None,
-):
+    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)
@@ -418,12 +420,12 @@ def compute_action_match(
 
 
 def compute_total_action_number(
-    solution_str=None,
-    ground_truth=None,
-    gold_sentences=None,
-    data_source=None,
-    extra_info=None,
-):
+    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)
@@ -434,7 +436,7 @@ def compute_total_action_number(
 # define reward functions for evaluation
 
 
-def compute_scores(answer, ground_truth):
+def compute_scores(answer: str, ground_truth: str) -> float:
     parsed_answer = extract_answer(answer)
     if parsed_answer is None:
         return -0.1

examples/rag/wiki_retriever_mcp/wiki_retriever_mcp.py

@@ -1,5 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+# type: ignore
+
 import pickle
 
 import faiss

examples/spider/spider_eval/convert_dataset.py

@@ -1,32 +1,35 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+# type: ignore
+
 import os
 
 import pandas as pd
 
-data_dir = "data"
-target_data_dir = "data"
-columns = ["db_id", "question", "query"]
+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)
+    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)
+    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)
+    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 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)
+    # 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/sql_agent.py

@@ -12,14 +12,14 @@ import re
 import shutil
 import tempfile
 import time
-from typing import Any, Literal, Optional
+from typing import Any, Dict, Literal, Optional, cast
 
 import dotenv
 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 BaseMessage, HumanMessage, SystemMessage
+from langchain_core.messages import AnyMessage, HumanMessage
 from langchain_core.prompts import ChatPromptTemplate
 from langgraph.graph import END, START, MessagesState, StateGraph
 from langgraph.graph.state import CompiledStateGraph
@@ -176,7 +176,7 @@ class State(MessagesState):
     answer: str
     feedback: str
     num_turns: int
-    messages: list[BaseMessage]
+    messages: list[AnyMessage]
 
 
 class SQLAgent:
@@ -188,18 +188,18 @@ class SQLAgent:
         debug: bool = False,
         db_schema: str | None = None,
         endpoint: str | None = None,
-        verl_replacement: dict | None = None,
+        verl_replacement: Dict[str, Any] | None = None,
         table_info_truncate: int = 2048,
         execution_truncate: int = 2048,
     ):
-        self.db = SQLDatabase.from_uri(db)
+        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 = verl_replacement["model"]
+            self.model_name: str = verl_replacement["model"]  # type: ignore
             assert endpoint is not None
             self.llm = init_chat_model(
                 self.model_name,
@@ -211,7 +211,7 @@ class SQLAgent:
                 max_tokens=2048,
             )
         else:
-            self.model_name = os.environ.get("MODEL", "gpt-4.1-mini")
+            self.model_name: str = os.environ.get("MODEL", "gpt-4.1-mini")
             self.llm = init_chat_model(
                 self.model_name,
                 model_provider="openai",
@@ -237,7 +237,7 @@ class SQLAgent:
                 return self.db_schema
             return "No schema available."
 
-    def invoke_prompt(self, prompt: Any) -> BaseMessage:
+    def invoke_prompt(self, prompt: Any) -> AnyMessage:
         if self.debug:
             for message in prompt.messages:
                 termcolor.cprint(message.pretty_repr(), "blue")
@@ -252,7 +252,7 @@ class SQLAgent:
         if self.debug:
             termcolor.cprint(result.pretty_repr(), "green")
 
-        return result
+        return result  # type: ignore
 
     def truncate_execuion(self, execution: str) -> str:
         """Truncate the execution result to a reasonable length."""
@@ -260,28 +260,28 @@ class SQLAgent:
             return execution[: self.execution_truncate] + "\n... (truncated)"
         return execution
 
-    def parse_query(self, message: BaseMessage) -> str | None:
-        result = None
-        for match in re.finditer(r".*```\w*\n(.*?)\n```.*", message.content, re.DOTALL):
-            result = match.group(1).strip()
-        return result
+    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):
+    def write_query(self, state: State) -> State:
         """Generate SQL query to fetch information."""
-        prompt = WRITE_QUERY_PROMPT.invoke(
+        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)
+        result = self.invoke_prompt(prompt)  # type: ignore
 
-        query = self.parse_query(result) or result.content
+        query = self.parse_query(result) or result.content  # type: ignore
 
-        return {
+        return {  # type: ignore
             **state,
-            "query": query,
+            "query": query,  # type: ignore
             "num_turns": 1,
             "messages": [*prompt.messages, result],
         }
@@ -289,7 +289,7 @@ class SQLAgent:
     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"])
+        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)
@@ -299,7 +299,7 @@ class SQLAgent:
 
     def check_query(self, state: State) -> State:
         """Check the SQL query for correctness."""
-        prompt = CHECK_QUERY_PROMPT.invoke(
+        prompt: Any = CHECK_QUERY_PROMPT.invoke(  # type: ignore
             {
                 "dialect": self.db.dialect,
                 "input": state["question"],
@@ -308,18 +308,18 @@ class SQLAgent:
                 "table_info": self.get_table_info(),
             }
         )
-        result = self.invoke_prompt(prompt)
+        result = self.invoke_prompt(prompt)  # type: ignore
 
-        res = {
+        res = {  # type: ignore
             **state,
-            "feedback": result.content,
+            "feedback": result.content,  # type: ignore
             "messages": [*state.get("messages", []), *prompt.messages, result],
         }
-        return res
+        return res  # type: ignore
 
     def rewrite_query(self, state: State) -> State:
         """Rewrite SQL query if necessary."""
-        prompt = REWRITE_QUERY_PROMPT.invoke(
+        prompt: Any = REWRITE_QUERY_PROMPT.invoke(  # type: ignore
             {
                 "dialect": self.db.dialect,
                 "input": state["question"],
@@ -329,9 +329,9 @@ class SQLAgent:
                 "table_info": self.get_table_info(),
             }
         )
-        result = self.invoke_prompt(prompt)
+        result = self.invoke_prompt(prompt)  # type: ignore
 
-        rewritten_query = self.parse_query(result)
+        rewritten_query = self.parse_query(result)  # type: ignore
 
         return {
             **state,
@@ -342,14 +342,14 @@ class SQLAgent:
 
     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):
+        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:
-                if "THE QUERY IS INCORRECT" in last_message.content:
+            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")
-                    incorrect_index = last_message.content.rfind("THE QUERY IS INCORRECT")
+                    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:
@@ -362,21 +362,21 @@ class SQLAgent:
 
     def graph(self) -> CompiledStateGraph[State]:
         builder = StateGraph(State)
-        builder.add_node(self.write_query)
-        builder.add_node(self.execute_query)
-        builder.add_node(self.check_query)
-        builder.add_node(self.rewrite_query)
+        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,
+            self.should_continue,  # type: ignore
         )
         builder.add_edge("rewrite_query", "execute_query")
 
-        return builder.compile()
+        return builder.compile()  # type: ignore
 
 
 def evaluate_query(query: str, ground_truth: str, database: str, raise_on_error: bool = True) -> float:
@@ -411,7 +411,7 @@ def evaluate_query(query: str, ground_truth: str, database: str, raise_on_error:
             return 0.0
 
 
-class LitSQLAgent(agentlightning.LitAgent):
+class LitSQLAgent(agentlightning.LitAgent[Any]):
 
     def __init__(
         self,
@@ -433,7 +433,7 @@ class LitSQLAgent(agentlightning.LitAgent):
     ) -> float | None:
         question = sample["question"]
         start_time = time.time()
-        llm: agentlightning.LLM = resources["main_llm"]
+        llm: agentlightning.LLM = cast(agentlightning.LLM, resources["main_llm"])
 
         if is_training:
             original_db_path = os.path.join(self.spider_dir, "database", sample["db_id"], sample["db_id"] + ".sqlite")
@@ -484,9 +484,9 @@ class LitSQLAgent(agentlightning.LitAgent):
                 ),
             ).graph()
             try:
-                result = agent.invoke(
-                    {"question": question},
-                    {"callbacks": [self.tracer.get_langchain_callback_handler()], "recursion_limit": 100},
+                result = agent.invoke(  # type: ignore
+                    {"question": question},  # type: ignore
+                    {"callbacks": [self.tracer.get_langchain_callback_handler()], "recursion_limit": 100},  # type: ignore
                 )
             except Exception as e:
                 logger.exception(f"[Rollout {rollout_id}] Error during agent invocation: {e}")
@@ -512,10 +512,10 @@ class LitSQLAgent(agentlightning.LitAgent):
 
         return reward
 
-    def training_rollout(self, task: Any, rollout_id: str, resources: agentlightning.NamedResources) -> Any:
+    def training_rollout(self, task: Any, rollout_id: str, resources: agentlightning.NamedResources) -> Any:  # type: ignore
         return self._execute_rollout(task, resources=resources, rollout_id=rollout_id, is_training=True)
 
-    def validation_rollout(self, task: Any, rollout_id: str, resources: agentlightning.NamedResources) -> Any:
+    def validation_rollout(self, task: Any, rollout_id: str, resources: agentlightning.NamedResources) -> Any:  # type: ignore
         return self._execute_rollout(task, resources=resources, rollout_id=rollout_id, is_training=False)
 
 
@@ -526,7 +526,7 @@ def spider_dev_data():
     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)
+    df = pd.read_parquet(spider_dev_data_path)  # type: ignore
     if "OPENAI_API_BASE" not in os.environ:
         logger.warning(
             "Environment variable OPENAI_API_BASE is not set. Using default value 'https://api.openai.com/v1'."
@@ -544,7 +544,7 @@ def spider_dev_data():
             },
         )
     }
-    return agentlightning.DevTaskLoader(df.head(10).to_dict(orient="records"), resource)
+    return agentlightning.DevTaskLoader(df.head(10).to_dict(orient="records"), resource)  # type: ignore
 
 
 if __name__ == "__main__":

pyproject.toml

@@ -26,6 +26,7 @@ dev = [
   "pytest-rerunfailures",
   "black",
   "isort",
+  "pyright",
   "mkdocs",
   "mkdocs-material",
   "mkdocstrings[python]",

pyrightconfig.fast.json

@@ -0,0 +1,22 @@
+{
+  "include": ["agentlightning"],
+  "exclude": [
+    "**/data",
+    "**/assets",
+    "agentlightning/verl",
+    "agentlightning/instrumentation",
+    "agentlightning/algorithm/verl",
+    "agentlightning/cli/vllm.py"
+  ],
+
+  "pythonVersion": "3.12",
+
+  "typeCheckingMode": "strict",
+
+  "reportMissingTypeStubs": "none",
+  "reportUnknownMemberType": "error",
+  "reportUnknownVariableType": "error",
+  "reportMissingImports": "error",
+  "reportMissingModuleSource": "error",
+  "reportOptionalMemberAccess": "error"
+}

pyrightconfig.json

@@ -0,0 +1,19 @@
+{
+  // Paths to check
+  "include": ["agentlightning", "tests", "examples"],
+  "exclude": ["**/data", "**/assets"],
+
+  // Lock Python version for consistent semantics
+  "pythonVersion": "3.12",
+
+  // Start strict; downgrade only if noisy
+  "typeCheckingMode": "strict",
+
+  // reporting tweaks
+  "reportMissingTypeStubs": "none",
+  "reportUnknownMemberType": "error",
+  "reportUnknownVariableType": "error",
+  "reportMissingImports": "error",
+  "reportMissingModuleSource": "error",
+  "reportOptionalMemberAccess": "error"
+}

scripts/setup_stable_full.sh

@@ -0,0 +1,10 @@
+set -ex
+python -m pip install --upgrade --no-cache-dir pip
+
+# CPU version full installation
+
+pip install --no-cache-dir packaging ninja numpy pandas ipython ipykernel gdown wheel setuptools
+pip install --no-cache-dir vllm  # pytorch auto installed when installing vllm
+pip install --no-cache-dir verl==0.5.0
+
+pip install --no-cache-dir -e .[dev,agent]

tests/test_client.py

@@ -1,5 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+# type: ignore
+
 import asyncio
 import time
 from typing import Any, AsyncGenerator, Dict

tests/test_config.py

@@ -1,5 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+# type: ignore
+
 """
 This file is not carefully reviewed.
 It is to ensure the *somewhat* correctness of the code in agentlightning/config.py.
@@ -107,7 +109,7 @@ class OptionalNoDefaultConfig:
         ("NoneValue", "NoneValue"),  # Should not convert if not exact keyword
     ],
 )
-def test_nullable_str(input_val, expected_output):
+def test_nullable_str(input_val: str, expected_output: Optional[str]) -> None:
     """Tests the nullable_str function for various inputs."""
     assert config.nullable_str(input_val) == expected_output
 
@@ -142,13 +144,13 @@ def test_nullable_str(input_val, expected_output):
         (False, False),  # Direct bool pass-through
     ],
 )
-def test_str_to_bool_valid(input_val, expected_output):
+def test_str_to_bool_valid(input_val: Union[str, bool], expected_output: bool) -> None:
     """Tests _str_to_bool with valid boolean string representations."""
     assert config._str_to_bool(input_val) == expected_output
 
 
 @pytest.mark.parametrize("invalid_input", ["maybe", "2", "", " ", "trueish", "falsey"])
-def test_str_to_bool_invalid(invalid_input):
+def test_str_to_bool_invalid(invalid_input: str) -> None:
     """Tests _str_to_bool with invalid inputs, expecting ArgumentTypeError."""
     with pytest.raises(argparse.ArgumentTypeError):
         config._str_to_bool(invalid_input)
@@ -180,7 +182,9 @@ def test_str_to_bool_invalid(invalid_input):
         (Optional[List[Any]], List[Any], True, True),
     ],
 )
-def test_get_param_type_details(annotation, expected_core_type, expected_is_optional, expected_is_list):
+def test_get_param_type_details(
+    annotation: Any, expected_core_type: Any, expected_is_optional: bool, expected_is_list: bool
+) -> None:
     """Tests _get_param_type_details for various type annotations."""
     core_type, is_optional, is_list = config._get_param_type_details(annotation)
 
@@ -219,7 +223,9 @@ def test_get_param_type_details(annotation, expected_core_type, expected_is_opti
         (List[Dict[str, int]], True, {"nargs": "*", "type": str}),
     ],
 )
-def test_determine_argparse_type_and_nargs(core_param_type, is_param_list, expected_kwargs):
+def test_determine_argparse_type_and_nargs(
+    core_param_type: Any, is_param_list: bool, expected_kwargs: Dict[str, Any]
+) -> None:
     """Tests _determine_argparse_type_and_nargs for type and nargs mapping."""
     assert config._determine_argparse_type_and_nargs(core_param_type, is_param_list) == expected_kwargs
 
@@ -246,7 +252,9 @@ def test_determine_argparse_type_and_nargs(core_param_type, is_param_list, expec
         ("C", "p_empty", inspect.Parameter.empty, False, False, "For C: 'p_empty'. Inferred type: Any."),
     ],
 )
-def test_build_help_string(cls_name, param_name, core_type, is_optional, is_list, expected_help):
+def test_build_help_string(
+    cls_name: str, param_name: str, core_type: Any, is_optional: bool, is_list: bool, expected_help: str
+) -> None:
     """Tests _build_help_string for generating correct help messages."""
     assert config._build_help_string(cls_name, param_name, core_type, is_optional, is_list) == expected_help
 
@@ -349,7 +357,7 @@ def test_add_argument_for_parameter(
 
 
 # --- Tests for _add_arguments_for_class ---
-def test_add_arguments_for_class(mock_parser):
+def test_add_arguments_for_class(mock_parser: Any) -> None:
     """Tests _add_arguments_for_class by checking calls to _add_argument_for_parameter."""
     class_arg_configs_maps = {}
     with mock.patch("agentlightning.config._add_argument_for_parameter") as mock_add_param_func:
@@ -377,7 +385,7 @@ def test_add_arguments_for_class(mock_parser):
         assert mock_add_param_func.call_count == expected_calls
 
 
-def test_add_arguments_for_class_no_init_params(mock_parser):
+def test_add_arguments_for_class_no_init_params(mock_parser: Any) -> None:
     """Tests _add_arguments_for_class with a class having no __init__ parameters."""
     class_arg_configs_maps = {}
     with mock.patch("agentlightning.config._add_argument_for_parameter") as mock_add_param_func:
@@ -386,7 +394,7 @@ def test_add_arguments_for_class_no_init_params(mock_parser):
         assert class_arg_configs_maps[NoInitParamsConfig] == {}
 
 
-def test_create_argument_parser():
+def test_create_argument_parser() -> None:
     """Tests _create_argument_parser for basic parser properties."""
     parser = config._create_argument_parser()
     assert isinstance(parser, argparse.ArgumentParser)
@@ -395,7 +403,7 @@ def test_create_argument_parser():
     assert type(parser.formatter_class) == type(argparse.ArgumentDefaultsHelpFormatter)
 
 
-def test_instantiate_classes():
+def test_instantiate_classes() -> None:
     """Tests _instantiate_classes with various argument types and defaults."""
     parsed_args = argparse.Namespace(
         simpleconfig_name="TestName",

tests/test_decorator_executability.py

@@ -1,5 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+# type: ignore
+
 """Test that @llm_rollout and @rollout decorators preserve function executability."""
 
 import inspect

tests/test_litagent_hooks.py

@@ -1,6 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 from contextlib import contextmanager
+from typing import Any, Iterator, List, Optional
 
 import pytest
 
@@ -11,67 +12,67 @@ from agentlightning.tracer import BaseTracer, TripletExporter
 
 class DummyTracer(BaseTracer):
     @contextmanager
-    def trace_context(self, name=None):
+    def trace_context(self, name: Optional[str] = None) -> Iterator[None]:
         yield
 
-    def get_last_trace(self):
+    def get_last_trace(self) -> List[Any]:
         return []
 
 
 class DummyClient:
-    def __init__(self):
+    def __init__(self) -> None:
         self.posted = None
         self.polled = False
 
-    def poll_next_task(self):
+    def poll_next_task(self) -> Optional[Task]:
         if self.polled:
             return None
         self.polled = True
         return Task(rollout_id="1", input={}, mode="train", resources_id=None)
 
-    def get_latest_resources(self):
+    def get_latest_resources(self) -> ResourcesUpdate:
         return ResourcesUpdate(resources_id="r", resources={})
 
-    def post_rollout(self, rollout):
+    def post_rollout(self, rollout: Any) -> None:
         self.posted = rollout
 
 
 class DummyAsyncClient:
-    def __init__(self):
+    def __init__(self) -> None:
         self.posted = None
         self.polled = False
 
-    async def poll_next_task_async(self):
+    async def poll_next_task_async(self) -> Optional[Task]:
         if self.polled:
             return None
         self.polled = True
         return Task(rollout_id="1", input={}, mode="train", resources_id=None)
 
-    async def get_latest_resources_async(self):
+    async def get_latest_resources_async(self) -> ResourcesUpdate:
         return ResourcesUpdate(resources_id="r", resources={})
 
-    async def post_rollout_async(self, rollout):
+    async def post_rollout_async(self, rollout: Any) -> None:
         self.posted = rollout
 
 
-class HookAgent(LitAgent):
+class HookAgent(LitAgent[Any]):
     def __init__(self):
         super().__init__()
         self.start_called = False
         self.end_called = False
         self.end_rollout = None
 
-    def training_rollout(self, task, resources, rollout):
+    def training_rollout(self, task: Any, resources: Any, rollout: Any) -> float:
         return 0.5
 
-    async def training_rollout_async(self, task, resources, rollout):
+    async def training_rollout_async(self, task: Any, resources: Any, rollout: Any) -> float:
         return 0.5
 
-    def on_rollout_start(self, task, runner, tracer):
+    def on_rollout_start(self, task: Any, runner: Any, tracer: Any) -> None:
         self.start_called = True
         self.start_task = task
 
-    def on_rollout_end(self, task, rollout, runner, tracer):
+    def on_rollout_end(self, task: Any, rollout: Any, runner: Any, tracer: Any) -> None:
         self.end_called = True
         self.end_rollout = rollout
 
@@ -80,12 +81,12 @@ def test_runner_calls_hooks():
     agent = HookAgent()
     client = DummyClient()
     tracer = DummyTracer()
-    runner = AgentRunner(agent, client, tracer, TripletExporter())
+    runner = AgentRunner(agent, client, tracer, TripletExporter())  # type: ignore
 
     assert runner.run() is True
     assert agent.start_called
     assert agent.end_called
-    assert agent.end_rollout.final_reward == 0.5
+    assert agent.end_rollout.final_reward == 0.5  # type: ignore
 
 
 @pytest.mark.asyncio
@@ -93,9 +94,9 @@ async def test_runner_calls_hooks_async():
     agent = HookAgent()
     client = DummyAsyncClient()
     tracer = DummyTracer()
-    runner = AgentRunner(agent, client, tracer, TripletExporter())
+    runner = AgentRunner(agent, client, tracer, TripletExporter())  # type: ignore
 
     assert await runner.run_async() is True
     assert agent.start_called
     assert agent.end_called
-    assert agent.end_rollout.final_reward == 0.5
+    assert agent.end_rollout.final_reward == 0.5  # type: ignore

tests/test_tracer.py

@@ -1,5 +1,7 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+# type: ignore
+
 """
 Integration tests for various agent frameworks with AgentLightning.
 
@@ -26,7 +28,7 @@ import re
 import threading
 import time
 from contextlib import asynccontextmanager, contextmanager
-from typing import Any, Dict, List, Literal, Optional
+from typing import Any, Callable, Dict, Iterator, List, Literal, Optional, Tuple
 from unittest.mock import AsyncMock, MagicMock, Mock, patch
 
 import agentops
@@ -95,7 +97,7 @@ class MockOpenAICompatibleServer:
     Now supports replaying from prompt caches.
     """
 
-    def __init__(self, host="127.0.0.1", port=8000):
+    def __init__(self, host: str = "127.0.0.1", port: int = 8000) -> None:
         self.host = host
         self.port = port
         self.app = FastAPI()
@@ -116,13 +118,13 @@ class MockOpenAICompatibleServer:
                         continue
         return caches
 
-    def _find_best_cache_match(self, request_dict):
+    def _find_best_cache_match(self, request_dict: Dict[str, Any]) -> Tuple[Optional[Dict[str, Any]], float]:
         """
         Find the cached request with the highest similarity to the incoming request.
         Returns (response, similarity_score) or (None, 0.0) if not found.
         """
 
-        def normalize_messages(msgs):
+        def normalize_messages(msgs: List[Dict[str, Any]]) -> str:
             # Flatten messages to a string for comparison
             if not msgs:
                 return ""
@@ -180,14 +182,14 @@ class MockOpenAICompatibleServer:
 
         return self
 
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         if self.server:
             self.server.should_exit = True
         if self.server_thread and self.server_thread.is_alive():
             self.server_thread.join(timeout=5)
 
 
-async def run_agent(agent_func):
+async def run_agent(agent_func: Callable[[], Any]) -> None:
     """
     Run an agent function with mock server, handling both sync and async functions.
 
@@ -217,7 +219,7 @@ async def run_agent(agent_func):
             return agent_func()
 
 
-def agent_pure_openai():
+def agent_pure_openai() -> None:
     """A simple agent using the `openai` library."""
     client = OpenAI(base_url=OPENAI_BASE_URL, api_key=OPENAI_API_KEY)
     response = client.chat.completions.create(
@@ -226,7 +228,7 @@ def agent_pure_openai():
     assert "Paris" in response.choices[0].message.content
 
 
-def agent_litellm():
+def agent_litellm() -> None:
     """Agent using `litellm` to call the mock server."""
     response = litellm.completion(
         model="openai/" + OPENAI_MODEL,
@@ -237,7 +239,7 @@ def agent_litellm():
     assert "4" in response.choices[0].message.content
 
 
-def agent_langchain():
+def agent_langchain() -> None:
     """A simple LangChain agent."""
     llm = ChatOpenAI(model=OPENAI_MODEL, openai_api_base=OPENAI_BASE_URL, openai_api_key=OPENAI_API_KEY)
     prompt = ChatPromptTemplate.from_messages([("human", "{input}")])
@@ -246,7 +248,7 @@ def agent_langchain():
     assert "Paris" in result
 
 
-def agent_langchain_tooluse():
+def agent_langchain_tooluse() -> None:
     """A LangChain agent that uses a calculator tool."""
 
     @tool
@@ -272,14 +274,14 @@ def agent_langchain_tooluse():
     assert "504" in result["output"]
 
 
-def agent_langgraph():
+def agent_langgraph() -> None:
     """An agent built with LangGraph for stateful, cyclical workflows."""
     llm = init_chat_model("openai:" + OPENAI_MODEL, openai_api_base=OPENAI_BASE_URL, openai_api_key=OPENAI_API_KEY)
     db = SQLDatabase.from_uri("sqlite:///" + os.path.join(os.path.dirname(__file__), "assets/chinook.db"))
     toolkit = SQLDatabaseToolkit(db=db, llm=llm)
     tools = toolkit.get_tools()
 
-    def get_tool(name):
+    def get_tool(name: str) -> Any:
         return next(t for t in tools if t.name == name)
 
     get_schema_tool = next(tool for tool in tools if tool.name == "sql_db_schema")
@@ -369,7 +371,7 @@ def agent_langgraph():
     assert len(result["messages"]) > 5
 
 
-async def agent_autogen_multiagent():
+async def agent_autogen_multiagent() -> None:
     """A multi-agent conversation with AutoGen."""
 
     model_client = OpenAIChatCompletionClient(
@@ -401,7 +403,7 @@ async def agent_autogen_multiagent():
     assert "critic" in sources
 
 
-async def agent_autogen_mcp():
+async def agent_autogen_mcp() -> None:
     """An AutoGen agent using the Multi-agent Conversation Platform (MCP) and a tool (fixed usage)."""
     calculator_mcp_server = StdioServerParams(command="uvx", args=["mcp-server-calculator"])
 
@@ -417,14 +419,14 @@ async def agent_autogen_mcp():
         assert "504" in response.messages[-1].content
 
 
-def openai_agents_sdk_run_config():
+def openai_agents_sdk_run_config() -> RunConfig:
     return RunConfig(
         model=OPENAI_MODEL,
         model_provider=OpenAIProvider(api_key=OPENAI_API_KEY, base_url=OPENAI_BASE_URL, use_responses=False),
     )
 
 
-async def openai_agents_sdk_eval_hook_and_guardrail():
+async def openai_agents_sdk_eval_hook_and_guardrail() -> None:
     class HomeworkOutput(BaseModel):
         is_homework: bool
         reasoning: str
@@ -473,7 +475,7 @@ async def openai_agents_sdk_eval_hook_and_guardrail():
     assert hasattr(result, "final_output")
 
 
-async def openai_agents_sdk_mcp_tool_use():
+async def openai_agents_sdk_mcp_tool_use() -> None:
     async with MCPServerStdio(params={"command": "uvx", "args": ["mcp-server-calculator"]}) as mcp_server:
         agent = Agent(
             name="MCP Tool Agent",
@@ -487,7 +489,7 @@ async def openai_agents_sdk_mcp_tool_use():
         assert "2451" in result.final_output_as(str)
 
 
-async def openai_agents_sdk_handoff_tool_output_type_and_reward():
+async def openai_agents_sdk_handoff_tool_output_type_and_reward() -> None:
 
     class MathOutput(BaseModel):
         answer: int
@@ -608,7 +610,7 @@ AGENTOPS_EXPECTED_REWARDS = {
 }
 
 
-def assert_expected_pairs_in_tree(root_tuple, expected_pairs):
+def assert_expected_pairs_in_tree(root_tuple: Tuple[str, List[Any]], expected_pairs: List[Tuple[str, str]]) -> None:
     """
     Assert that every (ancestor_name, child_name) pair in `expected_pairs`
     occurs somewhere in the tree produced by TraceTree.names_tuple().
@@ -647,7 +649,7 @@ def assert_expected_pairs_in_tree(root_tuple, expected_pairs):
             )
 
 
-def iterate_over_agents():
+def iterate_over_agents() -> Iterator[Callable[[], Any]]:
     yield from [
         agent_pure_openai,
         agent_litellm,
@@ -662,16 +664,16 @@ def iterate_over_agents():
     ]
 
 
-def run_one(agent_func):
+def run_one(agent_func: Callable[[], Any]) -> None:
     asyncio.get_event_loop().run_until_complete(run_agent(agent_func))
 
 
-def run_all():
+def run_all() -> None:
     for agent_func in iterate_over_agents():
         run_one(agent_func)
 
 
-def run_with_agentops_tracer():
+def run_with_agentops_tracer() -> None:
     tracer = AgentOpsTracer()
     tracer.init()
     tracer.init_worker(0)
@@ -724,7 +726,7 @@ def run_with_agentops_tracer():
     tracer.teardown()
 
 
-def run_with_http_tracer():
+def run_with_http_tracer() -> None:
     import httpdbg.hooks.all
 
     @contextmanager
@@ -754,7 +756,7 @@ def run_with_http_tracer():
     tracer.teardown()
 
 
-def create_prompt_caches():
+def create_prompt_caches() -> None:
     """Create prompt caches for the agent frameworks.
     This should only be run once to populate the caches.
     """

tests/test_tracer_http.py

@@ -1,6 +1,5 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-#!/usr/bin/env python3
 """
 Functional tests for HttpTracer with real HTTP requests.
 
@@ -9,8 +8,6 @@ that the HttpTracer correctly captures HTTP traffic in both
 normal and subprocess modes.
 """
 
-import asyncio
-
 import aiohttp
 import pytest
 import requests
@@ -68,6 +65,7 @@ def test_normal_mode_sync_requests():
 
     # Verify span attributes
     for span in spans:
+        assert span.attributes is not None
         assert "http.method" in span.attributes
         assert "http.url" in span.attributes
         assert "http.status_code" in span.attributes
@@ -155,6 +153,7 @@ def test_span_attributes_detailed():
 
     for span in spans:
         # Basic HTTP attributes
+        assert span.attributes is not None
         assert "http.method" in span.attributes
         assert "http.url" in span.attributes
         assert "http.target" in span.attributes
@@ -182,6 +181,7 @@ def test_span_attributes_minimal():
 
     for span in spans:
         # Basic HTTP attributes should still be present
+        assert span.attributes is not None
         assert "http.method" in span.attributes
         assert "http.url" in span.attributes
         assert "http.status_code" in span.attributes