From cd9dca47407c2884571dcff4f5bd3f1505868259 Mon Sep 17 00:00:00 2001
From: Eric Zhu <ekzhu@users.noreply.github.com>
Date: Fri, 31 Jan 2025 15:07:14 -0800
Subject: [PATCH] feat: add o3 to model info; update chess example (#5311)

Because.
---
 .../src/autogen_core/models/_model_client.py          |  3 ++-
 .../src/autogen_ext/models/openai/_model_info.py      |  7 +++++++
 python/samples/agentchat_chess_game/.gitignore        |  1 +
 python/samples/agentchat_chess_game/README.md         | 11 ++++++++++-
 python/samples/agentchat_chess_game/main.py           |  2 +-
 5 files changed, 21 insertions(+), 3 deletions(-)

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 62c26c51a677..d76b7191b925 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
@@ -20,12 +20,13 @@ class ModelFamily:
 
     GPT_4O = "gpt-4o"
     O1 = "o1"
+    O3 = "o3"
     GPT_4 = "gpt-4"
     GPT_35 = "gpt-35"
     R1 = "r1"
     UNKNOWN = "unknown"
 
-    ANY: TypeAlias = Literal["gpt-4o", "o1", "gpt-4", "gpt-35", "r1", "unknown"]
+    ANY: TypeAlias = Literal["gpt-4o", "o1", "o3", "gpt-4", "gpt-35", "r1", "unknown"]
 
     def __new__(cls, *args: Any, **kwargs: Any) -> ModelFamily:
         raise TypeError(f"{cls.__name__} is a namespace class and cannot be instantiated.")
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 c69ed87d1d55..316725f0bfd7 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
@@ -5,6 +5,7 @@
 # Based on: https://platform.openai.com/docs/models/continuous-model-upgrades
 # This is a moving target, so correctness is checked by the model value returned by openai against expected values at runtime``
 _MODEL_POINTERS = {
+    "o3-mini": "o3-mini-2025-01-31",
     "o1": "o1-2024-12-17",
     "o1-preview": "o1-preview-2024-09-12",
     "o1-mini": "o1-mini-2024-09-12",
@@ -19,6 +20,12 @@
 }
 
 _MODEL_INFO: Dict[str, ModelInfo] = {
+    "o3-mini-2025-01-31": {
+        "vision": False,
+        "function_calling": True,
+        "json_output": True,
+        "family": ModelFamily.O3,
+    },
     "o1-2024-12-17": {
         "vision": False,
         "function_calling": False,
diff --git a/python/samples/agentchat_chess_game/.gitignore b/python/samples/agentchat_chess_game/.gitignore
index 189b1a838595..74565a19d2ce 100644
--- a/python/samples/agentchat_chess_game/.gitignore
+++ b/python/samples/agentchat_chess_game/.gitignore
@@ -1 +1,2 @@
 model_config.yml
+model_config.yaml
\ No newline at end of file
diff --git a/python/samples/agentchat_chess_game/README.md b/python/samples/agentchat_chess_game/README.md
index 00251bcfd7c1..7dee8a5cf889 100644
--- a/python/samples/agentchat_chess_game/README.md
+++ b/python/samples/agentchat_chess_game/README.md
@@ -27,7 +27,16 @@ For example, to use `gpt-4o` model from OpenAI, you can use the following config
 provider: autogen_ext.models.openai.OpenAIChatCompletionClient
 config:
   model: gpt-4o
-  api_key: REPLACE_WITH_YOUR_API_KEY
+  api_key: replace with your API key or skip it if you have environment variable OPENAI_API_KEY set
+```
+
+To use `o3-mini-2025-01-31` model from OpenAI, you can use the following configuration:
+
+```yaml
+provider: autogen_ext.models.openai.OpenAIChatCompletionClient
+config:
+  model: o3-mini-2025-01-31
+  api_key: replace with your API key or skip it if you have environment variable OPENAI_API_KEY set
 ```
 
 To use a locally hosted DeepSeek-R1:8b model using Ollama throught its compatibility endpoint,
diff --git a/python/samples/agentchat_chess_game/main.py b/python/samples/agentchat_chess_game/main.py
index 51870a9135cb..e9b2f314f0f6 100644
--- a/python/samples/agentchat_chess_game/main.py
+++ b/python/samples/agentchat_chess_game/main.py
@@ -12,7 +12,7 @@
 
 def create_ai_player() -> AssistantAgent:
     # Load the model client from config.
-    with open("model_config.yml", "r") as f:
+    with open("model_config.yaml", "r") as f:
         model_config = yaml.safe_load(f)
     model_client = ChatCompletionClient.load_component(model_config)
     # Create an agent that can use the model client.