diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 8dd6287..589e66d 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -5,11 +5,6 @@ repos: hooks: - id: ruff args: [--fix, --exit-non-zero-on-fix] - - repo: https://github.com/psf/black - rev: 22.3.0 - hooks: - - id: black - language_version: python3 - repo: https://github.com/pre-commit/pre-commit-hooks rev: v4.3.0 hooks: diff --git a/genai_impact/tracer.py b/genai_impact/tracer.py index 39d3e82..662eb8b 100644 --- a/genai_impact/tracer.py +++ b/genai_impact/tracer.py @@ -15,7 +15,7 @@ def init_instruments() -> None: def init_openai_instrumentor() -> None: - if importlib.util.find_spec('openai') is not None: + if importlib.util.find_spec("openai") is not None: from genai_impact.tracers.openai_tracer import OpenAIInstrumentor instrumentor = OpenAIInstrumentor() @@ -23,7 +23,7 @@ def init_openai_instrumentor() -> None: def init_anthropic_instrumentor() -> None: - if importlib.util.find_spec('anthropic') is not None: + if importlib.util.find_spec("anthropic") is not None: from genai_impact.tracers.anthropic_tracer import AnthropicInstrumentor instrumentor = AnthropicInstrumentor() @@ -31,7 +31,7 @@ def init_anthropic_instrumentor() -> None: def init_mistralai_instrumentor() -> None: - if importlib.util.find_spec('mistralai') is not None: + if importlib.util.find_spec("mistralai") is not None: from genai_impact.tracers.mistralai_tracer import MistralAIInstrumentor instrumentor = MistralAIInstrumentor() diff --git a/genai_impact/tracers/anthropic_tracer.py b/genai_impact/tracers/anthropic_tracer.py index 37b6d66..6015846 100644 --- a/genai_impact/tracers/anthropic_tracer.py +++ b/genai_impact/tracers/anthropic_tracer.py @@ -4,7 +4,6 @@ from genai_impact.compute_impacts import Impacts, compute_llm_impact - try: from anthropic import Anthropic as _Anthropic from anthropic.types import Message as _Message @@ -23,7 +22,7 @@ class Message(_Message): impacts: Impacts -def _set_impacts(response): +def _set_impacts(response: Message) -> Impacts: model_size = _MODEL_SIZES.get(response.model) output_tokens = response.usage.output_tokens impacts = compute_llm_impact( @@ -41,7 +40,7 @@ def anthropic_chat_wrapper( class AnthropicInstrumentor: - def __init__(self): + def __init__(self) -> None: self.wrapped_methods = [ { "module": "anthropic.resources", @@ -50,7 +49,7 @@ def __init__(self): }, ] - def instrument(self): + def instrument(self) -> None: for wrapper in self.wrapped_methods: wrap_function_wrapper( wrapper["module"], diff --git a/genai_impact/tracers/mistralai_tracer.py b/genai_impact/tracers/mistralai_tracer.py index a482dc4..7bdeac5 100644 --- a/genai_impact/tracers/mistralai_tracer.py +++ b/genai_impact/tracers/mistralai_tracer.py @@ -39,7 +39,7 @@ def mistralai_chat_wrapper( class MistralAIInstrumentor: - def __init__(self): + def __init__(self) -> None: self.wrapped_methods = [ { "module": "mistralai.client", @@ -48,7 +48,7 @@ def __init__(self): }, ] - def instrument(self): + def instrument(self) -> None: for wrapper in self.wrapped_methods: wrap_function_wrapper( wrapper["module"], diff --git a/genai_impact/tracers/openai_tracer.py b/genai_impact/tracers/openai_tracer.py index 97c4352..4bb1b9c 100644 --- a/genai_impact/tracers/openai_tracer.py +++ b/genai_impact/tracers/openai_tracer.py @@ -44,7 +44,7 @@ def openai_chat_wrapper( class OpenAIInstrumentor: - def __init__(self): + def __init__(self) -> None: self.wrapped_methods = [ { "module": "openai.resources.chat.completions", @@ -53,7 +53,7 @@ def __init__(self): }, ] - def instrument(self): + def instrument(self) -> None: for wrapper in self.wrapped_methods: wrap_function_wrapper( wrapper["module"],