Repository: oumi-ai/oumi
File: src/oumi/builders/inference_engines.py (Python)

Code before:
from typing import Optional

from oumi.core.configs import (
    GenerationParams,
    InferenceEngineType,
    ModelParams,
    RemoteParams,
)
from oumi.core.inference import BaseInferenceEngine
from oumi.inference import (
    AnthropicInferenceEngine,
    DeepSeekInferenceEngine,
    GoogleGeminiInferenceEngine,
    GoogleVertexInferenceEngine,
    LlamaCppInferenceEngine,
    NativeTextInferenceEngine,
    OpenAIInferenceEngine,
    ParasailInferenceEngine,
    RemoteInferenceEngine,
    RemoteVLLMInferenceEngine,
    SGLangInferenceEngine,
    TogetherInferenceEngine,
    VLLMInferenceEngine,
)

ENGINE_MAP: dict[InferenceEngineType, type[BaseInferenceEngine]] = {
    InferenceEngineType.ANTHROPIC: AnthropicInferenceEngine,
    InferenceEngineType.DEEPSEEK: DeepSeekInferenceEngine,
    InferenceEngineType.GOOGLE_GEMINI: GoogleGeminiInferenceEngine,
    InferenceEngineType.GOOGLE_VERTEX: GoogleVertexInferenceEngine,
    InferenceEngineType.LLAMACPP: LlamaCppInferenceEngine,
    InferenceEngineType.NATIVE: NativeTextInferenceEngine,
    InferenceEngineType.OPENAI: OpenAIInferenceEngine,
    InferenceEngineType.PARASAIL: ParasailInferenceEngine,
    InferenceEngineType.REMOTE_VLLM: RemoteVLLMInferenceEngine,
    InferenceEngineType.REMOTE: RemoteInferenceEngine,
    InferenceEngineType.SGLANG: SGLangInferenceEngine,
    InferenceEngineType.TOGETHER: TogetherInferenceEngine,
    InferenceEngineType.VLLM: VLLMInferenceEngine,
}


def build_inference_engine(
    engine_type: InferenceEngineType,
    model_params: ModelParams,
    remote_params: Optional[RemoteParams] = None,
    generation_params: Optional[GenerationParams] = None,
) -> BaseInferenceEngine:
    """Returns the inference engine based on the provided config.

    Args:

Review comment (suggestion):
consider using `Mapping[...]` type annotation with `types.MappingProxyType`
https://github.com/oumi-ai/oumi/blob/3b7e07cea72c3e0e7411cd98a0afe5e648e2678c/src/oumi/core/configs/internal/supported_models.py#L275
to make the map is immutable

Code after:
from types import MappingProxyType
from typing import Optional

from oumi.core.configs import (
    GenerationParams,
    InferenceEngineType,
    ModelParams,
    RemoteParams,
)
from oumi.core.inference import BaseInferenceEngine
from oumi.inference import (
    AnthropicInferenceEngine,
    DeepSeekInferenceEngine,
    GoogleGeminiInferenceEngine,
    GoogleVertexInferenceEngine,
    LlamaCppInferenceEngine,
    NativeTextInferenceEngine,
    OpenAIInferenceEngine,
    ParasailInferenceEngine,
    RemoteInferenceEngine,
    RemoteVLLMInferenceEngine,
    SGLangInferenceEngine,
    TogetherInferenceEngine,
    VLLMInferenceEngine,
)

ENGINE_MAP: MappingProxyType[InferenceEngineType, type[BaseInferenceEngine]] = (
    MappingProxyType(
        {
            InferenceEngineType.ANTHROPIC: AnthropicInferenceEngine,
            InferenceEngineType.DEEPSEEK: DeepSeekInferenceEngine,
            InferenceEngineType.GOOGLE_GEMINI: GoogleGeminiInferenceEngine,
            InferenceEngineType.GOOGLE_VERTEX: GoogleVertexInferenceEngine,
            InferenceEngineType.LLAMACPP: LlamaCppInferenceEngine,
            InferenceEngineType.NATIVE: NativeTextInferenceEngine,
            InferenceEngineType.OPENAI: OpenAIInferenceEngine,
            InferenceEngineType.PARASAIL: ParasailInferenceEngine,
            InferenceEngineType.REMOTE_VLLM: RemoteVLLMInferenceEngine,
            InferenceEngineType.REMOTE: RemoteInferenceEngine,
            InferenceEngineType.SGLANG: SGLangInferenceEngine,
            InferenceEngineType.TOGETHER: TogetherInferenceEngine,
            InferenceEngineType.VLLM: VLLMInferenceEngine,
        }
    )
)


def build_inference_engine(
    engine_type: InferenceEngineType,
    model_params: ModelParams,
    remote_params: Optional[RemoteParams] = None,