Skip to content

Commit 1b49c07

Browse files
fix: add OCI tracer to init
1 parent adb54b8 commit 1b49c07

File tree

1 file changed

+22
-3
lines changed

1 file changed

+22
-3
lines changed

src/openlayer/lib/__init__.py

Lines changed: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
"trace_async_openai",
1212
"trace_async",
1313
"trace_bedrock",
14+
"trace_oci",
1415
]
1516

1617
# ---------------------------------- Tracing --------------------------------- #
@@ -95,11 +96,29 @@ def trace_bedrock(client):
9596
try:
9697
import boto3
9798
except ImportError:
98-
raise ImportError("boto3 is required for Bedrock tracing. Install with: pip install boto3")
99+
raise ImportError(
100+
"boto3 is required for Bedrock tracing. Install with: pip install boto3"
101+
)
99102

100103
from .integrations import bedrock_tracer
101104

102105
# Check if it's a boto3 client for bedrock-runtime service
103-
if not hasattr(client, "_service_model") or client._service_model.service_name != "bedrock-runtime":
104-
raise ValueError("Invalid client. Please provide a boto3 bedrock-runtime client.")
106+
if (
107+
not hasattr(client, "_service_model")
108+
or client._service_model.service_name != "bedrock-runtime"
109+
):
110+
raise ValueError(
111+
"Invalid client. Please provide a boto3 bedrock-runtime client."
112+
)
105113
return bedrock_tracer.trace_bedrock(client)
114+
115+
116+
def trace_oci(client):
117+
"""Trace OCI queries."""
118+
# pylint: disable=import-outside-toplevel
119+
import oci
120+
from .integrations import oci_tracer
121+
122+
if not isinstance(client, oci.generative_ai_inference.GenerativeAiInferenceClient):
123+
raise ValueError("Invalid client. Please provide an OCI Generative AI client.")
124+
return oci_tracer.trace_oci_genai(client)

0 commit comments

Comments
 (0)