Skip to content

Instantly share code, notes, and snippets.

@gyliu513
Last active November 15, 2023 17:42
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save gyliu513/303573c537b96da654bdab4e3349b381 to your computer and use it in GitHub Desktop.
Save gyliu513/303573c537b96da654bdab4e3349b381 to your computer and use it in GitHub Desktop.
WatsonxInstrumentor
class WatsonxInstrumentor(BaseInstrumentor):
"""An instrumentor for Watsonx's client library."""
def instrumentation_dependencies(self) -> Collection[str]:
return _instruments
def _instrument(self, **kwargs):
print("calling instrument")
tracer_provider = kwargs.get("tracer_provider")
tracer = get_tracer(__name__, __version__, tracer_provider)
wrapped_methods = (
WRAPPED_METHODS_VERSION_1
)
for wrapped_method in wrapped_methods:
wrap_module = wrapped_method.get("module")
wrap_object = wrapped_method.get("object")
wrap_method = wrapped_method.get("method")
wrap_function_wrapper(
wrap_module,
f"{wrap_object}.{wrap_method}",
_wrap(tracer, wrapped_method),
)
def _uninstrument(self, **kwargs):
wrapped_methods = (
WRAPPED_METHODS_VERSION_1
)
for wrapped_method in wrapped_methods:
wrap_object = wrapped_method.get("object")
unwrap(f"openai.{wrap_object}", wrapped_method.get("method"))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment