ag2 0.4.1__py3-none-any.whl → 0.4.2b1__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Potentially problematic release.
This version of ag2 might be problematic. Click here for more details.
- ag2-0.4.2b1.dist-info/METADATA +19 -0
- ag2-0.4.2b1.dist-info/RECORD +6 -0
- ag2-0.4.2b1.dist-info/top_level.txt +1 -0
- ag2-0.4.1.dist-info/METADATA +0 -500
- ag2-0.4.1.dist-info/RECORD +0 -158
- ag2-0.4.1.dist-info/top_level.txt +0 -1
- autogen/__init__.py +0 -17
- autogen/_pydantic.py +0 -116
- autogen/agentchat/__init__.py +0 -42
- autogen/agentchat/agent.py +0 -142
- autogen/agentchat/assistant_agent.py +0 -85
- autogen/agentchat/chat.py +0 -306
- autogen/agentchat/contrib/__init__.py +0 -0
- autogen/agentchat/contrib/agent_builder.py +0 -788
- autogen/agentchat/contrib/agent_eval/agent_eval.py +0 -107
- autogen/agentchat/contrib/agent_eval/criterion.py +0 -47
- autogen/agentchat/contrib/agent_eval/critic_agent.py +0 -47
- autogen/agentchat/contrib/agent_eval/quantifier_agent.py +0 -42
- autogen/agentchat/contrib/agent_eval/subcritic_agent.py +0 -48
- autogen/agentchat/contrib/agent_eval/task.py +0 -43
- autogen/agentchat/contrib/agent_optimizer.py +0 -450
- autogen/agentchat/contrib/capabilities/__init__.py +0 -0
- autogen/agentchat/contrib/capabilities/agent_capability.py +0 -21
- autogen/agentchat/contrib/capabilities/generate_images.py +0 -297
- autogen/agentchat/contrib/capabilities/teachability.py +0 -406
- autogen/agentchat/contrib/capabilities/text_compressors.py +0 -72
- autogen/agentchat/contrib/capabilities/transform_messages.py +0 -92
- autogen/agentchat/contrib/capabilities/transforms.py +0 -565
- autogen/agentchat/contrib/capabilities/transforms_util.py +0 -120
- autogen/agentchat/contrib/capabilities/vision_capability.py +0 -217
- autogen/agentchat/contrib/captainagent/tools/__init__.py +0 -0
- autogen/agentchat/contrib/captainagent/tools/data_analysis/calculate_correlation.py +0 -41
- autogen/agentchat/contrib/captainagent/tools/data_analysis/calculate_skewness_and_kurtosis.py +0 -29
- autogen/agentchat/contrib/captainagent/tools/data_analysis/detect_outlier_iqr.py +0 -29
- autogen/agentchat/contrib/captainagent/tools/data_analysis/detect_outlier_zscore.py +0 -29
- autogen/agentchat/contrib/captainagent/tools/data_analysis/explore_csv.py +0 -22
- autogen/agentchat/contrib/captainagent/tools/data_analysis/shapiro_wilk_test.py +0 -31
- autogen/agentchat/contrib/captainagent/tools/information_retrieval/arxiv_download.py +0 -26
- autogen/agentchat/contrib/captainagent/tools/information_retrieval/arxiv_search.py +0 -55
- autogen/agentchat/contrib/captainagent/tools/information_retrieval/extract_pdf_image.py +0 -54
- autogen/agentchat/contrib/captainagent/tools/information_retrieval/extract_pdf_text.py +0 -39
- autogen/agentchat/contrib/captainagent/tools/information_retrieval/get_wikipedia_text.py +0 -22
- autogen/agentchat/contrib/captainagent/tools/information_retrieval/get_youtube_caption.py +0 -35
- autogen/agentchat/contrib/captainagent/tools/information_retrieval/image_qa.py +0 -61
- autogen/agentchat/contrib/captainagent/tools/information_retrieval/optical_character_recognition.py +0 -62
- autogen/agentchat/contrib/captainagent/tools/information_retrieval/perform_web_search.py +0 -48
- autogen/agentchat/contrib/captainagent/tools/information_retrieval/scrape_wikipedia_tables.py +0 -34
- autogen/agentchat/contrib/captainagent/tools/information_retrieval/transcribe_audio_file.py +0 -22
- autogen/agentchat/contrib/captainagent/tools/information_retrieval/youtube_download.py +0 -36
- autogen/agentchat/contrib/captainagent/tools/math/calculate_circle_area_from_diameter.py +0 -22
- autogen/agentchat/contrib/captainagent/tools/math/calculate_day_of_the_week.py +0 -19
- autogen/agentchat/contrib/captainagent/tools/math/calculate_fraction_sum.py +0 -29
- autogen/agentchat/contrib/captainagent/tools/math/calculate_matrix_power.py +0 -32
- autogen/agentchat/contrib/captainagent/tools/math/calculate_reflected_point.py +0 -17
- autogen/agentchat/contrib/captainagent/tools/math/complex_numbers_product.py +0 -26
- autogen/agentchat/contrib/captainagent/tools/math/compute_currency_conversion.py +0 -24
- autogen/agentchat/contrib/captainagent/tools/math/count_distinct_permutations.py +0 -28
- autogen/agentchat/contrib/captainagent/tools/math/evaluate_expression.py +0 -29
- autogen/agentchat/contrib/captainagent/tools/math/find_continuity_point.py +0 -35
- autogen/agentchat/contrib/captainagent/tools/math/fraction_to_mixed_numbers.py +0 -40
- autogen/agentchat/contrib/captainagent/tools/math/modular_inverse_sum.py +0 -23
- autogen/agentchat/contrib/captainagent/tools/math/simplify_mixed_numbers.py +0 -37
- autogen/agentchat/contrib/captainagent/tools/math/sum_of_digit_factorials.py +0 -16
- autogen/agentchat/contrib/captainagent/tools/math/sum_of_primes_below.py +0 -16
- autogen/agentchat/contrib/captainagent/tools/requirements.txt +0 -10
- autogen/agentchat/contrib/captainagent/tools/tool_description.tsv +0 -34
- autogen/agentchat/contrib/captainagent.py +0 -490
- autogen/agentchat/contrib/gpt_assistant_agent.py +0 -545
- autogen/agentchat/contrib/graph_rag/__init__.py +0 -0
- autogen/agentchat/contrib/graph_rag/document.py +0 -30
- autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py +0 -111
- autogen/agentchat/contrib/graph_rag/falkor_graph_rag_capability.py +0 -81
- autogen/agentchat/contrib/graph_rag/graph_query_engine.py +0 -56
- autogen/agentchat/contrib/graph_rag/graph_rag_capability.py +0 -64
- autogen/agentchat/contrib/img_utils.py +0 -390
- autogen/agentchat/contrib/llamaindex_conversable_agent.py +0 -123
- autogen/agentchat/contrib/llava_agent.py +0 -176
- autogen/agentchat/contrib/math_user_proxy_agent.py +0 -471
- autogen/agentchat/contrib/multimodal_conversable_agent.py +0 -128
- autogen/agentchat/contrib/qdrant_retrieve_user_proxy_agent.py +0 -325
- autogen/agentchat/contrib/retrieve_assistant_agent.py +0 -56
- autogen/agentchat/contrib/retrieve_user_proxy_agent.py +0 -705
- autogen/agentchat/contrib/society_of_mind_agent.py +0 -203
- autogen/agentchat/contrib/swarm_agent.py +0 -463
- autogen/agentchat/contrib/text_analyzer_agent.py +0 -76
- autogen/agentchat/contrib/tool_retriever.py +0 -120
- autogen/agentchat/contrib/vectordb/__init__.py +0 -0
- autogen/agentchat/contrib/vectordb/base.py +0 -243
- autogen/agentchat/contrib/vectordb/chromadb.py +0 -326
- autogen/agentchat/contrib/vectordb/mongodb.py +0 -559
- autogen/agentchat/contrib/vectordb/pgvectordb.py +0 -958
- autogen/agentchat/contrib/vectordb/qdrant.py +0 -334
- autogen/agentchat/contrib/vectordb/utils.py +0 -126
- autogen/agentchat/contrib/web_surfer.py +0 -305
- autogen/agentchat/conversable_agent.py +0 -2908
- autogen/agentchat/groupchat.py +0 -1668
- autogen/agentchat/user_proxy_agent.py +0 -109
- autogen/agentchat/utils.py +0 -207
- autogen/browser_utils.py +0 -291
- autogen/cache/__init__.py +0 -10
- autogen/cache/abstract_cache_base.py +0 -78
- autogen/cache/cache.py +0 -182
- autogen/cache/cache_factory.py +0 -85
- autogen/cache/cosmos_db_cache.py +0 -150
- autogen/cache/disk_cache.py +0 -109
- autogen/cache/in_memory_cache.py +0 -61
- autogen/cache/redis_cache.py +0 -128
- autogen/code_utils.py +0 -745
- autogen/coding/__init__.py +0 -22
- autogen/coding/base.py +0 -113
- autogen/coding/docker_commandline_code_executor.py +0 -262
- autogen/coding/factory.py +0 -45
- autogen/coding/func_with_reqs.py +0 -203
- autogen/coding/jupyter/__init__.py +0 -22
- autogen/coding/jupyter/base.py +0 -32
- autogen/coding/jupyter/docker_jupyter_server.py +0 -164
- autogen/coding/jupyter/embedded_ipython_code_executor.py +0 -182
- autogen/coding/jupyter/jupyter_client.py +0 -224
- autogen/coding/jupyter/jupyter_code_executor.py +0 -161
- autogen/coding/jupyter/local_jupyter_server.py +0 -168
- autogen/coding/local_commandline_code_executor.py +0 -410
- autogen/coding/markdown_code_extractor.py +0 -44
- autogen/coding/utils.py +0 -57
- autogen/exception_utils.py +0 -46
- autogen/extensions/__init__.py +0 -0
- autogen/formatting_utils.py +0 -76
- autogen/function_utils.py +0 -362
- autogen/graph_utils.py +0 -148
- autogen/io/__init__.py +0 -15
- autogen/io/base.py +0 -105
- autogen/io/console.py +0 -43
- autogen/io/websockets.py +0 -213
- autogen/logger/__init__.py +0 -11
- autogen/logger/base_logger.py +0 -140
- autogen/logger/file_logger.py +0 -287
- autogen/logger/logger_factory.py +0 -29
- autogen/logger/logger_utils.py +0 -42
- autogen/logger/sqlite_logger.py +0 -459
- autogen/math_utils.py +0 -356
- autogen/oai/__init__.py +0 -33
- autogen/oai/anthropic.py +0 -428
- autogen/oai/bedrock.py +0 -606
- autogen/oai/cerebras.py +0 -270
- autogen/oai/client.py +0 -1148
- autogen/oai/client_utils.py +0 -167
- autogen/oai/cohere.py +0 -453
- autogen/oai/completion.py +0 -1216
- autogen/oai/gemini.py +0 -469
- autogen/oai/groq.py +0 -281
- autogen/oai/mistral.py +0 -279
- autogen/oai/ollama.py +0 -582
- autogen/oai/openai_utils.py +0 -811
- autogen/oai/together.py +0 -343
- autogen/retrieve_utils.py +0 -487
- autogen/runtime_logging.py +0 -163
- autogen/token_count_utils.py +0 -259
- autogen/types.py +0 -20
- autogen/version.py +0 -7
- {ag2-0.4.1.dist-info → ag2-0.4.2b1.dist-info}/LICENSE +0 -0
- {ag2-0.4.1.dist-info → ag2-0.4.2b1.dist-info}/NOTICE.md +0 -0
- {ag2-0.4.1.dist-info → ag2-0.4.2b1.dist-info}/WHEEL +0 -0
autogen/coding/jupyter/base.py
DELETED
|
@@ -1,32 +0,0 @@
|
|
|
1
|
-
# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
|
|
2
|
-
#
|
|
3
|
-
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
-
#
|
|
5
|
-
# Portions derived from https://github.com/microsoft/autogen are under the MIT License.
|
|
6
|
-
# SPDX-License-Identifier: MIT
|
|
7
|
-
from dataclasses import dataclass
|
|
8
|
-
from typing import Optional, Protocol, runtime_checkable
|
|
9
|
-
|
|
10
|
-
|
|
11
|
-
@dataclass
|
|
12
|
-
class JupyterConnectionInfo:
|
|
13
|
-
"""(Experimental)"""
|
|
14
|
-
|
|
15
|
-
host: str
|
|
16
|
-
"""`str` - Host of the Jupyter gateway server"""
|
|
17
|
-
use_https: bool
|
|
18
|
-
"""`bool` - Whether to use HTTPS"""
|
|
19
|
-
port: Optional[int] = None
|
|
20
|
-
"""`Optional[int]` - Port of the Jupyter gateway server. If None, the default port is used"""
|
|
21
|
-
token: Optional[str] = None
|
|
22
|
-
"""`Optional[str]` - Token for authentication. If None, no token is used"""
|
|
23
|
-
|
|
24
|
-
|
|
25
|
-
@runtime_checkable
|
|
26
|
-
class JupyterConnectable(Protocol):
|
|
27
|
-
"""(Experimental)"""
|
|
28
|
-
|
|
29
|
-
@property
|
|
30
|
-
def connection_info(self) -> JupyterConnectionInfo:
|
|
31
|
-
"""Return the connection information for this connectable."""
|
|
32
|
-
pass
|
|
@@ -1,164 +0,0 @@
|
|
|
1
|
-
# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
|
|
2
|
-
#
|
|
3
|
-
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
-
#
|
|
5
|
-
# Portions derived from https://github.com/microsoft/autogen are under the MIT License.
|
|
6
|
-
# SPDX-License-Identifier: MIT
|
|
7
|
-
from __future__ import annotations
|
|
8
|
-
|
|
9
|
-
import atexit
|
|
10
|
-
import io
|
|
11
|
-
import logging
|
|
12
|
-
import secrets
|
|
13
|
-
import sys
|
|
14
|
-
import uuid
|
|
15
|
-
from pathlib import Path
|
|
16
|
-
from types import TracebackType
|
|
17
|
-
from typing import Dict, Optional, Type, Union
|
|
18
|
-
|
|
19
|
-
import docker
|
|
20
|
-
|
|
21
|
-
from ..docker_commandline_code_executor import _wait_for_ready
|
|
22
|
-
|
|
23
|
-
if sys.version_info >= (3, 11):
|
|
24
|
-
from typing import Self
|
|
25
|
-
else:
|
|
26
|
-
from typing_extensions import Self
|
|
27
|
-
|
|
28
|
-
|
|
29
|
-
from .base import JupyterConnectable, JupyterConnectionInfo
|
|
30
|
-
from .jupyter_client import JupyterClient
|
|
31
|
-
|
|
32
|
-
|
|
33
|
-
class DockerJupyterServer(JupyterConnectable):
|
|
34
|
-
DEFAULT_DOCKERFILE = """FROM quay.io/jupyter/docker-stacks-foundation
|
|
35
|
-
|
|
36
|
-
SHELL ["/bin/bash", "-o", "pipefail", "-c"]
|
|
37
|
-
|
|
38
|
-
USER ${NB_UID}
|
|
39
|
-
RUN mamba install --yes jupyter_kernel_gateway ipykernel && \
|
|
40
|
-
mamba clean --all -f -y && \
|
|
41
|
-
fix-permissions "${CONDA_DIR}" && \
|
|
42
|
-
fix-permissions "/home/${NB_USER}"
|
|
43
|
-
|
|
44
|
-
ENV TOKEN="UNSET"
|
|
45
|
-
CMD python -m jupyter kernelgateway --KernelGatewayApp.ip=0.0.0.0 \
|
|
46
|
-
--KernelGatewayApp.port=8888 \
|
|
47
|
-
--KernelGatewayApp.auth_token="${TOKEN}" \
|
|
48
|
-
--JupyterApp.answer_yes=true \
|
|
49
|
-
--JupyterWebsocketPersonality.list_kernels=true
|
|
50
|
-
|
|
51
|
-
EXPOSE 8888
|
|
52
|
-
|
|
53
|
-
WORKDIR "${HOME}"
|
|
54
|
-
"""
|
|
55
|
-
|
|
56
|
-
class GenerateToken:
|
|
57
|
-
pass
|
|
58
|
-
|
|
59
|
-
def __init__(
|
|
60
|
-
self,
|
|
61
|
-
*,
|
|
62
|
-
custom_image_name: Optional[str] = None,
|
|
63
|
-
container_name: Optional[str] = None,
|
|
64
|
-
auto_remove: bool = True,
|
|
65
|
-
stop_container: bool = True,
|
|
66
|
-
docker_env: Dict[str, str] = {},
|
|
67
|
-
token: Union[str, GenerateToken] = GenerateToken(),
|
|
68
|
-
):
|
|
69
|
-
"""Start a Jupyter kernel gateway server in a Docker container.
|
|
70
|
-
|
|
71
|
-
Args:
|
|
72
|
-
custom_image_name (Optional[str], optional): Custom image to use. If this is None,
|
|
73
|
-
then the bundled image will be built and used. The default image is based on
|
|
74
|
-
quay.io/jupyter/docker-stacks-foundation and extended to include jupyter_kernel_gateway
|
|
75
|
-
container_name (Optional[str], optional): Name of the container to start.
|
|
76
|
-
A name will be generated if None.
|
|
77
|
-
auto_remove (bool, optional): If true the Docker container will be deleted
|
|
78
|
-
when it is stopped.
|
|
79
|
-
stop_container (bool, optional): If true the container will be stopped,
|
|
80
|
-
either by program exit or using the context manager
|
|
81
|
-
docker_env (Dict[str, str], optional): Extra environment variables to pass
|
|
82
|
-
to the running Docker container.
|
|
83
|
-
token (Union[str, GenerateToken], optional): Token to use for authentication.
|
|
84
|
-
If GenerateToken is used, a random token will be generated. Empty string
|
|
85
|
-
will be unauthenticated.
|
|
86
|
-
"""
|
|
87
|
-
if container_name is None:
|
|
88
|
-
container_name = f"autogen-jupyterkernelgateway-{uuid.uuid4()}"
|
|
89
|
-
|
|
90
|
-
client = docker.from_env()
|
|
91
|
-
if custom_image_name is None:
|
|
92
|
-
image_name = "autogen-jupyterkernelgateway"
|
|
93
|
-
# Make sure the image exists
|
|
94
|
-
try:
|
|
95
|
-
client.images.get(image_name)
|
|
96
|
-
except docker.errors.ImageNotFound:
|
|
97
|
-
# Build the image
|
|
98
|
-
# Get this script directory
|
|
99
|
-
here = Path(__file__).parent
|
|
100
|
-
dockerfile = io.BytesIO(self.DEFAULT_DOCKERFILE.encode("utf-8"))
|
|
101
|
-
logging.info(f"Image {image_name} not found. Building it now.")
|
|
102
|
-
client.images.build(path=here, fileobj=dockerfile, tag=image_name)
|
|
103
|
-
logging.info(f"Image {image_name} built successfully.")
|
|
104
|
-
else:
|
|
105
|
-
image_name = custom_image_name
|
|
106
|
-
# Check if the image exists
|
|
107
|
-
try:
|
|
108
|
-
client.images.get(image_name)
|
|
109
|
-
except docker.errors.ImageNotFound:
|
|
110
|
-
raise ValueError(f"Custom image {image_name} does not exist")
|
|
111
|
-
|
|
112
|
-
if isinstance(token, DockerJupyterServer.GenerateToken):
|
|
113
|
-
self._token = secrets.token_hex(32)
|
|
114
|
-
else:
|
|
115
|
-
self._token = token
|
|
116
|
-
|
|
117
|
-
# Run the container
|
|
118
|
-
env = {"TOKEN": self._token}
|
|
119
|
-
env.update(docker_env)
|
|
120
|
-
container = client.containers.run(
|
|
121
|
-
image_name,
|
|
122
|
-
detach=True,
|
|
123
|
-
auto_remove=auto_remove,
|
|
124
|
-
environment=env,
|
|
125
|
-
publish_all_ports=True,
|
|
126
|
-
name=container_name,
|
|
127
|
-
)
|
|
128
|
-
_wait_for_ready(container)
|
|
129
|
-
container_ports = container.ports
|
|
130
|
-
self._port = int(container_ports["8888/tcp"][0]["HostPort"])
|
|
131
|
-
self._container_id = container.id
|
|
132
|
-
|
|
133
|
-
def cleanup() -> None:
|
|
134
|
-
try:
|
|
135
|
-
inner_container = client.containers.get(container.id)
|
|
136
|
-
inner_container.stop()
|
|
137
|
-
except docker.errors.NotFound:
|
|
138
|
-
pass
|
|
139
|
-
|
|
140
|
-
atexit.unregister(cleanup)
|
|
141
|
-
|
|
142
|
-
if stop_container:
|
|
143
|
-
atexit.register(cleanup)
|
|
144
|
-
|
|
145
|
-
self._cleanup_func = cleanup
|
|
146
|
-
self._stop_container = stop_container
|
|
147
|
-
|
|
148
|
-
@property
|
|
149
|
-
def connection_info(self) -> JupyterConnectionInfo:
|
|
150
|
-
return JupyterConnectionInfo(host="127.0.0.1", use_https=False, port=self._port, token=self._token)
|
|
151
|
-
|
|
152
|
-
def stop(self) -> None:
|
|
153
|
-
self._cleanup_func()
|
|
154
|
-
|
|
155
|
-
def get_client(self) -> JupyterClient:
|
|
156
|
-
return JupyterClient(self.connection_info)
|
|
157
|
-
|
|
158
|
-
def __enter__(self) -> Self:
|
|
159
|
-
return self
|
|
160
|
-
|
|
161
|
-
def __exit__(
|
|
162
|
-
self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
|
|
163
|
-
) -> None:
|
|
164
|
-
self.stop()
|
|
@@ -1,182 +0,0 @@
|
|
|
1
|
-
# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
|
|
2
|
-
#
|
|
3
|
-
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
-
#
|
|
5
|
-
# Portions derived from https://github.com/microsoft/autogen are under the MIT License.
|
|
6
|
-
# SPDX-License-Identifier: MIT
|
|
7
|
-
import base64
|
|
8
|
-
import json
|
|
9
|
-
import os
|
|
10
|
-
import re
|
|
11
|
-
import uuid
|
|
12
|
-
from pathlib import Path
|
|
13
|
-
from queue import Empty
|
|
14
|
-
from typing import Any, ClassVar, List
|
|
15
|
-
|
|
16
|
-
from jupyter_client import KernelManager # type: ignore[attr-defined]
|
|
17
|
-
from jupyter_client.kernelspec import KernelSpecManager
|
|
18
|
-
from pydantic import BaseModel, Field, field_validator
|
|
19
|
-
|
|
20
|
-
from ...agentchat.agent import LLMAgent
|
|
21
|
-
from ..base import CodeBlock, CodeExtractor, IPythonCodeResult
|
|
22
|
-
from ..markdown_code_extractor import MarkdownCodeExtractor
|
|
23
|
-
|
|
24
|
-
__all__ = "EmbeddedIPythonCodeExecutor"
|
|
25
|
-
|
|
26
|
-
|
|
27
|
-
class EmbeddedIPythonCodeExecutor(BaseModel):
|
|
28
|
-
"""(Experimental) A code executor class that executes code statefully using an embedded
|
|
29
|
-
IPython kernel managed by this class.
|
|
30
|
-
|
|
31
|
-
**This will execute LLM generated code on the local machine.**
|
|
32
|
-
|
|
33
|
-
Each execution is stateful and can access variables created from previous
|
|
34
|
-
executions in the same session. The kernel must be installed before using
|
|
35
|
-
this class. The kernel can be installed using the following command:
|
|
36
|
-
`python -m ipykernel install --user --name {kernel_name}`
|
|
37
|
-
where `kernel_name` is the name of the kernel to install.
|
|
38
|
-
|
|
39
|
-
Args:
|
|
40
|
-
timeout (int): The timeout for code execution, by default 60.
|
|
41
|
-
kernel_name (str): The kernel name to use. Make sure it is installed.
|
|
42
|
-
By default, it is "python3".
|
|
43
|
-
output_dir (str): The directory to save output files, by default ".".
|
|
44
|
-
"""
|
|
45
|
-
|
|
46
|
-
timeout: int = Field(default=60, ge=1, description="The timeout for code execution.")
|
|
47
|
-
kernel_name: str = Field(default="python3", description="The kernel name to use. Make sure it is installed.")
|
|
48
|
-
output_dir: str = Field(default=".", description="The directory to save output files.")
|
|
49
|
-
|
|
50
|
-
@field_validator("output_dir")
|
|
51
|
-
@classmethod
|
|
52
|
-
def _output_dir_must_exist(cls, value: str) -> str:
|
|
53
|
-
if not os.path.exists(value):
|
|
54
|
-
raise ValueError(f"Output directory {value} does not exist.")
|
|
55
|
-
return value
|
|
56
|
-
|
|
57
|
-
def __init__(self, **kwargs: Any):
|
|
58
|
-
super().__init__(**kwargs)
|
|
59
|
-
# Check if the kernel is installed.
|
|
60
|
-
if self.kernel_name not in KernelSpecManager().find_kernel_specs():
|
|
61
|
-
raise ValueError(
|
|
62
|
-
f"Kernel {self.kernel_name} is not installed. "
|
|
63
|
-
"Please first install it with "
|
|
64
|
-
f"`python -m ipykernel install --user --name {self.kernel_name}`."
|
|
65
|
-
)
|
|
66
|
-
self._kernel_manager = KernelManager(kernel_name=self.kernel_name)
|
|
67
|
-
self._kernel_manager.start_kernel()
|
|
68
|
-
self._kernel_client = self._kernel_manager.client()
|
|
69
|
-
self._kernel_client.start_channels()
|
|
70
|
-
self._timeout = self.timeout
|
|
71
|
-
self._kernel_name = self.kernel_name
|
|
72
|
-
self._output_dir = Path(self.output_dir)
|
|
73
|
-
|
|
74
|
-
@property
|
|
75
|
-
def code_extractor(self) -> CodeExtractor:
|
|
76
|
-
"""(Experimental) Export a code extractor that can be used by an agent."""
|
|
77
|
-
return MarkdownCodeExtractor()
|
|
78
|
-
|
|
79
|
-
def execute_code_blocks(self, code_blocks: List[CodeBlock]) -> IPythonCodeResult:
|
|
80
|
-
"""(Experimental) Execute a list of code blocks and return the result.
|
|
81
|
-
|
|
82
|
-
This method executes a list of code blocks as cells in an IPython kernel
|
|
83
|
-
managed by this class.
|
|
84
|
-
See: https://jupyter-client.readthedocs.io/en/stable/messaging.html
|
|
85
|
-
for the message protocol.
|
|
86
|
-
|
|
87
|
-
Args:
|
|
88
|
-
code_blocks (List[CodeBlock]): A list of code blocks to execute.
|
|
89
|
-
|
|
90
|
-
Returns:
|
|
91
|
-
IPythonCodeResult: The result of the code execution.
|
|
92
|
-
"""
|
|
93
|
-
self._kernel_client.wait_for_ready()
|
|
94
|
-
outputs = []
|
|
95
|
-
output_files = []
|
|
96
|
-
for code_block in code_blocks:
|
|
97
|
-
code = self._process_code(code_block.code)
|
|
98
|
-
self._kernel_client.execute(code, store_history=True)
|
|
99
|
-
while True:
|
|
100
|
-
try:
|
|
101
|
-
msg = self._kernel_client.get_iopub_msg(timeout=self._timeout)
|
|
102
|
-
msg_type = msg["msg_type"]
|
|
103
|
-
content = msg["content"]
|
|
104
|
-
if msg_type in ["execute_result", "display_data"]:
|
|
105
|
-
for data_type, data in content["data"].items():
|
|
106
|
-
if data_type == "text/plain":
|
|
107
|
-
# Output is a text.
|
|
108
|
-
outputs.append(data)
|
|
109
|
-
elif data_type.startswith("image/"):
|
|
110
|
-
# Output is an image.
|
|
111
|
-
path = self._save_image(data)
|
|
112
|
-
outputs.append(f"Image data saved to {path}")
|
|
113
|
-
output_files.append(path)
|
|
114
|
-
elif data_type == "text/html":
|
|
115
|
-
# Output is an html.
|
|
116
|
-
path = self._save_html(data)
|
|
117
|
-
outputs.append(f"HTML data saved to {path}")
|
|
118
|
-
output_files.append(path)
|
|
119
|
-
else:
|
|
120
|
-
# Output raw data.
|
|
121
|
-
outputs.append(json.dumps(data))
|
|
122
|
-
elif msg_type == "stream":
|
|
123
|
-
# Output is a text.
|
|
124
|
-
outputs.append(content["text"])
|
|
125
|
-
elif msg_type == "error":
|
|
126
|
-
# Output is an error.
|
|
127
|
-
return IPythonCodeResult(
|
|
128
|
-
exit_code=1,
|
|
129
|
-
output=f"ERROR: {content['ename']}: {content['evalue']}\n{content['traceback']}",
|
|
130
|
-
)
|
|
131
|
-
if msg_type == "status" and content["execution_state"] == "idle":
|
|
132
|
-
break
|
|
133
|
-
# handle time outs.
|
|
134
|
-
except Empty:
|
|
135
|
-
return IPythonCodeResult(
|
|
136
|
-
exit_code=1,
|
|
137
|
-
output=f"ERROR: Timeout waiting for output from code block: {code_block.code}",
|
|
138
|
-
)
|
|
139
|
-
# We return the full output.
|
|
140
|
-
return IPythonCodeResult(
|
|
141
|
-
exit_code=0, output="\n".join([str(output) for output in outputs]), output_files=output_files
|
|
142
|
-
)
|
|
143
|
-
|
|
144
|
-
def restart(self) -> None:
|
|
145
|
-
"""(Experimental) Restart a new session."""
|
|
146
|
-
self._kernel_client.stop_channels()
|
|
147
|
-
self._kernel_manager.shutdown_kernel()
|
|
148
|
-
self._kernel_manager = KernelManager(kernel_name=self.kernel_name)
|
|
149
|
-
self._kernel_manager.start_kernel()
|
|
150
|
-
self._kernel_client = self._kernel_manager.client()
|
|
151
|
-
self._kernel_client.start_channels()
|
|
152
|
-
|
|
153
|
-
def _save_image(self, image_data_base64: str) -> str:
|
|
154
|
-
"""Save image data to a file."""
|
|
155
|
-
image_data = base64.b64decode(image_data_base64)
|
|
156
|
-
# Randomly generate a filename.
|
|
157
|
-
filename = f"{uuid.uuid4().hex}.png"
|
|
158
|
-
path = os.path.join(self.output_dir, filename)
|
|
159
|
-
with open(path, "wb") as f:
|
|
160
|
-
f.write(image_data)
|
|
161
|
-
return os.path.abspath(path)
|
|
162
|
-
|
|
163
|
-
def _save_html(self, html_data: str) -> str:
|
|
164
|
-
"""Save html data to a file."""
|
|
165
|
-
# Randomly generate a filename.
|
|
166
|
-
filename = f"{uuid.uuid4().hex}.html"
|
|
167
|
-
path = os.path.join(self.output_dir, filename)
|
|
168
|
-
with open(path, "w") as f:
|
|
169
|
-
f.write(html_data)
|
|
170
|
-
return os.path.abspath(path)
|
|
171
|
-
|
|
172
|
-
def _process_code(self, code: str) -> str:
|
|
173
|
-
"""Process code before execution."""
|
|
174
|
-
# Find lines that start with `! pip install` and make sure "-qqq" flag is added.
|
|
175
|
-
lines = code.split("\n")
|
|
176
|
-
for i, line in enumerate(lines):
|
|
177
|
-
# use regex to find lines that start with `! pip install` or `!pip install`.
|
|
178
|
-
match = re.search(r"^! ?pip install", line)
|
|
179
|
-
if match is not None:
|
|
180
|
-
if "-qqq" not in line:
|
|
181
|
-
lines[i] = line.replace(match.group(0), match.group(0) + " -qqq")
|
|
182
|
-
return "\n".join(lines)
|
|
@@ -1,224 +0,0 @@
|
|
|
1
|
-
# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
|
|
2
|
-
#
|
|
3
|
-
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
-
#
|
|
5
|
-
# Portions derived from https://github.com/microsoft/autogen are under the MIT License.
|
|
6
|
-
# SPDX-License-Identifier: MIT
|
|
7
|
-
from __future__ import annotations
|
|
8
|
-
|
|
9
|
-
import sys
|
|
10
|
-
from dataclasses import dataclass
|
|
11
|
-
from types import TracebackType
|
|
12
|
-
from typing import Any, Dict, List, Optional, Type, cast
|
|
13
|
-
|
|
14
|
-
if sys.version_info >= (3, 11):
|
|
15
|
-
from typing import Self
|
|
16
|
-
else:
|
|
17
|
-
from typing_extensions import Self
|
|
18
|
-
|
|
19
|
-
import datetime
|
|
20
|
-
import json
|
|
21
|
-
import uuid
|
|
22
|
-
|
|
23
|
-
import requests
|
|
24
|
-
import websocket
|
|
25
|
-
from requests.adapters import HTTPAdapter, Retry
|
|
26
|
-
from websocket import WebSocket
|
|
27
|
-
|
|
28
|
-
from .base import JupyterConnectionInfo
|
|
29
|
-
|
|
30
|
-
|
|
31
|
-
class JupyterClient:
|
|
32
|
-
def __init__(self, connection_info: JupyterConnectionInfo):
|
|
33
|
-
"""(Experimental) A client for communicating with a Jupyter gateway server.
|
|
34
|
-
|
|
35
|
-
Args:
|
|
36
|
-
connection_info (JupyterConnectionInfo): Connection information
|
|
37
|
-
"""
|
|
38
|
-
self._connection_info = connection_info
|
|
39
|
-
self._session = requests.Session()
|
|
40
|
-
retries = Retry(total=5, backoff_factor=0.1)
|
|
41
|
-
self._session.mount("http://", HTTPAdapter(max_retries=retries))
|
|
42
|
-
|
|
43
|
-
def _get_headers(self) -> Dict[str, str]:
|
|
44
|
-
if self._connection_info.token is None:
|
|
45
|
-
return {}
|
|
46
|
-
return {"Authorization": f"token {self._connection_info.token}"}
|
|
47
|
-
|
|
48
|
-
def _get_api_base_url(self) -> str:
|
|
49
|
-
protocol = "https" if self._connection_info.use_https else "http"
|
|
50
|
-
port = f":{self._connection_info.port}" if self._connection_info.port else ""
|
|
51
|
-
return f"{protocol}://{self._connection_info.host}{port}"
|
|
52
|
-
|
|
53
|
-
def _get_ws_base_url(self) -> str:
|
|
54
|
-
port = f":{self._connection_info.port}" if self._connection_info.port else ""
|
|
55
|
-
return f"ws://{self._connection_info.host}{port}"
|
|
56
|
-
|
|
57
|
-
def list_kernel_specs(self) -> Dict[str, Dict[str, str]]:
|
|
58
|
-
response = self._session.get(f"{self._get_api_base_url()}/api/kernelspecs", headers=self._get_headers())
|
|
59
|
-
return cast(Dict[str, Dict[str, str]], response.json())
|
|
60
|
-
|
|
61
|
-
def list_kernels(self) -> List[Dict[str, str]]:
|
|
62
|
-
response = self._session.get(f"{self._get_api_base_url()}/api/kernels", headers=self._get_headers())
|
|
63
|
-
return cast(List[Dict[str, str]], response.json())
|
|
64
|
-
|
|
65
|
-
def start_kernel(self, kernel_spec_name: str) -> str:
|
|
66
|
-
"""Start a new kernel.
|
|
67
|
-
|
|
68
|
-
Args:
|
|
69
|
-
kernel_spec_name (str): Name of the kernel spec to start
|
|
70
|
-
|
|
71
|
-
Returns:
|
|
72
|
-
str: ID of the started kernel
|
|
73
|
-
"""
|
|
74
|
-
|
|
75
|
-
response = self._session.post(
|
|
76
|
-
f"{self._get_api_base_url()}/api/kernels",
|
|
77
|
-
headers=self._get_headers(),
|
|
78
|
-
json={"name": kernel_spec_name},
|
|
79
|
-
)
|
|
80
|
-
return cast(str, response.json()["id"])
|
|
81
|
-
|
|
82
|
-
def delete_kernel(self, kernel_id: str) -> None:
|
|
83
|
-
response = self._session.delete(
|
|
84
|
-
f"{self._get_api_base_url()}/api/kernels/{kernel_id}", headers=self._get_headers()
|
|
85
|
-
)
|
|
86
|
-
response.raise_for_status()
|
|
87
|
-
|
|
88
|
-
def restart_kernel(self, kernel_id: str) -> None:
|
|
89
|
-
response = self._session.post(
|
|
90
|
-
f"{self._get_api_base_url()}/api/kernels/{kernel_id}/restart", headers=self._get_headers()
|
|
91
|
-
)
|
|
92
|
-
response.raise_for_status()
|
|
93
|
-
|
|
94
|
-
def get_kernel_client(self, kernel_id: str) -> JupyterKernelClient:
|
|
95
|
-
ws_url = f"{self._get_ws_base_url()}/api/kernels/{kernel_id}/channels"
|
|
96
|
-
ws = websocket.create_connection(ws_url, header=self._get_headers())
|
|
97
|
-
return JupyterKernelClient(ws)
|
|
98
|
-
|
|
99
|
-
|
|
100
|
-
class JupyterKernelClient:
|
|
101
|
-
"""(Experimental) A client for communicating with a Jupyter kernel."""
|
|
102
|
-
|
|
103
|
-
@dataclass
|
|
104
|
-
class ExecutionResult:
|
|
105
|
-
@dataclass
|
|
106
|
-
class DataItem:
|
|
107
|
-
mime_type: str
|
|
108
|
-
data: str
|
|
109
|
-
|
|
110
|
-
is_ok: bool
|
|
111
|
-
output: str
|
|
112
|
-
data_items: List[DataItem]
|
|
113
|
-
|
|
114
|
-
def __init__(self, websocket: WebSocket):
|
|
115
|
-
self._session_id: str = uuid.uuid4().hex
|
|
116
|
-
self._websocket: WebSocket = websocket
|
|
117
|
-
|
|
118
|
-
def __enter__(self) -> Self:
|
|
119
|
-
return self
|
|
120
|
-
|
|
121
|
-
def __exit__(
|
|
122
|
-
self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType]
|
|
123
|
-
) -> None:
|
|
124
|
-
self.stop()
|
|
125
|
-
|
|
126
|
-
def stop(self) -> None:
|
|
127
|
-
self._websocket.close()
|
|
128
|
-
|
|
129
|
-
def _send_message(self, *, content: Dict[str, Any], channel: str, message_type: str) -> str:
|
|
130
|
-
timestamp = datetime.datetime.now().isoformat()
|
|
131
|
-
message_id = uuid.uuid4().hex
|
|
132
|
-
message = {
|
|
133
|
-
"header": {
|
|
134
|
-
"username": "autogen",
|
|
135
|
-
"version": "5.0",
|
|
136
|
-
"session": self._session_id,
|
|
137
|
-
"msg_id": message_id,
|
|
138
|
-
"msg_type": message_type,
|
|
139
|
-
"date": timestamp,
|
|
140
|
-
},
|
|
141
|
-
"parent_header": {},
|
|
142
|
-
"channel": channel,
|
|
143
|
-
"content": content,
|
|
144
|
-
"metadata": {},
|
|
145
|
-
"buffers": {},
|
|
146
|
-
}
|
|
147
|
-
self._websocket.send_text(json.dumps(message))
|
|
148
|
-
return message_id
|
|
149
|
-
|
|
150
|
-
def _receive_message(self, timeout_seconds: Optional[float]) -> Optional[Dict[str, Any]]:
|
|
151
|
-
self._websocket.settimeout(timeout_seconds)
|
|
152
|
-
try:
|
|
153
|
-
data = self._websocket.recv()
|
|
154
|
-
if isinstance(data, bytes):
|
|
155
|
-
data = data.decode("utf-8")
|
|
156
|
-
return cast(Dict[str, Any], json.loads(data))
|
|
157
|
-
except websocket.WebSocketTimeoutException:
|
|
158
|
-
return None
|
|
159
|
-
|
|
160
|
-
def wait_for_ready(self, timeout_seconds: Optional[float] = None) -> bool:
|
|
161
|
-
message_id = self._send_message(content={}, channel="shell", message_type="kernel_info_request")
|
|
162
|
-
while True:
|
|
163
|
-
message = self._receive_message(timeout_seconds)
|
|
164
|
-
# This means we timed out with no new messages.
|
|
165
|
-
if message is None:
|
|
166
|
-
return False
|
|
167
|
-
if (
|
|
168
|
-
message.get("parent_header", {}).get("msg_id") == message_id
|
|
169
|
-
and message["msg_type"] == "kernel_info_reply"
|
|
170
|
-
):
|
|
171
|
-
return True
|
|
172
|
-
|
|
173
|
-
def execute(self, code: str, timeout_seconds: Optional[float] = None) -> ExecutionResult:
|
|
174
|
-
message_id = self._send_message(
|
|
175
|
-
content={
|
|
176
|
-
"code": code,
|
|
177
|
-
"silent": False,
|
|
178
|
-
"store_history": True,
|
|
179
|
-
"user_expressions": {},
|
|
180
|
-
"allow_stdin": False,
|
|
181
|
-
"stop_on_error": True,
|
|
182
|
-
},
|
|
183
|
-
channel="shell",
|
|
184
|
-
message_type="execute_request",
|
|
185
|
-
)
|
|
186
|
-
|
|
187
|
-
text_output = []
|
|
188
|
-
data_output = []
|
|
189
|
-
while True:
|
|
190
|
-
message = self._receive_message(timeout_seconds)
|
|
191
|
-
if message is None:
|
|
192
|
-
return JupyterKernelClient.ExecutionResult(
|
|
193
|
-
is_ok=False, output="ERROR: Timeout waiting for output from code block.", data_items=[]
|
|
194
|
-
)
|
|
195
|
-
|
|
196
|
-
# Ignore messages that are not for this execution.
|
|
197
|
-
if message.get("parent_header", {}).get("msg_id") != message_id:
|
|
198
|
-
continue
|
|
199
|
-
|
|
200
|
-
msg_type = message["msg_type"]
|
|
201
|
-
content = message["content"]
|
|
202
|
-
if msg_type in ["execute_result", "display_data"]:
|
|
203
|
-
for data_type, data in content["data"].items():
|
|
204
|
-
if data_type == "text/plain":
|
|
205
|
-
text_output.append(data)
|
|
206
|
-
elif data_type.startswith("image/") or data_type == "text/html":
|
|
207
|
-
data_output.append(self.ExecutionResult.DataItem(mime_type=data_type, data=data))
|
|
208
|
-
else:
|
|
209
|
-
text_output.append(json.dumps(data))
|
|
210
|
-
elif msg_type == "stream":
|
|
211
|
-
text_output.append(content["text"])
|
|
212
|
-
elif msg_type == "error":
|
|
213
|
-
# Output is an error.
|
|
214
|
-
return JupyterKernelClient.ExecutionResult(
|
|
215
|
-
is_ok=False,
|
|
216
|
-
output=f"ERROR: {content['ename']}: {content['evalue']}\n{content['traceback']}",
|
|
217
|
-
data_items=[],
|
|
218
|
-
)
|
|
219
|
-
if msg_type == "status" and content["execution_state"] == "idle":
|
|
220
|
-
break
|
|
221
|
-
|
|
222
|
-
return JupyterKernelClient.ExecutionResult(
|
|
223
|
-
is_ok=True, output="\n".join([str(output) for output in text_output]), data_items=data_output
|
|
224
|
-
)
|