Resolve types issues in coding

.github/workflows/type-check.yml

@@ -20,6 +20,7 @@ jobs:
       - run: pip install ".[jupyter-executor]" mypy
       # As more modules are type check clean, add them here
       - run: |
-          mypy --install-types --non-interactive \
+          mypy \
             autogen/logger \
-            autogen/exception_utils.py
+            autogen/exception_utils.py \
+            autogen/coding

autogen/code_utils.py

@@ -14,6 +14,8 @@
 
 import docker
 
+from .types import UserMessageImageContentPart, UserMessageTextContentPart
+
 SENTINEL = object()
 DEFAULT_MODEL = "gpt-4"
 FAST_MODEL = "gpt-3.5-turbo"
@@ -37,8 +39,8 @@
 logger = logging.getLogger(__name__)
 
 
-def content_str(content: Union[str, List[Dict[str, Any]], None]) -> str:
-    """Converts `content` into a string format.
+def content_str(content: Union[str, List[Union[UserMessageTextContentPart, UserMessageImageContentPart]], None]) -> str:
+    """Converts the `content` field of an OpenAI merssage into a string format.
 
     This function processes content that may be a string, a list of mixed text and image URLs, or None,
     and converts it into a string. Text is directly appended to the result string, while image URLs are
@@ -241,7 +243,7 @@ def get_powershell_command():
         raise PermissionError("No permission to run powershell.") from e
 
 
-def _cmd(lang):
+def _cmd(lang: str) -> str:
     if lang.startswith("python") or lang in ["bash", "sh"]:
         return lang
     if lang in ["shell"]:

autogen/coding/__init__.py

@@ -1,7 +1,7 @@
 from .base import CodeBlock, CodeExecutor, CodeExtractor, CodeResult
 from .factory import CodeExecutorFactory
 from .markdown_code_extractor import MarkdownCodeExtractor
-from .local_commandline_code_executor import LocalCommandLineCodeExecutor, CommandLineCodeResult
+from .local_commandline_code_executor import LocalCommandLineCodeExecutor
 from .docker_commandline_code_executor import DockerCommandLineCodeExecutor
 
 __all__ = (
@@ -12,6 +12,5 @@
     "CodeExecutorFactory",
     "MarkdownCodeExtractor",
     "LocalCommandLineCodeExecutor",
-    "CommandLineCodeResult",
     "DockerCommandLineCodeExecutor",
 )

autogen/coding/base.py

@@ -1,10 +1,12 @@
-from typing import Any, Dict, List, Optional, Protocol, Union, runtime_checkable
+from __future__ import annotations
+from typing import Any, List, Literal, Mapping, Optional, Protocol, TypedDict, Union, runtime_checkable
 
 from pydantic import BaseModel, Field
 
 from ..agentchat.agent import LLMAgent
+from ..types import UserMessageImageContentPart, UserMessageTextContentPart
 
-__all__ = ("CodeBlock", "CodeResult", "CodeExtractor", "CodeExecutor")
+__all__ = ("CodeBlock", "CodeResult", "CodeExtractor", "CodeExecutor", "CodeExecutionConfig")
 
 
 class CodeBlock(BaseModel):
@@ -26,7 +28,9 @@ class CodeResult(BaseModel):
 class CodeExtractor(Protocol):
     """(Experimental) A code extractor class that extracts code blocks from a message."""
 
-    def extract_code_blocks(self, message: Union[str, List[Dict[str, Any]], None]) -> List[CodeBlock]:
+    def extract_code_blocks(
+        self, message: Union[str, List[Union[UserMessageTextContentPart, UserMessageImageContentPart]], None]
+    ) -> List[CodeBlock]:
         """(Experimental) Extract code blocks from a message.
 
         Args:
@@ -79,6 +83,21 @@ class IPythonCodeResult(CodeResult):
     )
 
 
+CodeExecutionConfig = TypedDict(
+    "CodeExecutionConfig",
+    {
+        "executor": Union[Literal["ipython-embedded", "commandline-local"], CodeExecutor],
+        "last_n_messages": Union[int, Literal["auto"]],
+        "timeout": int,
+        "use_docker": Union[bool, str, List[str]],
+        "work_dir": str,
+        "ipython-embedded": Mapping[str, Any],
+        "commandline-local": Mapping[str, Any],
+    },
+    total=False,
+)
+
+
 class CommandLineCodeResult(CodeResult):
     """(Experimental) A code result class for command line code executor."""
 

autogen/coding/docker_commandline_code_executor.py

@@ -6,9 +6,8 @@
 from time import sleep
 from types import TracebackType
 import uuid
-from typing import List, Optional, Type, Union
+from typing import Any, List, Optional, Type, Union
 import docker
-from docker.models.containers import Container
 from docker.errors import ImageNotFound
 
 from .utils import _get_file_name_from_content
@@ -25,8 +24,8 @@
     from typing_extensions import Self
 
 
-def _wait_for_ready(container: Container, timeout: int = 60, stop_time: int = 0.1) -> None:
-    elapsed_time = 0
+def _wait_for_ready(container: Any, timeout: int = 60, stop_time: float = 0.1) -> None:
+    elapsed_time = 0.0
     while container.status != "running" and elapsed_time < timeout:
         sleep(stop_time)
         elapsed_time += stop_time
@@ -114,7 +113,7 @@ def __init__(
 
         _wait_for_ready(self._container)
 
-        def cleanup():
+        def cleanup() -> None:
             try:
                 container = client.containers.get(container_name)
                 container.stop()

autogen/coding/factory.py

@@ -1,6 +1,4 @@
-from typing import Any, Dict
-
-from .base import CodeExecutor
+from .base import CodeExecutor, CodeExecutionConfig
 
 __all__ = ("CodeExecutorFactory",)
 
@@ -9,7 +7,7 @@ class CodeExecutorFactory:
     """(Experimental) A factory class for creating code executors."""
 
     @staticmethod
-    def create(code_execution_config: Dict[str, Any]) -> CodeExecutor:
+    def create(code_execution_config: CodeExecutionConfig) -> CodeExecutor:
         """(Experimental) Get a code executor based on the code execution config.
 
         Args:

autogen/coding/jupyter/docker_jupyter_server.py

@@ -123,7 +123,7 @@ def __init__(
         self._port = int(container_ports["8888/tcp"][0]["HostPort"])
         self._container_id = container.id
 
-        def cleanup():
+        def cleanup() -> None:
             try:
                 inner_container = client.containers.get(container.id)
                 inner_container.stop()
@@ -142,7 +142,7 @@ def cleanup():
     def connection_info(self) -> JupyterConnectionInfo:
         return JupyterConnectionInfo(host="127.0.0.1", use_https=False, port=self._port, token=self._token)
 
-    def stop(self):
+    def stop(self) -> None:
         self._cleanup_func()
 
     def get_client(self) -> JupyterClient:

autogen/coding/local_commandline_code_executor.py

@@ -170,12 +170,12 @@ def restart(self) -> None:
 
 # From stack overflow: https://stackoverflow.com/a/52087847/2214524
 class _DeprecatedClassMeta(type):
-    def __new__(cls, name, bases, classdict, *args, **kwargs):
+    def __new__(cls, name, bases, classdict, *args, **kwargs):  # type: ignore[no-untyped-def]
         alias = classdict.get("_DeprecatedClassMeta__alias")
 
         if alias is not None:
 
-            def new(cls, *args, **kwargs):
+            def new(cls, *args, **kwargs):  # type: ignore[no-untyped-def]
                 alias = getattr(cls, "_DeprecatedClassMeta__alias")
 
                 if alias is not None:
@@ -209,14 +209,14 @@ def new(cls, *args, **kwargs):
             if b not in fixed_bases:
                 fixed_bases.append(b)
 
-        fixed_bases = tuple(fixed_bases)
+        fixed_bases = tuple(fixed_bases)  # type: ignore[assignment]
 
-        return super().__new__(cls, name, fixed_bases, classdict, *args, **kwargs)
+        return super().__new__(cls, name, fixed_bases, classdict, *args, **kwargs)  # type: ignore[call-overload]
 
-    def __instancecheck__(cls, instance):
-        return any(cls.__subclasscheck__(c) for c in {type(instance), instance.__class__})
+    def __instancecheck__(cls, instance):  # type: ignore[no-untyped-def]
+        return any(cls.__subclasscheck__(c) for c in {type(instance), instance.__class__})  # type: ignore[no-untyped-call]
 
-    def __subclasscheck__(cls, subclass):
+    def __subclasscheck__(cls, subclass):  # type: ignore[no-untyped-def]
         if subclass is cls:
             return True
         else:

autogen/coding/markdown_code_extractor.py

@@ -2,15 +2,18 @@
 from typing import Any, Dict, List, Optional, Union
 
 from ..code_utils import CODE_BLOCK_PATTERN, UNKNOWN, content_str, infer_lang
-from .base import CodeBlock
+from .base import CodeBlock, CodeExtractor
+from ..types import UserMessageImageContentPart, UserMessageTextContentPart
 
 __all__ = ("MarkdownCodeExtractor",)
 
 
-class MarkdownCodeExtractor:
+class MarkdownCodeExtractor(CodeExtractor):
     """(Experimental) A class that extracts code blocks from a message using Markdown syntax."""
 
-    def extract_code_blocks(self, message: Union[str, List[Dict[str, Any]], None]) -> List[CodeBlock]:
+    def extract_code_blocks(
+        self, message: Union[str, List[Union[UserMessageTextContentPart, UserMessageImageContentPart]], None]
+    ) -> List[CodeBlock]:
         """(Experimental) Extract code blocks from a message. If no code blocks are found,
         return an empty list.
 

autogen/types.py

@@ -0,0 +1,12 @@
+from typing import Dict, List, Literal, TypedDict, Union
+
+
+class UserMessageTextContentPart(TypedDict):
+    type: Literal["text"]
+    text: str
+
+
+class UserMessageImageContentPart(TypedDict):
+    type: Literal["image_url"]
+    # Ignoring the other "detail param for now"
+    image_url: Dict[Literal["url"], str]