wisent 0.1.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of wisent might be problematic. Click here for more details.

Files changed (34) hide show
  1. wisent-0.1.1/LICENSE +21 -0
  2. wisent-0.1.1/PKG-INFO +142 -0
  3. wisent-0.1.1/README.md +112 -0
  4. wisent-0.1.1/pyproject.toml +22 -0
  5. wisent-0.1.1/setup.cfg +4 -0
  6. wisent-0.1.1/setup.py +31 -0
  7. wisent-0.1.1/tests/test_activations.py +350 -0
  8. wisent-0.1.1/tests/test_client.py +49 -0
  9. wisent-0.1.1/tests/test_control_vector.py +427 -0
  10. wisent-0.1.1/tests/test_examples.py +343 -0
  11. wisent-0.1.1/tests/test_inference.py +511 -0
  12. wisent-0.1.1/wisent/__init__.py +8 -0
  13. wisent-0.1.1/wisent/activations/__init__.py +9 -0
  14. wisent-0.1.1/wisent/activations/client.py +97 -0
  15. wisent-0.1.1/wisent/activations/extractor.py +251 -0
  16. wisent-0.1.1/wisent/activations/models.py +95 -0
  17. wisent-0.1.1/wisent/client.py +45 -0
  18. wisent-0.1.1/wisent/control_vector/__init__.py +9 -0
  19. wisent-0.1.1/wisent/control_vector/client.py +85 -0
  20. wisent-0.1.1/wisent/control_vector/manager.py +168 -0
  21. wisent-0.1.1/wisent/control_vector/models.py +70 -0
  22. wisent-0.1.1/wisent/inference/__init__.py +9 -0
  23. wisent-0.1.1/wisent/inference/client.py +103 -0
  24. wisent-0.1.1/wisent/inference/inferencer.py +250 -0
  25. wisent-0.1.1/wisent/inference/models.py +66 -0
  26. wisent-0.1.1/wisent/utils/__init__.py +3 -0
  27. wisent-0.1.1/wisent/utils/auth.py +30 -0
  28. wisent-0.1.1/wisent/utils/http.py +228 -0
  29. wisent-0.1.1/wisent/version.py +3 -0
  30. wisent-0.1.1/wisent.egg-info/PKG-INFO +142 -0
  31. wisent-0.1.1/wisent.egg-info/SOURCES.txt +32 -0
  32. wisent-0.1.1/wisent.egg-info/dependency_links.txt +1 -0
  33. wisent-0.1.1/wisent.egg-info/requires.txt +7 -0
  34. wisent-0.1.1/wisent.egg-info/top_level.txt +1 -0
