Skip to content

Commit d521c4b

Browse files
feat(tracer): enhance conditional imports and type hinting for OpenAI integration
- Improved conditional import handling for the `openai` library, ensuring graceful degradation when the library is not installed. - Enhanced type hints using forward references for `openai` types to improve code clarity and maintainability. - Added an informative error message when the `openai` library is missing, guiding users on how to install it. - This update ensures better compatibility and user experience when working with optional dependencies in the OpenAI integration.
1 parent 7639007 commit d521c4b

1 file changed

Lines changed: 16 additions & 4 deletions

File tree

src/openlayer/lib/integrations/async_openai_tracer.py

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,16 @@
44
import logging
55
import time
66
from functools import wraps
7-
from typing import Any, AsyncIterator, Optional, Union
7+
from typing import Any, AsyncIterator, Optional, Union, TYPE_CHECKING
88

9-
import openai
9+
try:
10+
import openai
11+
HAVE_OPENAI = True
12+
except ImportError:
13+
HAVE_OPENAI = False
14+
15+
if TYPE_CHECKING:
16+
import openai
1017

1118
from .openai_tracer import (
1219
get_model_parameters,
@@ -19,8 +26,8 @@
1926

2027

2128
def trace_async_openai(
22-
client: Union[openai.AsyncOpenAI, openai.AsyncAzureOpenAI],
23-
) -> Union[openai.AsyncOpenAI, openai.AsyncAzureOpenAI]:
29+
client: Union["openai.AsyncOpenAI", "openai.AsyncAzureOpenAI"],
30+
) -> Union["openai.AsyncOpenAI", "openai.AsyncAzureOpenAI"]:
2431
"""Patch the AsyncOpenAI or AsyncAzureOpenAI client to trace chat completions.
2532
2633
The following information is collected for each chat completion:
@@ -47,6 +54,11 @@ def trace_async_openai(
4754
Union[openai.AsyncOpenAI, openai.AsyncAzureOpenAI]
4855
The patched AsyncOpenAI client.
4956
"""
57+
if not HAVE_OPENAI:
58+
raise ImportError(
59+
"OpenAI library is not installed. Please install it with: pip install openai"
60+
)
61+
5062
is_azure_openai = isinstance(client, openai.AsyncAzureOpenAI)
5163
create_func = client.chat.completions.create
5264

0 commit comments

Comments
 (0)