diff --git a/python/packages/autogen-core/src/autogen_core/models/_model_client.py b/python/packages/autogen-core/src/autogen_core/models/_model_client.py index e70bacdcf57d..e74e83ad8fb4 100644 --- a/python/packages/autogen-core/src/autogen_core/models/_model_client.py +++ b/python/packages/autogen-core/src/autogen_core/models/_model_client.py @@ -50,6 +50,13 @@ class ModelFamily: MISTRAL = "mistral" MINISTRAL = "ministral" PIXTRAL = "pixtral" + GROK_CODE_FAST_1 = "grok-code-fast-1" + GROK_4_FAST_REASONING = "grok-4-fast-reasoning" + GROK_4_FAST_NON_REASONING = "grok-4-fast-non-reasoning" + GROK_4 = "grok-4" + GROK_3_MINI = "grok-3-mini" + GROK_3 = "grok-3" + UNKNOWN = "unknown" ANY: TypeAlias = Literal[ @@ -90,6 +97,13 @@ class ModelFamily: "mistral", "ministral", "pixtral", + # grok_models + "grok-code-fast-1", + "grok-4-fast-reasoning", + "grok-4-fast-non-reasoning", + "grok-4", + "grok-3-mini", + "grok-3", # unknown "unknown", ] @@ -152,6 +166,17 @@ def is_mistral(family: str) -> bool: ModelFamily.MINISTRAL, ModelFamily.PIXTRAL, ) + + @staticmethod + def is_grok(family: str) -> bool: + return family in ( + ModelFamily.GROK_CODE_FAST_1, + ModelFamily.GROK_4_FAST_REASONING, + ModelFamily.GROK_4_FAST_NON_REASONING, + ModelFamily.GROK_4, + ModelFamily.GROK_3_MINI, + ModelFamily.GROK_3, + ) @deprecated("Use the ModelInfo class instead ModelCapabilities.") diff --git a/python/packages/autogen-ext/src/autogen_ext/models/openai/_model_info.py b/python/packages/autogen-ext/src/autogen_ext/models/openai/_model_info.py index 91e86c1c67e4..f5f9f4518d90 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/openai/_model_info.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/openai/_model_info.py @@ -46,6 +46,13 @@ "llama-3.3-70b": "Llama-3.3-70B-Instruct", "llama-4-scout": "Llama-4-Scout-17B-16E-Instruct-FP8", "llama-4-maverick": "Llama-4-Maverick-17B-128E-Instruct-FP8", + # Grok models + "grok-code-fast-1": "grok-code-fast-1", + "grok-4-fast-reasoning": "grok-4-fast-reasoning", + "grok-4-fast-non-reasoning": "grok-4-fast-non-reasoning", + "grok-4": "grok-4", + "grok-3-mini": "grok-3-mini", + "grok-3": "grok-3", } _MODEL_INFO: Dict[str, ModelInfo] = { @@ -441,6 +448,48 @@ "structured_output": True, "multiple_system_messages": True, }, + "grok-code-fast-1": { + "vision": False, + "function_calling": True, + "json_output": True, + "family": ModelFamily.GROK_CODE_FAST_1, + "structured_output": True + }, + "grok-4-fast-reasoning": { + "vision": False, + "function_calling": True, + "json_output": True, + "family": ModelFamily.GROK_4_FAST_REASONING, + "structured_output": True + }, + "grok-4-fast-non-reasoning": { + "vision": False, + "function_calling": True, + "json_output": True, + "family": ModelFamily.GROK_4_FAST_NON_REASONING, + "structured_output": True + }, + "grok-4": { + "vision": False, + "function_calling": True, + "json_output": True, + "family": ModelFamily.GROK_4, + "structured_output": True + }, + "grok-3-mini": { + "vision": False, + "function_calling": True, + "json_output": True, + "family": ModelFamily.GROK_3_MINI, + "structured_output": True + }, + "grok-3": { + "vision": False, + "function_calling": True, + "json_output": True, + "family": ModelFamily.GROK_3, + "structured_output": True + }, } _MODEL_TOKEN_LIMITS: Dict[str, int] = { @@ -491,11 +540,18 @@ "Llama-3.3-70B-Instruct": 128000, "Llama-4-Scout-17B-16E-Instruct-FP8": 128000, "Llama-4-Maverick-17B-128E-Instruct-FP8": 128000, + "grok-code-fast-1": 256000, + "grok-4-fast-reasoning": 2000000, + "grok-4-fast-non-reasoning": 2000000, + "grok-4": 256000, + "grok-3-mini": 131072, + "grok-3": 131072, } GEMINI_OPENAI_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/openai/" ANTHROPIC_OPENAI_BASE_URL = "https://api.anthropic.com/v1/" LLAMA_API_BASE_URL = "https://api.llama.com/compat/v1/" +GROK_API_BASE_URL = "https://api.x.ai/v1" def resolve_model(model: str) -> str: diff --git a/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py b/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py index a80e912534ab..4b41f700fabe 100644 --- a/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py +++ b/python/packages/autogen-ext/src/autogen_ext/models/openai/_openai_client.py @@ -1480,7 +1480,12 @@ def __init__(self, **kwargs: Unpack[OpenAIClientConfiguration]): copied_args["base_url"] = _model_info.LLAMA_API_BASE_URL if "api_key" not in copied_args and "LLAMA_API_KEY" in os.environ: copied_args["api_key"] = os.environ["LLAMA_API_KEY"] - + if copied_args["model"].startswith("Grok-"): + if "base_url" not in copied_args: + copied_args["base_url"] = _model_info.GROK_API_BASE_URL + if "api_key" not in copied_args and "GROK_API_KEY" in os.environ: + copied_args["api_key"] = os.environ["GROK_API_KEY"] + client = _openai_client_from_config(copied_args) create_args = _create_args_from_config(copied_args)