|
11 | 11 | "trace_async_openai",
|
12 | 12 | "trace_async",
|
13 | 13 | "trace_bedrock",
|
| 14 | + "trace_oci", |
14 | 15 | ]
|
15 | 16 |
|
16 | 17 | # ---------------------------------- Tracing --------------------------------- #
|
@@ -95,11 +96,40 @@ def trace_bedrock(client):
|
95 | 96 | try:
|
96 | 97 | import boto3
|
97 | 98 | 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 | + ) |
99 | 102 |
|
100 | 103 | from .integrations import bedrock_tracer
|
101 | 104 |
|
102 | 105 | # 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 | + ) |
105 | 113 | return bedrock_tracer.trace_bedrock(client)
|
| 114 | + |
| 115 | + |
| 116 | + |
| 117 | +def trace_oci_genai(client, estimate_tokens: bool = True): |
| 118 | + """Trace OCI GenAI chat completions. |
| 119 | + |
| 120 | + Args: |
| 121 | + client: OCI GenAI client. |
| 122 | + estimate_tokens: Whether to estimate tokens when not available. Defaults to True. |
| 123 | + """ |
| 124 | + # pylint: disable=import-outside-toplevel |
| 125 | + try: |
| 126 | + import oci |
| 127 | + except ImportError: |
| 128 | + raise ImportError("oci is required for OCI GenAI tracing. Install with: pip install oci") |
| 129 | + |
| 130 | + from .integrations import oci_tracer |
| 131 | + |
| 132 | + if not isinstance(client, oci.generative_ai_inference.GenerativeAiInferenceClient): |
| 133 | + raise ValueError("Invalid client. Please provide an OCI GenAI client.") |
| 134 | + |
| 135 | + return oci_tracer.trace_oci_genai(client, estimate_tokens=estimate_tokens) |
0 commit comments