diff --git a/py/src/braintrust/env.py b/py/src/braintrust/env.py index 90ab2a44b..f08afb668 100644 --- a/py/src/braintrust/env.py +++ b/py/src/braintrust/env.py @@ -230,6 +230,11 @@ class BraintrustEnv: DISABLE_ATEXIT_FLUSH = EnvVar("BRAINTRUST_DISABLE_ATEXIT_FLUSH", EnvParser.BOOL) CAPTURE_USER_AUDIO_ATTACHMENTS = EnvVar("BRAINTRUST_CAPTURE_USER_AUDIO_ATTACHMENTS", EnvParser.BOOL) CAPTURE_AGENT_AUDIO_ATTACHMENTS = EnvVar("BRAINTRUST_CAPTURE_AGENT_AUDIO_ATTACHMENTS", EnvParser.BOOL) + MAX_GENERATOR_ITEMS = EnvVar("BRAINTRUST_MAX_GENERATOR_ITEMS", EnvParser.INT) + PROMPT_CACHE_MEMORY_MAX_SIZE = EnvVar("BRAINTRUST_PROMPT_CACHE_MEMORY_MAX_SIZE", EnvParser.INT) + PROMPT_CACHE_DISK_MAX_SIZE = EnvVar("BRAINTRUST_PROMPT_CACHE_DISK_MAX_SIZE", EnvParser.INT) + PARAMETERS_CACHE_MEMORY_MAX_SIZE = EnvVar("BRAINTRUST_PARAMETERS_CACHE_MEMORY_MAX_SIZE", EnvParser.INT) + PARAMETERS_CACHE_DISK_MAX_SIZE = EnvVar("BRAINTRUST_PARAMETERS_CACHE_DISK_MAX_SIZE", EnvParser.INT) OTEL_COMPAT = EnvVar("BRAINTRUST_OTEL_COMPAT", EnvParser.BOOL) # Opt out of the default OpenTelemetry-compatible hex span/trace IDs and use # legacy UUID-based IDs (and V3 span-component export) instead. diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index bf7df50b0..a63487d01 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -626,27 +626,23 @@ def default_get_api_conn(): self.reset_login_info() self._prompt_cache = PromptCache( - memory_cache=LRUCache( - max_size=int(os.environ.get("BRAINTRUST_PROMPT_CACHE_MEMORY_MAX_SIZE", str(1 << 10))) - ), + memory_cache=LRUCache(max_size=BraintrustEnv.PROMPT_CACHE_MEMORY_MAX_SIZE.get(1 << 10)), disk_cache=DiskCache( cache_dir=os.environ.get( "BRAINTRUST_PROMPT_CACHE_DIR", f"{os.environ.get('HOME')}/.braintrust/prompt_cache" ), - max_size=int(os.environ.get("BRAINTRUST_PROMPT_CACHE_DISK_MAX_SIZE", str(1 << 20))), + max_size=BraintrustEnv.PROMPT_CACHE_DISK_MAX_SIZE.get(1 << 20), serializer=lambda x: x.as_dict(), deserializer=PromptSchema.from_dict_deep, ), ) self._parameters_cache = ParametersCache( - memory_cache=LRUCache( - max_size=int(os.environ.get("BRAINTRUST_PARAMETERS_CACHE_MEMORY_MAX_SIZE", str(1 << 10))) - ), + memory_cache=LRUCache(max_size=BraintrustEnv.PARAMETERS_CACHE_MEMORY_MAX_SIZE.get(1 << 10)), disk_cache=DiskCache( cache_dir=os.environ.get( "BRAINTRUST_PARAMETERS_CACHE_DIR", f"{os.environ.get('HOME')}/.braintrust/parameters_cache" ), - max_size=int(os.environ.get("BRAINTRUST_PARAMETERS_CACHE_DISK_MAX_SIZE", str(1 << 20))), + max_size=BraintrustEnv.PARAMETERS_CACHE_DISK_MAX_SIZE.get(1 << 20), serializer=lambda x: x.as_dict(), deserializer=RemoteEvalParameters.from_dict_deep, ), @@ -2951,7 +2947,7 @@ async def wrapper_async_gen(*f_args, **f_kwargs): _try_log_input(span, f_sig, f_args, f_kwargs) # Get max items from environment or default - max_items = int(os.environ.get("BRAINTRUST_MAX_GENERATOR_ITEMS", "1000")) + max_items = BraintrustEnv.MAX_GENERATOR_ITEMS.get(1000) if trace_io and max_items != 0: # Collect output up to limit @@ -2992,7 +2988,7 @@ def wrapper_sync_gen(*f_args, **f_kwargs): _try_log_input(span, f_sig, f_args, f_kwargs) # Get max items from environment or default - max_items = int(os.environ.get("BRAINTRUST_MAX_GENERATOR_ITEMS", "1000")) + max_items = BraintrustEnv.MAX_GENERATOR_ITEMS.get(1000) if trace_io and max_items != 0: # Collect output up to limit diff --git a/py/src/braintrust/test_logger.py b/py/src/braintrust/test_logger.py index 9745c17be..1749f8d73 100644 --- a/py/src/braintrust/test_logger.py +++ b/py/src/braintrust/test_logger.py @@ -2838,6 +2838,47 @@ async def large_async_generator(): os.environ["BRAINTRUST_MAX_GENERATOR_ITEMS"] = original +@pytest.mark.parametrize("value", ["", " ", "not-a-number"]) +def test_traced_generators_ignore_invalid_max_items_env(with_memory_logger, monkeypatch, value): + """An empty or non-numeric BRAINTRUST_MAX_GENERATOR_ITEMS falls back to the default instead of raising.""" + init_test_logger(__name__) + monkeypatch.setenv("BRAINTRUST_MAX_GENERATOR_ITEMS", value) + + @logger.traced + def sync_generator(): + yield from range(3) + + @logger.traced + async def async_generator(): + for i in range(3): + yield i + + async def collect(): + return [v async for v in async_generator()] + + assert list(sync_generator()) == [0, 1, 2] + assert asyncio.run(collect()) == [0, 1, 2] + + logs = with_memory_logger.pop() + assert [log.get("output") for log in logs] == [[0, 1, 2], [0, 1, 2]] + + +@pytest.mark.parametrize( + "name", + [ + "BRAINTRUST_PROMPT_CACHE_MEMORY_MAX_SIZE", + "BRAINTRUST_PROMPT_CACHE_DISK_MAX_SIZE", + "BRAINTRUST_PARAMETERS_CACHE_MEMORY_MAX_SIZE", + "BRAINTRUST_PARAMETERS_CACHE_DISK_MAX_SIZE", + ], +) +@pytest.mark.parametrize("value", ["", "not-a-number"]) +def test_state_ignores_invalid_cache_size_env(monkeypatch, name, value): + """An empty or non-numeric cache size variable keeps the default instead of failing state creation.""" + monkeypatch.setenv(name, value) + logger.BraintrustState() + + def test_traced_sync_generator_zero_limit_drops_output(with_memory_logger): """Test sync generator with limit=0 drops all output but still yields values.""" init_test_logger(__name__)