-
Notifications
You must be signed in to change notification settings - Fork 2
feat: AICORE_SERVICE_KEY parsing extended + Authentication refactor #103
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
yamaceay
wants to merge
3
commits into
main
Choose a base branch
from
feat/aicore-service-key-2
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Open
Changes from all commits
Commits
Show all changes
3 commits
Select commit
Hold shift + click to select a range
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,290 +1,67 @@ | ||
| from __future__ import annotations | ||
|
|
||
| from typing import Any, Dict, Final, List, Optional, Callable, Tuple | ||
| import json | ||
| import os | ||
| import pathlib | ||
|
|
||
| from dataclasses import dataclass | ||
|
|
||
| from gen_ai_hub.evaluations.constants import ( | ||
| AI_CORE_PREFIX, | ||
| AUTH_ENDPOINT_SUFFIX, | ||
| ENV_VAR_AICORE_CONFIG_FILE, | ||
| ENV_VAR_AICORE_PROFILE, | ||
| VCAP_AICORE_SERVICE_NAME, | ||
| ENV_VAR_VCAP_SERVICES, | ||
| ENV_VAR_AICORE_HOME_PATH, | ||
| DEFAULT_HOME_PATH, | ||
| from typing import Dict, Final, List | ||
|
|
||
| from ai_core_sdk.credentials import ( | ||
| CORE_CREDENTIAL_VALUES, | ||
| CredentialsValue, | ||
| Service, | ||
| Source, | ||
| VCAPEnvironment, | ||
| extract_credentials as _extract_core_credentials, | ||
| fetch_credentials as _fetch_core_credentials, | ||
| get_nested_value, | ||
| init_conf, | ||
| resolve_credentials as _resolve_core_credentials, | ||
| resolve_resource_group, | ||
| validate_credentials, | ||
| ) | ||
| from gen_ai_hub.evaluations.helpers.logging import get_logger | ||
|
|
||
| logger = get_logger() | ||
|
|
||
|
|
||
| def get_home() -> str: | ||
| return os.environ.get(ENV_VAR_AICORE_HOME_PATH, DEFAULT_HOME_PATH) | ||
|
|
||
|
|
||
| def get_nested_value(data_dict, keys: List[str]): | ||
| """ | ||
| Retrieve a nested value from a dictionary using a list of strings. | ||
|
|
||
| :param data_dict: The dictionary to search. | ||
| :param keys: A list of strings representing nested keys. | ||
| :return: The value associated with the nested keys, or None if not found. | ||
| """ | ||
| current_value = data_dict | ||
| for key in keys: | ||
| current_value = current_value[key] | ||
| return current_value | ||
|
|
||
|
|
||
| @dataclass | ||
| class VCAPEnvironment: | ||
| services: List[Service] | ||
|
|
||
| @classmethod | ||
| def from_env(cls, env_var: Optional[str] = None): | ||
| env_var = env_var or ENV_VAR_VCAP_SERVICES | ||
| env = json.loads(os.environ.get(env_var, '{}')) | ||
| return cls.from_dict(env) | ||
|
|
||
| @classmethod | ||
| def from_dict(cls, env: Dict[str, Any]): | ||
| services = [Service(service) for services in env.values() for service in services] | ||
| return cls(services=services) | ||
|
|
||
| def __getitem__(self, name) -> Service: | ||
| return self.get_service(name, exactly_one=True) | ||
|
|
||
| def get_service(self, label, exactly_one: bool = True) -> Service: | ||
| services = [s for s in self.services if s.label == label] | ||
| if exactly_one: | ||
| if len(services) == 0: | ||
| raise KeyError(f"No service found with label '{label}'.") | ||
| return services[0] | ||
| else: | ||
| return services | ||
|
|
||
| def get_service_by_name(self, name, exactly_one: bool = True) -> Service: | ||
| services = [s for s in self.services if s.name == name] | ||
| if exactly_one: | ||
| if len(services) == 0: | ||
| raise KeyError(f"No service found with name '{name}'.") | ||
| return services[0] | ||
| else: | ||
| return services | ||
|
|
||
|
|
||
| NoDefault = object() | ||
|
|
||
|
|
||
| class Service: | ||
|
|
||
| def __init__(self, env: Dict[str, Any]): | ||
| self._env = env | ||
|
|
||
| @property | ||
| def label(self) -> Optional[str]: | ||
| return self._env.get('label') | ||
|
|
||
| @property | ||
| def name(self) -> Optional[str]: | ||
| return self._env.get('name') | ||
|
|
||
| def __getitem__(self, key): | ||
| return self.get(key) | ||
|
|
||
| def get(self, key, default=NoDefault): | ||
| if isinstance(key, str): | ||
| key_splitted = key.split('.') | ||
| else: | ||
| key_splitted = key | ||
| try: | ||
| return get_nested_value(self._env, key_splitted) or default | ||
| except KeyError: | ||
| if default is NoDefault: | ||
| raise KeyError(f"Key '{key}' not found in service '{self.name}'.") | ||
| return default | ||
|
|
||
|
|
||
| @dataclass | ||
| class CredentialsValue: | ||
| name: str | ||
| vcap_key: Optional[Tuple[str, ...]] = None | ||
| transform_fn: Optional[Callable] = None | ||
|
|
||
|
|
||
| @dataclass | ||
| class Source: | ||
| name: str | ||
| get: Callable[[CredentialsValue], Optional[str]] | ||
|
|
||
|
|
||
| CREDENTIAL_VALUES: Final[List[CredentialsValue]] = [ | ||
| CredentialsValue(name='client_id', vcap_key=('credentials', 'clientid')), | ||
| CredentialsValue(name='client_secret', vcap_key=('credentials', 'clientsecret')), | ||
| CredentialsValue(name='auth_url', | ||
| vcap_key=('credentials', 'url'), | ||
| transform_fn=lambda url: url.rstrip('/') + | ||
| ('' if url.endswith(AUTH_ENDPOINT_SUFFIX) else AUTH_ENDPOINT_SUFFIX)), | ||
| CredentialsValue(name='base_url', | ||
| vcap_key=('credentials', 'serviceurls', 'AI_API_URL'), | ||
| transform_fn=lambda url: url.rstrip('/') + ('' if url.endswith('/v2') else '/v2')), | ||
| CredentialsValue(name='resource_group'), | ||
| CredentialsValue(name='cert_url', vcap_key=('credentials', 'certurl'), | ||
| transform_fn=lambda url: url.rstrip('/') + | ||
| ('' if url.endswith(AUTH_ENDPOINT_SUFFIX) else AUTH_ENDPOINT_SUFFIX)), | ||
| # Even though the certificate and key in VCAP_SERVICES are not file paths, the names are defined this way in order | ||
| # to keep it compatible with the config names. It'll be handled in fetch_credentials function. | ||
| CredentialsValue(name='cert_file_path'), | ||
| CredentialsValue(name='key_file_path'), | ||
| CredentialsValue(name='cert_str', vcap_key=('credentials', 'certificate'), | ||
| transform_fn=lambda cert_str: cert_str.replace('\\n', '\n')), | ||
| CredentialsValue(name='key_str', vcap_key=('credentials', 'key'), | ||
| transform_fn=lambda key_str: key_str.replace('\\n', '\n')), | ||
| # Currently supporting only the AWS creds, would need to extend to other hyperscalers in future. | ||
| CredentialsValue(name='aws_access_key_id'), | ||
| CredentialsValue(name='aws_secret_access_key'), | ||
| CredentialsValue(name='orchestration_url'), | ||
| CredentialsValue(name='input_object_store_secret_name'), | ||
| from ai_core_sdk.helpers import get_home | ||
|
|
||
| # Re-exported for backward compatibility: these are generic and fully reused from ai_core_sdk.credentials. | ||
| __all__ = [ | ||
| "CredentialsValue", | ||
| "Service", | ||
| "Source", | ||
| "VCAPEnvironment", | ||
| "CREDENTIAL_VALUES", | ||
| "EVAL_CREDENTIAL_VALUES", | ||
| "extract_credentials", | ||
| "fetch_credentials", | ||
| "get_home", | ||
| "get_nested_value", | ||
| "init_conf", | ||
| "resolve_credentials", | ||
| "resolve_resource_group", | ||
| "validate_credentials", | ||
| ] | ||
|
|
||
| # Extends the core credential values with evaluation-specific ones. | ||
| # Currently supporting only the AWS creds, would need to extend to other hyperscalers in future. | ||
| EVAL_CREDENTIAL_VALUES: Final[List[CredentialsValue]] = CORE_CREDENTIAL_VALUES + [ | ||
| CredentialsValue(name='aws_access_key_id', vcap_key=('credentials', 'aws_access_key_id')), | ||
| CredentialsValue(name='aws_secret_access_key', vcap_key=('credentials', 'aws_secret_access_key')), | ||
| CredentialsValue(name='orchestration_url', vcap_key=('credentials', 'orchestration_url')), | ||
| CredentialsValue(name='input_object_store_secret_name', vcap_key=('credentials', 'input_object_store_secret_name')), | ||
| ] | ||
|
|
||
| def init_conf(profile: str = None): | ||
| # Read configuration from ${AICORE_HOME}/config_<profile>.json. | ||
| home = pathlib.Path(get_home()) | ||
| profile = profile or os.environ.get(ENV_VAR_AICORE_PROFILE) | ||
| profile_config_file = f'config_{profile}.json' | ||
| direct_config_file = pathlib.Path(os.getenv(ENV_VAR_AICORE_CONFIG_FILE)) if os.getenv(ENV_VAR_AICORE_CONFIG_FILE) else None | ||
| path_to_config = (direct_config_file or | ||
| (home / ('config.json' if profile in ('default', '', None) else profile_config_file))) | ||
| config = {} | ||
| if path_to_config.exists(): | ||
| logger.debug('Config file path %s', path_to_config) | ||
| try: | ||
| with path_to_config.open(encoding='utf-8') as f: | ||
| return json.load(f) | ||
| except json.decoder.JSONDecodeError: | ||
| raise KeyError(f'{path_to_config} is not a valid json file. Please fix or remove it!') | ||
| except PermissionError as e: | ||
| logger.warning("Permission denied when trying to read config file '%s'. File ignored.", path_to_config) | ||
| return config | ||
| elif profile: | ||
| raise FileNotFoundError(f"Unable to locate profile config file '{profile_config_file}' " | ||
| f"in AICORE_HOME '{home}')") | ||
| return config | ||
|
|
||
| CREDENTIAL_VALUES: Final[List[CredentialsValue]] = EVAL_CREDENTIAL_VALUES | ||
|
|
||
| def extract_credentials(source: Source, exclude: List[str] = None) -> Dict[str, str]: | ||
| """Extract all credentials from a source.""" | ||
| exclude = exclude or [] | ||
| credentials = {} | ||
| for cv in CREDENTIAL_VALUES: | ||
| if cv.name in exclude: | ||
| continue | ||
| if value := source.get(cv): | ||
| credentials[cv.name] = cv.transform_fn(value) if cv.transform_fn else value | ||
| return credentials | ||
| """Extract all evaluation credentials from a source.""" | ||
| return _extract_core_credentials(source, credential_values=EVAL_CREDENTIAL_VALUES, exclude=exclude) | ||
|
|
||
|
|
||
| def resolve_credentials(sources: List[Source]) -> Dict[str, str]: | ||
| """Extract credentials from the first source that has any defined.""" | ||
| for source in sources: | ||
| if credentials := extract_credentials(source, exclude=['resource_group']): | ||
| logger.debug(f"Using credentials from: {source.name}") | ||
| return credentials | ||
| raise ValueError("No credentials found in any source") | ||
|
|
||
|
|
||
| def resolve_resource_group(sources: List[Source]) -> Optional[str]: | ||
| """Find resource_group from the first source that defines it.""" | ||
| rg_cred = CredentialsValue(name='resource_group') | ||
| for source in sources: | ||
| if value := source.get(rg_cred): | ||
| logger.debug("Using resource_group '%s' from: %s", value, source.name) | ||
| return value | ||
| logger.debug("No resource_group found in any source") | ||
| return None | ||
|
|
||
|
|
||
| def validate_credentials(credentials: Dict[str, str]) -> None: | ||
| """Validate that we have a complete authentication method.""" | ||
| required_base = {'client_id', 'auth_url', 'base_url'} | ||
|
|
||
| # Check which auth method we have | ||
| has_client_secret = 'client_secret' in credentials | ||
| has_cert_files = 'cert_file_path' in credentials and 'key_file_path' in credentials | ||
| has_cert_strings = 'cert_str' in credentials and 'key_str' in credentials | ||
|
|
||
| # Must have exactly one auth method | ||
| auth_methods = sum([has_client_secret, has_cert_files, has_cert_strings]) | ||
|
|
||
| if auth_methods == 0: | ||
| raise ValueError( | ||
| "No authentication method found. Must provide one of:\n" | ||
| "1. client_secret\n" | ||
| "2. cert_file_path AND key_file_path\n" | ||
| "3. cert_str AND key_str" | ||
| ) | ||
|
|
||
| if auth_methods > 1: | ||
| raise ValueError( | ||
| "Multiple authentication methods found. Please provide only one of:\n" | ||
| "1. client_secret\n" | ||
| "2. cert_file_path AND key_file_path\n" | ||
| "3. cert_str AND key_str" | ||
| ) | ||
|
|
||
| # Check required base fields | ||
| missing = required_base - set(credentials.keys()) | ||
| if missing: | ||
| raise ValueError(f"Missing required credentials: {missing}") | ||
|
|
||
|
|
||
| def _str_or_none(value) -> Optional[str]: | ||
| return str(value) if value else None | ||
| """Extract evaluation credentials from the first source that has any defined.""" | ||
| return _resolve_core_credentials(sources, credential_values=EVAL_CREDENTIAL_VALUES) | ||
|
|
||
|
|
||
| def fetch_credentials(profile: str = None, **kwargs) -> Dict[str, str]: | ||
| """ | ||
| Fetch credentials from a single source based on precedence. | ||
|
|
||
| Precedence order: kwargs > environment variables > config file > VCAP service | ||
| Fetch evaluation credentials from a single source based on precedence. | ||
|
|
||
| Once a source is selected (first one with any credential), all credentials | ||
| come from that source only. Resource group is an exception and follows | ||
| precedence independently. | ||
| Precedence order: kwargs > AICORE_SERVICE_KEY > environment variables > config file > VCAP service | ||
| (see ai_core_sdk.credentials.fetch_credentials for the full behavior). | ||
| """ | ||
| config = init_conf(profile=profile) | ||
|
|
||
| try: | ||
| vcap_service = VCAPEnvironment.from_env()[VCAP_AICORE_SERVICE_NAME] | ||
| except KeyError: | ||
| vcap_service = None | ||
|
|
||
| sources = [ | ||
| Source("kwargs", | ||
| lambda cv: _str_or_none(kwargs.get(cv.name))), | ||
| Source("environment variables", | ||
| lambda cv: _str_or_none(os.environ.get(f'{AI_CORE_PREFIX}_{cv.name.upper()}'))), | ||
| Source("config file", | ||
| lambda cv: _str_or_none(config.get(f'{AI_CORE_PREFIX}_{cv.name.upper()}'))), | ||
| Source("VCAP service", | ||
| lambda cv: _str_or_none(vcap_service.get(cv.vcap_key, None) if vcap_service and cv.vcap_key else None)), | ||
| ] | ||
|
|
||
| credentials = resolve_credentials(sources) | ||
|
|
||
| # Use cert_url as auth_url if present (VCAP provides cert_url for certificate auth) | ||
| if 'cert_url' in credentials: | ||
| credentials['auth_url'] = credentials.pop('cert_url') | ||
|
|
||
| validate_credentials(credentials) | ||
|
|
||
| resource_group = resolve_resource_group(sources) | ||
| if resource_group: | ||
| credentials['resource_group'] = resource_group | ||
|
|
||
| return credentials | ||
| return _fetch_core_credentials(profile=profile, credential_values=EVAL_CREDENTIAL_VALUES, **kwargs) | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.