diff --git a/packages/base/ai_api_client_sdk/ai_api_v2_client.py b/packages/base/ai_api_client_sdk/ai_api_v2_client.py index 9de3eebf..17470c14 100644 --- a/packages/base/ai_api_client_sdk/ai_api_v2_client.py +++ b/packages/base/ai_api_client_sdk/ai_api_v2_client.py @@ -1,7 +1,9 @@ import os -from typing import Callable +from typing import Callable, List from ai_api_client_sdk.exception import AIAPIAuthenticatorException +from ai_api_client_sdk.models.model import Model +from ai_api_client_sdk.models.status import Status from ai_api_client_sdk.helpers import _is_all_none from ai_api_client_sdk.helpers.authenticator import Authenticator from ai_api_client_sdk.helpers.constants import SKIP_AUTH_ENV_VAR, Timeouts @@ -122,3 +124,98 @@ def __init__(self, base_url: str, auth_url: str = None, client_id: str = None, c self.meta: MetaClient = MetaClient(rest_client=self.rest_client) admin_rest_client = RestClient(base_url=base_url[:-3], get_token=token_creator, resource_group=resource_group) self.resource_groups: ResourceGroupsClient = ResourceGroupsClient(rest_client=admin_rest_client) + + def get_models(self, resource_group: str = None) -> List[Model]: + """Lists the available foundation models. + + :param resource_group: Resource Group which the request should be sent on behalf. Either this or a default + resource group in the :class:`ai_api_client_sdk.ai_api_v2_client.AIAPIV2Client` should be specified, + defaults to None + :type resource_group: str + :return: List of available models + :rtype: List[class:`ai_api_client_sdk.models.model.Model`] + """ + return self.model.query(resource_group=resource_group).resources + + def get_foundation_model_deployment_id(self, model_name: str, executable_id: str, + resource_group: str = None) -> str: + """Resolves the deployment ID for a running foundation model. + + :param model_name: Model name (e.g. ``'gpt-4o'``) + :type model_name: str + :param executable_id: Executable / provider ID (e.g. ``'azure-openai'``) + :type executable_id: str + :param resource_group: Resource Group which the request should be sent on behalf. Either this or a default + resource group in the :class:`ai_api_client_sdk.ai_api_v2_client.AIAPIV2Client` should be specified, + defaults to None + :type resource_group: str + :raises ValueError: if no matching running deployment is found + :return: The deployment ID + :rtype: str + """ + deployments = self.deployment.query( + scenario_id='foundation-models', + executable_ids=[executable_id], + status=Status.RUNNING, + resource_group=resource_group, + ).resources + for d in deployments: + details = getattr(d, 'details', None) or {} + backend_model = ( + details + .get('resources', {}) + .get('backend_details', {}) + .get('model', {}) + ) + if backend_model.get('name') == model_name: + return d.id + raise ValueError( + f"No running deployment found for model '{model_name}' " + f"with executable '{executable_id}'." + ) + + def get_orchestration_deployment_id(self, resource_group: str = None) -> str: + """Resolves the deployment ID for the orchestration scenario. + + :param resource_group: Resource Group which the request should be sent on behalf. Either this or a default + resource group in the :class:`ai_api_client_sdk.ai_api_v2_client.AIAPIV2Client` should be specified, + defaults to None + :type resource_group: str + :raises ValueError: if no running orchestration deployment is found + :return: The deployment ID + :rtype: str + """ + deployments = self.deployment.query( + scenario_id='orchestration', + status=Status.RUNNING, + resource_group=resource_group, + ).resources + if not deployments: + raise ValueError('No running orchestration deployment found.') + return deployments[0].id + + def get_orchestration_deployment_url(self, resource_group: str = None) -> str: + """Resolves the deployment URL for the orchestration scenario. + + :param resource_group: Resource Group which the request should be sent on behalf. Either this or a default + resource group in the :class:`ai_api_client_sdk.ai_api_v2_client.AIAPIV2Client` should be specified, + defaults to None + :type resource_group: str + :raises ValueError: if no running orchestration deployment is found or it has no URL + :return: The deployment URL + :rtype: str + """ + deployments = self.deployment.query( + scenario_id='orchestration', + status=Status.RUNNING, + resource_group=resource_group, + ).resources + if not deployments: + raise ValueError('No running orchestration deployment found.') + url = getattr(deployments[0], 'deployment_url', None) + if not url: + raise ValueError( + 'Orchestration deployment has no deployment URL. ' + 'Ensure the deployment is running.' + ) + return url diff --git a/packages/base/ai_api_client_sdk/helpers/llm_helper.py b/packages/base/ai_api_client_sdk/helpers/llm_helper.py index 56b3b555..721bdf52 100644 --- a/packages/base/ai_api_client_sdk/helpers/llm_helper.py +++ b/packages/base/ai_api_client_sdk/helpers/llm_helper.py @@ -1,6 +1,9 @@ +import copy +from typing import Dict, List from functools import partial from ai_api_client_sdk.helpers.constants import SCENARIO_LABEL_NAME_PATTERN +from ai_api_client_sdk.models.model import Model def get_attr(obj, attr): @@ -41,3 +44,67 @@ def filter_for_llm_scenarios(response_dict): res_dict['count'] = len(filtered_scenarios) res_dict['resources'] = filtered_scenarios return res_dict + + +def filter_foundation_models_by_capability(models: List[Model], capability: str) -> List[Model]: + """Filters models by a model capability (e.g. ``'vision'``). + + :param models: The list of models to filter + :type models: List[class:`ai_api_client_sdk.models.model.Model`] + :param capability: The capability to filter by + :type capability: str + :return: Models where at least one non-deprecated version declares the given capability + :rtype: List[class:`ai_api_client_sdk.models.model.Model`] + """ + return [ + m + for m in models + if any( + not v.deprecated + and v.capabilities + and capability in v.capabilities + for v in (m.versions or []) + ) + ] + + +def filter_foundation_models_non_deprecated(models: List[Model]) -> List[Model]: + """Removes deprecated versions from each model. Models with no remaining active versions are excluded. + + :param models: The list of models to filter + :type models: List[class:`ai_api_client_sdk.models.model.Model`] + :return: Models with deprecated versions removed + :rtype: List[class:`ai_api_client_sdk.models.model.Model`] + """ + result: List[Model] = [] + for m in models: + active = [v for v in (m.versions or []) if not v.deprecated] + if active: + filtered = copy.copy(m) + filtered.versions = active + result.append(filtered) + return result + + +def get_foundation_model_latest_versions(models: List[Model]) -> List[Dict]: + """Gets the latest version(s) of each model. If a version is marked ``is_latest=True``, only that version is + returned for the model. If no version is marked as latest, all versions are returned. + + :param models: The list of models + :type models: List[class:`ai_api_client_sdk.models.model.Model`] + :return: A list of dicts with keys ``model``, ``executable_id``, and ``latest_versions`` + :rtype: List[Dict] + """ + result = [] + for m in models: + versions = m.versions or [] + marked = [v for v in versions if v.is_latest] + latest = marked if marked else versions + result.append( + { + "model": m.model, + "executable_id": m.executable_id, + "latest_versions": latest, + } + ) + return result diff --git a/packages/base/ai_api_client_sdk/resource_clients/model_client.py b/packages/base/ai_api_client_sdk/resource_clients/model_client.py index a3e039c9..0415cd11 100644 --- a/packages/base/ai_api_client_sdk/resource_clients/model_client.py +++ b/packages/base/ai_api_client_sdk/resource_clients/model_client.py @@ -1,3 +1,6 @@ +from typing import List + +from ai_api_client_sdk.models.model import Model from ai_api_client_sdk.models.model_query_response import ModelQueryResponse from ai_api_client_sdk.resource_clients.base_client import BaseClient diff --git a/packages/base/tests/helpers/test_model_discovery.py b/packages/base/tests/helpers/test_model_discovery.py new file mode 100644 index 00000000..2700db5e --- /dev/null +++ b/packages/base/tests/helpers/test_model_discovery.py @@ -0,0 +1,161 @@ +import unittest + +from ai_api_client_sdk.helpers.llm_helper import ( + filter_foundation_models_by_capability, + filter_foundation_models_non_deprecated, + get_foundation_model_latest_versions, +) +from ai_api_client_sdk.models.model import Model +from ai_api_client_sdk.models.model_version import ModelVersion + + +def _make_version(name, is_latest=False, deprecated=False, capabilities=None): + return ModelVersion( + name=name, + is_latest=is_latest, + deprecated=deprecated, + capabilities=capabilities, + ) + + +def _make_model(model_name, executable_id, versions): + return Model( + executable_id=executable_id, + model=model_name, + description=f"{model_name} description", + versions=versions, + ) + + +class TestFilterFoundationModelsByCapability(unittest.TestCase): + def setUp(self): + self.models = [ + _make_model( + "gpt-4o", + "azure-openai", + [ + _make_version("2024-11-20", is_latest=True, capabilities=["vision", "tools"]), + _make_version("2024-05-13", deprecated=True, capabilities=["vision"]), + ], + ), + _make_model( + "text-embedding-3-large", + "azure-openai", + [ + _make_version("1", is_latest=True, capabilities=["embedding"]), + ], + ), + ] + + def test_returns_models_with_capability(self): + result = filter_foundation_models_by_capability(self.models, "vision") + self.assertEqual(len(result), 1) + self.assertEqual(result[0].model, "gpt-4o") + + def test_returns_empty_when_no_match(self): + result = filter_foundation_models_by_capability(self.models, "nonexistent") + self.assertEqual(len(result), 0) + + def test_ignores_deprecated_versions(self): + # Only deprecated version has "vision" → should be excluded + models = [ + _make_model( + "old-model", + "azure-openai", + [_make_version("1.0", deprecated=True, capabilities=["vision"])], + ), + ] + result = filter_foundation_models_by_capability(models, "vision") + self.assertEqual(len(result), 0) + + +class TestFilterFoundationModelsNonDeprecated(unittest.TestCase): + def setUp(self): + self.models = [ + _make_model( + "gpt-4o", + "azure-openai", + [ + _make_version("2024-11-20", is_latest=True), + _make_version("2024-05-13", deprecated=True), + ], + ), + _make_model( + "text-embedding-3-large", + "azure-openai", + [_make_version("1", is_latest=True)], + ), + ] + + def test_removes_deprecated_versions(self): + result = filter_foundation_models_non_deprecated(self.models) + gpt4o = next(m for m in result if m.model == "gpt-4o") + self.assertEqual(len(gpt4o.versions), 1) + self.assertEqual(gpt4o.versions[0].name, "2024-11-20") + + def test_excludes_models_with_all_deprecated(self): + models = [ + _make_model( + "old-model", + "azure-openai", + [_make_version("1", deprecated=True)], + ), + ] + result = filter_foundation_models_non_deprecated(models) + self.assertEqual(len(result), 0) + + def test_keeps_models_with_active_versions(self): + result = filter_foundation_models_non_deprecated(self.models) + self.assertEqual(len(result), 2) + + def test_does_not_mutate_original(self): + original_len = len(self.models[0].versions) + filter_foundation_models_non_deprecated(self.models) + self.assertEqual(len(self.models[0].versions), original_len) + + +class TestGetFoundationModelLatestVersions(unittest.TestCase): + def test_returns_one_entry_per_model(self): + models = [ + _make_model("a", "x", [_make_version("1", is_latest=True)]), + _make_model("b", "y", [_make_version("1", is_latest=True)]), + ] + result = get_foundation_model_latest_versions(models) + self.assertEqual(len(result), 2) + + def test_returns_marked_latest(self): + models = [ + _make_model( + "gpt-4o", + "azure-openai", + [ + _make_version("2024-11-20", is_latest=True), + _make_version("2024-05-13"), + ], + ), + ] + result = get_foundation_model_latest_versions(models) + self.assertEqual(len(result[0]["latest_versions"]), 1) + self.assertEqual(result[0]["latest_versions"][0].name, "2024-11-20") + + def test_returns_all_when_none_marked_latest(self): + models = [ + _make_model( + "no-latest", + "azure-openai", + [ + _make_version("1"), + _make_version("2"), + ], + ), + ] + result = get_foundation_model_latest_versions(models) + self.assertEqual(len(result[0]["latest_versions"]), 2) + + def test_includes_model_and_executable_id(self): + models = [ + _make_model("gpt-4o", "azure-openai", [_make_version("1", is_latest=True)]), + ] + result = get_foundation_model_latest_versions(models) + self.assertEqual(result[0]["model"], "gpt-4o") + self.assertEqual(result[0]["executable_id"], "azure-openai") diff --git a/packages/base/tests/test_ai_api_v2_client_convenience.py b/packages/base/tests/test_ai_api_v2_client_convenience.py new file mode 100644 index 00000000..d7984ade --- /dev/null +++ b/packages/base/tests/test_ai_api_v2_client_convenience.py @@ -0,0 +1,144 @@ +import os +import unittest +from unittest.mock import MagicMock, patch + +from ai_api_client_sdk.ai_api_v2_client import AIAPIV2Client +from ai_api_client_sdk.helpers.constants import SKIP_AUTH_ENV_VAR +from ai_api_client_sdk.models.model import Model +from ai_api_client_sdk.models.model_version import ModelVersion + + +def _make_client(): + """Create an AIAPIV2Client with auth skipped.""" + os.environ[SKIP_AUTH_ENV_VAR] = 'true' + try: + return AIAPIV2Client(base_url='https://test.example.com/v2/lm') + finally: + del os.environ[SKIP_AUTH_ENV_VAR] + + +class TestAIAPIV2ClientGetModels(unittest.TestCase): + def test_get_models_delegates_to_model_query(self): + client = _make_client() + sentinel = [MagicMock(spec=Model)] + mock_response = MagicMock() + mock_response.resources = sentinel + client.model = MagicMock() + client.model.query.return_value = mock_response + + result = client.get_models() + client.model.query.assert_called_once_with(resource_group=None) + self.assertIs(result, sentinel) + + def test_get_models_passes_resource_group(self): + client = _make_client() + sentinel = [MagicMock(spec=Model)] + mock_response = MagicMock() + mock_response.resources = sentinel + client.model = MagicMock() + client.model.query.return_value = mock_response + + result = client.get_models(resource_group='my-group') + client.model.query.assert_called_once_with(resource_group='my-group') + + +class TestAIAPIV2ClientGetFoundationModelDeploymentId(unittest.TestCase): + def test_resolves_deployment_id(self): + client = _make_client() + mock_deployment = MagicMock() + mock_deployment.id = 'dep-123' + mock_deployment.details = { + 'resources': { + 'backend_details': { + 'model': {'name': 'gpt-4o', 'version': 'latest'} + } + } + } + mock_response = MagicMock() + mock_response.resources = [mock_deployment] + client.deployment = MagicMock() + client.deployment.query.return_value = mock_response + + result = client.get_foundation_model_deployment_id('gpt-4o', 'azure-openai') + self.assertEqual(result, 'dep-123') + + def test_raises_when_no_match(self): + client = _make_client() + mock_deployment = MagicMock() + mock_deployment.id = 'dep-123' + mock_deployment.details = { + 'resources': { + 'backend_details': { + 'model': {'name': 'other-model', 'version': 'latest'} + } + } + } + mock_response = MagicMock() + mock_response.resources = [mock_deployment] + client.deployment = MagicMock() + client.deployment.query.return_value = mock_response + + with self.assertRaises(ValueError) as ctx: + client.get_foundation_model_deployment_id('gpt-4o', 'azure-openai') + self.assertIn('gpt-4o', str(ctx.exception)) + + +class TestAIAPIV2ClientGetOrchestrationDeploymentId(unittest.TestCase): + def test_resolves_orchestration_deployment_id(self): + client = _make_client() + mock_deployment = MagicMock() + mock_deployment.id = 'orch-dep-1' + mock_response = MagicMock() + mock_response.resources = [mock_deployment] + client.deployment = MagicMock() + client.deployment.query.return_value = mock_response + + result = client.get_orchestration_deployment_id() + self.assertEqual(result, 'orch-dep-1') + + def test_raises_when_no_deployment(self): + client = _make_client() + mock_response = MagicMock() + mock_response.resources = [] + client.deployment = MagicMock() + client.deployment.query.return_value = mock_response + + with self.assertRaises(ValueError): + client.get_orchestration_deployment_id() + + +class TestAIAPIV2ClientGetOrchestrationDeploymentUrl(unittest.TestCase): + def test_resolves_orchestration_deployment_url(self): + client = _make_client() + mock_deployment = MagicMock() + mock_deployment.deployment_url = 'https://api.example.com/v2/inference/deployments/orch-1' + mock_response = MagicMock() + mock_response.resources = [mock_deployment] + client.deployment = MagicMock() + client.deployment.query.return_value = mock_response + + result = client.get_orchestration_deployment_url() + self.assertIn('orch-1', result) + + def test_raises_when_no_deployment(self): + client = _make_client() + mock_response = MagicMock() + mock_response.resources = [] + client.deployment = MagicMock() + client.deployment.query.return_value = mock_response + + with self.assertRaises(ValueError): + client.get_orchestration_deployment_url() + + def test_raises_when_no_url(self): + client = _make_client() + mock_deployment = MagicMock() + mock_deployment.deployment_url = None + mock_response = MagicMock() + mock_response.resources = [mock_deployment] + client.deployment = MagicMock() + client.deployment.query.return_value = mock_response + + with self.assertRaises(ValueError) as ctx: + client.get_orchestration_deployment_url() + self.assertIn('no deployment URL', str(ctx.exception))