diff --git a/tests/openai/evals/conftest.py b/tests/openai/evals/conftest.py index e263076..6039834 100644 --- a/tests/openai/evals/conftest.py +++ b/tests/openai/evals/conftest.py @@ -1,13 +1,39 @@ # conftest.py from openai import AsyncOpenAI, OpenAI import pytest +import os + +try: + import braintrust + + wrap_openai = braintrust.wrap_openai +except ImportError: + + def wrap_openai(x): + return x @pytest.fixture(scope="session") def client(): - yield OpenAI() + if os.environ.get("BRAINTRUST_API_KEY"): + yield wrap_openai( + OpenAI( + api_key=os.environ["BRAINTRUST_API_KEY"], + base_url="https://braintrustproxy.com/v1", + ) + ) + else: + yield OpenAI() @pytest.fixture(scope="session") def aclient(): - yield AsyncOpenAI() + if os.environ.get("BRAINTRUST_API_KEY"): + yield wrap_openai( + AsyncOpenAI( + api_key=os.environ["BRAINTRUST_API_KEY"], + base_url="https://braintrustproxy.com/v1", + ) + ) + else: + yield AsyncOpenAI()