makes POST /activation/all support pagination

apps/inference/neuronpedia_inference/endpoints/activation/all.py

@@ -3,7 +3,7 @@
 from typing import Any
 
 import torch
-from fastapi import APIRouter
+from fastapi import APIRouter, HTTPException
 from fastapi.responses import JSONResponse
 from neuronpedia_inference_client.models.activation_all_post200_response import (
     ActivationAllPost200Response,
@@ -87,6 +87,9 @@ async def activation_all(
         logger.info("Activations result processed successfully")
 
         return result
+    except HTTPException:
+        # Let FastAPI handle client errors we raise explicitly
+        raise
     except Exception as e:
         logger.error(f"Error processing activations: {str(e)}")
         import traceback
@@ -117,16 +120,28 @@ def process_activations(
             request.prompt, prepend_bos, max_layer
         )
         # ensure sort_by_token_indexes doesn't have any out of range indexes
-        # TODO: return a better error for this (currently returns a 500 error)
         for token_index in request.sort_by_token_indexes:
             if token_index >= len(str_tokens) or token_index < 0:
-                raise ValueError(
-                    f"Sort by token index {token_index} is out of range for the given prompt, which only has {len(str_tokens)} tokens."
+                raise HTTPException(
+                    status_code=400,
+                    detail=(
+                        f"Sort by token index {token_index} is out of range for the "
+                        f"given prompt, which only has {len(str_tokens)} tokens."
+                    ),
                 )
 
         source_activations = self._process_sources(request, cache)
 
         sorted_activations = self._sort_and_filter_results(source_activations, request)
+
+        # get pyright checks to pass
+        assert request.num_results is not None
+        assert request.offset is not None
+
+        start = request.offset
+        end = start + request.num_results
+        sorted_activations = sorted_activations[start:end]
+
         feature_activations = self._format_result_and_calculate_dfa(
             sorted_activations, cache, request
         )
@@ -262,7 +277,7 @@ def _sort_and_filter_results(
         if request.ignore_bos and Model.get_instance().cfg.default_prepend_bos:
             sorted_activations = sorted_activations[sorted_activations[:, 3] != 0]
 
-        return sorted_activations[: request.num_results].tolist()
+        return sorted_activations.tolist()
 
     def _format_result_and_calculate_dfa(
         self,

apps/inference/tests/unit/test_all.py

@@ -278,8 +278,13 @@ def test_process_activations_invalid_token_index(
     mock_config_class.get_instance.return_value = mock_config
     # Set invalid token index (beyond token length)
     sample_request.sort_by_token_indexes = [10]  # tokens only go 0-4
-    with pytest.raises(ValueError, match="Sort by token index .* is out of range"):
+    from fastapi import HTTPException
+
+    with pytest.raises(HTTPException) as exc_info:
         processor.process_activations(sample_request)
+    assert exc_info.value.status_code == 400
+    assert "Sort by token index" in exc_info.value.detail
+    assert "is out of range" in exc_info.value.detail
 
 
 @patch("neuronpedia_inference.endpoints.activation.all.get_activations_by_index")

packages/python/neuronpedia-inference-client/README.md

@@ -4,7 +4,7 @@ No description provided (generated by Openapi Generator https://github.com/opena
 This Python package is automatically generated by the [OpenAPI Generator](https://openapi-generator.tech) project:
 
 - API version: 1.1.0
-- Package version: 1.3.0
+- Package version: 1.4.0
 - Generator version: 7.12.0
 - Build package: org.openapitools.codegen.languages.PythonClientCodegen
 

packages/python/neuronpedia-inference-client/docs/ActivationAllPostRequest.md

@@ -14,6 +14,7 @@ Name | Type | Description | Notes
 **ignore_bos** | **bool** | Whether or not to include features whose highest activation value is the BOS token. | [default to True]
 **feature_filter** | **List[int]** | Optional. If specified, will only return features that match the indexes specified. Can only be used if we're testing just one SAE (\"selected_sources\" length = 1). | [optional] 
 **num_results** | **int** | Optional. The number of top features to return. | [optional] [default to 25]
+**offset** | **int** | Skip this many top-ranked results before returning the next page. | [optional] [default to 0]
 
 ## Example
 

packages/python/neuronpedia-inference-client/neuronpedia_inference_client/__init__.py

@@ -15,7 +15,7 @@
 """  # noqa: E501
 
 
-__version__ = "1.3.0"
+__version__ = "1.4.0"
 
 # import apis into sdk package
 from neuronpedia_inference_client.api.default_api import DefaultApi

packages/python/neuronpedia-inference-client/neuronpedia_inference_client/api_client.py

@@ -91,7 +91,7 @@ def __init__(
             self.default_headers[header_name] = header_value
         self.cookie = cookie
         # Set default User-Agent.
-        self.user_agent = 'OpenAPI-Generator/1.3.0/python'
+        self.user_agent = 'OpenAPI-Generator/1.4.0/python'
         self.client_side_validation = configuration.client_side_validation
 
     def __enter__(self):

packages/python/neuronpedia-inference-client/neuronpedia_inference_client/configuration.py

@@ -532,7 +532,7 @@ def to_debug_report(self) -> str:
                "OS: {env}\n"\
                "Python Version: {pyversion}\n"\
                "Version of the API: 1.1.0\n"\
-               "SDK Package Version: 1.3.0".\
+               "SDK Package Version: 1.4.0".\
                format(env=sys.platform, pyversion=sys.version)
 
     def get_host_settings(self) -> List[HostSetting]:

packages/python/neuronpedia-inference-client/neuronpedia_inference_client/models/activation_all_post_request.py

@@ -35,7 +35,8 @@ class ActivationAllPostRequest(BaseModel):
     ignore_bos: StrictBool = Field(description="Whether or not to include features whose highest activation value is the BOS token.")
     feature_filter: Optional[List[StrictInt]] = Field(default=None, description="Optional. If specified, will only return features that match the indexes specified. Can only be used if we're testing just one SAE (\"selected_sources\" length = 1).")
     num_results: Optional[StrictInt] = Field(default=25, description="Optional. The number of top features to return.")
-    __properties: ClassVar[List[str]] = ["prompt", "model", "source_set", "selected_sources", "sort_by_token_indexes", "ignore_bos", "feature_filter", "num_results"]
+    offset: Optional[StrictInt] = Field(default=0, description="Skip this many top-ranked results before returning the next page.")
+    __properties: ClassVar[List[str]] = ["prompt", "model", "source_set", "selected_sources", "sort_by_token_indexes", "ignore_bos", "feature_filter", "num_results", "offset"]
 
     model_config = ConfigDict(
         populate_by_name=True,
@@ -95,7 +96,8 @@ def from_dict(cls, obj: Optional[Dict[str, Any]]) -> Optional[Self]:
             "sort_by_token_indexes": obj.get("sort_by_token_indexes"),
             "ignore_bos": obj.get("ignore_bos") if obj.get("ignore_bos") is not None else True,
             "feature_filter": obj.get("feature_filter"),
-            "num_results": obj.get("num_results") if obj.get("num_results") is not None else 25
+            "num_results": obj.get("num_results") if obj.get("num_results") is not None else 25,
+            "offset": obj.get("offset") if obj.get("offset") is not None else 0
         })
         return _obj
 

packages/python/neuronpedia-inference-client/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "neuronpedia_inference_client"
-version = "1.3.0"
+version = "1.4.0"
 description = "Neuronpedia - Inference Server"
 authors = ["OpenAPI Generator Community <johnny@neuronpedia.org>"]
 license = "MIT"

packages/python/neuronpedia-inference-client/setup.py

@@ -22,7 +22,7 @@
 # prerequisite: setuptools
 # http://pypi.python.org/pypi/setuptools
 NAME = "neuronpedia-inference-client"
-VERSION = "1.3.0"
+VERSION = "1.4.0"
 PYTHON_REQUIRES = ">= 3.8"
 REQUIRES = [
     "urllib3 >= 1.25.3, < 3.0.0",

packages/python/neuronpedia-inference-client/test/test_activation_all_post_request.py

@@ -49,7 +49,8 @@ def make_instance(self, include_optional) -> ActivationAllPostRequest:
                 feature_filter = [
                     56
                     ],
-                num_results = 56
+                num_results = 56,
+                offset = 56
             )
         else:
             return ActivationAllPostRequest(

packages/typescript/neuronpedia-inference-client/README.md

@@ -1,4 +1,4 @@
-## neuronpedia-inference-client@1.3.0
+## neuronpedia-inference-client@1.4.0
 
 This generator creates TypeScript/JavaScript client that utilizes [Fetch API](https://fetch.spec.whatwg.org/). The generated Node module can be used in the following environments:
 
@@ -36,7 +36,7 @@ navigate to the folder of your consuming project and run one of the following co
 _published:_
 
 ```
-npm install neuronpedia-inference-client@1.3.0 --save
+npm install neuronpedia-inference-client@1.4.0 --save
 ```
 
 _unPublished (not recommended):_

packages/typescript/neuronpedia-inference-client/package.json

@@ -1,6 +1,6 @@
 {
   "name": "neuronpedia-inference-client",
-  "version": "1.3.1",
+  "version": "1.4.0",
   "description": "OpenAPI client for neuronpedia-inference-client",
   "author": "OpenAPI-Generator",
   "repository": {

packages/typescript/neuronpedia-inference-client/src/models/ActivationAllPostRequest.ts

@@ -67,6 +67,12 @@ export interface ActivationAllPostRequest {
      * @memberof ActivationAllPostRequest
      */
     numResults?: number;
+    /**
+     * Skip this many top-ranked results before returning the next page.
+     * @type {number}
+     * @memberof ActivationAllPostRequest
+     */
+    offset?: number;
 }
 
 /**
@@ -100,6 +106,7 @@ export function ActivationAllPostRequestFromJSONTyped(json: any, ignoreDiscrimin
         'ignoreBos': json['ignore_bos'],
         'featureFilter': json['feature_filter'] == null ? undefined : json['feature_filter'],
         'numResults': json['num_results'] == null ? undefined : json['num_results'],
+        'offset': json['offset'] == null ? undefined : json['offset'],
     };
 }
 
@@ -122,6 +129,7 @@ export function ActivationAllPostRequestToJSONTyped(value?: ActivationAllPostReq
         'ignore_bos': value['ignoreBos'],
         'feature_filter': value['featureFilter'],
         'num_results': value['numResults'],
+        'offset': value['offset'],
     };
 }
 

schemas/openapi/inference/paths/activation/all.yaml

@@ -49,6 +49,10 @@ post:
               type: integer
               default: 25
               description: Optional. The number of top features to return.
+            offset:
+              type: integer
+              default: 0
+              description: Skip this many top-ranked results before returning the next page.
   responses:
     "200":
       description: Successfully retrieved results