feat(api): add metric filtering and retrieval for features

- Introduced new endpoints to retrieve available metrics and count features based on specified filters.
- Enhanced existing feature retrieval to support metric filters, allowing for more granular data access.
- Updated the `FeatureRecord` model to include an optional `metric` field for better data representation.
- Implemented frontend logic to manage metric filters and display filtering options in the UI.

Author: dest1n1s

server/app.py

@@ -1,4 +1,5 @@
 import io
+import json
 import os
 from functools import lru_cache
 from typing import Any, Optional
@@ -166,11 +167,60 @@ def get_image(dataset_name: str, context_idx: int, image_idx: int, shard_idx: in
     return Response(content=img_byte_arr.getvalue(), media_type="image/png")
 
 
+@app.get("/dictionaries/{name}/metrics")
+def get_available_metrics(name: str):
+    """Get available metrics for a dictionary.
+
+    Args:
+        name: Name of the dictionary/SAE
+
+    Returns:
+        List of available metric names
+    """
+    metrics = client.get_available_metrics(name, sae_series=sae_series)
+    return {"metrics": metrics}
+
+
+@app.get("/dictionaries/{name}/features/count")
+def count_features_with_filters(
+    name: str,
+    feature_analysis_name: str | None = None,
+    metric_filters: str | None = None,
+):
+    """Count features that match the given filters.
+
+    Args:
+        name: Name of the dictionary/SAE
+        feature_analysis_name: Optional analysis name
+        metric_filters: Optional JSON string of metric filters
+
+    Returns:
+        Count of features matching the filters
+    """
+    # Parse metric filters if provided
+    parsed_metric_filters = None
+    if metric_filters:
+        try:
+            parsed_metric_filters = json.loads(metric_filters)
+        except (json.JSONDecodeError, TypeError):
+            return Response(
+                content=f"Invalid metric_filters format: {metric_filters}",
+                status_code=400,
+            )
+
+    count = client.count_features_with_filters(
+        sae_name=name, sae_series=sae_series, name=feature_analysis_name, metric_filters=parsed_metric_filters
+    )
+
+    return {"count": count}
+
+
 @app.get("/dictionaries/{name}/features/{feature_index}")
 def get_feature(
     name: str,
     feature_index: str | int,
     feature_analysis_name: str | None = None,
+    metric_filters: str | None = None,
 ):
     # Parse feature_index if it's a string
     if isinstance(feature_index, str) and feature_index != "random":
@@ -182,9 +232,22 @@ def get_feature(
                 status_code=400,
             )
 
+    # Parse metric filters if provided
+    parsed_metric_filters = None
+    if metric_filters:
+        try:
+            parsed_metric_filters = json.loads(metric_filters)
+        except (json.JSONDecodeError, TypeError):
+            return Response(
+                content=f"Invalid metric_filters format: {metric_filters}",
+                status_code=400,
+            )
+
     # Get feature data
     feature = (
-        client.get_random_alive_feature(sae_name=name, sae_series=sae_series, name=feature_analysis_name)
+        client.get_random_alive_feature(
+            sae_name=name, sae_series=sae_series, name=feature_analysis_name, metric_filters=parsed_metric_filters
+        )
         if feature_index == "random"
         else client.get_feature(sae_name=name, sae_series=sae_series, index=feature_index)
     )

src/lm_saes/database.py

@@ -55,6 +55,7 @@ class FeatureRecord(BaseModel):
     index: int
     analyses: list[FeatureAnalysis] = []
     interpretation: Optional[dict[str, Any]] = None
+    metric: Optional[dict[str, float]] = None
 
 
 class AnalysisRecord(BaseModel):
@@ -235,14 +236,19 @@ def get_sae(self, sae_name: str, sae_series: str) -> Optional[SAERecord]:
         return SAERecord.model_validate(sae)
 
     def get_random_alive_feature(
-        self, sae_name: str, sae_series: str, name: str | None = None
+        self,
+        sae_name: str,
+        sae_series: str,
+        name: str | None = None,
+        metric_filters: Optional[dict[str, dict[str, float]]] = None,
     ) -> Optional[FeatureRecord]:
         """Get a random feature that has non-zero activation.
 
         Args:
             sae_name: Name of the SAE model
             sae_series: Series of the SAE model
             name: Name of the analysis
+            metric_filters: Optional dict of metric filters in the format {"metric_name": {"$gte": value, "$lte": value}}
 
         Returns:
             A random feature record with non-zero activation, or None if no such feature exists
@@ -251,14 +257,19 @@ def get_random_alive_feature(
         if name is not None:
             elem_match["name"] = name
 
+        match_filter: dict[str, Any] = {
+            "sae_name": sae_name,
+            "sae_series": sae_series,
+            "analyses": {"$elemMatch": elem_match},
+        }
+
+        # Add metric filters if provided
+        if metric_filters:
+            for metric_name, filters in metric_filters.items():
+                match_filter[f"metric.{metric_name}"] = filters
+
         pipeline = [
-            {
-                "$match": {
-                    "sae_name": sae_name,
-                    "sae_series": sae_series,
-                    "analyses": {"$elemMatch": elem_match},
-                }
-            },
+            {"$match": match_filter},
             {"$sample": {"size": 1}},
         ]
         feature = next(self.feature_collection.aggregate(pipeline), None)
@@ -590,3 +601,60 @@ def get_bookmark_count(self, sae_name: Optional[str] = None, sae_series: Optiona
             query["sae_series"] = sae_series
 
         return self.bookmark_collection.count_documents(query)
+
+    def get_available_metrics(self, sae_name: str, sae_series: str) -> list[str]:
+        """Get available metrics for an SAE by checking the first feature.
+
+        Args:
+            sae_name: Name of the SAE model
+            sae_series: Series of the SAE model
+
+        Returns:
+            List of available metric names
+        """
+        # Use projection to avoid loading large arrays from analyses[0].samplings
+        projection = {
+            "metric": 1,
+        }
+
+        first_feature = self.feature_collection.find_one({"sae_name": sae_name, "sae_series": sae_series}, projection)
+
+        if first_feature is None or first_feature.get("metric") is None:
+            return []
+
+        return list(first_feature["metric"].keys())
+
+    def count_features_with_filters(
+        self,
+        sae_name: str,
+        sae_series: str,
+        name: str | None = None,
+        metric_filters: Optional[dict[str, dict[str, float]]] = None,
+    ) -> int:
+        """Count features that match the given filters.
+
+        Args:
+            sae_name: Name of the SAE model
+            sae_series: Series of the SAE model
+            name: Name of the analysis
+            metric_filters: Optional dict of metric filters in the format {"metric_name": {"$gte": value, "$lte": value}}
+
+        Returns:
+            Number of features matching the filters
+        """
+        elem_match: dict[str, Any] = {"max_feature_acts": {"$gt": 0}}
+        if name is not None:
+            elem_match["name"] = name
+
+        match_filter: dict[str, Any] = {
+            "sae_name": sae_name,
+            "sae_series": sae_series,
+            "analyses": {"$elemMatch": elem_match},
+        }
+
+        # Add metric filters if provided
+        if metric_filters:
+            for metric_name, filters in metric_filters.items():
+                match_filter[f"metric.{metric_name}"] = filters
+
+        return self.feature_collection.count_documents(match_filter)