maas-base/gpustack/codegen/templates/clientset.py.jinja

import base64
from urllib.parse import urlparse
from .generated_http_client import HTTPClient
from typing import Optional
{% for class_name in class_names %}
from .generated_{{ class_name | to_snake_case }}_client import {{ class_name }}Client
{%- endfor %}

from gpustack.utils.network import use_proxy_env_for_url


class ClientSet:
    def __init__(
        self,
        base_url: str,
        api_key: Optional[str] = None,
        username: Optional[str] = None,
        password: Optional[str] = None,
        headers: Optional[dict] = None,
        timeout: Optional[float] = 60.0,
        enable_cache: bool = True,
    ):
        if headers is None:
            headers = {}

        if api_key:
            headers["Authorization"] = f"Bearer {api_key}"
        elif username and password:
            base64_credentials = base64.b64encode(
                f"{username}:{password}".encode()
            ).decode()
            headers["Authorization"] = f"Basic {base64_credentials}"

        self.base_url = base_url
        self.headers = headers

        verify = None
        parsed_url = urlparse(base_url)
        if parsed_url.hostname == "127.0.0.1" and parsed_url.scheme == "https":
            verify = False

        use_proxy_env = use_proxy_env_for_url(base_url)
        http_client = (
            HTTPClient(
                base_url=base_url,
                verify_ssl=verify,
                httpx_args={"trust_env": use_proxy_env},
            )
            .with_headers(headers)
            .with_timeout(timeout)
        )
        self.http_client = http_client
{% for class_name in class_names %}
        self.{{ class_name | to_underscore_plural }} = {{ class_name }}Client(
            http_client,
            enable_cache=enable_cache,
        )
{%- endfor %}{{ "\n" }}