From 0f7f177c4ccef3d259eba147d7a5b2e330d3bffc Mon Sep 17 00:00:00 2001 From: m-peko Date: Wed, 27 Aug 2025 09:11:00 +0200 Subject: [PATCH 1/2] Support fetching model and benchmark by ID --- src/atlas/resources/benchmarks/benchmarks.py | 53 ++++++++++++++++++- .../resources/evaluations/evaluations.py | 8 +-- src/atlas/resources/models/models.py | 44 +++++++++++++++ tests/resources/test_benchmarks.py | 29 ++++++++++ tests/resources/test_models_resource.py | 29 ++++++++++ 5 files changed, 158 insertions(+), 5 deletions(-) diff --git a/src/atlas/resources/benchmarks/benchmarks.py b/src/atlas/resources/benchmarks/benchmarks.py index d6239ab..fb4d843 100644 --- a/src/atlas/resources/benchmarks/benchmarks.py +++ b/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 @@ -53,6 +58,28 @@ 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) + class AsyncBenchmarks(AsyncAPIResource): async def get( @@ -97,3 +124,27 @@ 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) diff --git a/src/atlas/resources/evaluations/evaluations.py b/src/atlas/resources/evaluations/evaluations.py index c97162e..d1ea851 100644 --- a/src/atlas/resources/evaluations/evaluations.py +++ b/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, ) diff --git a/src/atlas/resources/models/models.py b/src/atlas/resources/models/models.py index d4d3ad3..66eb51c 100644 --- a/src/atlas/resources/models/models.py +++ b/src/atlas/resources/models/models.py @@ -62,6 +62,28 @@ 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) + class AsyncModels(AsyncAPIResource): async def get( @@ -115,3 +137,25 @@ 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) diff --git a/tests/resources/test_benchmarks.py b/tests/resources/test_benchmarks.py index 5bf0aa0..aabaae1 100644 --- a/tests/resources/test_benchmarks.py +++ b/tests/resources/test_benchmarks.py @@ -229,6 +229,35 @@ 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 + class TestBenchmarksErrorHandling: """Test error handling in Benchmarks resource.""" diff --git a/tests/resources/test_models_resource.py b/tests/resources/test_models_resource.py index 48a5ab6..ab93f46 100644 --- a/tests/resources/test_models_resource.py +++ b/tests/resources/test_models_resource.py @@ -299,6 +299,35 @@ def test_get_models_custom_model_attributes(self, models_resource, mock_custom_m assert custom_model.disabled is False assert custom_model.api_url == "https://api.example.com/v1/chat" + def test_get_by_id_custom_benchmark(self, models_resource, sample_custom_model_data): + """get_by_id returns CustomModel when organization_id is present.""" + sample_custom_model_data["organization_id"] = "org-123" # Required for type detection + models_resource._get.return_value = {"data": sample_custom_model_data} + + result = models_resource.get_by_id("custom-model-456") + + assert isinstance(result, CustomModel) + assert result.id == "custom-model-456" + assert result.name == "My Custom Model" + + def test_get_by_id_public_benchmark(self, models_resource, sample_model_data): + """get_by_id returns PublicModel when organization_id is missing.""" + models_resource._get.return_value = {"data": sample_model_data} + + result = models_resource.get_by_id("model-123") + + assert isinstance(result, PublicModel) + assert result.id == "model-123" + assert result.name == "GPT-4" + + def test_get_by_id_invalid_response(self, models_resource): + """get_by_id returns None for invalid responses.""" + models_resource._get.return_value = "not-a-dict" + + result = models_resource.get_by_id("invalid") + + assert result is None + class TestModelsErrorHandling: """Test error handling in Models resource.""" From e5d48ff8b5ffb035f0468734db45b4a8d87003a9 Mon Sep 17 00:00:00 2001 From: m-peko Date: Wed, 27 Aug 2025 09:52:43 +0200 Subject: [PATCH 2/2] Support fetching models and benchmarks by key --- src/atlas/resources/benchmarks/benchmarks.py | 44 +++++++++++++++++++- src/atlas/resources/models/models.py | 44 +++++++++++++++++++- tests/resources/test_benchmarks.py | 40 ++++++++++++++++++ tests/resources/test_models_resource.py | 40 ++++++++++++++++++ 4 files changed, 164 insertions(+), 4 deletions(-) diff --git a/src/atlas/resources/benchmarks/benchmarks.py b/src/atlas/resources/benchmarks/benchmarks.py index fb4d843..c694f16 100644 --- a/src/atlas/resources/benchmarks/benchmarks.py +++ b/src/atlas/resources/benchmarks/benchmarks.py @@ -21,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, @@ -80,6 +83,23 @@ def get_by_id(self, id: str, *, timeout: float | httpx.Timeout | None = DEFAULT_ 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( @@ -88,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, @@ -148,3 +171,20 @@ async def get_by_id( 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 diff --git a/src/atlas/resources/models/models.py b/src/atlas/resources/models/models.py index 66eb51c..7e0cd70 100644 --- a/src/atlas/resources/models/models.py +++ b/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: @@ -84,6 +87,23 @@ def get_by_id(self, id: str, *, timeout: float | httpx.Timeout | None = DEFAULT_ 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( @@ -92,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, @@ -101,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: @@ -159,3 +182,20 @@ async def get_by_id(self, id: str, *, timeout: float | httpx.Timeout | None = DE 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 diff --git a/tests/resources/test_benchmarks.py b/tests/resources/test_benchmarks.py index aabaae1..e65c14b 100644 --- a/tests/resources/test_benchmarks.py +++ b/tests/resources/test_benchmarks.py @@ -258,6 +258,46 @@ def test_get_by_id_invalid_response(self, benchmarks_resource): 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.""" diff --git a/tests/resources/test_models_resource.py b/tests/resources/test_models_resource.py index ab93f46..47e8d56 100644 --- a/tests/resources/test_models_resource.py +++ b/tests/resources/test_models_resource.py @@ -328,6 +328,46 @@ def test_get_by_id_invalid_response(self, models_resource): assert result is None + def test_get_by_key_custom_model(self, models_resource, sample_custom_model_data): + """get_by_key returns CustomModel when key matches and organization_id is present.""" + sample_custom_model_data["organization_id"] = "org-123" + custom_model = CustomModel(**sample_custom_model_data) + models_resource.get = Mock(return_value=[custom_model]) + + result = models_resource.get_by_key(key="my-model") + + assert isinstance(result, CustomModel) + assert result.key == "my-model" + assert result.name == "My Custom Model" + + def test_get_by_key_public_model(self, models_resource, sample_model_data): + """get_by_key returns PublicModel when key matches and organization_id is missing.""" + public_model = PublicModel(**sample_model_data) + models_resource.get = Mock(return_value=[public_model]) + + result = models_resource.get_by_key(key="gpt-4") + + assert isinstance(result, PublicModel) + assert result.key == "gpt-4" + assert result.name == "GPT-4" + + def test_get_by_key_no_match(self, models_resource, sample_model_data): + """get_by_key returns None if no model has the exact key.""" + public_model = PublicModel(**sample_model_data) + models_resource.get = Mock(return_value=[public_model]) + + result = models_resource.get_by_key(key="nonexistent-key") + + assert result is None + + def test_get_by_key_invalid_response(self, models_resource): + """get_by_key returns None when get() returns None or invalid type.""" + models_resource.get = Mock(return_value=None) + + result = models_resource.get_by_key(key="some-key") + + assert result is None + class TestModelsErrorHandling: """Test error handling in Models resource."""