Merge branch 'ag2ai:main' into main

autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py

@@ -28,8 +28,8 @@ def __init__(
         name: str,
         host: str = "127.0.0.1",
         port: int = 6379,
-        username: str | None = None,
-        password: str | None = None,
+        username: Optional[str] = None,
+        password: Optional[str] = None,
         model: Optional["GenerativeModel"] = None,
         ontology: Optional["Ontology"] = None,
     ):

autogen/agentchat/contrib/graph_rag/graph_query_engine.py

@@ -28,7 +28,7 @@ class GraphQueryEngine(Protocol):
     This interface defines the basic methods for graph-based RAG.
     """
 
-    def init_db(self, input_doc: list[Document] | None = None):
+    def init_db(self, input_doc: Optional[list[Document]] = None):
         """This method initializes graph database with the input documents or records.
         Usually, it takes the following steps,
         1. connecting to a graph database.

autogen/agentchat/contrib/graph_rag/neo4j_graph_query_engine.py

@@ -94,7 +94,7 @@ def __init__(
         self.schema = schema
         self.strict = strict
 
-    def init_db(self, input_doc: list[Document] | None = None):
+    def init_db(self, input_doc: Optional[list[Document]] = None):
         """Build the knowledge graph with input documents."""
         self.documents = self._load_doc(input_doc)
 

autogen/agentchat/contrib/vectordb/pgvectordb.py

@@ -10,7 +10,6 @@
 from typing import Callable, Optional, Union
 
 import numpy as np
-from sentence_transformers import SentenceTransformer
 
 from ....import_utils import optional_import_block, require_optional_import
 from .base import Document, ItemID, QueryResults, VectorDB
@@ -20,12 +19,13 @@
     import pgvector  # noqa: F401
     import psycopg
     from pgvector.psycopg import register_vector
+    from sentence_transformers import SentenceTransformer
 
 PGVECTOR_MAX_BATCH_SIZE = os.environ.get("PGVECTOR_MAX_BATCH_SIZE", 40000)
 logger = get_logger(__name__)
 
 
-@require_optional_import("psycopg", "retrievechat-pgvector")
+@require_optional_import(["psycopg", "sentence_transformers"], "retrievechat-pgvector")
 class Collection:
     """A Collection object for PGVector.
 
@@ -539,7 +539,7 @@ def create_collection(
         cursor.close()
 
 
-@require_optional_import(["pgvector", "psycopg"], "retrievechat-pgvector")
+@require_optional_import(["pgvector", "psycopg", "sentence_transformers"], "retrievechat-pgvector")
 class PGVectorDB(VectorDB):
     """A vector database that uses PGVector as the backend."""
 

autogen/coding/jupyter/__init__.py

@@ -4,13 +4,6 @@
 #
 # Original portions of this file are derived from https://github.com/microsoft/autogen under the MIT License.
 # SPDX-License-Identifier: MIT
-from .helpers import is_jupyter_kernel_gateway_installed
-
-if not is_jupyter_kernel_gateway_installed():
-    raise ImportError(
-        "jupyter-kernel-gateway is required for JupyterCodeExecutor, please install it with `pip install ag2[jupyter-executor]`"
-    )
-
 
 from .base import JupyterConnectable, JupyterConnectionInfo
 from .docker_jupyter_server import DockerJupyterServer

autogen/coding/jupyter/docker_jupyter_server.py

@@ -24,9 +24,11 @@
 
 from ..docker_commandline_code_executor import _wait_for_ready
 from .base import JupyterConnectable, JupyterConnectionInfo
+from .import_utils import require_jupyter_kernel_gateway_installed
 from .jupyter_client import JupyterClient
 
 
+@require_jupyter_kernel_gateway_installed()
 class DockerJupyterServer(JupyterConnectable):
     DEFAULT_DOCKERFILE = """FROM quay.io/jupyter/docker-stacks-foundation
 

autogen/coding/jupyter/embedded_ipython_code_executor.py

@@ -18,6 +18,7 @@
 from ...import_utils import optional_import_block, require_optional_import
 from ..base import CodeBlock, CodeExtractor, IPythonCodeResult
 from ..markdown_code_extractor import MarkdownCodeExtractor
+from .import_utils import require_jupyter_kernel_gateway_installed
 
 with optional_import_block():
     from jupyter_client import KernelManager  # type: ignore[attr-defined]
@@ -27,6 +28,7 @@
 
 
 @require_optional_import("jupyter_client", "jupyter-executor")
+@require_jupyter_kernel_gateway_installed()
 class EmbeddedIPythonCodeExecutor(BaseModel):
     """(Experimental) A code executor class that executes code statefully using an embedded
     IPython kernel managed by this class.

autogen/coding/jupyter/import_utils.py

@@ -0,0 +1,78 @@
+# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import subprocess
+from functools import lru_cache
+from logging import getLogger
+from typing import Callable, TypeVar
+
+from ...import_utils import patch_object
+
+logger = getLogger(__name__)
+
+__all__ = ["require_jupyter_kernel_gateway_installed", "skip_on_missing_jupyter_kernel_gateway"]
+
+
+@lru_cache()
+def is_jupyter_kernel_gateway_installed() -> bool:
+    """Check if jupyter-kernel-gateway is installed."""
+    try:
+        subprocess.run(
+            ["jupyter", "kernelgateway", "--version"],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            check=True,
+        )
+        return True
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        logger.warning(
+            "jupyter-kernel-gateway is required for JupyterCodeExecutor, please install it with `pip install ag2[jupyter-executor]`"
+        )
+        return False
+
+
+T = TypeVar("T")
+
+
+def require_jupyter_kernel_gateway_installed() -> Callable[[T], T]:
+    """Decorator to handle optional module dependencies
+
+    Args:
+        modules: Module name or list of module names required
+        dep_target: Target name for pip installation (e.g. 'test' in pip install ag2[test])
+    """
+    if is_jupyter_kernel_gateway_installed():
+
+        def decorator(o: T) -> T:
+            return o
+    else:
+
+        def decorator(o: T) -> T:
+            return patch_object(o, missing_modules=[], dep_target="jupyter-executor")
+
+    return decorator
+
+
+def skip_on_missing_jupyter_kernel_gateway() -> Callable[[T], T]:
+    """Decorator to skip a test if an optional module is missing
+
+    Args:
+        module: Module name
+        dep_target: Target name for pip installation (e.g. 'test' in pip install ag2[test])
+    """
+
+    if is_jupyter_kernel_gateway_installed():
+
+        def decorator(o: T) -> T:
+            return o
+    else:
+
+        def decorator(o: T) -> T:
+            import pytest
+
+            return pytest.mark.skip(
+                reason="jupyter-kernel-gateway is required for JupyterCodeExecutor, please install it with `pip install ag2[jupyter-executor]`"
+            )(o)  # type: ignore[return-value]
+
+    return decorator

autogen/coding/jupyter/local_jupyter_server.py

@@ -20,9 +20,11 @@
     from typing_extensions import Self
 
 from .base import JupyterConnectable, JupyterConnectionInfo
+from .import_utils import require_jupyter_kernel_gateway_installed
 from .jupyter_client import JupyterClient
 
 
+@require_jupyter_kernel_gateway_installed()
 class LocalJupyterServer(JupyterConnectable):
     class GenerateToken:
         pass

autogen/formatting_utils.py

@@ -7,7 +7,7 @@
 from __future__ import annotations
 
 from collections.abc import Iterable
-from typing import Literal
+from typing import Literal, Optional
 
 from .import_utils import optional_import_block
 
@@ -70,12 +70,12 @@
 
     def colored(
         text: object,
-        color: Color | None = None,
-        on_color: Highlight | None = None,
-        attrs: Iterable[Attribute] | None = None,
+        color: Optional[Color] = None,
+        on_color: Optional[Highlight] = None,
+        attrs: Optional[Iterable[Attribute]] = None,
         *,
-        no_color: bool | None = None,
-        force_color: bool | None = None,
+        no_color: Optional[bool] = None,
+        force_color: Optional[bool] = None,
     ) -> str:
         return str(text)