arize-phoenix 0.1.1__tar.gz → 1.0.0__tar.gz

{arize_phoenix-0.1.1 → arize_phoenix-1.0.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: arize-phoenix
-Version: 0.1.1
+Version: 1.0.0
 Summary: ML Observability in your notebook
 Project-URL: Documentation, https://docs.arize.com/phoenix/
 Project-URL: Issues, https://github.com/Arize-ai/phoenix/issues
@@ -38,15 +38,15 @@ Requires-Dist: arize[autoembeddings,llm-evaluation]; extra == 'dev'
 Requires-Dist: gcsfs; extra == 'dev'
 Requires-Dist: hatch; extra == 'dev'
 Requires-Dist: jupyter; extra == 'dev'
-Requires-Dist: langchain>=0.0.324; extra == 'dev'
-Requires-Dist: llama-index>=0.8.29; extra == 'dev'
+Requires-Dist: langchain>=0.0.334; extra == 'dev'
+Requires-Dist: llama-index>=0.8.64; extra == 'dev'
 Requires-Dist: nbqa; extra == 'dev'
 Requires-Dist: pandas-stubs<=2.0.2.230605; extra == 'dev'
 Requires-Dist: pre-commit; extra == 'dev'
 Requires-Dist: pytest; extra == 'dev'
 Requires-Dist: pytest-cov; extra == 'dev'
 Requires-Dist: pytest-lazy-fixture; extra == 'dev'
-Requires-Dist: ruff==0.1.3; extra == 'dev'
+Requires-Dist: ruff==0.1.5; extra == 'dev'
 Requires-Dist: strawberry-graphql[debug-server]==0.208.2; extra == 'dev'
 Provides-Extra: experimental
 Requires-Dist: tenacity; extra == 'experimental'
@@ -101,6 +101,7 @@ Phoenix provides MLOps and LLMOps insights at lightning speed with zero-config o
     -   [Exportable Clusters](#exportable-clusters)
 -   [Retrieval-Augmented Generation Analysis](#retrieval-augmented-generation-analysis)
 -   [Structured Data Analysis](#structured-data-analysis)
+-   [Breaking Changes](#breaking-changes)
 -   [Community](#community)
 -   [Thanks](#thanks)
 -   [Copyright, Patent, and License](#copyright-patent-and-license)
@@ -418,6 +419,10 @@ train_ds = px.Dataset(dataframe=train_df, schema=schema, name="training")
 session = px.launch_app(primary=prod_ds, reference=train_ds)
 ```
 
+## Breaking Changes
+
+-   **v1.0.0** - Phoenix now exclusively supports the `openai>=1.0.0` sdk. If you are using an older version of the OpenAI SDK, you can continue to use `arize-phoenix==0.1.1`. However, we recommend upgrading to the latest version of the OpenAI SDK as it contains many improvements. If you are using Phoenix with LlamaIndex and and LangChain, you will have to upgrade to the versions of these packages that support the OpenAI `1.0.0` SDK as well (`llama-index>=0.8.64`, `langchain>=0.0.334`)
+
 ## Community
 
 Join our community to connect with thousands of machine learning practitioners and ML observability enthusiasts.

{arize_phoenix-0.1.1 → arize_phoenix-1.0.0}/README.md

@@ -47,6 +47,7 @@ Phoenix provides MLOps and LLMOps insights at lightning speed with zero-config o
     -   [Exportable Clusters](#exportable-clusters)
 -   [Retrieval-Augmented Generation Analysis](#retrieval-augmented-generation-analysis)
 -   [Structured Data Analysis](#structured-data-analysis)
+-   [Breaking Changes](#breaking-changes)
 -   [Community](#community)
 -   [Thanks](#thanks)
 -   [Copyright, Patent, and License](#copyright-patent-and-license)
@@ -364,6 +365,10 @@ train_ds = px.Dataset(dataframe=train_df, schema=schema, name="training")
 session = px.launch_app(primary=prod_ds, reference=train_ds)
 ```
 
+## Breaking Changes
+
+-   **v1.0.0** - Phoenix now exclusively supports the `openai>=1.0.0` sdk. If you are using an older version of the OpenAI SDK, you can continue to use `arize-phoenix==0.1.1`. However, we recommend upgrading to the latest version of the OpenAI SDK as it contains many improvements. If you are using Phoenix with LlamaIndex and and LangChain, you will have to upgrade to the versions of these packages that support the OpenAI `1.0.0` SDK as well (`llama-index>=0.8.64`, `langchain>=0.0.334`)
+
 ## Community
 
 Join our community to connect with thousands of machine learning practitioners and ML observability enthusiasts.

{arize_phoenix-0.1.1 → arize_phoenix-1.0.0}/pyproject.toml

@@ -47,7 +47,7 @@ dev = [
   "hatch",
   "jupyter",
   "nbqa",
-  "ruff==0.1.3",
+  "ruff==0.1.5",
   "pandas-stubs<=2.0.2.230605",  # version 2.0.3.230814 is causing a dependency conflict.
   "pytest",
   "pytest-cov",
@@ -55,8 +55,8 @@ dev = [
   "strawberry-graphql[debug-server]==0.208.2",
   "pre-commit",
   "arize[AutoEmbeddings, LLM_Evaluation]",
-  "llama-index>=0.8.29",
-  "langchain>=0.0.324",
+  "llama-index>=0.8.64",
+  "langchain>=0.0.334",
 ]
 experimental = [
   "tenacity",
@@ -91,9 +91,9 @@ dependencies = [
   "pytest-cov",
   "pytest-lazy-fixture",
   "arize",
-  "langchain>=0.0.324",
-  "llama-index>=0.8.29",
-  "openai<1.0.0",
+  "langchain>=0.0.334",
+  "llama-index>=0.8.63.post2",
+  "openai>=1.0.0",
   "tenacity",
   "nltk==3.8.1",
   "sentence-transformers==2.2.2",
@@ -103,24 +103,26 @@ dependencies = [
   "responses",
   "tiktoken",
   "typing-extensions<4.6.0",  # for Colab
+  "httpx", # For OpenAI testing
+  "respx", # For OpenAI testing
 ]
 
 [tool.hatch.envs.type]
 dependencies = [
   "mypy==1.5.1",
-  "llama-index>=0.8.29",
+  "llama-index>=0.8.64",
   "pandas-stubs<=2.0.2.230605",  # version 2.0.3.230814 is causing a dependency conflict.
   "types-psutil",
   "types-tqdm",
   "types-requests",
   "types-protobuf",
-  "openai<1.0.0",
+  "openai>=1.0.0",
 ]
 
 [tool.hatch.envs.style]
 detached = true
 dependencies = [
-  "ruff~=0.1.3",
+  "ruff~=0.1.5",
 ]
 
 [tool.hatch.envs.notebooks]

{arize_phoenix-0.1.1 → arize_phoenix-1.0.0}/src/phoenix/__init__.py

@@ -5,7 +5,7 @@ from .session.session import Session, active_session, close_app, launch_app
 from .trace.fixtures import load_example_traces
 from .trace.trace_dataset import TraceDataset
 
-__version__ = "0.1.1"
+__version__ = "1.0.0"
 
 # module level doc-string
 __doc__ = """

{arize_phoenix-0.1.1 → arize_phoenix-1.0.0}/src/phoenix/experimental/evals/models/openai.py

@@ -1,7 +1,18 @@
 import logging
 import os
-from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+from dataclasses import dataclass, field, fields
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Tuple,
+    Union,
+    get_args,
+    get_origin,
+)
 
 from phoenix.experimental.evals.models.base import BaseEvalModel
 
@@ -9,7 +20,7 @@ if TYPE_CHECKING:
     from tiktoken import Encoding
 
 OPENAI_API_KEY_ENVVAR_NAME = "OPENAI_API_KEY"
-MINIMUM_OPENAI_VERSION = "0.26.4"
+MINIMUM_OPENAI_VERSION = "1.0.0"
 MODEL_TOKEN_LIMIT_MAPPING = {
     "gpt-3.5-turbo-instruct": 4096,
     "gpt-3.5-turbo-0301": 4096,
@@ -24,17 +35,31 @@ LEGACY_COMPLETION_API_MODELS = ("gpt-3.5-turbo-instruct",)
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class AzureOptions:
+    api_version: str
+    azure_endpoint: str
+    azure_deployment: Optional[str]
+    azure_ad_token: Optional[str]
+    azure_ad_token_provider: Optional[Callable[[], str]]
+
+
 @dataclass
 class OpenAIModel(BaseEvalModel):
-    openai_api_type: Optional[str] = field(default=None)
-    openai_api_version: Optional[str] = field(default=None)
-    openai_api_key: Optional[str] = field(repr=False, default=None)
-    openai_api_base: Optional[str] = field(repr=False, default=None)
-    openai_organization: Optional[str] = field(repr=False, default=None)
-    engine: str = ""
-    """Azure engine (the Deployment Name of your model)"""
+    api_key: Optional[str] = field(repr=False, default=None)
+    """Your OpenAI key. If not provided, will be read from the environment variable"""
+    organization: Optional[str] = field(repr=False, default=None)
+    """
+    The organization to use for the OpenAI API. If not provided, will default
+    to what's configured in OpenAI
+    """
+    base_url: Optional[str] = field(repr=False, default=None)
+    """
+    An optional base URL to use for the OpenAI API. If not provided, will default
+    to what's configured in OpenAI
+    """
     model_name: str = "gpt-4"
-    """Model name to use."""
+    """Model name to use. In of azure, this is the deployment name such as gpt-35-instant"""
     temperature: float = 0.0
     """What sampling temperature to use."""
     max_tokens: int = 256
@@ -63,6 +88,18 @@ class OpenAIModel(BaseEvalModel):
     retry_max_seconds: int = 60
     """Maximum number of seconds to wait when retrying."""
 
+    # Azure options
+    api_version: Optional[str] = field(default=None)
+    """https://learn.microsoft.com/en-us/azure/ai-services/openai/reference#rest-api-versioning"""
+    azure_endpoint: Optional[str] = field(default=None)
+    """
+    The endpoint to use for azure openai. Available in the azure portal.
+    https://learn.microsoft.com/en-us/azure/cognitive-services/openai/how-to/create-resource?pivots=web-portal#create-a-resource
+    """
+    azure_deployment: Optional[str] = field(default=None)
+    azure_ad_token: Optional[str] = field(default=None)
+    azure_ad_token_provider: Optional[Callable[[], str]] = field(default=None)
+
     def __post_init__(self) -> None:
         self._init_environment()
         self._init_open_ai()
@@ -71,12 +108,10 @@ class OpenAIModel(BaseEvalModel):
     def _init_environment(self) -> None:
         try:
             import openai
-            import openai.util
-            from openai import error as openai_error
+            import openai._utils as openai_util
 
             self._openai = openai
-            self._openai_error = openai_error
-            self._openai_util = openai.util
+            self._openai_util = openai_util
         except ImportError:
             self._raise_import_error(
                 package_display_name="OpenAI",
@@ -93,52 +128,76 @@ class OpenAIModel(BaseEvalModel):
             )
 
     def _init_open_ai(self) -> None:
+        # For Azure, you need to provide the endpoint and the endpoint
+        self._is_azure = bool(self.azure_endpoint)
+
         self._model_uses_legacy_completion_api = self.model_name.startswith(
             LEGACY_COMPLETION_API_MODELS
         )
-        if self.openai_api_key is None:
+        if self.api_key is None:
             api_key = os.getenv(OPENAI_API_KEY_ENVVAR_NAME)
             if api_key is None:
                 # TODO: Create custom AuthenticationError
                 raise RuntimeError(
-                    "OpenAI's API key not provided. Pass it as an argument to 'openai_api_key' "
+                    "OpenAI's API key not provided. Pass it as an argument to 'api_key' "
                     "or set it in your environment: 'export OPENAI_API_KEY=sk-****'"
                 )
-            self.openai_api_key = api_key
-        self.openai_api_base = self.openai_api_base or self._openai.api_base
-        self.openai_api_type = self.openai_api_type or self._openai.api_type
-        self.openai_api_version = self.openai_api_version or self._openai.api_version
-        self.openai_organization = self.openai_organization or self._openai.organization
-        # use enum to validate api type
-        self._openai_util.ApiType.from_str(self.openai_api_type)  # type: ignore
-        self._is_azure = self.openai_api_type.lower().startswith("azure")
+            self.api_key = api_key
 
+        # Set the version, organization, and base_url - default to openAI
+        self.api_version = self.api_version or self._openai.api_version
+        self.organization = self.organization or self._openai.organization
+
+        # Initialize specific clients depending on the API backend
+        # Set the type first
+        self._client: Union[self._openai.OpenAI, self._openai.AzureOpenAI]  # type: ignore
         if self._is_azure:
-            if not self.engine:
-                raise ValueError(
-                    "You must provide the deployment name in the 'engine' parameter "
-                    "to access the Azure OpenAI service"
-                )
-            self._openai_api_model_name = self.engine
-        elif self.model_name in MODEL_TOKEN_LIMIT_MAPPING.keys():
-            self._openai_api_model_name = self.model_name
-        elif "gpt-3.5-turbo" in self.model_name:
-            self._openai_api_model_name = "gpt-3.5-turbo-0613"
-        elif "gpt-4" in self.model_name:
-            self._openai_api_model_name = "gpt-4-0613"
-        else:
-            raise NotImplementedError(
-                f"openai_api_model_name not available for model {self.model_name}. "
+            # Validate the azure options and construct a client
+            azure_options = self._get_azure_options()
+            self._client = self._openai.AzureOpenAI(
+                azure_endpoint=azure_options.azure_endpoint,
+                azure_deployment=azure_options.azure_deployment,
+                api_version=azure_options.api_version,
+                azure_ad_token=azure_options.azure_ad_token,
+                azure_ad_token_provider=azure_options.azure_ad_token_provider,
+                api_key=self.api_key,
+                organization=self.organization,
             )
+            # return early since we don't need to check the model
+            return
+
+        # The client is not azure, so it must be openai
+        self._client = self._openai.OpenAI(
+            api_key=self.api_key,
+            organization=self.organization,
+            base_url=(self.base_url or self._openai.base_url),
+        )
 
     def _init_tiktoken(self) -> None:
         try:
-            encoding = self._tiktoken.encoding_for_model(self.openai_api_model_name)
+            encoding = self._tiktoken.encoding_for_model(self.model_name)
         except KeyError:
-            logger.warning("Warning: model not found. Using cl100k_base encoding.")
             encoding = self._tiktoken.get_encoding("cl100k_base")
         self._tiktoken_encoding = encoding
 
+    def _get_azure_options(self) -> AzureOptions:
+        options = {}
+        for option in fields(AzureOptions):
+            if (value := getattr(self, option.name)) is not None:
+                options[option.name] = value
+            else:
+                # raise ValueError if field is not optional
+                # See if the field is optional - e.g. get_origin(Optional[T])  = typing.Union
+                option_is_optional = get_origin(option.type) is Union and type(None) in get_args(
+                    option.type
+                )
+                if not option_is_optional:
+                    raise ValueError(
+                        f"Option '{option.name}' must be set when using Azure OpenAI API"
+                    )
+                options[option.name] = None
+        return AzureOptions(**options)
+
     @staticmethod
     def _build_messages(
         prompt: str, system_instruction: Optional[str] = None
@@ -173,11 +232,11 @@ class OpenAIModel(BaseEvalModel):
     def _generate_with_retry(self, **kwargs: Any) -> Any:
         """Use tenacity to retry the completion call."""
         openai_retry_errors = [
-            self._openai_error.Timeout,
-            self._openai_error.APIError,
-            self._openai_error.APIConnectionError,
-            self._openai_error.RateLimitError,
-            self._openai_error.ServiceUnavailableError,
+            self._openai.APITimeoutError,
+            self._openai.APIError,
+            self._openai.APIConnectionError,
+            self._openai.RateLimitError,
+            self._openai.InternalServerError,
         ]
 
         @self.retry(
@@ -193,17 +252,22 @@ class OpenAIModel(BaseEvalModel):
                         (message.get("content") or "")
                         for message in (kwargs.pop("messages", None) or ())
                     )
-                return self._openai.Completion.create(**kwargs)
-            return self._openai.ChatCompletion.create(**kwargs)
+                # OpenAI 1.0.0 API responses are pydantic objects, not dicts
+                # We must dump the model to get the dict
+                return self._client.completions.create(**kwargs).model_dump()
+            return self._client.chat.completions.create(**kwargs).model_dump()
 
         return _completion_with_retry(**kwargs)
 
     @property
     def max_context_size(self) -> int:
-        model_name = self.openai_api_model_name
+        model_name = self.model_name
         # handling finetuned models
         if "ft-" in model_name:
             model_name = self.model_name.split(":")[0]
+        if model_name == "gpt-4":
+            # Map gpt-4 to the current default
+            model_name = "gpt-4-0613"
 
         context_size = MODEL_TOKEN_LIMIT_MAPPING.get(model_name, None)
 
@@ -219,7 +283,7 @@ class OpenAIModel(BaseEvalModel):
     @property
     def public_invocation_params(self) -> Dict[str, Any]:
         return {
-            **({"engine": self.engine} if self._is_azure else {"model": self.model_name}),
+            **({"model": self.model_name}),
             **self._default_params,
             **self.model_kwargs,
         }
@@ -228,18 +292,6 @@ class OpenAIModel(BaseEvalModel):
     def invocation_params(self) -> Dict[str, Any]:
         return {
             **self.public_invocation_params,
-            **self._credentials,
-        }
-
-    @property
-    def _credentials(self) -> Dict[str, Any]:
-        """Get the default parameters for calling OpenAI API."""
-        return {
-            "api_key": self.openai_api_key,
-            "api_base": self.openai_api_base,
-            "api_type": self.openai_api_type,
-            "api_version": self.openai_api_version,
-            "organization": self.openai_organization,
         }
 
     @property
@@ -252,13 +304,9 @@ class OpenAIModel(BaseEvalModel):
             "presence_penalty": self.presence_penalty,
             "top_p": self.top_p,
             "n": self.n,
-            "request_timeout": self.request_timeout,
+            "timeout": self.request_timeout,
         }
 
-    @property
-    def openai_api_model_name(self) -> str:
-        return self._openai_api_model_name
-
     @property
     def encoder(self) -> "Encoding":
         return self._tiktoken_encoding
@@ -268,7 +316,7 @@ class OpenAIModel(BaseEvalModel):
 
         Official documentation: https://github.com/openai/openai-cookbook/blob/main/examples/How_to_format_inputs_to_ChatGPT_models.ipynb
         """  # noqa
-        model_name = self.openai_api_model_name
+        model_name = self.model_name
         if model_name == "gpt-3.5-turbo-0301":
             tokens_per_message = 4  # every message follows <|start|>{role/name}\n{content}<|end|>\n
             tokens_per_name = -1  # if there's a name, the role is omitted
@@ -297,10 +345,10 @@ class OpenAIModel(BaseEvalModel):
     def supports_function_calling(self) -> bool:
         if (
             self._is_azure
-            and self.openai_api_version
+            and self.api_version
             # The first api version supporting function calling is 2023-07-01-preview.
             # See https://github.com/Azure/azure-rest-api-specs/blob/58e92dd03733bc175e6a9540f4bc53703b57fcc9/specification/cognitiveservices/data-plane/AzureOpenAI/inference/preview/2023-07-01-preview/inference.json#L895 # noqa E501
-            and self.openai_api_version[:10] < "2023-07-01"
+            and self.api_version[:10] < "2023-07-01"
         ):
             return False
         if self._model_uses_legacy_completion_api:

{arize_phoenix-0.1.1 → arize_phoenix-1.0.0}/src/phoenix/experimental/evals/retrievals.py

@@ -75,19 +75,22 @@ def classify_relevance(query: str, document: str, model_name: str) -> Optional[b
             unparseable output.
     """
 
-    from openai import ChatCompletion
+    from openai import OpenAI
+
+    client = OpenAI()
 
     prompt = _QUERY_CONTEXT_PROMPT_TEMPLATE.format(
         query=query,
         reference=document,
     )
-    response = ChatCompletion.create(  # type: ignore
+    response = client.chat.completions.create(
         messages=[
             {"role": "system", "content": _EVALUATION_SYSTEM_MESSAGE},
             {"role": "user", "content": prompt},
         ],
         model=model_name,
     )
-    raw_response_text = str(response["choices"][0]["message"]["content"]).strip()
+
+    raw_response_text = str(response.choices[0].message.content).strip()
     relevance_classification = {"relevant": True, "irrelevant": False}.get(raw_response_text)
     return relevance_classification

{arize_phoenix-0.1.1 → arize_phoenix-1.0.0}/src/phoenix/trace/openai/instrumentor.py

@@ -10,9 +10,10 @@ from typing import (
     List,
     Mapping,
     Optional,
-    cast,
 )
 
+from typing_extensions import TypeGuard
+
 from phoenix.trace.schemas import (
     SpanAttributes,
     SpanEvent,
@@ -44,7 +45,7 @@ from phoenix.trace.utils import get_stacktrace, import_package
 from ..tracer import Tracer
 
 if TYPE_CHECKING:
-    from openai.openai_response import OpenAIResponse
+    from openai.types.chat import ChatCompletion
 
 
 Parameters = Mapping[str, Any]
@@ -75,21 +76,21 @@ class OpenAIInstrumentor:
         """
         openai = import_package("openai")
         is_instrumented = hasattr(
-            openai.api_requestor.APIRequestor.request,
+            openai.OpenAI,
             INSTRUMENTED_ATTRIBUTE_NAME,
         )
         if not is_instrumented:
-            openai.api_requestor.APIRequestor.request = _wrap_openai_api_requestor(
-                openai.api_requestor.APIRequestor.request, self._tracer
+            openai.OpenAI.request = _wrapped_openai_client_request_function(
+                openai.OpenAI.request, self._tracer
             )
             setattr(
-                openai.api_requestor.APIRequestor.request,
+                openai.OpenAI,
                 INSTRUMENTED_ATTRIBUTE_NAME,
                 True,
             )
 
 
-def _wrap_openai_api_requestor(
+def _wrapped_openai_client_request_function(
     request_fn: Callable[..., Any], tracer: Tracer
 ) -> Callable[..., Any]:
     """Wraps the OpenAI APIRequestor.request method to create spans for each API call.
@@ -105,9 +106,10 @@ def _wrap_openai_api_requestor(
     def wrapped(*args: Any, **kwargs: Any) -> Any:
         call_signature = signature(request_fn)
         bound_arguments = call_signature.bind(*args, **kwargs)
-        parameters = bound_arguments.arguments["params"]
-        is_streaming = parameters.get("stream", False)
-        url = bound_arguments.arguments["url"]
+        is_streaming = bound_arguments.arguments["stream"]
+        options = bound_arguments.arguments["options"]
+        parameters = options.json_data
+        url = options.url
         current_status_code = SpanStatusCode.UNSET
         events: List[SpanEvent] = []
         attributes: SpanAttributes = dict()
@@ -118,13 +120,13 @@ def _wrap_openai_api_requestor(
             ) in _PARAMETER_ATTRIBUTE_FUNCTIONS.items():
                 if (attribute_value := get_parameter_attribute_fn(parameters)) is not None:
                     attributes[attribute_name] = attribute_value
-            outputs = None
+            response = None
             try:
                 start_time = datetime.now()
-                outputs = request_fn(*args, **kwargs)
+                response = request_fn(*args, **kwargs)
                 end_time = datetime.now()
                 current_status_code = SpanStatusCode.OK
-                return outputs
+                return response
             except Exception as error:
                 end_time = datetime.now()
                 current_status_code = SpanStatusCode.ERROR
@@ -138,16 +140,17 @@ def _wrap_openai_api_requestor(
                 )
                 raise
             finally:
-                if outputs:
-                    response = outputs[0]
+                if _is_chat_completion(response):
                     for (
                         attribute_name,
-                        get_response_attribute_fn,
-                    ) in _RESPONSE_ATTRIBUTE_FUNCTIONS.items():
-                        if (attribute_value := get_response_attribute_fn(response)) is not None:
+                        get_chat_completion_attribute_fn,
+                    ) in _CHAT_COMPLETION_ATTRIBUTE_FUNCTIONS.items():
+                        if (
+                            attribute_value := get_chat_completion_attribute_fn(response)
+                        ) is not None:
                             attributes[attribute_name] = attribute_value
                 tracer.create_span(
-                    name="openai.ChatCompletion.create",
+                    name="OpenAI Chat Completion",
                     span_kind=SpanKind.LLM,
                     start_time=start_time,
                     end_time=end_time,
@@ -182,48 +185,46 @@ def _llm_invocation_parameters(
     return json.dumps(parameters)
 
 
-def _output_value(response: "OpenAIResponse") -> str:
-    return json.dumps(response.data)
+def _output_value(chat_completion: "ChatCompletion") -> str:
+    return chat_completion.json()
 
 
 def _output_mime_type(_: Any) -> MimeType:
     return MimeType.JSON
 
 
-def _llm_output_messages(response: "OpenAIResponse") -> List[OpenInferenceMessage]:
+def _llm_output_messages(chat_completion: "ChatCompletion") -> List[OpenInferenceMessage]:
     return [
-        _to_openinference_message(choice["message"], expects_name=False)
-        for choice in response.data["choices"]
+        _to_openinference_message(choice.message.dict(), expects_name=False)
+        for choice in chat_completion.choices
     ]
 
 
-def _llm_token_count_prompt(response: "OpenAIResponse") -> Optional[int]:
-    if token_usage := response.data.get("usage"):
-        return cast(int, token_usage["prompt_tokens"])
+def _llm_token_count_prompt(chat_completion: "ChatCompletion") -> Optional[int]:
+    if completion_usage := chat_completion.usage:
+        return completion_usage.prompt_tokens
     return None
 
 
-def _llm_token_count_completion(response: "OpenAIResponse") -> Optional[int]:
-    if token_usage := response.data.get("usage"):
-        return cast(int, token_usage["completion_tokens"])
+def _llm_token_count_completion(chat_completion: "ChatCompletion") -> Optional[int]:
+    if completion_usage := chat_completion.usage:
+        return completion_usage.completion_tokens
     return None
 
 
-def _llm_token_count_total(response: "OpenAIResponse") -> Optional[int]:
-    if token_usage := response.data.get("usage"):
-        return cast(int, token_usage["total_tokens"])
+def _llm_token_count_total(chat_completion: "ChatCompletion") -> Optional[int]:
+    if completion_usage := chat_completion.usage:
+        return completion_usage.total_tokens
     return None
 
 
 def _llm_function_call(
-    response: "OpenAIResponse",
+    chat_completion: "ChatCompletion",
 ) -> Optional[str]:
-    choices = response.data["choices"]
+    choices = chat_completion.choices
     choice = choices[0]
-    if choice.get("finish_reason") == "function_call" and (
-        function_call_data := choice["message"].get("function_call")
-    ):
-        return json.dumps(function_call_data)
+    if choice.finish_reason == "function_call" and (function_call := choice.message.function_call):
+        return function_call.json()
     return None
 
 
@@ -274,7 +275,7 @@ _PARAMETER_ATTRIBUTE_FUNCTIONS: Dict[str, Callable[[Parameters], Any]] = {
     LLM_INPUT_MESSAGES: _llm_input_messages,
     LLM_INVOCATION_PARAMETERS: _llm_invocation_parameters,
 }
-_RESPONSE_ATTRIBUTE_FUNCTIONS: Dict[str, Callable[["OpenAIResponse"], Any]] = {
+_CHAT_COMPLETION_ATTRIBUTE_FUNCTIONS: Dict[str, Callable[["ChatCompletion"], Any]] = {
     OUTPUT_VALUE: _output_value,
     OUTPUT_MIME_TYPE: _output_mime_type,
     LLM_OUTPUT_MESSAGES: _llm_output_messages,
@@ -283,3 +284,11 @@ _RESPONSE_ATTRIBUTE_FUNCTIONS: Dict[str, Callable[["OpenAIResponse"], Any]] = {
     LLM_TOKEN_COUNT_TOTAL: _llm_token_count_total,
     LLM_FUNCTION_CALL: _llm_function_call,
 }
+
+
+def _is_chat_completion(response: Any) -> TypeGuard["ChatCompletion"]:
+    """
+    Type guard for ChatCompletion.
+    """
+    openai = import_package("openai")
+    return isinstance(response, openai.types.chat.ChatCompletion)