From d63d789e533b9cbc1d4b14b00fe098cf20bc4cdf Mon Sep 17 00:00:00 2001
From: Kelsey Wong <kelsey.wong@cleanlab.ai>
Date: Mon, 3 Mar 2025 11:49:24 -0800
Subject: [PATCH] include cleanlab-codex version in headers for codex backend
 requests

---
 src/cleanlab_codex/utils/analytics.py |  8 +++++++-
 tests/test_project.py                 |  4 +++-
 tests/utils/test_analytics.py         | 22 +++++++++++++++++-----
 3 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/src/cleanlab_codex/utils/analytics.py b/src/cleanlab_codex/utils/analytics.py
index ed4157a..8139cc0 100644
--- a/src/cleanlab_codex/utils/analytics.py
+++ b/src/cleanlab_codex/utils/analytics.py
@@ -1,13 +1,19 @@
 from __future__ import annotations
 
+from cleanlab_codex.__about__ import __version__ as package_version
+
 
 class AnalyticsMetadata:
-    def __init__(self, *, integration_type: str, source: str = "codex-python-sdk"):
+    def __init__(
+        self, *, integration_type: str, package_version: str = package_version, source: str = "cleanlab-codex-python"
+    ):
         self._integration_type = integration_type
+        self._package_version = package_version
         self._source = source
 
     def to_headers(self) -> dict[str, str]:
         return {
             "X-Integration-Type": self._integration_type,
+            "X-Client-Library-Version": self._package_version,
             "X-Source": self._source,
         }
diff --git a/tests/test_project.py b/tests/test_project.py
index 1458183..e5fcf8e 100644
--- a/tests/test_project.py
+++ b/tests/test_project.py
@@ -8,6 +8,7 @@
 from codex.types.projects.access_key_retrieve_project_id_response import AccessKeyRetrieveProjectIDResponse
 from codex.types.projects.entry import Entry as SDKEntry
 
+from cleanlab_codex.__about__ import __version__ as package_version
 from cleanlab_codex.project import MissingProjectError, Project
 from cleanlab_codex.types.entry import EntryCreate
 
@@ -137,7 +138,8 @@ def test_query_read_only(mock_client_from_access_key: MagicMock) -> None:
     res = project.query("What is the capital of France?", read_only=True)
     expected_headers = {
         "X-Integration-Type": "backup",
-        "X-Source": "codex-python-sdk",
+        "X-Source": "cleanlab-codex-python",
+        "X-Client-Library-Version": package_version,
     }
     mock_client_from_access_key.projects.entries.query.assert_called_once_with(
         FAKE_PROJECT_ID, question="What is the capital of France?", extra_headers=expected_headers
diff --git a/tests/utils/test_analytics.py b/tests/utils/test_analytics.py
index ebeb030..2b1f8b8 100644
--- a/tests/utils/test_analytics.py
+++ b/tests/utils/test_analytics.py
@@ -1,16 +1,28 @@
+from __future__ import annotations
+
 import pytest
 
+from cleanlab_codex.__about__ import __version__ as package_version
 from cleanlab_codex.utils.analytics import AnalyticsMetadata
 
 
-def test_analytics_metadata_to_headers_uses_default_source() -> None:
+def test_analytics_metadata_to_headers_uses_defaults() -> None:
     metadata = AnalyticsMetadata(integration_type="backup")
-    assert metadata.to_headers() == {"X-Integration-Type": "backup", "X-Source": "codex-python-sdk"}
+
+    assert metadata.to_headers() == {
+        "X-Integration-Type": "backup",
+        "X-Source": "cleanlab-codex-python",
+        "X-Client-Library-Version": package_version,
+    }
 
 
-def test_analytics_metadata_to_headers_uses_integration_type() -> None:
-    metadata = AnalyticsMetadata(integration_type="tool-call", source="test")
-    assert metadata.to_headers() == {"X-Integration-Type": "tool-call", "X-Source": "test"}
+def test_analytics_metadata_to_headers_uses_custom_values() -> None:
+    metadata = AnalyticsMetadata(integration_type="tool-call", source="test", package_version="2.0.0")
+    assert metadata.to_headers() == {
+        "X-Integration-Type": "tool-call",
+        "X-Source": "test",
+        "X-Client-Library-Version": "2.0.0",
+    }
 
 
 def test_analytics_metadata_requires_integration_type() -> None: