From 33e24284cf1b2adb2596e44fd87b96f77d394b80 Mon Sep 17 00:00:00 2001
From: Jack Gerrits <jack@jackgerrits.com>
Date: Fri, 10 Jan 2025 19:46:10 -0500
Subject: [PATCH] fix

---
 .../autogen-ext/src/autogen_ext/teams/magentic_one.py    | 9 +++++++--
 .../magentic-one-cli/src/magentic_one_cli/_m1.py         | 8 +++-----
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one.py b/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one.py
index 904c486426c0..fc2e4f6b9129 100644
--- a/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one.py
+++ b/python/packages/autogen-ext/src/autogen_ext/teams/magentic_one.py
@@ -1,9 +1,10 @@
 import warnings
-from typing import List
+from typing import Awaitable, Callable, List, Optional, Union
 
 from autogen_agentchat.agents import CodeExecutorAgent, UserProxyAgent
 from autogen_agentchat.base import ChatAgent
 from autogen_agentchat.teams import MagenticOneGroupChat
+from autogen_core import CancellationToken
 from autogen_core.models import ChatCompletionClient
 
 from autogen_ext.agents.file_surfer import FileSurfer
@@ -12,6 +13,10 @@
 from autogen_ext.code_executors.local import LocalCommandLineCodeExecutor
 from autogen_ext.models.openai._openai_client import BaseOpenAIChatCompletionClient
 
+SyncInputFunc = Callable[[str], str]
+AsyncInputFunc = Callable[[str, Optional[CancellationToken]], Awaitable[str]]
+InputFuncType = Union[SyncInputFunc, AsyncInputFunc]
+
 
 class MagenticOne(MagenticOneGroupChat):
     """
@@ -120,7 +125,7 @@ def __init__(
         self,
         client: ChatCompletionClient,
         hil_mode: bool = False,
-        input_func: UserProxyAgent.InputFuncType | None = None,
+        input_func: InputFuncType | None = None,
     ):
         self.client = client
         self._validate_client_capabilities(client)
diff --git a/python/packages/magentic-one-cli/src/magentic_one_cli/_m1.py b/python/packages/magentic-one-cli/src/magentic_one_cli/_m1.py
index ce5c3ac86b00..e7a3f2ed1e89 100644
--- a/python/packages/magentic-one-cli/src/magentic_one_cli/_m1.py
+++ b/python/packages/magentic-one-cli/src/magentic_one_cli/_m1.py
@@ -1,10 +1,10 @@
 import argparse
 import asyncio
-from typing import Optional
 import warnings
+from typing import Optional
 
-from autogen_agentchat.ui import Console
-from autogen_agentchat.ui import UserInputManager
+from aioconsole import ainput  # type: ignore
+from autogen_agentchat.ui import Console, UserInputManager
 from autogen_core import CancellationToken
 from autogen_ext.models.openai import OpenAIChatCompletionClient
 from autogen_ext.teams.magentic_one import MagenticOne
@@ -13,8 +13,6 @@
 warnings.filterwarnings(action="ignore", message="unclosed", category=ResourceWarning)
 
 
-
-
 async def cancellable_input(prompt: str, cancellation_token: Optional[CancellationToken]) -> str:
     task: asyncio.Task[str] = asyncio.create_task(ainput(prompt))  # type: ignore
     if cancellation_token is not None: