sync_api.py

"""Sync I/O protocol implementation for the LM Studio remote access API."""

import itertools
import time
import weakref

from abc import abstractmethod
from concurrent.futures import Future as SyncFuture, ThreadPoolExecutor, as_completed
from contextlib import (
    contextmanager,
    ExitStack,
)
from types import EllipsisType, TracebackType
from typing import (
    Any,
    ContextManager,
    Generator,
    Iterable,
    Iterator,
    Callable,
    Generic,
    Sequence,
    Type,
    TypeAlias,
    TypeVar,
)
from typing_extensions import (
    # Native in 3.11+
    Self,
    # Native in 3.13+
    TypeIs,
)

import httpx

# Synchronous API still uses an async websocket (just in a background thread)
from httpx_ws import AsyncWebSocketSession

from .sdk_api import (
    LMStudioRuntimeError,
    LMStudioValueError,
    sdk_callback_invocation,
    sdk_public_api,
)
from .schemas import AnyLMStudioStruct, DictObject
from .history import (
    AssistantResponse,
    Chat,
    ChatHistoryDataDict,
    FileHandle,
    LocalFileInput,
    _LocalFileData,
    ToolCallRequest,
    ToolResultMessage,
)
from .json_api import (
    ActResult,
    AnyLoadConfig,
    AnyModelSpecifier,
    AvailableModelBase,
    ChannelEndpoint,
    ChannelHandler,
    ChatResponseEndpoint,
    ClientBase,
    ClientSession,
    CompletionEndpoint,
    DEFAULT_TTL,
    DownloadedModelBase,
    DownloadFinalizedCallback,
    DownloadProgressCallback,
    EmbeddingLoadModelConfig,
    EmbeddingLoadModelConfigDict,
    EmbeddingModelInfo,
    GetOrLoadEndpoint,
    LlmInfo,
    LlmLoadModelConfig,
    LlmLoadModelConfigDict,
    LlmPredictionConfig,
    LlmPredictionConfigDict,
    LlmPredictionFragment,
    LMStudioCancelledError,
    LMStudioClientError,
    LMStudioPredictionError,
    LMStudioTimeoutError,
    LMStudioWebsocket,
    LMStudioWebsocketError,
    LoadModelEndpoint,
    ModelDownloadOptionBase,
    ModelHandleBase,
    ModelInstanceInfo,
    ModelLoadingCallback,
    ModelSessionTypes,
    ModelTypesEmbedding,
    ModelTypesLlm,
    PredictionEndpoint,
    PredictionFirstTokenCallback,
    PredictionFragmentCallback,
    PredictionFragmentEvent,
    PredictionMessageCallback,
    PredictionResult,
    PredictionRoundResult,
    PredictionRxEvent,
    PredictionStreamBase,
    PredictionToolCallEvent,
    PromptProcessingCallback,
    RemoteCallHandler,
    ResponseSchema,
    TModelInfo,
    ToolDefinition,
    check_model_namespace,
    load_struct,
    _model_spec_to_api_dict,
    _redact_json,
)
from ._ws_impl import SyncFutureTimeout
from ._ws_thread import AsyncWebsocketThread, SyncToAsyncWebsocketBridge
from ._kv_config import TLoadConfig, TLoadConfigDict, parse_server_config
from ._sdk_models import (
    EmbeddingRpcCountTokensParameter,
    EmbeddingRpcEmbedStringParameter,
    EmbeddingRpcTokenizeParameter,
    LlmApplyPromptTemplateOpts,
    LlmApplyPromptTemplateOptsDict,
    LlmRpcApplyPromptTemplateParameter,
    ModelCompatibilityType,
)

from ._logging import new_logger, LogEventContext

# Only the sync API itself is published from
# this module. Anything needed for type hints
# and similar tasks is published from `json_api`.
# Bypassing the high level API, and working more
# directly with the underlying websocket(s) is
# supported (hence the public names), but they're
# not exported via the top-level `lmstudio` API.
__all__ = [
    "AnyDownloadedModel",
    "Client",
    "DownloadedEmbeddingModel",
    "DownloadedLlm",
    "EmbeddingModel",
    "LLM",
    "SyncModelHandle",
    "PredictionStream",
    "configure_default_client",
    "get_default_client",
    "get_sync_api_timeout",
    "embedding_model",
    "list_downloaded_models",
    "list_loaded_models",
    "llm",
    "prepare_image",
    "set_sync_api_timeout",
]

#
_DEFAULT_TIMEOUT: float | None = 60.0


@sdk_public_api()
def get_sync_api_timeout() -> float | None:
    """Return the current default sync API timeout when waiting for server messages."""
    return _DEFAULT_TIMEOUT


@sdk_public_api()
def set_sync_api_timeout(timeout: float | None) -> None:
    """Set the default sync API timeout when waiting for server messages."""
    global _DEFAULT_TIMEOUT
    if timeout is not None:
        timeout = float(timeout)
    _DEFAULT_TIMEOUT = timeout


T = TypeVar("T")
CallWithTimeout: TypeAlias = Callable[[float | None], Any]
TimeoutOption: TypeAlias = float | None | EllipsisType


class SyncChannel(Generic[T]):
    """Communication subchannel over multiplexed async websocket."""

    def __init__(
        self,
        channel_id: int,
        get_message: CallWithTimeout,
        endpoint: ChannelEndpoint[T, Any, Any],
        send_json: Callable[[DictObject], None],
        log_context: LogEventContext,
        timeout: TimeoutOption = ...,
    ) -> None:
        """Initialize synchronous websocket streaming channel."""
        self._is_finished = False
        self._get_message = get_message
        self._send_json = send_json
        self._timeout = timeout
        self._api_channel = ChannelHandler(channel_id, endpoint, log_context)

    def get_creation_message(self) -> DictObject:
        """Get the message to send to create this channel."""
        return self._api_channel.get_creation_message()

    def send_message(self, message: DictObject) -> None:
        """Send given message on this channel."""
        wrapped_message = self._api_channel.wrap_message(message)
        self._send_json(wrapped_message)

    def cancel(self) -> None:
        """Cancel the channel."""
        if self._is_finished:
            return
        cancel_message = self._api_channel.get_cancel_message()
        self._send_json(cancel_message)

    @property
    def timeout(self) -> float | None:
        """Permitted time between received messages for this channel."""
        timeout = self._timeout
        if timeout is ...:
            return _DEFAULT_TIMEOUT
        return timeout

    def rx_stream(
        self,
    ) -> Iterator[DictObject | None]:
        """Stream received channel messages until channel is closed by server."""
        while not self._is_finished:
            with sdk_public_api():
                # Avoid emitting tracebacks that delve into supporting libraries
                # (we can't easily suppress the SDK's own frames for iterators)
                try:
                    message = self._get_message(self.timeout)
                except SyncFutureTimeout:
                    raise LMStudioTimeoutError from None
                if message is None:
                    raise LMStudioWebsocketError("Client unexpectedly disconnected.")
                contents = self._api_channel.handle_rx_message(message)
            if contents is None:
                self._is_finished = True
                break
            yield contents

    def wait_for_result(self) -> T:
        """Wait for the channel to finish and return the result."""
        endpoint = self._api_channel.endpoint
        for contents in self.rx_stream():
            endpoint.handle_message_events(contents)
            if endpoint.is_finished:
                break
        return endpoint.result()


class SyncRemoteCall:
    """Remote procedure call over multiplexed async websocket."""

    def __init__(
        self,
        call_id: int,
        get_message: CallWithTimeout,
        log_context: LogEventContext,
        notice_prefix: str = "RPC",
        timeout: TimeoutOption = ...,
    ) -> None:
        """Initialize synchronous remote procedure call."""
        self._get_message = get_message
        self._timeout = timeout
        self._rpc = RemoteCallHandler(call_id, log_context, notice_prefix)
        self._logger = logger = new_logger(type(self).__name__)
        logger.update_context(log_context, call_id=call_id)

    def get_rpc_message(
        self, endpoint: str, params: AnyLMStudioStruct | None
    ) -> DictObject:
        """Get the message to send to initiate this remote procedure call."""
        return self._rpc.get_rpc_message(endpoint, params)

    @property
    def timeout(self) -> float | None:
        """Permitted time to wait for a reply to this call."""
        timeout = self._timeout
        if timeout is ...:
            return _DEFAULT_TIMEOUT
        return timeout

    def receive_result(self) -> Any:
        """Receive call response on the receive queue."""
        try:
            message = self._get_message(self.timeout)
        except SyncFutureTimeout:
            raise LMStudioTimeoutError from None
        if message is None:
            raise LMStudioWebsocketError("Client unexpectedly disconnected.")
        return self._rpc.handle_rx_message(message)


class SyncLMStudioWebsocket(LMStudioWebsocket[SyncToAsyncWebsocketBridge]):
    """Synchronous websocket client that handles demultiplexing of reply messages."""

    def __init__(
        self,
        ws_thread: AsyncWebsocketThread,
        ws_url: str,
        auth_details: DictObject,
        log_context: LogEventContext | None = None,
    ) -> None:
        """Initialize synchronous websocket client."""
        super().__init__(ws_url, auth_details, log_context)
        self._ws_thread = ws_thread

    @property
    def _httpx_ws(self) -> AsyncWebSocketSession | None:
        # Underlying HTTPX session is accessible for testing purposes
        ws_thread = self._ws
        if ws_thread is None:
            return None
        return ws_thread._ws

    def __enter__(self) -> Self:
        # Handle reentrancy the same way files do:
        # allow nested use as a CM, but close on the first exit
        if self._ws is None:
            self.connect()
        return self

    def __exit__(self, *args: Any) -> None:
        self.disconnect()

    def connect(self) -> Self:
        """Connect to and authenticate with the LM Studio API."""
        self._fail_if_connected("Attempted to connect already connected websocket")
        ws = SyncToAsyncWebsocketBridge(
            self._ws_thread,
            self._ws_url,
            self._auth_details,
            self._logger.event_context,
        )
        if not ws.connect():
            if ws._connection_failure is not None:
                raise self._get_connection_failure_error(ws._connection_failure)
            if ws._auth_failure is not None:
                raise self._get_auth_failure_error(ws._auth_failure)
            self._logger.error("Connection failed, but no failure reason reported.")
            raise self._get_connection_failure_error()
        self._ws = ws
        return self

    def disconnect(self) -> None:
        """Drop the LM Studio API connection."""
        ws = self._ws
        self._ws = None
        if ws is not None:
            self._logger.debug(f"Disconnecting websocket session ({self._ws_url})")
            ws.notify_client_termination_threadsafe()
            ws.disconnect()
        self._logger.info(f"Websocket session disconnected ({self._ws_url})")

    close = disconnect

    def _send_json(self, message: DictObject) -> None:
        # Callers are expected to call `_ensure_connected` before this method
        ws = self._ws
        assert ws is not None
        # Background thread handles the exception conversion
        ws.send_json(message)

    def _connect_to_endpoint(self, channel: SyncChannel[Any]) -> None:
        """Connect channel to specified endpoint."""
        self._ensure_connected("open channel endpoints")
        create_message = channel.get_creation_message()
        self._logger.debug("Connecting channel endpoint", json=create_message)
        self._send_json(create_message)

    @contextmanager
    def open_channel(
        self,
        endpoint: ChannelEndpoint[T, Any, Any],
    ) -> Generator[SyncChannel[T], None, None]:
        """Open a streaming channel over the websocket."""
        ws = self._ws
        assert ws is not None
        with ws.open_channel() as (channel_id, getter):
            channel = SyncChannel(
                channel_id,
                getter,
                endpoint,
                self._send_json,
                self._logger.event_context,
            )
            self._connect_to_endpoint(channel)
            yield channel

    def _send_call(
        self,
        rpc: SyncRemoteCall,
        endpoint: str,
        params: AnyLMStudioStruct | None = None,
    ) -> None:
        """Initiate remote call to specified endpoint."""
        self._ensure_connected("send remote procedure call")
        call_message = rpc.get_rpc_message(endpoint, params)
        # TODO: Improve logging for large requests (such as file uploads)
        #       without requiring explicit special casing here
        logged_message: DictObject
        if call_message.get("endpoint") == "uploadFileBase64":
            logged_message = _redact_json(call_message)
        else:
            logged_message = call_message
        self._logger.debug("Sending RPC request", json=logged_message)
        self._send_json(call_message)

    def remote_call(
        self,
        endpoint: str,
        params: AnyLMStudioStruct | None,
        notice_prefix: str = "RPC",
    ) -> Any:
        """Make a remote procedure call over the websocket."""
        ws = self._ws
        assert ws is not None
        with ws.start_call() as (call_id, getter):
            rpc = SyncRemoteCall(
                call_id, getter, self._logger.event_context, notice_prefix
            )
            self._send_call(rpc, endpoint, params)
            return rpc.receive_result()


class _SyncSession(ClientSession["Client", SyncLMStudioWebsocket]):
    """Sync client session interfaces applicable to all API namespaces."""

    def __init__(self, client: "Client") -> None:
        """Initialize synchronous API client session."""
        super().__init__(client)
        self._resources = ExitStack()

    def _ensure_connected(self) -> None:
        # Allow lazy connection of the session websocket
        if self._lmsws is None:
            self.connect()

    def __enter__(self) -> Self:
        # Handle reentrancy the same way files do:
        # allow nested use as a CM, but close on the first exit
        self._ensure_connected()
        return self

    def __exit__(self, *args: Any) -> None:
        self.disconnect()

    @sdk_public_api()
    def connect(self) -> SyncLMStudioWebsocket:
        """Connect the client session."""
        self._fail_if_connected("Attempted to connect already connected session")
        api_host = self._client.api_host
        namespace = self.API_NAMESPACE
        if namespace is None:
            raise LMStudioClientError(
                f"No API namespace defined for {type(self).__name__}"
            )
        session_url = f"ws://{api_host}/{namespace}"
        resources = self._resources
        self._lmsws = lmsws = resources.enter_context(
            SyncLMStudioWebsocket(
                self._client._ws_thread, session_url, self._client._auth_details
            )
        )
        return lmsws

    @sdk_public_api()
    def disconnect(self) -> None:
        """Disconnect the client session."""
        self._lmsws = None
        self._resources.close()

    close = disconnect

    # To allow for client level management of the session lifecycles
    # without requiring network I/O on property access, we implicitly
    # connect the websocket (if necessary) when sending requests

    @contextmanager
    def _create_channel(
        self,
        endpoint: ChannelEndpoint[T, Any, Any],
    ) -> Generator[SyncChannel[T], None, None]:
        """Connect a channel to an LM Studio streaming endpoint."""
        self._ensure_connected()
        lmsws = self._get_lmsws("create channels")
        with lmsws.open_channel(endpoint) as channel:
            yield channel

    @sdk_public_api()
    def remote_call(
        self,
        endpoint: str,
        params: AnyLMStudioStruct | None = None,
        notice_prefix: str = "RPC",
    ) -> Any:
        """Send a remote call to the given RPC endpoint and wait for the result."""
        self._ensure_connected()
        lmsws = self._get_lmsws("make remote calls")
        return lmsws.remote_call(endpoint, params, notice_prefix)


TSyncSessionModel = TypeVar(
    "TSyncSessionModel", bound="_SyncSessionModel[Any, Any, Any, Any]"
)
TModelHandle = TypeVar("TModelHandle", bound="SyncModelHandle[Any]")


class DownloadedModel(
    Generic[
        TModelInfo,
        TSyncSessionModel,
        TLoadConfig,
        TLoadConfigDict,
        TModelHandle,
    ],
    DownloadedModelBase[TModelInfo, TSyncSessionModel],
):
    @sdk_public_api()
    def load_new_instance(
        self,
        *,
        ttl: int | None = DEFAULT_TTL,
        instance_identifier: str | None = None,
        config: TLoadConfig | TLoadConfigDict | None = None,
        on_load_progress: ModelLoadingCallback | None = None,
    ) -> TModelHandle:
        """Load this model with the given identifier and configuration.

        Note: details of configuration fields may change in SDK feature releases.
        """
        handle: TModelHandle = self._session._load_new_instance(
            self.model_key, instance_identifier, ttl, config, on_load_progress
        )
        return handle

    @sdk_public_api()
    def model(
        self,
        *,
        ttl: int | None = DEFAULT_TTL,
        config: TLoadConfig | TLoadConfigDict | None = None,
        on_load_progress: ModelLoadingCallback | None = None,
    ) -> TModelHandle:
        """Retrieve model with default identifier, or load it with given configuration.

        Note: configuration of retrieved model is NOT checked against the given config.
        Note: details of configuration fields may change in SDK feature releases.
        """
        # Call _get_or_load directly, since we have a model identifier
        handle: TModelHandle = self._session._get_or_load(
            self.model_key, ttl, config, on_load_progress
        )
        return handle


class DownloadedEmbeddingModel(
    DownloadedModel[
        EmbeddingModelInfo,
        "_SyncSessionEmbedding",
        EmbeddingLoadModelConfig,
        EmbeddingLoadModelConfigDict,
        "EmbeddingModel",
    ],
):
    """Download listing for an embedding model."""

    def __init__(
        self, model_info: DictObject, session: "_SyncSessionEmbedding"
    ) -> None:
        """Initialize downloaded embedding model details."""
        super().__init__(EmbeddingModelInfo, model_info, session)


class DownloadedLlm(
    DownloadedModel[
        LlmInfo,
        "_SyncSessionLlm",
        LlmLoadModelConfig,
        LlmLoadModelConfigDict,
        "LLM",
    ]
):
    """Download listing for an LLM."""

    def __init__(self, model_info: DictObject, session: "_SyncSessionLlm") -> None:
        """Initialize downloaded embedding model details."""
        super().__init__(LlmInfo, model_info, session)


AnyDownloadedModel: TypeAlias = DownloadedModel[Any, Any, Any, Any, Any]


class _SyncSessionSystem(_SyncSession):
    """Sync client session for the system namespace."""

    API_NAMESPACE = "system"

    @sdk_public_api()
    def list_downloaded_models(self) -> Sequence[AnyDownloadedModel]:
        """Get the list of all downloaded models that are available for loading."""
        # The list of downloaded models is only available via the system API namespace
        models = self.remote_call("listDownloadedModels")
        return [self._process_download_listing(m) for m in models]

    def _process_download_listing(self, model_info: DictObject) -> AnyDownloadedModel:
        model_type = model_info.get("type")
        if model_type is None:
            raise LMStudioClientError(
                f"No 'type' field in download listing: {model_info}"
            )
        match model_type:
            case "embedding":
                return DownloadedEmbeddingModel(model_info, self._client.embedding)
            case "llm":
                return DownloadedLlm(model_info, self._client.llm)
        raise LMStudioClientError(
            f"Unknown model type {model_type!r} in download listing: {model_info}"
        )


class _SyncSessionFiles(_SyncSession):
    """Sync client session for the files namespace."""

    API_NAMESPACE = "files"

    def _fetch_file_handle(self, file_data: _LocalFileData) -> FileHandle:
        handle = self.remote_call("uploadFileBase64", file_data._as_fetch_param())
        # Returned dict provides the handle identifier, file type, and size in bytes
        # Add the extra fields needed for a FileHandle (aka ChatMessagePartFileData)
        handle["name"] = file_data.name
        handle["type"] = "file"
        return load_struct(handle, FileHandle)

    # Not yet implemented (server API only supports the same file types as prepare_image)
    # @sdk_public_api()
    def _prepare_file(self, src: LocalFileInput, name: str | None = None) -> FileHandle:
        """Add a file to the server. Returns a file handle for use in prediction requests."""
        file_data = _LocalFileData(src, name)
        return self._fetch_file_handle(file_data)

    @sdk_public_api()
    def prepare_image(self, src: LocalFileInput, name: str | None = None) -> FileHandle:
        """Add an image to the server. Returns a file handle for use in prediction requests."""
        file_data = _LocalFileData(src, name)
        return self._fetch_file_handle(file_data)


class ModelDownloadOption(ModelDownloadOptionBase[_SyncSession]):
    """A single download option for a model search result."""

    @sdk_public_api()
    def download(
        self,
        on_progress: DownloadProgressCallback | None = None,
        on_finalize: DownloadFinalizedCallback | None = None,
    ) -> str:
        """Download a model and get its path for loading."""
        endpoint = self._get_download_endpoint(on_progress, on_finalize)
        with self._session._create_channel(endpoint) as channel:
            return channel.wait_for_result()


class AvailableModel(AvailableModelBase[_SyncSession]):
    """A model available for download from the model repository."""

    @sdk_public_api()
    def get_download_options(
        self,
    ) -> Sequence[ModelDownloadOption]:
        """Get the download options for the specified model."""
        params = self._get_download_query_params()
        options = self._session.remote_call("getModelDownloadOptions", params)
        final = []
        for m in options["results"]:
            final.append(ModelDownloadOption(m, self._session))
        return final


class _SyncSessionRepository(_SyncSession):
    """Sync client session for the repository namespace."""

    API_NAMESPACE = "repository"

    @sdk_public_api()
    def search_models(
        self,
        search_term: str | None = None,
        limit: int | None = None,
        compatibility_types: list[ModelCompatibilityType] | None = None,
    ) -> Sequence[AvailableModel]:
        """Search for downloadable models satisfying a search query."""
        params = self._get_model_search_params(search_term, limit, compatibility_types)
        models = self.remote_call("searchModels", params)
        return [AvailableModel(m, self) for m in models["results"]]


TDownloadedModel = TypeVar("TDownloadedModel", bound=AnyDownloadedModel)


