diff --git a/genai_impact/exceptions.py b/genai_impact/exceptions.py index 567eac0..689c8f0 100644 --- a/genai_impact/exceptions.py +++ b/genai_impact/exceptions.py @@ -1,6 +1,7 @@ class GenAIImpactError(Exception): pass + class TracerInitializationError(GenAIImpactError): - "Tracer is initialized twice" + """Tracer is initialized twice""" pass diff --git a/genai_impact/tracer.py b/genai_impact/tracer.py index 3cf3070..7f56bcc 100644 --- a/genai_impact/tracer.py +++ b/genai_impact/tracer.py @@ -19,6 +19,7 @@ def init_instruments() -> None: init_anthropic_instrumentor() init_mistralai_instrumentor() + def init_openai_instrumentor() -> None: if importlib.util.find_spec("openai") is not None: from genai_impact.tracers.openai_tracer import OpenAIInstrumentor @@ -26,6 +27,7 @@ def init_openai_instrumentor() -> None: instrumentor = OpenAIInstrumentor() instrumentor.instrument() + def init_anthropic_instrumentor() -> None: if importlib.util.find_spec("anthropic") is not None: from genai_impact.tracers.anthropic_tracer import AnthropicInstrumentor @@ -33,6 +35,7 @@ def init_anthropic_instrumentor() -> None: instrumentor = AnthropicInstrumentor() instrumentor.instrument() + def init_mistralai_instrumentor() -> None: if importlib.util.find_spec("mistralai") is not None: from genai_impact.tracers.mistralai_tracer import MistralAIInstrumentor diff --git a/genai_impact/tracers/mistralai_tracer.py b/genai_impact/tracers/mistralai_tracer.py index a930883..f48bc16 100644 --- a/genai_impact/tracers/mistralai_tracer.py +++ b/genai_impact/tracers/mistralai_tracer.py @@ -44,7 +44,7 @@ def mistralai_chat_wrapper( async def mistralai_async_chat_wrapper( wrapped: Callable, - instance: _MistralAsyncClient, # noqa: ARG001 + instance: _MistralAsyncClient, # noqa: ARG001 args: Any, kwargs: Any, ) -> ChatCompletionResponse: diff --git a/genai_impact/tracers/openai_tracer.py b/genai_impact/tracers/openai_tracer.py index dbc5b15..93b0bf9 100644 --- a/genai_impact/tracers/openai_tracer.py +++ b/genai_impact/tracers/openai_tracer.py @@ -25,12 +25,14 @@ def compute_impacts_and_return_response(response: Any) -> ChatCompletion: ) return ChatCompletion(**response.model_dump(), impacts=impacts) + def openai_chat_wrapper( wrapped: Callable, instance: Completions, args: Any, kwargs: Any # noqa: ARG001 ) -> ChatCompletion: response = wrapped(*args, **kwargs) return compute_impacts_and_return_response(response) + async def openai_async_chat_wrapper( wrapped: Callable, instance: AsyncCompletions, # noqa: ARG001 @@ -41,7 +43,6 @@ async def openai_async_chat_wrapper( return compute_impacts_and_return_response(response) - class OpenAIInstrumentor: def __init__(self) -> None: self.wrapped_methods = [ diff --git a/tests/test_tracer.py b/tests/test_tracer.py index 2dcc6c9..7a0a5e8 100644 --- a/tests/test_tracer.py +++ b/tests/test_tracer.py @@ -2,6 +2,7 @@ from genai_impact import Tracer from genai_impact.exceptions import GenAIImpactError + def test_double_init(tracer_init): with pytest.raises(GenAIImpactError) as e: - Tracer.init() # Second initialization + Tracer.init() # Second initialization