From e255706790d4df0f003c2e569744fe7806b308ee Mon Sep 17 00:00:00 2001 From: CHEN Juntong Date: Thu, 13 Mar 2025 10:05:56 +0000 Subject: [PATCH] perf: Optimize has_more_than_n_keys function and BatchTraceProcessor Co-Authored-By: Juntong CHEN --- src/agents/strict_schema.py | 7 +------ src/agents/tracing/processors.py | 15 +++++++-------- tests/test_has_more_than_n_keys.py | 14 ++++++++++++++ 3 files changed, 22 insertions(+), 14 deletions(-) create mode 100644 tests/test_has_more_than_n_keys.py diff --git a/src/agents/strict_schema.py b/src/agents/strict_schema.py index 910ad85f..28292cb4 100644 --- a/src/agents/strict_schema.py +++ b/src/agents/strict_schema.py @@ -159,9 +159,4 @@ def is_list(obj: object) -> TypeGuard[list[object]]: def has_more_than_n_keys(obj: dict[str, object], n: int) -> bool: - i = 0 - for _ in obj.keys(): - i += 1 - if i > n: - return True - return False + return len(obj) > n diff --git a/src/agents/tracing/processors.py b/src/agents/tracing/processors.py index 308adf2a..76f6aa42 100644 --- a/src/agents/tracing/processors.py +++ b/src/agents/tracing/processors.py @@ -224,23 +224,22 @@ def _export_batches(self, force: bool = False): """ while True: items_to_export: list[Span[Any] | Trace] = [] - - # Gather a batch of spans up to max_batch_size - while not self._queue.empty() and ( - force or len(items_to_export) < self._max_batch_size - ): + max_items = self._max_batch_size if not force else float('inf') + # Try to get up to max_items from the queue without checking empty() repeatedly + for _ in range(int(max_items)): try: items_to_export.append(self._queue.get_nowait()) except queue.Empty: - # Another thread might have emptied the queue between checks + # Queue is empty, no need to continue the loop break - # If we collected nothing, we're done if not items_to_export: break - # Export the batch self._exporter.export(items_to_export) + # If not forcing and queue is likely empty, we can stop + if not force and len(items_to_export) < max_items: + break # Create a shared global instance: diff --git a/tests/test_has_more_than_n_keys.py b/tests/test_has_more_than_n_keys.py new file mode 100644 index 00000000..8672da94 --- /dev/null +++ b/tests/test_has_more_than_n_keys.py @@ -0,0 +1,14 @@ +from agents.strict_schema import has_more_than_n_keys + + +def test_has_more_than_n_keys(): + # Test with empty dict + assert has_more_than_n_keys({}, 0) is False + # Test with dict having exactly n keys + assert has_more_than_n_keys({"a": 1}, 1) is False + # Test with dict having more than n keys + assert has_more_than_n_keys({"a": 1, "b": 2}, 1) is True + # Test with large dict + large_dict = {str(i): object() for i in range(1000)} + assert has_more_than_n_keys(large_dict, 500) is True + assert has_more_than_n_keys(large_dict, 1000) is False