class _SyncSessionModel(
    _SyncSession,
    Generic[TModelHandle, TLoadConfig, TLoadConfigDict, TDownloadedModel],
):
    """Sync client session for a model (LLM/embedding) namespace."""

    _API_TYPES: Type[ModelSessionTypes[TLoadConfig]]

    @property
    def _system_session(self) -> _SyncSessionSystem:
        return self._client.system

    @property
    def _files_session(self) -> _SyncSessionFiles:
        return self._client.files

    def _get_load_config(self, model_specifier: AnyModelSpecifier) -> AnyLoadConfig:
        """Get the model load config for the specified model."""
        # Note that the configuration reported here uses the *server* config names,
        # not the attributes used to set the configuration in the client SDK
        params = self._API_TYPES.REQUEST_LOAD_CONFIG._from_api_dict(
            {
                "specifier": _model_spec_to_api_dict(model_specifier),
            }
        )
        config = self.remote_call("getLoadConfig", params)
        result_type = self._API_TYPES.MODEL_LOAD_CONFIG
        return result_type._from_any_api_dict(parse_server_config(config))

    def _get_api_model_info(self, model_specifier: AnyModelSpecifier) -> Any:
        """Get the raw model info (if any) for a model matching the given criteria."""
        params = self._API_TYPES.REQUEST_MODEL_INFO._from_api_dict(
            {
                "specifier": _model_spec_to_api_dict(model_specifier),
                "throwIfNotFound": True,
            }
        )
        return self.remote_call("getModelInfo", params)

    @sdk_public_api()
    def get_model_info(self, model_specifier: AnyModelSpecifier) -> ModelInstanceInfo:
        """Get the model info (if any) for a model matching the given criteria."""
        response = self._get_api_model_info(model_specifier)
        model_info = self._API_TYPES.MODEL_INSTANCE_INFO._from_any_api_dict(response)
        return model_info

    def _get_context_length(self, model_specifier: AnyModelSpecifier) -> int:
        """Get the context length of the specified model."""
        raw_model_info = self._get_api_model_info(model_specifier)
        return int(raw_model_info.get("contextLength", -1))

    def _count_tokens(self, model_specifier: AnyModelSpecifier, input: str) -> int:
        params = EmbeddingRpcCountTokensParameter._from_api_dict(
            {
                "specifier": _model_spec_to_api_dict(model_specifier),
                "inputString": input,
            }
        )
        response = self.remote_call("countTokens", params)
        return int(response["tokenCount"])

    # Private helper method to allow the main API to easily accept iterables
    def _tokenize_text(
        self, model_specifier: AnyModelSpecifier, input: str
    ) -> Sequence[int]:
        params = EmbeddingRpcTokenizeParameter._from_api_dict(
            {
                "specifier": _model_spec_to_api_dict(model_specifier),
                "inputString": input,
            }
        )
        response = self.remote_call("tokenize", params)
        return response.get("tokens", []) if response else []

    # Alas, type hints don't properly support distinguishing str vs Iterable[str]:
    #     https://github.com/python/typing/issues/256
    def _tokenize(
        self, model_specifier: AnyModelSpecifier, input: str | Iterable[str]
    ) -> Sequence[int] | Sequence[Sequence[int]]:
        """Tokenize the input string(s) using the specified model."""
        if isinstance(input, str):
            return self._tokenize_text(model_specifier, input)
        return [self._tokenize_text(model_specifier, i) for i in input]

    @abstractmethod
    def _create_handle(self, model_identifier: str) -> TModelHandle:
        """Get a symbolic handle to the specified model."""
        ...

    @sdk_public_api()
    def model(
        self,
        model_key: str | None = None,
        /,
        *,
        ttl: int | None = DEFAULT_TTL,
        config: TLoadConfig | TLoadConfigDict | None = None,
        on_load_progress: ModelLoadingCallback | None = None,
    ) -> TModelHandle:
        """Get a handle to the specified model (loading it if necessary).

        Note: configuration of retrieved model is NOT checked against the given config.
        Note: details of configuration fields may change in SDK feature releases.
        """
        if model_key is None:
            # Should this raise an error if a config is supplied?
            return self._get_any()
        return self._get_or_load(model_key, ttl, config, on_load_progress)

    @sdk_public_api()
    def list_loaded(self) -> Sequence[TModelHandle]:
        """Get the list of currently loaded models."""
        models = self.remote_call("listLoaded")
        return [self._create_handle(m["identifier"]) for m in models]

    @sdk_public_api()
    def unload(self, model_identifier: str) -> None:
        """Unload the specified model."""
        params = self._API_TYPES.REQUEST_UNLOAD(identifier=model_identifier)
        self.remote_call("unloadModel", params)

    # N.B. Canceling a load from the UI doesn't update the load process for a while.
    # Fortunately, this is not our fault. The server just delays in broadcasting it.
    @sdk_public_api()
    def load_new_instance(
        self,
        model_key: str,
        instance_identifier: str | None = None,
        *,
        ttl: int | None = DEFAULT_TTL,
        config: TLoadConfig | TLoadConfigDict | None = None,
        on_load_progress: ModelLoadingCallback | None = None,
    ) -> TModelHandle:
        """Load the specified model with the given identifier and configuration.

        Note: details of configuration fields may change in SDK feature releases.
        """
        return self._load_new_instance(
            model_key, instance_identifier, ttl, config, on_load_progress
        )

    def _load_new_instance(
        self,
        model_key: str,
        instance_identifier: str | None,
        ttl: int | None,
        config: TLoadConfig | TLoadConfigDict | None,
        on_load_progress: ModelLoadingCallback | None,
    ) -> TModelHandle:
        channel_type = self._API_TYPES.REQUEST_NEW_INSTANCE
        config_type = self._API_TYPES.MODEL_LOAD_CONFIG
        endpoint = LoadModelEndpoint(
            model_key,
            instance_identifier,
            ttl,
            channel_type,
            config_type,
            config,
            on_load_progress,
        )
        with self._create_channel(endpoint) as channel:
            result = channel.wait_for_result()
            return self._create_handle(result.identifier)

    def _get_or_load(
        self,
        model_key: str,
        ttl: int | None,
        config: TLoadConfig | TLoadConfigDict | None,
        on_load_progress: ModelLoadingCallback | None,
    ) -> TModelHandle:
        """Get the specified model if it is already loaded, otherwise load it."""
        channel_type = self._API_TYPES.REQUEST_GET_OR_LOAD
        config_type = self._API_TYPES.MODEL_LOAD_CONFIG
        endpoint = GetOrLoadEndpoint(
            model_key, ttl, channel_type, config_type, config, on_load_progress
        )
        with self._create_channel(endpoint) as channel:
            result = channel.wait_for_result()
            return self._create_handle(result.identifier)

    def _get_any(self) -> TModelHandle:
        """Get a handle to any loaded model."""
        loaded_models = self.list_loaded()
        if not loaded_models:
            raise LMStudioClientError(
                f"Could not get model handle in namespace {self.API_NAMESPACE} (no models are currently loaded)."
            )
        return self._create_handle(loaded_models[0].identifier)

    @classmethod
    def _is_relevant_model(cls, model: AnyDownloadedModel) -> TypeIs[TDownloadedModel]:
        return bool(model.type == cls.API_NAMESPACE)

    @sdk_public_api()
    def list_downloaded(self) -> Sequence[TDownloadedModel]:
        """Get the list of currently downloaded models that are available for loading."""
        models = self._system_session.list_downloaded_models()
        return [m for m in models if self._is_relevant_model(m)]

    def _fetch_file_handle(self, file_data: _LocalFileData) -> FileHandle:
        return self._files_session._fetch_file_handle(file_data)


SyncPredictionChannel: TypeAlias = SyncChannel[PredictionResult]
SyncPredictionCM: TypeAlias = ContextManager[SyncPredictionChannel]


class PredictionStream(PredictionStreamBase):
    """Sync context manager for an ongoing prediction process."""

    def __init__(
        self,
        channel_cm: SyncPredictionCM,
        endpoint: PredictionEndpoint,
    ) -> None:
        """Initialize a prediction process representation."""
        self._resources = ExitStack()
        self._channel_cm: SyncPredictionCM = channel_cm
        self._channel: SyncPredictionChannel | None = None
        super().__init__(endpoint)

    @sdk_public_api()
    def start(self) -> None:
        """Send the prediction request."""
        if self._is_finished:
            raise LMStudioRuntimeError("Prediction result has already been received.")
        if self._is_started:
            raise LMStudioRuntimeError("Prediction request has already been sent.")
        # The given channel context manager is set up to send the relevant request
        self._channel = self._resources.enter_context(self._channel_cm)
        self._mark_started()

    @sdk_public_api()
    def close(self) -> None:
        """Terminate the prediction processing (if not already terminated)."""
        # Cancel the prediction (if unfinished) and release acquired resources
        if self._is_started and not self._is_finished:
            self._set_error(
                LMStudioCancelledError(
                    "Prediction cancelled unexpectedly: please use .cancel()"
                )
            )
        self._channel = None
        self._resources.close()

    def __enter__(self) -> Self:
        if self._channel is None:
            self.start()
        return self

    def __exit__(
        self,
        _exc_type: Type[BaseException] | None,
        exc_val: BaseException | None,
        _exc_tb: TracebackType | None,
    ) -> None:
        if exc_val and not self._is_finished:
            self._set_error(exc_val)
        self.close()

    def __iter__(self) -> Iterator[LlmPredictionFragment]:
        for event in self._iter_events():
            if isinstance(event, PredictionFragmentEvent):
                yield event.arg

    def _iter_events(self) -> Iterator[PredictionRxEvent]:
        endpoint = self._endpoint
        with self:
            assert self._channel is not None
            for contents in self._channel.rx_stream():
                for event in endpoint.iter_message_events(contents):
                    endpoint.handle_rx_event(event)
                    yield event
                if endpoint.is_finished:
                    break
            self._mark_finished()

    @sdk_public_api()
    def wait_for_result(self) -> PredictionResult:
        """Wait for the result of the prediction."""
        for _ in self:
            pass
        return self.result()

    @sdk_public_api()
    def cancel(self) -> None:
        """Cancel the prediction process."""
        if not self._is_finished and self._channel:
            self._mark_cancelled()
            self._channel.cancel()


class _SyncSessionLlm(
    _SyncSessionModel[
        "LLM",
        LlmLoadModelConfig,
        LlmLoadModelConfigDict,
        DownloadedLlm,
    ]
):
    """Sync client session for LLM namespace."""

    API_NAMESPACE = "llm"
    _API_TYPES = ModelTypesLlm

    def __init__(self, client: "Client") -> None:
        """Initialize API client session for LLM interaction."""
        super().__init__(client)

    def _create_handle(self, model_identifier: str) -> "LLM":
        """Create a symbolic handle to the specified LLM model."""
        return LLM(model_identifier, self)

    def _complete_stream(
        self,
        model_specifier: AnyModelSpecifier,
        prompt: str,
        *,
        response_format: ResponseSchema | None = None,
        config: LlmPredictionConfig | LlmPredictionConfigDict | None = None,
        preset: str | None = None,
        on_message: PredictionMessageCallback | None = None,
        on_first_token: PredictionFirstTokenCallback | None = None,
        on_prediction_fragment: PredictionFragmentCallback | None = None,
        on_prompt_processing_progress: PromptProcessingCallback | None = None,
    ) -> PredictionStream:
        """Request a one-off prediction without any context and stream the generated tokens.

        Note: details of configuration fields may change in SDK feature releases.
        """
        endpoint = CompletionEndpoint(
            model_specifier,
            prompt,
            response_format,
            config,
            preset,
            on_message,
            on_first_token,
            on_prediction_fragment,
            on_prompt_processing_progress,
        )
        channel_cm = self._create_channel(endpoint)
        prediction_stream = PredictionStream(channel_cm, endpoint)
        return prediction_stream

    def _respond_stream(
        self,
        model_specifier: AnyModelSpecifier,
        history: Chat | ChatHistoryDataDict | str,
        *,
        response_format: ResponseSchema | None = None,
        config: LlmPredictionConfig | LlmPredictionConfigDict | None = None,
        preset: str | None = None,
        on_message: PredictionMessageCallback | None = None,
        on_first_token: PredictionFirstTokenCallback | None = None,
        on_prediction_fragment: PredictionFragmentCallback | None = None,
        on_prompt_processing_progress: PromptProcessingCallback | None = None,
    ) -> PredictionStream:
        """Request a response in an ongoing assistant chat session and stream the generated tokens.

        Note: details of configuration fields may change in SDK feature releases.
        """
        if not isinstance(history, Chat):
            history = Chat.from_history(history)
        endpoint = ChatResponseEndpoint(
            model_specifier,
            history,
            response_format,
            config,
            preset,
            on_message,
            on_first_token,
            on_prediction_fragment,
            on_prompt_processing_progress,
        )
        channel_cm = self._create_channel(endpoint)
        prediction_stream = PredictionStream(channel_cm, endpoint)
        return prediction_stream

    def _apply_prompt_template(
        self,
        model_specifier: AnyModelSpecifier,
        history: Chat | ChatHistoryDataDict | str,
        opts: LlmApplyPromptTemplateOpts | LlmApplyPromptTemplateOptsDict = {},
    ) -> str:
        """Apply a prompt template to the given history."""
        if not isinstance(history, Chat):
            history = Chat.from_history(history)
        if not isinstance(opts, LlmApplyPromptTemplateOpts):
            opts = LlmApplyPromptTemplateOpts.from_dict(opts)
        params = LlmRpcApplyPromptTemplateParameter._from_api_dict(
            {
                "specifier": _model_spec_to_api_dict(model_specifier),
                "history": history._get_history_for_prediction(),
                "predictionConfigStack": {"layers": []},
                "opts": opts.to_dict(),
            }
        )
        response = self.remote_call("applyPromptTemplate", params)
        return response.get("formatted", "") if response else ""


class _SyncSessionEmbedding(
    _SyncSessionModel[
        "EmbeddingModel",
        EmbeddingLoadModelConfig,
        EmbeddingLoadModelConfigDict,
        DownloadedEmbeddingModel,
    ]
):
    """Sync client session for embedding namespace."""

    API_NAMESPACE = "embedding"
    _API_TYPES = ModelTypesEmbedding

    def __init__(self, client: "Client") -> None:
        """Initialize API client session for embedding model interaction."""
        super().__init__(client)

    def _create_handle(self, model_identifier: str) -> "EmbeddingModel":
        """Create a symbolic handle to the specified embedding model."""
        return EmbeddingModel(model_identifier, self)

    # Private helper method to allow the main API to easily accept iterables
    def _embed_text(
        self, model_specifier: AnyModelSpecifier, input: str
    ) -> Sequence[float]:
        params = EmbeddingRpcEmbedStringParameter._from_api_dict(
            {
                "modelSpecifier": _model_spec_to_api_dict(model_specifier),
                "inputString": input,
            }
        )

        response = self.remote_call("embedString", params)
        return response.get("embedding", []) if response else []

    # Alas, type hints don't properly support distinguishing str vs Iterable[str]:
    #     https://github.com/python/typing/issues/256
    def _embed(
        self, model_specifier: AnyModelSpecifier, input: str | Iterable[str]
    ) -> Sequence[float] | Sequence[Sequence[float]]:
        """Request embedding vectors for the given input string(s)."""
        if isinstance(input, str):
            return self._embed_text(model_specifier, input)
        return [self._embed_text(model_specifier, i) for i in input]


class SyncModelHandle(ModelHandleBase[TSyncSessionModel]):
    """Reference to a loaded LM Studio model."""

    @sdk_public_api()
    def unload(self) -> None:
        """Unload this model."""
        self._session.unload(self.identifier)

    @sdk_public_api()
    def get_info(self) -> ModelInstanceInfo:
        """Get the model info for this model."""
        return self._session.get_model_info(self.identifier)

    @sdk_public_api()
    def get_load_config(self) -> AnyLoadConfig:
        """Get the model load config for this model."""
        return self._session._get_load_config(self.identifier)

    # Alas, type hints don't properly support distinguishing str vs Iterable[str]:
    #     https://github.com/python/typing/issues/256
    @sdk_public_api()
    def tokenize(
        self, input: str | Iterable[str]
    ) -> Sequence[int] | Sequence[Sequence[int]]:
        """Tokenize the input string(s) using this model."""
        return self._session._tokenize(self.identifier, input)

    @sdk_public_api()
    def count_tokens(self, input: str) -> int:
        """Report the number of tokens needed for the input string using this model."""
        return self._session._count_tokens(self.identifier, input)

    @sdk_public_api()
    def get_context_length(self) -> int:
        """Get the context length of this model."""
        return self._session._get_context_length(self.identifier)


AnySyncModel: TypeAlias = SyncModelHandle[Any]


class LLM(SyncModelHandle[_SyncSessionLlm]):
    """Reference to a loaded LLM model."""

    @sdk_public_api()
    def complete_stream(
        self,
        prompt: str,
        *,
        response_format: ResponseSchema | None = None,
        config: LlmPredictionConfig | LlmPredictionConfigDict | None = None,
        preset: str | None = None,
        on_message: PredictionMessageCallback | None = None,
        on_first_token: PredictionFirstTokenCallback | None = None,
        on_prediction_fragment: PredictionFragmentCallback | None = None,
        on_prompt_processing_progress: PromptProcessingCallback | None = None,
    ) -> PredictionStream:
        """Request a one-off prediction without any context and stream the generated tokens.

        Note: details of configuration fields may change in SDK feature releases.
        """
        return self._session._complete_stream(
            self.identifier,
            prompt,
            response_format=response_format,
            config=config,
            preset=preset,
            on_message=on_message,
            on_first_token=on_first_token,
            on_prediction_fragment=on_prediction_fragment,
            on_prompt_processing_progress=on_prompt_processing_progress,
        )

    @sdk_public_api()
    def complete(
        self,
        prompt: str,
        *,
        response_format: ResponseSchema | None = None,
        config: LlmPredictionConfig | LlmPredictionConfigDict | None = None,
        preset: str | None = None,
        on_message: PredictionMessageCallback | None = None,
        on_first_token: PredictionFirstTokenCallback | None = None,
        on_prediction_fragment: PredictionFragmentCallback | None = None,
        on_prompt_processing_progress: PromptProcessingCallback | None = None,
    ) -> PredictionResult:
        """Request a one-off prediction without any context.

        Note: details of configuration fields may change in SDK feature releases.
        """
        prediction_stream = self._session._complete_stream(
            self.identifier,
            prompt,
            response_format=response_format,
            config=config,
            preset=preset,
            on_message=on_message,
            on_first_token=on_first_token,
            on_prediction_fragment=on_prediction_fragment,
            on_prompt_processing_progress=on_prompt_processing_progress,
        )
        for _ in prediction_stream:
            # No yield in body means iterator reliably provides
            # prompt resource cleanup on coroutine cancellation
            pass
        return prediction_stream.result()

    @sdk_public_api()
    def respond_stream(
        self,
        history: Chat | ChatHistoryDataDict | str,
        *,
        response_format: ResponseSchema | None = None,
        config: LlmPredictionConfig | LlmPredictionConfigDict | None = None,
        preset: str | None = None,
        on_message: PredictionMessageCallback | None = None,
        on_first_token: PredictionFirstTokenCallback | None = None,
        on_prediction_fragment: PredictionFragmentCallback | None = None,
        on_prompt_processing_progress: PromptProcessingCallback | None = None,
    ) -> PredictionStream:
        """Request a response in an ongoing assistant chat session and stream the generated tokens.

        Note: details of configuration fields may change in SDK feature releases.
        """
        return self._session._respond_stream(
            self.identifier,
            history,
            response_format=response_format,
            config=config,
            preset=preset,
            on_message=on_message,
            on_first_token=on_first_token,
            on_prediction_fragment=on_prediction_fragment,
            on_prompt_processing_progress=on_prompt_processing_progress,
        )

    @sdk_public_api()
    def respond(
        self,
        history: Chat | ChatHistoryDataDict | str,
        *,
        response_format: ResponseSchema | None = None,
        config: LlmPredictionConfig | LlmPredictionConfigDict | None = None,
        preset: str | None = None,
        on_message: PredictionMessageCallback | None = None,
        on_first_token: PredictionFirstTokenCallback | None = None,
        on_prediction_fragment: PredictionFragmentCallback | None = None,
        on_prompt_processing_progress: PromptProcessingCallback | None = None,
    ) -> PredictionResult:
        """Request a response in an ongoing assistant chat session.

        Note: details of configuration fields may change in SDK feature releases.
        """
        prediction_stream = self._session._respond_stream(
            self.identifier,
            history,
            response_format=response_format,
            config=config,
            preset=preset,
            on_message=on_message,
            on_first_token=on_first_token,
            on_prediction_fragment=on_prediction_fragment,
            on_prompt_processing_progress=on_prompt_processing_progress,
        )
        for _ in prediction_stream:
            # No yield in body means iterator reliably provides
            # prompt resource cleanup on coroutine cancellation
            pass
        return prediction_stream.result()

    # TODO: Improve code sharing between sync and async multi-round predictions
    # TODO: Accept async tools in the tools iterable
    @sdk_public_api()
    def act(
        self,
        chat: Chat | ChatHistoryDataDict | str,
        tools: Iterable[ToolDefinition],
        *,
        max_prediction_rounds: int | None = None,
        max_parallel_tool_calls: int | None = 1,
        config: LlmPredictionConfig | LlmPredictionConfigDict | None = None,
        preset: str | None = None,
        on_message: Callable[[AssistantResponse | ToolResultMessage], Any]
        | None = None,
        on_first_token: Callable[[int], Any] | None = None,
        on_prediction_fragment: Callable[[LlmPredictionFragment, int], Any]
        | None = None,
        on_round_start: Callable[[int], Any] | None = None,
        on_round_end: Callable[[int], Any] | None = None,
        on_prediction_completed: Callable[[PredictionRoundResult], Any] | None = None,
        on_prompt_processing_progress: Callable[[float, int], Any] | None = None,
        handle_invalid_tool_request: Callable[
            [LMStudioPredictionError, ToolCallRequest | None], str | None
        ]
        | None = None,
    ) -> ActResult:
        """Request a response (with implicit tool use) in an ongoing agent chat session.

        Note: details of configuration fields may change in SDK feature releases.
        """
        start_time = time.perf_counter()
        # It is not yet possible to combine tool calling with requests for structured responses
        response_format = None
        agent_chat: Chat = Chat.from_history(chat)
        del chat  # Avoid any further access to the input chat history
        # Multiple rounds, until all tool calls are resolved or limit is reached
        round_counter: Iterable[int]
        if max_prediction_rounds is not None:
            if max_prediction_rounds < 1:
                raise LMStudioValueError(
                    f"Max prediction rounds must be at least 1 ({max_prediction_rounds!r} given)"
                )
            round_counter = range(max_prediction_rounds)
            final_round_index = max_prediction_rounds - 1
        else:
            # Do not force a final round when no limit is specified
            final_round_index = -1
            round_counter = itertools.count()
        llm_tool_args = ChatResponseEndpoint.parse_tools(tools)
        del tools
        # Supply the round index to any endpoint callbacks that expect one
        round_index: int
        on_first_token_for_endpoint: PredictionFirstTokenCallback | None = None
        if on_first_token is not None:

            def _wrapped_on_first_token() -> None:
                assert on_first_token is not None
                on_first_token(round_index)

            on_first_token_for_endpoint = _wrapped_on_first_token
        on_prediction_fragment_for_endpoint: PredictionFragmentCallback | None = None
        if on_prediction_fragment is not None:

            def _wrapped_on_prediction_fragment(
                fragment: LlmPredictionFragment,
            ) -> None:
                assert on_prediction_fragment is not None
                on_prediction_fragment(fragment, round_index)

            on_prediction_fragment_for_endpoint = _wrapped_on_prediction_fragment
        on_prompt_processing_for_endpoint: PromptProcessingCallback | None = None
        if on_prompt_processing_progress is not None:

            def _wrapped_on_prompt_processing_progress(progress: float) -> None:
                assert on_prompt_processing_progress is not None
                on_prompt_processing_progress(progress, round_index)

            on_prompt_processing_for_endpoint = _wrapped_on_prompt_processing_progress
        # TODO: Implementation to this point is common between the sync and async APIs
        # (aside from the allow_async flag when parsing the tool definitions)
        # Implementation past this point differs (as the async API uses the loop's executor)

        # Request predictions until no more tool call requests are received in response
        # (or the maximum number of prediction rounds is reached)
        with ThreadPoolExecutor(max_parallel_tool_calls) as pool:
            for round_index in round_counter:
                self._logger.debug(
                    "Starting .act() prediction round", round_index=round_index
                )
                if on_round_start is not None:
                    err_msg = f"Round start callback failed for {self!r}"
                    with sdk_callback_invocation(err_msg, self._logger):
                        on_round_start(round_index)
                # Update the endpoint definition on each iteration in order to:
                # * update the chat history with the previous round result
                # * be able to disallow tool use when the rounds are limited
                # TODO: Refactor endpoint API to avoid repeatedly performing the
                #       LlmPredictionConfig -> KvConfigStack transformation
                endpoint = ChatResponseEndpoint(
                    self.identifier,
                    agent_chat,
                    response_format,
                    config,
                    preset,
                    None,  # on_message is invoked directly
                    on_first_token_for_endpoint,
                    on_prediction_fragment_for_endpoint,
                    on_prompt_processing_for_endpoint,
                    handle_invalid_tool_request,
                    *(
                        llm_tool_args
                        if round_index != final_round_index
                        else (None, None)
                    ),
                )
                channel_cm = self._session._create_channel(endpoint)
                prediction_stream = PredictionStream(channel_cm, endpoint)
                tool_call_requests: list[ToolCallRequest] = []
                pending_tool_calls: dict[SyncFuture[Any], ToolCallRequest] = {}
                for event in prediction_stream._iter_events():
                    if isinstance(event, PredictionToolCallEvent):
                        tool_call_request = event.arg
                        tool_call_requests.append(tool_call_request)
                        # TODO: Also handle async tool calls here
                        tool_call = endpoint.request_tool_call(tool_call_request)
                        pending_tool_calls[pool.submit(tool_call)] = tool_call_request
                prediction = prediction_stream.result()
                self._logger.debug(
                    "Completed .act() prediction round", round_index=round_index
                )
                if on_prediction_completed:
                    round_result = PredictionRoundResult.from_result(
                        prediction, round_index
                    )
                    err_msg = f"Prediction completed callback failed for {self!r}"
                    with sdk_callback_invocation(err_msg, self._logger):
                        on_prediction_completed(round_result)
                if pending_tool_calls:

                    def _finish_tool_call(fut: SyncFuture[Any]) -> Any:
                        exc = fut.exception()
                        if exc is not None:
                            if not isinstance(exc, Exception):
                                # Don't allow base exceptions to be suppressed
                                raise exc
                            failed_request = pending_tool_calls[fut]
                            return endpoint._handle_failed_tool_request(
                                exc, failed_request
                            )
                        return fut.result()

                    tool_results = [
                        _finish_tool_call(fut)
                        for fut in as_completed(pending_tool_calls)
                    ]
                    requests_message = agent_chat.add_assistant_response(
                        prediction, tool_call_requests
                    )
                    results_message = agent_chat.add_tool_results(tool_results)
                    if on_message is not None:
                        err_msg = f"Tool request message callback failed for {self!r}"
                        with sdk_callback_invocation(err_msg, self._logger):
                            on_message(requests_message)
                        err_msg = f"Tool result message callback failed for {self!r}"
                        with sdk_callback_invocation(err_msg, self._logger):
                            on_message(results_message)
                elif on_message is not None:
                    err_msg = f"Final response message callback failed for {self!r}"
                    with sdk_callback_invocation(err_msg, self._logger):
                        on_message(agent_chat.add_assistant_response(prediction))
                if on_round_end is not None:
                    err_msg = f"Round end callback failed for {self!r}"
                    with sdk_callback_invocation(err_msg, self._logger):
                        on_round_end(round_index)
                if not tool_call_requests:
                    # No tool call requests -> we're done here
                    break
                if round_index == final_round_index:
                    # We somehow received at least one tool call request,
                    # even though tools are omitted on the final round
                    err_msg = "Model requested tool use on final prediction round."
                    endpoint._handle_invalid_tool_request(err_msg)
                    break
        num_rounds = round_index + 1
        duration = time.perf_counter() - start_time
        return ActResult(rounds=num_rounds, total_time_seconds=duration)

    @sdk_public_api()
    def apply_prompt_template(
        self,
        history: Chat | ChatHistoryDataDict | str,
        opts: LlmApplyPromptTemplateOpts | LlmApplyPromptTemplateOptsDict = {},
    ) -> str:
        """Apply a prompt template to the given history."""
        return self._session._apply_prompt_template(
            self.identifier,
            history,
            opts=opts,
        )


