diff --git a/any_llm_client/main.py b/any_llm_client/main.py
index 552bb40..d734e67 100644
--- a/any_llm_client/main.py
+++ b/any_llm_client/main.py
@@ -1,6 +1,8 @@
 import functools
 import typing
 
+import pydantic
+
 from any_llm_client.clients.mock import MockLLMClient, MockLLMConfig
 from any_llm_client.clients.openai import OpenAIClient, OpenAIConfig
 from any_llm_client.clients.yandexgpt import YandexGPTClient, YandexGPTConfig
@@ -8,7 +10,7 @@
 from any_llm_client.retry import RequestRetryConfig
 
 
-AnyLLMConfig = YandexGPTConfig | OpenAIConfig | MockLLMConfig
+AnyLLMConfig = typing.Annotated[YandexGPTConfig | OpenAIConfig | MockLLMConfig, pydantic.Discriminator("api_type")]
 
 
 if typing.TYPE_CHECKING: