Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0

### Changed

- Split Agno 2.x tool-call loops into ReAct step spans with one LLM span per
provider request, and populate streaming agent token usage from child LLM
spans when Agno run events do not include metrics.
- Propagate Agno run user and session identity to model and tool spans while
preferring ENTRY baggage identity over Agno framework-provided values.

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.

import importlib
import logging
from typing import Any, Collection

from wrapt import wrap_function_wrapper
Expand All @@ -30,8 +32,65 @@
_AGENT = "agno.agent"
_MODULE = "agno.models.base"
_TOOLKIT = "agno.tools.function"
_logger = logging.getLogger(__name__)
__all__ = ["AgnoInstrumentor"]

_INNER_MODEL_WRAPPERS = (
("Model._process_model_response", "process_model_response"),
("Model._aprocess_model_response", "aprocess_model_response"),
("Model.process_response_stream", "process_response_stream"),
("Model.aprocess_response_stream", "aprocess_response_stream"),
("Model.run_function_calls", "run_function_calls"),
("Model.arun_function_calls", "arun_function_calls"),
)

_FALLBACK_MODEL_WRAPPERS = (
("Model.response", "response"),
("Model.aresponse", "aresponse"),
("Model.response_stream", "response_stream"),
("Model.aresponse_stream", "aresponse_stream"),
)


def _has_wrap_target(module_name: str, name: str) -> bool:
try:
target = importlib.import_module(module_name)
for part in name.split("."):
target = getattr(target, part)
except Exception:
return False
return True


def _wrap_model_methods(model_wrapper: AgnoModelWrapper) -> None:
if all(
_has_wrap_target(_MODULE, name)
for name, _wrapper_name in _INNER_MODEL_WRAPPERS
):
for name, wrapper_name in _INNER_MODEL_WRAPPERS:
wrap_function_wrapper(
module=_MODULE,
name=name,
wrapper=getattr(model_wrapper, wrapper_name),
)
return

_logger.warning(
"Agno inner model hooks are unavailable; falling back to outer "
"Model.response wrappers."
)
for name, wrapper_name in _FALLBACK_MODEL_WRAPPERS:
wrap_function_wrapper(
module=_MODULE,
name=name,
wrapper=getattr(model_wrapper, wrapper_name),
)


def _unwrap_if_present(target: Any, name: str) -> None:
if hasattr(target, name):
unwrap(target, name)


class AgnoInstrumentor(BaseInstrumentor): # type: ignore
"""
Expand Down Expand Up @@ -91,27 +150,9 @@ def _instrument(self, **kwargs: Any) -> None:
wrapper=function_call_wrapper.aexecute,
)

# Wrap the model
wrap_function_wrapper(
module=_MODULE,
name="Model.response",
wrapper=model_wrapper.response,
)
wrap_function_wrapper(
module=_MODULE,
name="Model.aresponse",
wrapper=model_wrapper.aresponse,
)
wrap_function_wrapper(
module=_MODULE,
name="Model.response_stream",
wrapper=model_wrapper.response_stream,
)
wrap_function_wrapper(
module=_MODULE,
name="Model.aresponse_stream",
wrapper=model_wrapper.aresponse_stream,
)
# Wrap the model. Prefer Agno's per-request internals so a tool-call
# loop emits one LLM span per provider call.
_wrap_model_methods(model_wrapper)

def _uninstrument(self, **kwargs: Any) -> None:
# Unwrap the agent call function
Expand All @@ -129,8 +170,9 @@ def _uninstrument(self, **kwargs: Any) -> None:
# Unwrap the model
import agno.models.base # noqa: PLC0415

unwrap(agno.models.base.Model, "response")
unwrap(agno.models.base.Model, "aresponse")
unwrap(agno.models.base.Model, "response_stream")
unwrap(agno.models.base.Model, "aresponse_stream")
for name, _wrapper_name in (
*_INNER_MODEL_WRAPPERS,
*_FALLBACK_MODEL_WRAPPERS,
):
_unwrap_if_present(agno.models.base.Model, name.split(".", 1)[1])
self._handler = None
Loading
Loading