class EmbeddingModel(SyncModelHandle[_SyncSessionEmbedding]):
    """Reference to a loaded embedding model."""

    # Alas, type hints don't properly support distinguishing str vs Iterable[str]:
    #     https://github.com/python/typing/issues/256
    @sdk_public_api()
    def embed(
        self, input: str | Iterable[str]
    ) -> Sequence[float] | Sequence[Sequence[float]]:
        """Request embedding vectors for the given input string(s)."""
        return self._session._embed(self.identifier, input)


TSyncSession = TypeVar("TSyncSession", bound=_SyncSession)


class Client(ClientBase):
    """Synchronous SDK client interface."""

    def __init__(
        self, api_host: str | None = None, api_token: str | None = None
    ) -> None:
        """Initialize API client."""
        super().__init__(api_host, api_token)
        self._resources = rm = ExitStack()
        self._ws_thread = ws_thread = AsyncWebsocketThread(dict(client=repr(self)))
        ws_thread.start()
        rm.callback(ws_thread.terminate)
        self._sessions: dict[str, _SyncSession] = {}
        # Support GC-based resource management in the sync API by
        # finalizing at the client layer, and letting its resource
        # manager handle clearing up everything else
        rm.callback(self._sessions.clear)
        weakref.finalize(self, rm.close)

    def __enter__(self) -> Self:
        # Handle reentrancy the same way files do:
        # allow nested use as a CM, but close on the first exit
        with sdk_public_api():
            self._ensure_api_host_is_valid()
        return self

    def __exit__(self, *args: Any) -> None:
        self.close()

    def close(self) -> None:
        """Close any started client sessions."""
        self._resources.close()

    @staticmethod
    def _query_probe_url(url: str) -> httpx.Response:
        return httpx.get(url, timeout=1)

    @classmethod
    @sdk_public_api()
    def is_valid_api_host(cls, api_host: str) -> bool:
        """Report whether the given API host is running an API server instance."""
        probe_url = cls._get_probe_url(api_host)
        try:
            probe_response = cls._query_probe_url(probe_url)
        except (httpx.ConnectTimeout, httpx.ConnectError):
            return False
        return cls._check_probe_response(probe_response)

    @classmethod
    @sdk_public_api()
    def find_default_local_api_host(cls) -> str | None:
        """Query local ports for a running API server instance."""
        for api_host in cls._iter_default_api_hosts():
            if cls.is_valid_api_host(api_host):
                return api_host
        return None

    def _ensure_api_host_is_valid(self) -> None:
        specified_api_host = self._api_host
        if specified_api_host is None:
            api_host = self.find_default_local_api_host()
        elif self.is_valid_api_host(specified_api_host):
            api_host = specified_api_host
        else:
            api_host = None
        if api_host is None:
            raise self._get_probe_failure_error(specified_api_host)
        self._api_host = api_host

    # Doing network I/O in properties is generally considered undesirable.
    # The async API can't perform network I/O in properties at all.
    # Unlike the async API (which follows the principles of structured
    # concurrency), the sync API can resolve this by delaying the creation
    # of the underlying websockets until they're actually used.

    def _get_session(self, cls: Type[TSyncSession]) -> TSyncSession:
        """Get the client session of the given type."""
        namespace = cls.API_NAMESPACE
        assert namespace is not None
        try:
            session = self._sessions[namespace]
        except KeyError:
            pass
        else:
            # This *will* be an instance of the given type.
            # The assertion notifies typecheckers of that.
            assert isinstance(session, cls)
            return session
        # No session yet for this namespace, so create one
        session = cls(self)
        self._sessions[namespace] = session
        # Note: session starts itself on first use rather than immediately
        # We push the exit callback here so that only the client itself
        # needs to be explicitly managed as a context manager. Directly
        # managing the individual session is supported, but not required.
        self._resources.push(session)
        return session

    @property
    @sdk_public_api()
    def llm(self) -> _SyncSessionLlm:
        """Return the LLM API client session."""
        return self._get_session(_SyncSessionLlm)

    @property
    @sdk_public_api()
    def embedding(self) -> _SyncSessionEmbedding:
        """Return the embedding model API client session."""
        return self._get_session(_SyncSessionEmbedding)

    @property
    def system(self) -> _SyncSessionSystem:
        """Return the system API client session."""
        return self._get_session(_SyncSessionSystem)

    @property
    def files(self) -> _SyncSessionFiles:
        """Return the files API client session."""
        return self._get_session(_SyncSessionFiles)

    @property
    def repository(self) -> _SyncSessionRepository:
        """Return the repository API client session."""
        return self._get_session(_SyncSessionRepository)

    # Convenience methods
    # Not yet implemented (server API only supports the same file types as prepare_image)
    # @sdk_public_api()
    def _prepare_file(self, src: LocalFileInput, name: str | None = None) -> FileHandle:
        """Add a file to the server. Returns a file handle for use in prediction requests."""
        return self.files._prepare_file(src, name)

    @sdk_public_api()
    def prepare_image(self, src: LocalFileInput, name: str | None = None) -> FileHandle:
        """Add an image to the server. Returns a file handle for use in prediction requests."""
        return self.files.prepare_image(src, name)

    @sdk_public_api()
    def list_downloaded_models(
        self, namespace: str | None = None
    ) -> Sequence[AnyDownloadedModel]:
        """Get the list of downloaded models."""
        namespace_filter = check_model_namespace(namespace)
        if namespace_filter is None:
            return self.system.list_downloaded_models()
        if namespace_filter == "llm":
            return self.llm.list_downloaded()
        return self.embedding.list_downloaded()

    @sdk_public_api()
    def list_loaded_models(
        self, namespace: str | None = None
    ) -> Sequence[AnySyncModel]:
        """Get the list of loaded models using the default global client."""
        namespace_filter = check_model_namespace(namespace)
        loaded_models: list[AnySyncModel] = []
        if namespace_filter is None or namespace_filter == "llm":
            loaded_models.extend(self.llm.list_loaded())
        if namespace_filter is None or namespace_filter == "embedding":
            loaded_models.extend(self.embedding.list_loaded())
        return loaded_models


