feat: implement energy estimation

genai_impact/__init__.py

@@ -0,0 +1 @@
+from client_wrapper import OpenAI

genai_impact/__main__.py

@@ -1,12 +1,12 @@
-from genai_impact.gen_ai_impacts import GenAIImpacts
+from genai_impact.client_wrapper import OpenAI
 
-client = GenAIImpacts()
+client = OpenAI()
 
 response = client.chat.completions.create(
     model="gpt-3.5-turbo",
-    messages=[...]  # Your messages here
+    messages=[
+        {"role": "user", "content": "Hello World!"}
+    ]
 )
 
-impacts = client.get_impacts(response)
-print(impacts)
-# Outputs: Impacts(energy=42, energy_unit='kWh', gwp=15, gwp_unit='kgCO2eq')
+print(response.impacts)

genai_impact/client_wrapper.py

@@ -0,0 +1,43 @@
+from dataclasses import asdict
+from typing import Optional
+
+from openai import OpenAI as _OpenAI
+from wrapt import wrap_function_wrapper
+
+from genai_impact.compute_impacts import MODEL_SIZES, compute_llm_impact
+
+
+def chat_wrapper(wrapped, instance, args, kwargs):
+    response = wrapped(*args, **kwargs)
+    model_size = MODEL_SIZES.get(response.model)
+    output_tokens = response.usage.completion_tokens
+    impacts = compute_llm_impact(model_parameter_count=model_size, output_token_count=output_tokens)
+    response.impacts = asdict(impacts)
+    return response
+
+
+class OpenAI:
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        base_url: Optional[str] = None,
+        **kwargs
+    ):
+        self.__client = _OpenAI(
+            api_key=api_key,
+            base_url=base_url,
+            **kwargs
+        )
+
+    wrap_function_wrapper(
+        "openai.resources.chat.completions",
+        "Completions.create",
+        chat_wrapper
+    )
+
+    def __getattr__(self, name):
+        """
+        Redirect attribute access to the underlying openai client if the attribute
+        is not defined in this class.
+        """
+        return getattr(self.__client, name)

genai_impact/compute_impacts.py

@@ -0,0 +1,39 @@
+from dataclasses import dataclass
+
+
+ENERGY_PROFILE = 1.17e-4
+
+MODEL_SIZES = {
+    "gpt-4-0125-preview": None,
+    "gpt-4-turbo-preview": None,
+    "gpt-4-1106-preview": None,
+    "gpt-4-vision-preview": None,
+    "gpt-4": 200,
+    "gpt-4-0314": 200,
+    "gpt-4-0613": 200,
+    "gpt-4-32k": 200,
+    "gpt-4-32k-0314": 200,
+    "gpt-4-32k-0613": 200,
+    "gpt-3.5-turbo": 20,
+    "gpt-3.5-turbo-16k": 20,
+    "gpt-3.5-turbo-0301": 20,
+    "gpt-3.5-turbo-0613": 20,
+    "gpt-3.5-turbo-1106": 20,
+    "gpt-3.5-turbo-0125": 20,
+    "gpt-3.5-turbo-16k-0613": 20,
+}
+
+
+@dataclass
+class Impacts:
+    energy: float
+    energy_unit: str = 'Wh'
+
+
+def compute_llm_impact(
+    model_parameter_count: float,
+    output_token_count: int,
+) -> Impacts:
+    return Impacts(
+        energy=ENERGY_PROFILE * model_parameter_count * output_token_count
+    )

pyproject.toml

@@ -13,6 +13,7 @@ readme = "README.md"
 [tool.poetry.dependencies]
 python = ">=3.8,<4"
 openai = "^1.12.0"
+wrapt = "^1.16.0"
 
 
 [tool.poetry.group.dev.dependencies]