github-featured/microsoft--agent-lightning
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

Internal API update in preparation for Tinker integration (#226)

Browse Source
This commit is contained in:
Yuge Zhang
2025-10-28 11:26:49 +08:00
committed by GitHub
parent 8ce40a0410
commit d16538da96
6 changed files with 119 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
agentlightning/adapter/triplet.py
+47 -18
View File
@@ -3,6 +3,7 @@
from __future__ import annotations
import json
import logging
import re
from enum import Enum
from typing import Any, Dict, List, Optional, Tuple, Union, cast
@@ -14,6 +15,8 @@ from agentlightning.types import Span, SpanNames, Triplet
from .base import TraceAdapter
logger = logging.getLogger(__name__)
class Transition(BaseModel):
"""A single transition within a reinforcement learning trajectory.
@@ -505,6 +508,30 @@ class TraceTree:
return rewards
def span_to_triplet(self, span: Span, agent_name: str) -> Triplet:
"""Convert a span to a triplet.
Subclass can override this method to add more fields to the triplet,
such as chat messages and tool calls.
"""
prompt_token_ids = span.attributes.get("prompt_token_ids", []) # type: ignore
response_token_ids = span.attributes.get("response_token_ids", []) # type: ignore
response_id = span.attributes.get("gen_ai.response.id", None) # type: ignore
logprobs_content = span.attributes.get("logprobs.content", None) # type: ignore
if isinstance(logprobs_content, str):
logprobs_content = json.loads(logprobs_content)
response: Dict[str, Any] = {"token_ids": response_token_ids, "logprobs": logprobs_content}
else:
response = {"token_ids": response_token_ids}
return Triplet(
prompt={"token_ids": prompt_token_ids},
response=response,
reward=None,
metadata=dict(response_id=response_id, agent_name=agent_name),
)
def to_trajectory(
self,
llm_call_match: str = r"openai\.chat\.completion",
@@ -513,6 +540,7 @@ class TraceTree:
dedup_llm_call: bool = True,
reward_match: RewardMatchPolicy = RewardMatchPolicy.FIRST_OCCURRENCE,
final_reward: Optional[float] = None,
_skip_empty_token_spans: bool = False,
) -> List[Triplet]:
"""Convert the trace tree into a trajectory of [`Triplet`][agentlightning.Triplet] items.
@@ -537,25 +565,23 @@ class TraceTree:
within_llm_call=False if dedup_llm_call else None,
existing_llm_call_response_ids=set(),
)
id_transitions = [
(
llm_call.id,
Triplet(
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( # type: ignore
"gen_ai.response.id", None
), # it works at least for OpenAI
agent_name=agent_name,
),
),
)
for llm_call, agent_name in llm_calls
]
rewards = self.match_rewards(reward_match, [call for call, _ in llm_calls])
id_transitions: List[Tuple[str, Triplet]] = []
# We need to filter out the LLM calls with unrecorded token IDs
filtered_llm_calls: List[Tuple[TraceTree, str]] = []
for llm_call, agent_name in llm_calls:
triplet = self.span_to_triplet(llm_call.span, agent_name)
# This is a hot-fix for Tinker+CrewAI, which has some anonymous requests outside the trained agent.
# TODO: We might need to reconsider this.
if _skip_empty_token_spans and (
not triplet.prompt.get("token_ids") or not triplet.response.get("token_ids")
):
logger.warning(f"Skipping LLM call with unrecorded token IDs: {triplet}")
continue
filtered_llm_calls.append((llm_call, agent_name))
id_transitions.append((llm_call.id, triplet))
rewards = self.match_rewards(reward_match, [call for call, _ in filtered_llm_calls])
transitions = [
transition.model_copy(update={"reward": rewards.get(id, None)}) for id, transition in id_transitions
]
@@ -597,12 +623,14 @@ class TracerTraceToTriplet(TraceToTripletBase):
agent_match: Optional[str] = None,
exclude_llm_call_in_reward: bool = True,
reward_match: RewardMatchPolicy = RewardMatchPolicy.FIRST_OCCURRENCE,
_skip_empty_token_spans: bool = False,
):
self.repair_hierarchy = repair_hierarchy
self.llm_call_match = llm_call_match
self.agent_match = agent_match
self.exclude_llm_call_in_reward = exclude_llm_call_in_reward
self.reward_match = reward_match
self._skip_empty_token_spans = _skip_empty_token_spans
def visualize(
self,
@@ -654,6 +682,7 @@ class TracerTraceToTriplet(TraceToTripletBase):
agent_match=self.agent_match,
exclude_llm_call_in_reward=self.exclude_llm_call_in_reward,
reward_match=self.reward_match,
_skip_empty_token_spans=self._skip_empty_token_spans,
)
return trajectory
agentlightning/algorithm/fast.py
+4 -2
View File
@@ -128,11 +128,13 @@ class Baseline(FastAlgorithm):
# Attempts to adapt the spans using the adapter if provided
try:
adapter = self.get_adapter()
except ValueError:
logger.warning("No adapter set for MockAlgorithm. Skipping trace adaptation.")
adapter = None
if adapter is not None:
spans = await store.query_spans(rollout_id=rollout_id, attempt_id="latest")
transformed_data = adapter.adapt(spans)
logger.info(f"[Rollout {rollout_id}] Adapted data: {transformed_data}")
except ValueError:
logger.warning("No adapter set for MockAlgorithm. Skipping trace adaptation.")
async def _enqueue_rollouts(
self, dataset: Dataset[Any], train_indices: List[int], val_indices: List[int], resources_id: str
agentlightning/emitter/reward.py
+1
View File
@@ -149,6 +149,7 @@ def emit_reward(reward: float) -> ReadableSpan:
if not isinstance(reward, float):
raise ValueError(f"Reward must be a number, got: {type(reward)}")
# TODO: This should use the tracer from current context by tracer
tracer = get_tracer()
span = tracer.start_span(SpanNames.REWARD.value, attributes={"reward": reward})
# Do nothing; it's just a number
agentlightning/instrumentation/agentops.py
+52 -25
View File
@@ -2,12 +2,13 @@
from __future__ import annotations
import json
import logging
import multiprocessing
import signal
import socket
import time
from typing import Any, Callable
from typing import Any, Callable, no_type_check
import flask
import requests
@@ -40,41 +41,67 @@ def _patch_new_agentops():
_original_handle_chat_attributes = handle_chat_attributes # type: ignore
@no_type_check
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
attributes = _original_handle_chat_attributes(args=args, kwargs=kwargs, return_value=return_value, **kws)
if (
return_value is not None
and hasattr(return_value, "prompt_token_ids")
and return_value.prompt_token_ids is not None
):
attributes["prompt_token_ids"] = list(return_value.prompt_token_ids)
if (
return_value is not None
and hasattr(return_value, "response_token_ids")
and return_value.response_token_ids is not None
):
attributes["response_token_ids"] = list(return_value.response_token_ids[0])
# For LiteLLM Proxy (v0.2) with vLLM return_token_ids, response_token_ids now lives in choices
if (
not attributes.get("response_token_ids")
and return_value is not None
and hasattr(return_value, "choices") # type: ignore
and return_value.choices # type: ignore
and isinstance(return_value.choices, list) # type: ignore
return_value is not None
and hasattr(return_value, "choices")
and return_value.choices
and isinstance(return_value.choices, list)
and len(return_value.choices) > 0
):
first_choice = return_value.choices[0] # type: ignore
if hasattr(first_choice, "token_ids"): # type: ignore
attributes["response_token_ids"] = list(first_choice.token_ids) # type: ignore
# newer versions of OpenAI client SDK
elif hasattr(first_choice, "provider_specific_fields") and "token_ids" in first_choice.provider_specific_fields: # type: ignore
attributes["response_token_ids"] = list(first_choice.provider_specific_fields["token_ids"]) # type: ignore
first_choice = return_value.choices[0]
# Token IDs from "choices[0].token_ids"
if "response_token_ids" not in attributes:
if hasattr(first_choice, "token_ids") and first_choice.token_ids is not None:
attributes["response_token_ids"] = list(first_choice.token_ids)
# newer versions of OpenAI client SDK
elif (
hasattr(first_choice, "provider_specific_fields")
and first_choice.provider_specific_fields.get("token_ids") is not None
):
attributes["response_token_ids"] = list(first_choice.provider_specific_fields["token_ids"])
# log probability
# This is temporary. We need a unified convention for classifying and naming logprobs.
if hasattr(first_choice, "logprobs") and first_choice.logprobs is not None:
if hasattr(first_choice.logprobs, "content") and first_choice.logprobs.content is not None:
attributes["logprobs.content"] = json.dumps(
[logprob.model_dump() for logprob in first_choice.logprobs.content]
)
if hasattr(first_choice.logprobs, "refusal") and first_choice.logprobs.refusal is not None:
attributes["logprobs.refusal"] = json.dumps(
[logprob.model_dump() for logprob in first_choice.logprobs.refusal]
)
# For LiteLLM, response is a openai._legacy_response.LegacyAPIResponse
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
and hasattr(return_value, "http_response")
and return_value.http_response is not None
and hasattr(return_value.http_response, "json")
):
json_data = return_value.http_response.json() # type: ignore
json_data = return_value.http_response.json()
if isinstance(json_data, dict):
if "prompt_token_ids" in json_data:
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]) # type: ignore
if json_data.get("prompt_token_ids") is not None:
attributes["prompt_token_ids"] = list(json_data["prompt_token_ids"])
if json_data.get("response_token_ids") is not None:
attributes["response_token_ids"] = list(json_data["response_token_ids"][0])
return attributes
agentlightning/llm_proxy.py
+13 -2
View File
@@ -528,6 +528,7 @@ class LLMProxy:
host: str | None = None,
litellm_config: Dict[str, Any] | None = None,
num_retries: int = 0,
_add_return_token_ids: bool = True,
):
self.store = store
self.host = host or _get_default_ipv4_address()
@@ -544,6 +545,8 @@ class LLMProxy:
self._uvicorn_server = None
self._ready_event = threading.Event()
self._add_return_token_ids = _add_return_token_ids
def get_store(self) -> Optional[LightningStore]:
"""Get the store used by the proxy.
@@ -637,7 +640,7 @@ class LLMProxy:
logger.info("Adding a new middleware to the FastAPI app.")
app.add_middleware(RolloutAttemptMiddleware)
if not initialize_llm_callbacks():
if not initialize_llm_callbacks(self._add_return_token_ids):
# If it's not the first time to initialize the callbacks, also
# reset LiteLLM's logging worker so its asyncio.Queue binds to the new loop.
_reset_litellm_logging_worker()
@@ -827,13 +830,17 @@ def set_active_llm_proxy(proxy: LLMProxy) -> None:
_global_llm_proxy = proxy
def initialize_llm_callbacks() -> bool:
def initialize_llm_callbacks(_add_return_token_ids: bool = True) -> bool:
"""Restore `litellm.callbacks` to a state that is just initialized by agent-lightning.
When litellm is restarted multiple times in the same process, more and more callbacks
will be appended to `litellm.callbacks`, which may exceed the MAX_CALLBACKS limit.
This function remembers the initial state of `litellm.callbacks` and always restore to that state.
Args:
_add_return_token_ids: Whether to add the return token ids callback. Internal use only.
Ideally the callback should automatically be enabled when the backend supports it.
Returns:
Whether the callbacks are initialized for the first time.
"""
@@ -845,6 +852,10 @@ def initialize_llm_callbacks() -> bool:
AddReturnTokenIds(),
LightningOpenTelemetry(),
]
if _add_return_token_ids
else [
LightningOpenTelemetry(),
]
)
_callbacks_before_litellm_start = [*litellm.callbacks] # type: ignore
return True
examples/.gitignore
+2
View File
@@ -9,3 +9,5 @@ unsloth/models/
unsloth/unsloth_compiled_cache/
unsloth/unsloth_training_checkpoints/
apo/pomltrace/
tinker/logs/
tinker/crewai_*.html