feat: expose complexity consumption

src/kili/core/graphql/graphql_client.py

@@ -3,6 +3,7 @@
 import logging
 import os
 import threading
+import time
 from pathlib import Path
 from typing import Any, Dict, Optional, Union
 from urllib.parse import urlparse
@@ -80,6 +81,8 @@ def __init__(
         self.http_client = http_client
         self.verify = verify
         self.enable_schema_caching = enable_schema_caching
+        self.created_at = time.time()
+        self.complexity_consumed = 0
         self.graphql_schema_cache_dir = (
             Path(graphql_schema_cache_dir) if graphql_schema_cache_dir else None
         )
@@ -309,7 +312,7 @@ def _raw_execute(
     ) -> Dict[str, Any]:
         _limiter.try_acquire("GraphQLClient.execute")
         with _execute_lock:
-            return self._gql_client.execute(
+            res = self._gql_client.execute(
                 document=document,
                 variable_values=variables,
                 extra_args={
@@ -320,3 +323,9 @@ def _raw_execute(
                 },
                 **kwargs,
             )
+            transport = self._gql_client.transport
+            if transport:
+                headers = transport.response_headers  # pyright: ignore[reportGeneralTypeIssues]
+                returned_complexity = int(headers.get("x-complexity", 0)) if headers else 0
+                self.complexity_consumed += returned_complexity
+            return res

tests/e2e/test_e2e_models.py

@@ -44,7 +44,7 @@ def test_create_and_manage_project_and_model_resources(kili: Kili):
     project_id = project["id"]
 
     model_data = {
-        "credentials": {"api_key": "***", "endpoint": "your_open_ai_endpoint"},
+        "credentials": {"api_key": "***", "endpoint": "http://your_open_ai_endpoint.com"},
         "name": MODEL_NAME,
         "type": "OPEN_AI_SDK",
     }

tests/integration/entrypoints/client/test_client.py

@@ -140,3 +140,28 @@ def test_given_an_api_key_away_to_expiration_when_I_check_expiry_of_key_is_not_c
         warnings.simplefilter("error")  # checks that no warning is raised
         # When
         _ = Kili()
+
+
+@patch.dict(os.environ, {"KILI_API_KEY": "fake_key"})
+def test_complexity_increases_with_calls(
+    mocker: pytest_mock.MockerFixture,
+):
+    graphql_mock = mocker.MagicMock()
+    graphql_mock.execute.return_value = {"data": 1}
+    graphql_mock.transport.response_headers = {"x-complexity": "125"}
+
+    # Given
+    mocker.patch("kili.client.is_api_key_valid", return_value=True)
+    mocker.patch.object(ApiKeyUseCases, "check_expiry_of_key_is_close")
+    mocker.patch(
+        "kili.core.graphql.graphql_client.GraphQLClient._initizalize_graphql_client",
+        return_value=graphql_mock,
+    )
+    kili = Kili()
+
+    assert kili.graphql_client.complexity_consumed == 0
+    initialized_date = kili.graphql_client.created_at
+
+    kili.count_projects()
+    assert kili.graphql_client.complexity_consumed == 125
+    assert kili.graphql_client.created_at == initialized_date