diff --git a/src/atlas/resources/benchmarks/benchmarks.py b/src/atlas/resources/benchmarks/benchmarks.py index d6239ab..c694f16 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 @@ -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 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..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: @@ -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 diff --git a/tests/resources/test_benchmarks.py b/tests/resources/test_benchmarks.py index 5bf0aa0..e65c14b 100644 --- a/tests/resources/test_benchmarks.py +++ b/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.""" diff --git a/tests/resources/test_models_resource.py b/tests/resources/test_models_resource.py index 48a5ab6..47e8d56 100644 --- a/tests/resources/test_models_resource.py +++ b/tests/resources/test_models_resource.py @@ -299,6 +299,75 @@ 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 + + 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."""