test: Add Python test infrastructure and OpenAI wrapper tests

.github/workflows/py.yaml

@@ -44,7 +44,7 @@ jobs:
       - name: Test whether the Python SDK can be installed
         run: |
           python -m pip install -e ./core/py[all]
-          python -m pip install -e ./py[all]
+          python -m pip install -e ./py[all,dev]
       - name: Test whether the Python SDK can be imported
         run: |
           python -c 'import braintrust'
@@ -55,3 +55,6 @@ jobs:
       - name: Lint the Python SDK
         run: |
           pylint --errors-only $(git ls-files '*.py')
+      - name: Run Python tests
+        run: |
+          source env.sh && python -m pytest py/tests/ -v

Makefile

@@ -33,3 +33,10 @@ develop: ${VENV_PRE_COMMIT}
 
 fixup:
 	source env.sh && pre-commit run --all-files
+
+py-tests:
+	@if [ -z "$(tests)" ]; then \
+		source env.sh && python -m pytest py/tests/ -v; \
+	else \
+		source env.sh && python -m pytest $(tests) -v; \
+	fi

py/setup.py

@@ -33,8 +33,12 @@
         "flake8-isort",
         "IPython",
         "isort==5.10.1",
+        "openai",
         "pre-commit",
         "pytest",
+        "pytest-asyncio>=0.23.0",
+        "pytest-httpx<=0.22.0",
+        "responses",
         "twine",
     ],
     "doc": ["pydoc-markdown"],

py/src/braintrust/oai.py

@@ -121,7 +121,7 @@ def gen():
                                     }
                                 )
                                 first = False
-                            all_results.append(item if isinstance(item, dict) else item.dict())
+                            all_results.append(item if isinstance(item, dict) else item.model_dump())
                             yield item
 
                         span.log(**postprocess_streaming_results(all_results))
@@ -131,7 +131,7 @@ def gen():
                 should_end = False
                 return gen()
             else:
-                log_response = raw_response if isinstance(raw_response, dict) else raw_response.dict()
+                log_response = raw_response if isinstance(raw_response, dict) else raw_response.model_dump()
                 span.log(
                     metrics={
                         "time_to_first_token": time.time() - start,
@@ -179,7 +179,7 @@ async def gen():
                                     }
                                 )
                                 first = False
-                            all_results.append(item if isinstance(item, dict) else item.dict())
+                            all_results.append(item if isinstance(item, dict) else item.model_dump())
                             yield item
 
                         span.log(**postprocess_streaming_results(all_results))
@@ -189,7 +189,7 @@ async def gen():
                 should_end = False
                 return gen()
             else:
-                log_response = raw_response if isinstance(raw_response, dict) else raw_response.dict()
+                log_response = raw_response if isinstance(raw_response, dict) else raw_response.model_dump()
                 span.log(
                     metrics={
                         "tokens": log_response["usage"]["total_tokens"],
@@ -239,7 +239,7 @@ def create(self, *args, **kwargs):
             else:
                 raw_response = create_response
 
-            log_response = raw_response if isinstance(raw_response, dict) else raw_response.dict()
+            log_response = raw_response if isinstance(raw_response, dict) else raw_response.model_dump()
             span.log(
                 metrics={
                     "tokens": log_response["usage"]["total_tokens"],
@@ -263,7 +263,7 @@ async def acreate(self, *args, **kwargs):
                 log_headers(create_response, span)
             else:
                 raw_response = create_response
-            log_response = raw_response if isinstance(raw_response, dict) else raw_response.dict()
+            log_response = raw_response if isinstance(raw_response, dict) else raw_response.model_dump()
             span.log(
                 metrics={
                     "tokens": log_response["usage"]["total_tokens"],

py/tests/test_oai/test_wrap_openai.py

@@ -0,0 +1,103 @@
+import json
+import logging
+from inspect import iscoroutinefunction
+from typing import TYPE_CHECKING
+
+import httpx
+import openai
+import pytest
+from braintrust.oai import wrap_openai
+from openai.types import CompletionUsage
+from openai.types.chat import ChatCompletion
+from openai.types.chat.chat_completion import ChatCompletionMessage, Choice
+
+if TYPE_CHECKING:
+    reveal_type = print  # For type checking only
+else:
+
+    def reveal_type(obj):
+        pass  # No-op at runtime
+
+
+logging.basicConfig(level=logging.DEBUG)
+
+
+@pytest.fixture
+def openai_client():
+    return openai.OpenAI(api_key="sk-test", base_url="https://api.openai.com/v1/")
+
+
+@pytest.fixture
+def mock_completion():
+    return {
+        "id": "test-id",
+        "object": "chat.completion",
+        "created": 1677652288,
+        "model": "gpt-3.5-turbo",
+        "choices": [
+            {
+                "index": 0,
+                "message": {"role": "assistant", "content": "Hello, how can I help you?"},
+                "finish_reason": "stop",
+            }
+        ],
+        "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
+    }
+
+
+@pytest.fixture
+def setup_responses(mock_completion, httpx_mock):
+    httpx_mock.add_response(
+        method="POST",
+        url="https://api.openai.com/v1/chat/completions",
+        json=mock_completion,
+        headers={"Content-Type": "application/json"},
+        status_code=200,
+    )
+    return httpx_mock
+
+
+def test_wrap_openai_sync_types(openai_client):
+    wrapped = wrap_openai(openai_client)
+    reveal_type(wrapped)  # type: ignore  # Expected type: openai.OpenAI
+    reveal_type(wrapped.chat.completions)  # type: ignore  # Expected type: openai.resources.chat.completions.Completions
+    assert hasattr(wrapped.chat.completions, "create")
+    assert not hasattr(wrapped.chat.completions, "acreate")
+
+
+@pytest.mark.asyncio
+async def test_wrap_openai_async_types():
+    async_client = openai.AsyncOpenAI(
+        api_key="sk-test", base_url="https://api.openai.com/v1/", default_headers={"OpenAI-Version": "2020-10-01"}
+    )
+    wrapped = wrap_openai(async_client)
+    reveal_type(wrapped)  # type: ignore  # Expected type: openai.AsyncOpenAI
+    reveal_type(wrapped.chat.completions)  # type: ignore  # Expected type: openai.resources.chat.completions.AsyncCompletions
+    assert hasattr(wrapped.chat.completions, "create")
+    assert iscoroutinefunction(wrapped.chat.completions.create)
+
+
+def test_wrap_openai_sync_response_types(openai_client, mock_completion, setup_responses):
+    wrapped = wrap_openai(openai_client)
+    response = wrapped.chat.completions.create(model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hello"}])
+    reveal_type(response)  # type: ignore  # Expected type: openai.types.chat.ChatCompletion
+    reveal_type(response.choices[0])  # type: ignore  # Expected type: openai.types.chat.chat_completion.Choice
+    reveal_type(response.usage)  # type: ignore  # Expected type: openai.types.completion_usage.CompletionUsage
+    assert isinstance(response, ChatCompletion)
+    assert isinstance(response.choices[0], Choice)
+    assert isinstance(response.usage, CompletionUsage)
+
+
+@pytest.mark.asyncio
+async def test_wrap_openai_async_response_types(mock_completion, setup_responses):
+    async_client = openai.AsyncOpenAI(api_key="sk-test", base_url="https://api.openai.com/v1/")
+    wrapped = wrap_openai(async_client)
+    response = await wrapped.chat.completions.create(
+        model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hello"}]
+    )
+    reveal_type(response)  # type: ignore  # Expected type: openai.types.chat.ChatCompletion
+    reveal_type(response.choices[0])  # type: ignore  # Expected type: openai.types.chat.chat_completion.Choice
+    reveal_type(response.usage)  # type: ignore  # Expected type: openai.types.completion_usage.CompletionUsage
+    assert isinstance(response, ChatCompletion)
+    assert isinstance(response.choices[0], Choice)
+    assert isinstance(response.usage, CompletionUsage)

pyproject.toml

@@ -4,3 +4,21 @@ line-length = 119
 [tool.ruff]
 line-length = 119
 select = ["I"]
+
+[tool.pytest.ini_options]
+asyncio_mode = "strict"
+asyncio_default_fixture_loop_scope = "function"
+log_cli = true
+log_cli_level = "DEBUG"
+
+[tool.mypy]
+python_version = "3.8"
+warn_return_any = true
+warn_unused_configs = true
+disallow_untyped_defs = true
+check_untyped_defs = true
+warn_redundant_casts = true
+warn_unused_ignores = true
+warn_no_return = true
+warn_unreachable = true
+show_error_codes = true