Identify AI agents in API requests (#49)

Author: zvikagart

src/codeocean/client.py

@@ -13,15 +13,32 @@
 
 @dataclass
 class CodeOcean:
+    """
+    Code Ocean API client.
+
+    This class provides a unified interface to access Code Ocean's API endpoints
+    for managing capsules, pipelines, computations, and data assets.
+
+    Fields:
+        domain: The Code Ocean domain URL (e.g., 'https://codeocean.acme.com')
+        token: Code Ocean API access token
+        retries: Optional retry configuration for failed HTTP requests. Can be an integer
+                (number of retries) or a urllib3.util.Retry object for advanced
+                retry configuration. Defaults to 0 (no retries)
+        agent_id: Optional agent identifier for tracking AI agent API usage on behalf of users
+    """
 
     domain: str
     token: str
     retries: Optional[Retry | int] = 0
+    agent_id: Optional[str] = None
 
     def __post_init__(self):
         self.session = BaseUrlSession(base_url=f"{self.domain}/api/v1/")
         self.session.auth = (self.token, "")
         self.session.headers.update({"Content-Type": "application/json"})
+        if self.agent_id:
+            self.session.headers.update({"Agent-Id": self.agent_id})
         self.session.hooks["response"] = [
             lambda response, *args, **kwargs: response.raise_for_status()
         ]

tests/test_client.py

@@ -0,0 +1,81 @@
+import unittest
+from unittest.mock import patch
+from urllib3.util import Retry
+
+from codeocean.client import CodeOcean
+
+
+class TestClient(unittest.TestCase):
+    """Test cases for the CodeOcean client class."""
+
+    def setUp(self):
+        """Set up test fixtures."""
+        self.test_domain = "https://codeocean.acme.com"
+        self.test_token = "test_token_123"
+        self.test_agent_id = "test_agent_456"
+
+    def test_basic_init(self):
+        """Test a basic client initialization."""
+        client = CodeOcean(
+            domain=self.test_domain,
+            token=self.test_token,
+        )
+
+        # Verify base URL is set correctly
+        self.assertEqual(client.session.base_url, f"{self.test_domain}/api/v1/")
+
+        # Verify auth is correctly set
+        self.assertEqual(client.session.auth, (self.test_token, ""))
+
+        # Verify the session headers are correctly configured
+        headers = client.session.headers
+        self.assertIn("Content-Type", headers)
+        self.assertEqual(headers["Content-Type"], "application/json")
+
+    @patch("codeocean.client.TCPKeepAliveAdapter")
+    def test_retry_configuration_types(self, mock_adapter):
+        """Test that both integer and Retry object work for retries parameter."""
+        # Test with integer
+        CodeOcean(
+            domain=self.test_domain,
+            token=self.test_token,
+            retries=5,
+        )
+
+        # Test with Retry object
+        retry_obj = Retry(total=3, backoff_factor=0.3)
+        CodeOcean(
+            domain=self.test_domain,
+            token=self.test_token,
+            retries=retry_obj,
+        )
+
+        # Assert both configurations work
+        self.assertEqual(mock_adapter.call_count, 2)
+        mock_adapter.assert_any_call(max_retries=5)
+        mock_adapter.assert_any_call(max_retries=retry_obj)
+
+    def test_agent_id_header_set_when_provided(self):
+        """Test that Agent-Id header is set when agent_id is provided."""
+        client = CodeOcean(
+            domain=self.test_domain,
+            token=self.test_token,
+            agent_id=self.test_agent_id,
+        )
+
+        # Verify the session headers are correctly configured
+        headers = client.session.headers
+        self.assertIn("Agent-Id", headers)
+        self.assertEqual(headers["Agent-Id"], self.test_agent_id)
+
+    def test_agent_id_header_not_set_when_none(self):
+        """Test that Agent-Id header is not set when agent_id is None."""
+        client = CodeOcean(
+            domain=self.test_domain,
+            token=self.test_token,
+            agent_id=None,
+        )
+
+        # Verify the session headers are correctly configured
+        headers = client.session.headers
+        self.assertNotIn("Agent-Id", headers)