Support multiple functions in AgentSet agg method#2743
Conversation
Enhance the AgentSet.agg method to accept a list or tuple of functions in addition to a single function. This enables performing multiple aggregations in a single call, improving both convenience and efficiency when analyzing agent attributes.
Add test cases to verify that the AgentSet.agg method properly handles lists and tuples of aggregation functions.
|
Its amazing that you still trying to improve the |
|
Performance benchmarks:
|
It's so fun working on it, you can go from idea to complete PR in less than an hour.
Give it a try, it's the quickest way to get there ;) |
This comment was marked as resolved.
This comment was marked as resolved.
|
@quaquel thanks! Before merging, two questions.
|
|
|
Thanks for your perspective. Agree on both in some extend.
I will think about this a bit more. |
|
While I have no problem with this PR, here are some things I would like to know:
|
There was a problem hiding this comment.
Nice addition!
- I was initially considering allowing either multiple attributes or multiple methods, but not both. That said, I feel the case of multiple attributes is less common and provides limited performance benefits. So, for now, I think I’ll keep things as they are.
Personally, I’d still vote in favor of supporting multiple attributes as well—it gives the user more flexibility and freedom in how they want to use the function.
From an implementation perspective, you could use typing.overload to define valid combinations of parameters and return types. This way, the IDE will offer typing suggestions depending on the input.
from typing import overload, Collection, Any, Callable
from collections.abc import Iterable
@overload
def agg(
self,
attribute: str,
func: Callable
) -> Any:
...
@overload
def agg(
self,
attribute: Collection[str],
func: Callable
) -> dict[str, Any]:
...
@overload
def agg(
self,
attribute: str,
func: Iterable[Callable]
) -> tuple[Any, ...]:
...
@overload
def agg(
self,
attribute: Collection[str],
func: Iterable[Callable]
) -> dict[str, tuple[Any, ...]]:
...
def agg(
self,
attribute: str | Collection[str],
func: Callable | Iterable[Callable]
) -> Any | dict[str, Any] | tuple[Any, ...] | dict[str, tuple[Any, ...]]:
# implementation
|
Important Review skippedAuto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
Summary
This PR enhances the
AgentSet.aggmethod to accept lists and tuples of aggregation functions, enabling multiple aggregations in a single method call.Motive
Currently, when analyzing agent attributes, users need to make separate calls to
aggfor each aggregation function (min, max, sum, etc.). This creates performance overhead since the attribute values must be collected multiple times.This enhancement improves both performance and code readability by supporting multiple aggregation functions in a single call, allowing the attribute values to be collected just once and reused across all aggregation functions.
Implementation
aggmethod signature to accept a single callable or an iterable of callablesUsage Examples
Additional Notes
This enhancement is completely backward compatible, so existing code will continue to work without changes. The method will now return a tuple when multiple functions are provided, maintaining the same order as the input function list.
It might also reduce future datacollector complexity, because aggregation can be done efficiently outside it.
I also considered allowing multiple
attributevalues, but I decided against it since it adds complexity, its use case is less clear and the performance improvements are negligible.