Feature/filtering models and benchmarks

src/atlas/resources/benchmarks/benchmarks.py

@@ -4,7 +4,12 @@
 
 import httpx
 
-from ...models import Benchmark, CustomBenchmark, PublicBenchmark, BenchmarksResponse
+from ...models import (
+    Benchmark,
+    CustomBenchmark,
+    PublicBenchmark,
+    BenchmarksResponse,
+)
 from ..._resource import SyncAPIResource, AsyncAPIResource
 from ..._constants import DEFAULT_TIMEOUT
 
@@ -16,13 +21,16 @@ def get(
         timeout: float | httpx.Timeout | None = DEFAULT_TIMEOUT,
         type: Literal["custom", "public"] | None = None,
         name: Optional[str] = None,
+        key: Optional[str] = None,
     ) -> Optional[List[Benchmark]]:
         base_url = f"/organizations/{self._client.organization_id}/projects/{self._client.project_id}/benchmarks"
 
         def fetch(bench_type: str) -> BenchmarksResponse | None:
             params = {"type": bench_type}
             if name:
-                params["query"] = name
+                params["name"] = name
+            if key:
+                params["key"] = key
 
             resp = self._get(
                 base_url,
@@ -53,6 +61,45 @@ def cast_benchmark(b: Benchmark, bench_type: str) -> Benchmark:
 
         return benchmarks
 
+    def get_by_id(self, id: str, *, timeout: float | httpx.Timeout | None = DEFAULT_TIMEOUT) -> Optional[Benchmark]:
+        base_url = f"/organizations/{self._client.organization_id}/projects/{self._client.project_id}/benchmarks/{id}"
+
+        resp = self._get(
+            base_url,
+            timeout=timeout,
+            cast_to=dict,
+        )
+
+        if not isinstance(resp, dict):
+            return None
+
+        benchmark = resp.get("data")
+        if not isinstance(benchmark, dict):
+            return None
+
+        # Detect type dynamically: presence of "organization_id" means custom
+        if "organization_id" in benchmark:
+            return CustomBenchmark(**benchmark)
+        else:
+            return PublicBenchmark(**benchmark)
+
+    def get_by_key(
+        self,
+        key: str,
+        *,
+        timeout: float | httpx.Timeout | None = DEFAULT_TIMEOUT,
+    ) -> Optional[Benchmark]:
+        """Fetch a single benchmark by its unique key."""
+        benchmarks = self.get(timeout=timeout, key=key)
+
+        if not benchmarks:
+            return None
+
+        for benchmark in benchmarks:
+            if benchmark.key == key:
+                return benchmark
+        return None
+
 
 class AsyncBenchmarks(AsyncAPIResource):
     async def get(
@@ -61,13 +108,16 @@ async def get(
         timeout: float | httpx.Timeout | None = DEFAULT_TIMEOUT,
         type: Literal["custom", "public"] | None = None,
         name: Optional[str] = None,
+        key: Optional[str] = None,
     ) -> Optional[List[Benchmark]]:
         base_url = f"/organizations/{self._client.organization_id}/projects/{self._client.project_id}/benchmarks"
 
         async def fetch(bench_type: str) -> Optional[BenchmarksResponse]:
             params = {"type": bench_type}
             if name:
-                params["query"] = name
+                params["name"] = name
+            if key:
+                params["key"] = key
 
             resp = await self._get(
                 base_url,
@@ -97,3 +147,44 @@ def cast_benchmark(b: Benchmark, bench_type: str) -> Benchmark:
                 benchmarks.extend([cast_benchmark(b, type) for b in resp.data.benchmarks])
 
         return benchmarks
+
+    async def get_by_id(
+        self, id: str, *, timeout: float | httpx.Timeout | None = DEFAULT_TIMEOUT
+    ) -> Optional[Benchmark]:
+        base_url = f"/organizations/{self._client.organization_id}/projects/{self._client.project_id}/benchmarks/{id}"
+
+        resp = await self._get(
+            base_url,
+            timeout=timeout,
+            cast_to=dict,
+        )
+
+        if not isinstance(resp, dict):
+            return None
+
+        benchmark = resp.get("data")
+        if not isinstance(benchmark, dict):
+            return None
+
+        # Detect type dynamically: presence of "organization_id" means custom
+        if "organization_id" in benchmark:
+            return CustomBenchmark(**benchmark)
+        else:
+            return PublicBenchmark(**benchmark)
+
+    async def get_by_key(
+        self,
+        key: str,
+        *,
+        timeout: float | httpx.Timeout | None = DEFAULT_TIMEOUT,
+    ) -> Optional[Benchmark]:
+        """Fetch a single benchmark by its unique key."""
+        benchmarks = await self.get(timeout=timeout, key=key)
+
+        if not benchmarks:
+            return None
+
+        for benchmark in benchmarks:
+            if benchmark.key == key:
+                return benchmark
+        return None

src/atlas/resources/evaluations/evaluations.py

@@ -61,12 +61,12 @@ def get(
 
     def get_by_id(
         self,
-        evaluation_id: str,
+        id: str,
         *,
         timeout: float | httpx.Timeout | None = DEFAULT_TIMEOUT,
     ) -> Optional[Evaluation]:
         evaluation = self._get(
-            f"/evaluations/{evaluation_id}",
+            f"/evaluations/{id}",
             timeout=timeout,
             cast_to=Evaluation,
         )
@@ -195,12 +195,12 @@ async def get(
 
     async def get_by_id(
         self,
-        evaluation_id: str,
+        id: str,
         *,
         timeout: float | httpx.Timeout | None = DEFAULT_TIMEOUT,
     ) -> Optional[Evaluation]:
         evaluation = await self._get(
-            f"/evaluations/{evaluation_id}",
+            f"/evaluations/{id}",
             timeout=timeout,
             cast_to=Evaluation,
         )

src/atlas/resources/models/models.py

@@ -16,6 +16,7 @@ def get(
         timeout: float | httpx.Timeout | None = DEFAULT_TIMEOUT,
         type: Literal["custom", "public"] | None = None,
         name: Optional[str] = None,
+        key: Optional[str] = None,
         companies: Optional[List[str]] = None,
         regions: Optional[List[str]] = None,
         licenses: Optional[List[str]] = None,
@@ -25,7 +26,9 @@ def get(
         def fetch(model_type: str) -> ModelsResponse | None:
             params = {"type": model_type}
             if name:
-                params["query"] = name
+                params["name"] = name
+            if key:
+                params["key"] = key
             if companies:
                 params["companies"] = ",".join(companies)
             if regions:
@@ -62,6 +65,45 @@ def cast_model(m: Model, model_type: str) -> Model:
 
         return models
 
+    def get_by_id(self, id: str, *, timeout: float | httpx.Timeout | None = DEFAULT_TIMEOUT) -> Optional[Model]:
+        base_url = f"/organizations/{self._client.organization_id}/projects/{self._client.project_id}/models/{id}"
+
+        resp = self._get(
+            base_url,
+            timeout=timeout,
+            cast_to=dict,
+        )
+
+        if not isinstance(resp, dict):
+            return None
+
+        model = resp.get("data")
+        if not isinstance(model, dict):
+            return None
+
+        # Detect type dynamically: presence of "organization_id" means custom
+        if "organization_id" in model:
+            return CustomModel(**model)
+        else:
+            return PublicModel(**model)
+
+    def get_by_key(
+        self,
+        key: str,
+        *,
+        timeout: float | httpx.Timeout | None = DEFAULT_TIMEOUT,
+    ) -> Optional[Model]:
+        """Fetch a single model by its unique key."""
+        models = self.get(timeout=timeout, key=key)
+
+        if not models:
+            return None
+
+        for model in models:
+            if model.key == key:
+                return model
+        return None
+
 
 class AsyncModels(AsyncAPIResource):
     async def get(
@@ -70,6 +112,7 @@ async def get(
         timeout: float | httpx.Timeout | None = DEFAULT_TIMEOUT,
         type: Literal["custom", "public"] | None = None,
         name: Optional[str] = None,
+        key: Optional[str] = None,
         companies: Optional[List[str]] = None,
         regions: Optional[List[str]] = None,
         licenses: Optional[List[str]] = None,
@@ -79,7 +122,9 @@ async def get(
         async def fetch(model_type: str) -> ModelsResponse | None:
             params = {"type": model_type}
             if name:
-                params["query"] = name
+                params["name"] = name
+            if key:
+                params["key"] = key
             if companies:
                 params["companies"] = ",".join(companies)
             if regions:
@@ -115,3 +160,42 @@ def cast_model(m: Model, model_type: str) -> Model:
                 models.extend([cast_model(m, type) for m in resp.data.models])
 
         return models
+
+    async def get_by_id(self, id: str, *, timeout: float | httpx.Timeout | None = DEFAULT_TIMEOUT) -> Optional[Model]:
+        base_url = f"/organizations/{self._client.organization_id}/projects/{self._client.project_id}/models/{id}"
+
+        resp = await self._get(
+            base_url,
+            timeout=timeout,
+            cast_to=dict,
+        )
+
+        if not isinstance(resp, dict):
+            return None
+
+        model = resp.get("data")
+        if not isinstance(model, dict):
+            return None
+
+        # Detect type dynamically: presence of "organization_id" means custom
+        if "organization_id" in model:
+            return CustomModel(**model)
+        else:
+            return PublicModel(**model)
+
+    async def get_by_key(
+        self,
+        key: str,
+        *,
+        timeout: float | httpx.Timeout | None = DEFAULT_TIMEOUT,
+    ) -> Optional[Model]:
+        """Fetch a single model by its unique key."""
+        models = await self.get(timeout=timeout, key=key)
+
+        if not models:
+            return None
+
+        for model in models:
+            if model.key == key:
+                return model
+        return None

tests/resources/test_benchmarks.py

@@ -229,6 +229,75 @@ def test_get_benchmarks_with_none_timeout(self, benchmarks_resource, mock_benchm
         call_args = benchmarks_resource._get.call_args
         assert call_args.kwargs["timeout"] is None
 
+    def test_get_by_id_custom_benchmark(self, benchmarks_resource, sample_custom_benchmark_data):
+        """get_by_id returns CustomBenchmark when organization_id is present."""
+        sample_custom_benchmark_data["organization_id"] = "org-123"  # Required for type detection
+        benchmarks_resource._get.return_value = {"data": sample_custom_benchmark_data}
+
+        result = benchmarks_resource.get_by_id("custom-benchmark-456")
+
+        assert isinstance(result, CustomBenchmark)
+        assert result.id == "custom-benchmark-456"
+        assert result.name == "My Custom Benchmark"
+
+    def test_get_by_id_public_benchmark(self, benchmarks_resource, sample_benchmark_data):
+        """get_by_id returns PublicBenchmark when organization_id is missing."""
+        benchmarks_resource._get.return_value = {"data": sample_benchmark_data}
+
+        result = benchmarks_resource.get_by_id("benchmark-123")
+
+        assert isinstance(result, PublicBenchmark)
+        assert result.id == "benchmark-123"
+        assert result.name == "MMLU"
+
+    def test_get_by_id_invalid_response(self, benchmarks_resource):
+        """get_by_id returns None for invalid responses."""
+        benchmarks_resource._get.return_value = "not-a-dict"
+
+        result = benchmarks_resource.get_by_id("invalid")
+
+        assert result is None
+
+    def test_get_by_key_custom_benchmark(self, benchmarks_resource, sample_custom_benchmark_data):
+        """get_by_key returns CustomBenchmark when key matches and organization_id is present."""
+        sample_custom_benchmark_data["organization_id"] = "org-123"
+        custom_benchmark = CustomBenchmark(**sample_custom_benchmark_data)
+        benchmarks_resource.get = Mock(return_value=[custom_benchmark])
+
+        result = benchmarks_resource.get_by_key(key="my-benchmark")
+
+        assert isinstance(result, CustomBenchmark)
+        assert result.key == "my-benchmark"
+        assert result.name == "My Custom Benchmark"
+
+    def test_get_by_key_public_benchmark(self, benchmarks_resource, sample_benchmark_data):
+        """get_by_key returns PublicBenchmark when key matches and organization_id is missing."""
+        public_benchmark = PublicBenchmark(**sample_benchmark_data)
+        benchmarks_resource.get = Mock(return_value=[public_benchmark])
+
+        result = benchmarks_resource.get_by_key(key="mmlu")
+
+        assert isinstance(result, PublicBenchmark)
+        assert result.key == "mmlu"
+        assert result.name == "MMLU"
+
+    def test_get_by_key_no_match(self, benchmarks_resource, sample_benchmark_data):
+        """get_by_key returns None if no benchmark has the exact key."""
+        public_benchmark = PublicBenchmark(**sample_benchmark_data)
+        benchmarks_resource.get = Mock(return_value=[public_benchmark])
+
+        result = benchmarks_resource.get_by_key(key="nonexistent-key")
+
+        assert result is None
+
+    def test_get_by_key_invalid_response(self, benchmarks_resource):
+        """get_by_key returns None when get() returns None or invalid type."""
+        benchmarks_resource.get = Mock(return_value=None)
+
+        result = benchmarks_resource.get_by_key(key="some-key")
+
+        assert result is None
+
 
 class TestBenchmarksErrorHandling:
     """Test error handling in Benchmarks resource."""