# Convenience API
_default_api_host: str | None = None
_default_api_token: str | None = None
_default_client: Client | None = None


@sdk_public_api()
def configure_default_client(api_host: str, api_token: str | None = None) -> None:
    """Set the server API host for the default global client (without creating the client)."""
    global _default_api_host, _default_api_token
    if _default_client is not None:
        raise LMStudioClientError(
            "Default client is already created, cannot set its API host or token."
        )
    _default_api_host = api_host
    _default_api_token = api_token


@sdk_public_api()
def get_default_client(api_host: str | None = None) -> Client:
    """Get the default global client (creating it if necessary)."""
    # Note: call configure_default_client() explicitly to set the API token
    global _default_client
    if api_host is not None:
        # This will raise an exception if the client already exists
        configure_default_client(api_host)
    if _default_client is None:
        _default_client = Client(_default_api_host, _default_api_token)
        _default_client._ensure_api_host_is_valid()
    return _default_client


def _reset_default_client() -> None:
    # Allow the test suite to reset the client without
    # having to poke directly at the module's internals
    global _default_api_host, _default_api_token, _default_client
    previous_client = _default_client
    _default_api_host = _default_api_token = _default_client = None
    if previous_client is not None:
        previous_client.close()


@sdk_public_api()
def llm(
    model_key: str | None = None,
    /,
    *,
    ttl: int | None = DEFAULT_TTL,
    config: LlmLoadModelConfig | LlmLoadModelConfigDict | None = None,
) -> LLM:
    """Access an LLM using the default global client.

    Note: configuration of retrieved model is NOT checked against the given config.
    Note: details of configuration fields may change in SDK feature releases.
    """
    return get_default_client().llm.model(model_key, ttl=ttl, config=config)


@sdk_public_api()
def embedding_model(
    model_key: str | None = None,
    /,
    *,
    ttl: int | None = DEFAULT_TTL,
    config: EmbeddingLoadModelConfig | EmbeddingLoadModelConfigDict | None = None,
) -> EmbeddingModel:
    """Access an embedding model using the default global client.

    Note: configuration of retrieved model is NOT checked against the given config.
    Note: details of configuration fields may change in SDK feature releases.
    """
    return get_default_client().embedding.model(model_key, ttl=ttl, config=config)


# Not yet implemented (server API only supports the same file types as prepare_image)
# @sdk_public_api()
def _prepare_file(src: LocalFileInput, name: str | None = None) -> FileHandle:
    """Add a file to the server using the default global client."""
    return get_default_client()._prepare_file(src, name)


@sdk_public_api()
def prepare_image(src: LocalFileInput, name: str | None = None) -> FileHandle:
    """Add an image to the server using the default global client."""
    return get_default_client().prepare_image(src, name)


@sdk_public_api()
def list_downloaded_models(
    namespace: str | None = None,
) -> Sequence[AnyDownloadedModel]:
    """Get the list of downloaded models using the default global client."""
    return get_default_client().list_downloaded_models(namespace)


@sdk_public_api()
def list_loaded_models(namespace: str | None = None) -> Sequence[AnySyncModel]:
    """Get the list of loaded models using the default global client."""
    return get_default_client().list_loaded_models(namespace)