wisent-0.1.1/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2024 Wisent Team
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
wisent-0.1.1/PKG-INFO ADDED
@@ -0,0 +1,142 @@
1
+ Metadata-Version: 2.2
2
+ Name: wisent
3
+ Version: 0.1.1
4
+ Summary: Client library for interacting with the Wisent backend services
5
+ Home-page: https://github.com/wisent-ai/wisent
6
+ Author: Wisent Team
7
+ Author-email: info@wisent.ai
8
+ Classifier: Programming Language :: Python :: 3
9
+ Classifier: License :: OSI Approved :: MIT License
10
+ Classifier: Operating System :: OS Independent
11
+ Requires-Python: >=3.8
12
+ Description-Content-Type: text/markdown
13
+ License-File: LICENSE
14
+ Requires-Dist: requests>=2.25.0
15
+ Requires-Dist: pydantic>=2.0.0
16
+ Requires-Dist: aiohttp>=3.8.0
17
+ Requires-Dist: torch>=2.0.0
18
+ Requires-Dist: numpy>=1.20.0
19
+ Requires-Dist: tqdm>=4.60.0
20
+ Requires-Dist: transformers>=4.30.0
21
+ Dynamic: author
22
+ Dynamic: author-email
23
+ Dynamic: classifier
24
+ Dynamic: description
25
+ Dynamic: description-content-type
26
+ Dynamic: home-page
27
+ Dynamic: requires-dist
28
+ Dynamic: requires-python
29
+ Dynamic: summary
30
+
31
+ # Wisent
32
+
33
+ A Python client library for interacting with the Wisent backend services.
34
+
35
+ ## Installation
36
+
37
+ ```bash
38
+ pip install wisent
39
+ ```
40
+
41
+ ## Features
42
+
43
+ - **Activations**: Extract and send model activations to the Wisent backend
44
+ - **Control Vectors**: Retrieve and apply control vectors for model inference
45
+ - **Inference**: Utilities for applying control vectors during inference
46
+ - **Utilities**: Helper functions for common tasks
47
+
48
+ ## Quick Start
49
+
50
+ ```python
51
+ from wisent import WisentClient
52
+
53
+ # Initialize the client
54
+ client = WisentClient(api_key="your_api_key", base_url="https://api.wisent.ai")
55
+
56
+ # Extract activations from a model and send to backend
57
+ activations = client.activations.extract(
58
+ model_name="mistralai/Mistral-7B-Instruct-v0.1",
59
+ prompt="Tell me about quantum computing",
60
+ layers=[0, 12, 24]
61
+ )
62
+
63
+ # Get a control vector from the backend
64
+ control_vector = client.control_vector.get(
65
+ name="helpful",
66
+ model="mistralai/Mistral-7B-Instruct-v0.1"
67
+ )
68
+
69
+ # Apply a control vector during inference
70
+ response = client.inference.generate_with_control(
71
+ model_name="mistralai/Mistral-7B-Instruct-v0.1",
72
+ prompt="Tell me about quantum computing",
73
+ control_vectors={"helpful": 0.8, "concise": 0.5}
74
+ )
75
+
76
+ # Print the response
77
+ print(response.text)
78
+ ```
79
+
80
+ ## Advanced Usage
81
+
82
+ ### Extracting Activations
83
+
84
+ ```python
85
+ from wisent.activations import ActivationExtractor
86
+
87
+ # Create an extractor
88
+ extractor = ActivationExtractor(
89
+ model_name="mistralai/Mistral-7B-Instruct-v0.1",
90
+ device="cuda"
91
+ )
92
+
93
+ # Extract activations for a specific prompt
94
+ activations = extractor.extract(
95
+ prompt="Tell me about quantum computing",
96
+ layers=[0, 12, 24],
97
+ tokens_to_extract=[-10, -1] # Extract last 10 tokens and final token
98
+ )
99
+
100
+ # Send activations to the Wisent backend
101
+ from wisent import WisentClient
102
+ client = WisentClient(api_key="your_api_key")
103
+ client.activations.upload(activations)
104
+ ```
105
+
106
+ ### Working with Control Vectors
107
+
108
+ ```python
109
+ from wisent.control_vector import ControlVectorManager
110
+
111
+ # Initialize the manager
112
+ manager = ControlVectorManager(api_key="your_api_key")
113
+
114
+ # Get a control vector
115
+ helpful_vector = manager.get("helpful", model="mistralai/Mistral-7B-Instruct-v0.1")
116
+
117
+ # Combine multiple vectors
118
+ combined_vector = manager.combine(
119
+ vectors={
120
+ "helpful": 0.8,
121
+ "concise": 0.5
122
+ },
123
+ model="mistralai/Mistral-7B-Instruct-v0.1"
124
+ )
125
+
126
+ # Apply during inference
127
+ from wisent.inference import Inferencer
128
+ inferencer = Inferencer(model_name="mistralai/Mistral-7B-Instruct-v0.1")
129
+ response = inferencer.generate(
130
+ prompt="Tell me about quantum computing",
131
+ control_vector=combined_vector,
132
+ method="caa" # Context-Aware Addition
133
+ )
134
+ ```
135
+
136
+ ## Documentation
137
+
138
+ For full documentation, visit [docs.wisent.ai](https://docs.wisent.ai).
139
+
140
+ ## License
141
+
142
+ MIT
wisent-0.1.1/README.md ADDED
@@ -0,0 +1,112 @@
1
+ # Wisent
2
+
3
+ A Python client library for interacting with the Wisent backend services.
4
+
5
+ ## Installation
6
+
7
+ ```bash
8
+ pip install wisent
9
+ ```
10
+
11
+ ## Features
12
+
13
+ - **Activations**: Extract and send model activations to the Wisent backend
14
+ - **Control Vectors**: Retrieve and apply control vectors for model inference
15
+ - **Inference**: Utilities for applying control vectors during inference
16
+ - **Utilities**: Helper functions for common tasks
17
+
18
+ ## Quick Start
19
+
20
+ ```python
21
+ from wisent import WisentClient
22
+
23
+ # Initialize the client
24
+ client = WisentClient(api_key="your_api_key", base_url="https://api.wisent.ai")
25
+
26
+ # Extract activations from a model and send to backend
27
+ activations = client.activations.extract(
28
+ model_name="mistralai/Mistral-7B-Instruct-v0.1",
29
+ prompt="Tell me about quantum computing",
30
+ layers=[0, 12, 24]
31
+ )
32
+
33
+ # Get a control vector from the backend
34
+ control_vector = client.control_vector.get(
35
+ name="helpful",
36
+ model="mistralai/Mistral-7B-Instruct-v0.1"
37
+ )
38
+
39
+ # Apply a control vector during inference
40
+ response = client.inference.generate_with_control(
41
+ model_name="mistralai/Mistral-7B-Instruct-v0.1",
42
+ prompt="Tell me about quantum computing",
43
+ control_vectors={"helpful": 0.8, "concise": 0.5}
44
+ )
45
+
46
+ # Print the response
47
+ print(response.text)
48
+ ```
49
+
50
+ ## Advanced Usage
51
+
52
+ ### Extracting Activations
53
+
54
+ ```python
55
+ from wisent.activations import ActivationExtractor
56
+
57
+ # Create an extractor
58
+ extractor = ActivationExtractor(
59
+ model_name="mistralai/Mistral-7B-Instruct-v0.1",
60
+ device="cuda"
61
+ )
62
+
63
+ # Extract activations for a specific prompt
64
+ activations = extractor.extract(
65
+ prompt="Tell me about quantum computing",
66
+ layers=[0, 12, 24],
67
+ tokens_to_extract=[-10, -1] # Extract last 10 tokens and final token
68
+ )
69
+
70
+ # Send activations to the Wisent backend
71
+ from wisent import WisentClient
72
+ client = WisentClient(api_key="your_api_key")
73
+ client.activations.upload(activations)
74
+ ```
75
+
76
+ ### Working with Control Vectors
77
+
78
+ ```python
79
+ from wisent.control_vector import ControlVectorManager
80
+
81
+ # Initialize the manager
82
+ manager = ControlVectorManager(api_key="your_api_key")
83
+
84
+ # Get a control vector
85
+ helpful_vector = manager.get("helpful", model="mistralai/Mistral-7B-Instruct-v0.1")
86
+
87
+ # Combine multiple vectors
88
+ combined_vector = manager.combine(
89
+ vectors={
90
+ "helpful": 0.8,
91
+ "concise": 0.5
92
+ },
93
+ model="mistralai/Mistral-7B-Instruct-v0.1"
94
+ )
95
+
96
+ # Apply during inference
97
+ from wisent.inference import Inferencer
98
+ inferencer = Inferencer(model_name="mistralai/Mistral-7B-Instruct-v0.1")
99
+ response = inferencer.generate(
100
+ prompt="Tell me about quantum computing",
101
+ control_vector=combined_vector,
102
+ method="caa" # Context-Aware Addition
103
+ )
104
+ ```
105
+
106
+ ## Documentation
107
+
108
+ For full documentation, visit [docs.wisent.ai](https://docs.wisent.ai).
109
+
110
+ ## License
111
+
112
+ MIT
@@ -0,0 +1,22 @@
1
+ [build-system]
2
+ requires = ["setuptools>=42", "wheel"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [tool.black]
6
+ line-length = 88
7
+ target-version = ['py38']
8
+
9
+ [tool.isort]
10
+ profile = "black"
11
+ multi_line_output = 3
12
+
13
+ [tool.mypy]
14
+ python_version = "3.8"
15
+ warn_return_any = true
16
+ warn_unused_configs = true
17
+ disallow_untyped_defs = true
18
+ disallow_incomplete_defs = true
19
+
20
+ [tool.pytest.ini_options]
21
+ testpaths = ["tests"]
22
+ python_files = "test_*.py"
wisent-0.1.1/setup.cfg ADDED
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
wisent-0.1.1/setup.py ADDED
@@ -0,0 +1,31 @@
1
+ from setuptools import setup, find_packages
2
+
3
+ with open("README.md", "r", encoding="utf-8") as fh:
4
+ long_description = fh.read()
5
+
6
+ setup(
7
+ name="wisent",
8
+ version="0.1.1",
9
+ author="Wisent Team",
10
+ author_email="info@wisent.ai",
11
+ description="Client library for interacting with the Wisent backend services",
12
+ long_description=long_description,
13
+ long_description_content_type="text/markdown",
14
+ url="https://github.com/wisent-ai/wisent",
15
+ packages=find_packages(),
16
+ classifiers=[
17
+ "Programming Language :: Python :: 3",
18
+ "License :: OSI Approved :: MIT License",
19
+ "Operating System :: OS Independent",
20
+ ],
21
+ python_requires=">=3.8",
22
+ install_requires=[
23
+ "requests>=2.25.0",
24
+ "pydantic>=2.0.0",
25
+ "aiohttp>=3.8.0",
26
+ "torch>=2.0.0",
27
+ "numpy>=1.20.0",
28
+ "tqdm>=4.60.0",
29
+ "transformers>=4.30.0",
30
+ ],
31
+ )
@@ -0,0 +1,350 @@
1
+ """
2
+ Tests for the activations module.
3
+ """
4
+
5
+ import unittest
6
+ from unittest.mock import MagicMock, patch
7
+
8
+ import numpy as np
9
+ import torch
10
+
11
+ from wisent.activations import Activation, ActivationBatch, ActivationExtractor, ActivationsClient
12
+ from wisent.utils.auth import AuthManager
13
+
14
+
15
+ class TestActivation(unittest.TestCase):
16
+ """Tests for the Activation class."""
17
+
18
+ def test_init(self):
19
+ """Test initialization of an Activation."""
20
+ # Test with list values
21
+ values = [0.1, 0.2, 0.3]
22
+ activation = Activation(
23
+ model_name="test_model",
24
+ layer=0,
25
+ token_index=1,
26
+ values=values,
27
+ token_str="test"
28
+ )
29
+
30
+ self.assertEqual(activation.model_name, "test_model")
31
+ self.assertEqual(activation.layer, 0)
32
+ self.assertEqual(activation.token_index, 1)
33
+ self.assertEqual(activation.values, values)
34
+ self.assertEqual(activation.token_str, "test")
35
+
36
+ # Test with numpy array
37
+ values = np.array([0.1, 0.2, 0.3])
38
+ activation = Activation(
39
+ model_name="test_model",
40
+ layer=0,
41
+ token_index=1,
42
+ values=values
43
+ )
44
+
45
+ self.assertTrue(np.array_equal(activation.values, values))
46
+ self.assertIsNone(activation.token_str)
47
+
48
+ # Test with torch tensor
49
+ values = torch.tensor([0.1, 0.2, 0.3])
50
+ activation = Activation(
51
+ model_name="test_model",
52
+ layer=0,
53
+ token_index=1,
54
+ values=values
55
+ )
56
+
57
+ self.assertTrue(torch.equal(activation.values, values))
58
+
59
+ def test_to_dict(self):
60
+ """Test conversion to dictionary."""
61
+ # Test with list values
62
+ values = [0.1, 0.2, 0.3]
63
+ activation = Activation(
64
+ model_name="test_model",
65
+ layer=0,
66
+ token_index=1,
67
+ values=values,
68
+ token_str="test"
69
+ )
70
+
71
+ expected = {
72
+ "model_name": "test_model",
73
+ "layer": 0,
74
+ "token_index": 1,
75
+ "values": values,
76
+ "token_str": "test",
77
+ }
78
+
79
+ self.assertEqual(activation.to_dict(), expected)
80
+
81
+ # Test with numpy array
82
+ values = np.array([0.1, 0.2, 0.3])
83
+ activation = Activation(
84
+ model_name="test_model",
85
+ layer=0,
86
+ token_index=1,
87
+ values=values
88
+ )
89
+
90
+ expected = {
91
+ "model_name": "test_model",
92
+ "layer": 0,
93
+ "token_index": 1,
94
+ "values": values.tolist(),
95
+ "token_str": None,
96
+ }
97
+
98
+ self.assertEqual(activation.to_dict(), expected)
99
+
100
+ # Test with torch tensor
101
+ values = torch.tensor([0.1, 0.2, 0.3])
102
+ activation = Activation(
103
+ model_name="test_model",
104
+ layer=0,
105
+ token_index=1,
106
+ values=values
107
+ )
108
+
109
+ expected = {
110
+ "model_name": "test_model",
111
+ "layer": 0,
112
+ "token_index": 1,
113
+ "values": values.detach().cpu().numpy().tolist(),
114
+ "token_str": None,
115
+ }
116
+
117
+ self.assertEqual(activation.to_dict(), expected)
118
+
119
+
120
+ class TestActivationBatch(unittest.TestCase):
121
+ """Tests for the ActivationBatch class."""
122
+
123
+ def setUp(self):
124
+ """Set up test fixtures."""
125
+ self.activation1 = Activation(
126
+ model_name="test_model",
127
+ layer=0,
128
+ token_index=1,
129
+ values=[0.1, 0.2, 0.3],
130
+ token_str="test1"
131
+ )
132
+
133
+ self.activation2 = Activation(
134
+ model_name="test_model",
135
+ layer=1,
136
+ token_index=2,
137
+ values=[0.4, 0.5, 0.6],
138
+ token_str="test2"
139
+ )
140
+
141
+ def test_init(self):
142
+ """Test initialization of an ActivationBatch."""
143
+ batch = ActivationBatch(
144
+ model_name="test_model",
145
+ prompt="test prompt",
146
+ activations=[self.activation1, self.activation2],
147
+ metadata={"test": "metadata"}
148
+ )
149
+
150
+ self.assertEqual(batch.model_name, "test_model")
151
+ self.assertEqual(batch.prompt, "test prompt")
152
+ self.assertEqual(batch.activations, [self.activation1, self.activation2])
153
+ self.assertEqual(batch.metadata, {"test": "metadata"})
154
+
155
+ # Test with default metadata
156
+ batch = ActivationBatch(
157
+ model_name="test_model",
158
+ prompt="test prompt",
159
+ activations=[self.activation1, self.activation2]
160
+ )
161
+
162
+ self.assertEqual(batch.metadata, {})
163
+
164
+ def test_to_dict(self):
165
+ """Test conversion to dictionary."""
166
+ batch = ActivationBatch(
167
+ model_name="test_model",
168
+ prompt="test prompt",
169
+ activations=[self.activation1, self.activation2],
170
+ metadata={"test": "metadata"}
171
+ )
172
+
173
+ expected = {
174
+ "model_name": "test_model",
175
+ "prompt": "test prompt",
176
+ "activations": [
177
+ self.activation1.to_dict(),
178
+ self.activation2.to_dict()
179
+ ],
180
+ "metadata": {"test": "metadata"},
181
+ }
182
+
183
+ self.assertEqual(batch.to_dict(), expected)
184
+
185
+
186
+ class TestActivationsClient(unittest.TestCase):
187
+ """Tests for the ActivationsClient class."""
188
+
189
+ def setUp(self):
190
+ """Set up test fixtures."""
191
+ self.auth_manager = AuthManager("test_api_key")
192
+ self.base_url = "https://test.api.wisent.ai"
193
+ self.client = ActivationsClient(self.auth_manager, self.base_url)
194
+
195
+ @patch("wisent.activations.client.ActivationExtractor")
196
+ def test_extract(self, mock_extractor_class):
197
+ """Test extraction of activations."""
198
+ # Mock the extractor
199
+ mock_extractor = MagicMock()
200
+ mock_extractor_class.return_value = mock_extractor
201
+
202
+ # Mock the extract method
203
+ mock_batch = MagicMock()
204
+ mock_extractor.extract.return_value = mock_batch
205
+
206
+ # Call the method
207
+ result = self.client.extract(
208
+ model_name="test_model",
209
+ prompt="test prompt",
210
+ layers=[0, 1],
211
+ tokens_to_extract=[1, 2],
212
+ device="cpu"
213
+ )
214
+
215
+ # Check that the extractor was created correctly
216
+ mock_extractor_class.assert_called_once_with("test_model", device="cpu")
217
+
218
+ # Check that extract was called correctly
219
+ mock_extractor.extract.assert_called_once_with("test prompt", [0, 1], [1, 2])
220
+
221
+ # Check the result
222
+ self.assertEqual(result, mock_batch)
223
+
224
+ @patch("wisent.activations.client.HTTPClient")
225
+ def test_upload(self, mock_http_client_class):
226
+ """Test uploading of activations."""
227
+ # Mock the HTTP client
228
+ mock_http_client = MagicMock()
229
+ mock_http_client_class.return_value = mock_http_client
230
+
231
+ # Mock the post method
232
+ mock_response = {"id": "test_batch_id"}
233
+ mock_http_client.post.return_value = mock_response
234
+
235
+ # Create a client with the mocked HTTP client
236
+ client = ActivationsClient(self.auth_manager, self.base_url)
237
+
238
+ # Create a batch to upload
239
+ batch = ActivationBatch(
240
+ model_name="test_model",
241
+ prompt="test prompt",
242
+ activations=[
243
+ Activation(
244
+ model_name="test_model",
245
+ layer=0,
246
+ token_index=1,
247
+ values=[0.1, 0.2, 0.3],
248
+ token_str="test"
249
+ )
250
+ ]
251
+ )
252
+
253
+ # Call the method
254
+ result = client.upload(batch)
255
+
256
+ # Check that post was called correctly
257
+ mock_http_client.post.assert_called_once_with(
258
+ "/activations/upload",
259
+ json_data=batch.to_dict()
260
+ )
261
+
262
+ # Check the result
263
+ self.assertEqual(result, mock_response)
264
+
265
+ @patch("wisent.activations.client.HTTPClient")
266
+ def test_get(self, mock_http_client_class):
267
+ """Test getting activations by ID."""
268
+ # Mock the HTTP client
269
+ mock_http_client = MagicMock()
270
+ mock_http_client_class.return_value = mock_http_client
271
+
272
+ # Mock the get method
273
+ mock_response = {
274
+ "model_name": "test_model",
275
+ "prompt": "test prompt",
276
+ "activations": [
277
+ {
278
+ "model_name": "test_model",
279
+ "layer": 0,
280
+ "token_index": 1,
281
+ "values": [0.1, 0.2, 0.3],
282
+ "token_str": "test"
283
+ }
284
+ ],
285
+ "metadata": {"test": "metadata"}
286
+ }
287
+ mock_http_client.get.return_value = mock_response
288
+
289
+ # Create a client with the mocked HTTP client
290
+ client = ActivationsClient(self.auth_manager, self.base_url)
291
+
292
+ # Call the method
293
+ result = client.get("test_batch_id")
294
+
295
+ # Check that get was called correctly
296
+ mock_http_client.get.assert_called_once_with("/activations/test_batch_id")
297
+
298
+ # Check the result
299
+ self.assertEqual(result.model_name, "test_model")
300
+ self.assertEqual(result.prompt, "test prompt")
301
+ self.assertEqual(len(result.activations), 1)
302
+ self.assertEqual(result.activations[0].model_name, "test_model")
303
+ self.assertEqual(result.activations[0].layer, 0)
304
+ self.assertEqual(result.activations[0].token_index, 1)
305
+ self.assertEqual(result.activations[0].values, [0.1, 0.2, 0.3])
306
+ self.assertEqual(result.activations[0].token_str, "test")
307
+ self.assertEqual(result.metadata, {"test": "metadata"})
308
+
309
+ @patch("wisent.activations.client.HTTPClient")
310
+ def test_list(self, mock_http_client_class):
311
+ """Test listing activations."""
312
+ # Mock the HTTP client
313
+ mock_http_client = MagicMock()
314
+ mock_http_client_class.return_value = mock_http_client
315
+
316
+ # Mock the get method
317
+ mock_response = [
318
+ {"id": "batch1", "model_name": "test_model"},
319
+ {"id": "batch2", "model_name": "test_model"}
320
+ ]
321
+ mock_http_client.get.return_value = mock_response
322
+
323
+ # Create a client with the mocked HTTP client
324
+ client = ActivationsClient(self.auth_manager, self.base_url)
325
+
326
+ # Call the method
327
+ result = client.list(model_name="test_model", limit=10, offset=0)
328
+
329
+ # Check that get was called correctly
330
+ mock_http_client.get.assert_called_once_with(
331
+ "/activations",
332
+ params={"model_name": "test_model", "limit": 10, "offset": 0}
333
+ )
334
+
335
+ # Check the result
336
+ self.assertEqual(result, mock_response)
337
+
338
+ # Test without model_name
339
+ mock_http_client.get.reset_mock()
340
+ result = client.list(limit=10, offset=0)
341
+
342
+ # Check that get was called correctly
343
+ mock_http_client.get.assert_called_once_with(
344
+ "/activations",
345
+ params={"limit": 10, "offset": 0}
346
+ )
347
+
348
+
349
+ if __name__ == "__main__":
350
+ unittest.main()