zrb 1.0.0a2__py3-none-any.whl → 1.0.0a4__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.
- zrb/__init__.py +49 -40
- zrb/__main__.py +5 -3
- zrb/attr/type.py +2 -1
- zrb/builtin/__init__.py +42 -2
- zrb/builtin/base64.py +34 -0
- zrb/builtin/git.py +156 -0
- zrb/builtin/git_subtree.py +88 -0
- zrb/builtin/group.py +34 -0
- zrb/builtin/llm/llm_chat.py +47 -0
- zrb/builtin/llm/tool/cli.py +9 -0
- zrb/builtin/llm/tool/rag.py +189 -0
- zrb/builtin/llm/tool/web.py +74 -0
- zrb/builtin/md5.py +36 -0
- zrb/builtin/project/add/fastapp.py +72 -0
- zrb/builtin/project/add/fastapp_template/.gitignore +4 -0
- zrb/builtin/project/add/fastapp_template/README.md +7 -0
- zrb/builtin/project/add/fastapp_template/_zrb/config.py +17 -0
- zrb/builtin/project/add/fastapp_template/_zrb/group.py +16 -0
- zrb/builtin/project/add/fastapp_template/_zrb/helper.py +97 -0
- zrb/builtin/project/add/fastapp_template/_zrb/main.py +132 -0
- zrb/builtin/project/add/fastapp_template/_zrb/venv_task.py +22 -0
- zrb/builtin/project/add/fastapp_template/common/app.py +18 -0
- zrb/builtin/project/add/fastapp_template/common/db_engine.py +5 -0
- zrb/builtin/project/add/fastapp_template/common/db_repository.py +134 -0
- zrb/builtin/project/add/fastapp_template/common/error.py +8 -0
- zrb/builtin/project/add/fastapp_template/common/schema.py +5 -0
- zrb/builtin/project/add/fastapp_template/common/usecase.py +232 -0
- zrb/builtin/project/add/fastapp_template/config.py +29 -0
- zrb/builtin/project/add/fastapp_template/main.py +7 -0
- zrb/builtin/project/add/fastapp_template/migrate.py +3 -0
- zrb/builtin/project/add/fastapp_template/module/__init__.py +0 -0
- zrb/builtin/project/add/fastapp_template/module/auth/alembic.ini +117 -0
- zrb/builtin/project/add/fastapp_template/module/auth/client/api_client.py +7 -0
- zrb/builtin/project/add/fastapp_template/module/auth/client/base_client.py +27 -0
- zrb/builtin/project/add/fastapp_template/module/auth/client/direct_client.py +6 -0
- zrb/builtin/project/add/fastapp_template/module/auth/client/factory.py +9 -0
- zrb/builtin/project/add/fastapp_template/module/auth/migration/README +1 -0
- zrb/builtin/project/add/fastapp_template/module/auth/migration/env.py +108 -0
- zrb/builtin/project/add/fastapp_template/module/auth/migration/script.py.mako +26 -0
- zrb/builtin/project/add/fastapp_template/module/auth/migration/versions/3093c7336477_add_user_table.py +37 -0
- zrb/builtin/project/add/fastapp_template/module/auth/migration_metadata.py +6 -0
- zrb/builtin/project/add/fastapp_template/module/auth/route.py +22 -0
- zrb/builtin/project/add/fastapp_template/module/auth/service/__init__.py +0 -0
- zrb/builtin/project/add/fastapp_template/module/auth/service/user/__init__.py +0 -0
- zrb/builtin/project/add/fastapp_template/module/auth/service/user/repository/__init__.py +0 -0
- zrb/builtin/project/add/fastapp_template/module/auth/service/user/repository/db_repository.py +39 -0
- zrb/builtin/project/add/fastapp_template/module/auth/service/user/repository/factory.py +13 -0
- zrb/builtin/project/add/fastapp_template/module/auth/service/user/repository/repository.py +34 -0
- zrb/builtin/project/add/fastapp_template/module/auth/service/user/usecase.py +45 -0
- zrb/builtin/project/add/fastapp_template/module/gateway/alembic.ini +117 -0
- zrb/builtin/project/add/fastapp_template/module/gateway/migration/README +1 -0
- zrb/builtin/project/add/fastapp_template/module/gateway/migration/env.py +108 -0
- zrb/builtin/project/add/fastapp_template/module/gateway/migration/script.py.mako +26 -0
- zrb/builtin/project/add/fastapp_template/module/gateway/migration/versions/.gitkeep +0 -0
- zrb/builtin/project/add/fastapp_template/module/gateway/migration_metadata.py +3 -0
- zrb/builtin/project/add/fastapp_template/module/gateway/route.py +27 -0
- zrb/builtin/project/add/fastapp_template/requirements.txt +6 -0
- zrb/builtin/project/add/fastapp_template/schema/__init__.py +0 -0
- zrb/builtin/project/add/fastapp_template/schema/role.py +31 -0
- zrb/builtin/project/add/fastapp_template/schema/user.py +31 -0
- zrb/builtin/project/add/fastapp_template/template.env +2 -0
- zrb/builtin/project/create/__init__.py +0 -0
- zrb/builtin/project/create/create.py +41 -0
- zrb/builtin/project/create/project-template/README.md +3 -0
- zrb/builtin/project/create/project-template/zrb_init.py +7 -0
- zrb/builtin/python.py +11 -0
- zrb/builtin/shell/__init__.py +0 -5
- zrb/builtin/shell/autocomplete/__init__.py +0 -9
- zrb/builtin/shell/autocomplete/bash.py +5 -6
- zrb/builtin/shell/autocomplete/subcmd.py +7 -8
- zrb/builtin/shell/autocomplete/zsh.py +5 -6
- zrb/builtin/todo.py +219 -0
- zrb/callback/any_callback.py +1 -1
- zrb/callback/callback.py +5 -5
- zrb/cmd/cmd_val.py +2 -2
- zrb/config.py +16 -3
- zrb/content_transformer/any_content_transformer.py +1 -1
- zrb/content_transformer/content_transformer.py +2 -2
- zrb/context/any_context.py +1 -1
- zrb/context/any_shared_context.py +3 -3
- zrb/context/context.py +10 -8
- zrb/context/shared_context.py +9 -8
- zrb/env/__init__.py +0 -3
- zrb/env/any_env.py +1 -1
- zrb/env/env.py +3 -4
- zrb/env/env_file.py +4 -4
- zrb/env/env_map.py +2 -2
- zrb/group/__init__.py +0 -3
- zrb/group/any_group.py +3 -3
- zrb/group/group.py +7 -6
- zrb/input/any_input.py +1 -1
- zrb/input/base_input.py +4 -4
- zrb/input/bool_input.py +5 -5
- zrb/input/float_input.py +3 -3
- zrb/input/int_input.py +3 -3
- zrb/input/option_input.py +51 -0
- zrb/input/password_input.py +2 -2
- zrb/input/str_input.py +1 -1
- zrb/input/text_input.py +12 -10
- zrb/runner/cli.py +80 -45
- zrb/runner/web_app.py +150 -0
- zrb/runner/web_controller/__init__.py +0 -0
- zrb/runner/web_controller/group_info_ui/__init__.py +0 -0
- zrb/runner/{web_app → web_controller}/group_info_ui/controller.py +7 -8
- zrb/runner/{web_app → web_controller}/group_info_ui/view.html +2 -2
- zrb/runner/web_controller/home_page/__init__.py +0 -0
- zrb/runner/{web_app → web_controller}/home_page/controller.py +7 -6
- zrb/runner/{web_app → web_controller}/home_page/view.html +2 -2
- zrb/runner/web_controller/task_ui/__init__.py +0 -0
- zrb/runner/{web_app → web_controller}/task_ui/controller.py +8 -12
- zrb/runner/{web_app → web_controller}/task_ui/view.html +2 -2
- zrb/runner/web_util.py +5 -35
- zrb/session/any_session.py +13 -7
- zrb/session/session.py +78 -40
- zrb/session_state_log/session_state_log.py +7 -5
- zrb/session_state_logger/any_session_state_logger.py +1 -1
- zrb/session_state_logger/default_session_state_logger.py +2 -2
- zrb/session_state_logger/file_session_state_logger.py +19 -27
- zrb/task/any_task.py +4 -4
- zrb/task/base_task.py +33 -23
- zrb/task/base_trigger.py +11 -12
- zrb/task/cmd_task.py +72 -65
- zrb/task/http_check.py +13 -13
- zrb/task/llm_task.py +215 -0
- zrb/task/make_task.py +9 -9
- zrb/task/rsync_task.py +25 -25
- zrb/task/scaffolder.py +18 -15
- zrb/task/scheduler.py +6 -7
- zrb/task/task.py +1 -1
- zrb/task/tcp_check.py +11 -13
- zrb/util/attr.py +19 -3
- zrb/util/cli/style.py +71 -2
- zrb/util/cli/subcommand.py +2 -2
- zrb/util/codemod/__init__.py +0 -0
- zrb/util/codemod/add_code_to_class.py +35 -0
- zrb/util/codemod/add_code_to_function.py +36 -0
- zrb/util/codemod/add_code_to_method.py +55 -0
- zrb/util/codemod/add_key_to_dict.py +51 -0
- zrb/util/codemod/add_param_to_function_call.py +39 -0
- zrb/util/codemod/add_property_to_class.py +55 -0
- zrb/util/git.py +156 -0
- zrb/util/git_subtree.py +94 -0
- zrb/util/group.py +2 -2
- zrb/util/llm/tool.py +63 -0
- zrb/util/string/conversion.py +7 -0
- zrb/util/todo.py +259 -0
- {zrb-1.0.0a2.dist-info → zrb-1.0.0a4.dist-info}/METADATA +13 -5
- zrb-1.0.0a4.dist-info/RECORD +197 -0
- zrb/builtin/shell/_group.py +0 -9
- zrb/builtin/shell/autocomplete/_group.py +0 -6
- zrb/runner/web_app/any_request_handler.py +0 -24
- zrb/runner/web_server.py +0 -224
- zrb-1.0.0a2.dist-info/RECORD +0 -120
- /zrb/{runner/web_app → builtin/project}/__init__.py +0 -0
- /zrb/{runner/web_app/group_info_ui → builtin/project/add}/__init__.py +0 -0
- /zrb/{runner/web_app/home_page → builtin/project/add/fastapp_template}/__init__.py +0 -0
- /zrb/{runner/web_app/task_ui → builtin/project/add/fastapp_template/common}/__init__.py +0 -0
- /zrb/runner/{web_app → web_controller}/group_info_ui/partial/group_info.html +0 -0
- /zrb/runner/{web_app → web_controller}/group_info_ui/partial/group_li.html +0 -0
- /zrb/runner/{web_app → web_controller}/group_info_ui/partial/task_info.html +0 -0
- /zrb/runner/{web_app → web_controller}/group_info_ui/partial/task_li.html +0 -0
- /zrb/runner/{web_app → web_controller}/home_page/partial/group_info.html +0 -0
- /zrb/runner/{web_app → web_controller}/home_page/partial/group_li.html +0 -0
- /zrb/runner/{web_app → web_controller}/home_page/partial/task_info.html +0 -0
- /zrb/runner/{web_app → web_controller}/home_page/partial/task_li.html +0 -0
- /zrb/runner/{web_app → web_controller}/static/favicon-32x32.png +0 -0
- /zrb/runner/{web_app → web_controller}/static/pico.min.css +0 -0
- /zrb/runner/{web_app → web_controller}/task_ui/partial/common-util.js +0 -0
- /zrb/runner/{web_app → web_controller}/task_ui/partial/input.html +0 -0
- /zrb/runner/{web_app → web_controller}/task_ui/partial/main.js +0 -0
- /zrb/runner/{web_app → web_controller}/task_ui/partial/show-existing-session.js +0 -0
- /zrb/runner/{web_app → web_controller}/task_ui/partial/visualize-history.js +0 -0
- {zrb-1.0.0a2.dist-info → zrb-1.0.0a4.dist-info}/WHEEL +0 -0
- {zrb-1.0.0a2.dist-info → zrb-1.0.0a4.dist-info}/entry_points.txt +0 -0
zrb/task/cmd_task.py
CHANGED
@@ -2,17 +2,17 @@ import asyncio
|
|
2
2
|
import os
|
3
3
|
import sys
|
4
4
|
|
5
|
-
from
|
6
|
-
from
|
7
|
-
from
|
8
|
-
from
|
9
|
-
from
|
10
|
-
from
|
11
|
-
from
|
12
|
-
from
|
13
|
-
from
|
14
|
-
from .
|
15
|
-
from .
|
5
|
+
from zrb.attr.type import BoolAttr, IntAttr, StrAttr
|
6
|
+
from zrb.cmd.cmd_result import CmdResult
|
7
|
+
from zrb.cmd.cmd_val import AnyCmdVal, CmdVal, SingleCmdVal
|
8
|
+
from zrb.config import DEFAULT_SHELL
|
9
|
+
from zrb.context.any_context import AnyContext
|
10
|
+
from zrb.env.any_env import AnyEnv
|
11
|
+
from zrb.input.any_input import AnyInput
|
12
|
+
from zrb.task.any_task import AnyTask
|
13
|
+
from zrb.task.base_task import BaseTask
|
14
|
+
from zrb.util.attr import get_int_attr, get_str_attr
|
15
|
+
from zrb.util.cmd.remote import get_remote_cmd_script
|
16
16
|
|
17
17
|
|
18
18
|
class CmdTask(BaseTask):
|
@@ -23,25 +23,25 @@ class CmdTask(BaseTask):
|
|
23
23
|
icon: str | None = None,
|
24
24
|
description: str | None = None,
|
25
25
|
cli_only: bool = False,
|
26
|
-
input: list[AnyInput] | AnyInput | None = None,
|
27
|
-
env: list[AnyEnv] | AnyEnv | None = None,
|
26
|
+
input: list[AnyInput | None] | AnyInput | None = None,
|
27
|
+
env: list[AnyEnv | None] | AnyEnv | None = None,
|
28
28
|
shell: StrAttr | None = None,
|
29
|
-
|
29
|
+
render_shell: bool = True,
|
30
30
|
shell_flag: StrAttr | None = None,
|
31
|
-
|
31
|
+
render_shell_flag: bool = True,
|
32
32
|
remote_host: StrAttr | None = None,
|
33
|
-
|
33
|
+
render_remote_host: bool = True,
|
34
34
|
remote_port: IntAttr | None = None,
|
35
35
|
remote_user: StrAttr | None = None,
|
36
|
-
|
36
|
+
render_remote_user: bool = True,
|
37
37
|
remote_password: StrAttr | None = None,
|
38
|
-
|
38
|
+
render_remote_password: bool = True,
|
39
39
|
remote_ssh_key: StrAttr | None = None,
|
40
|
-
|
40
|
+
render_remote_ssh_key: bool = True,
|
41
41
|
cmd: CmdVal = "",
|
42
|
-
|
42
|
+
render_cmd: bool = True,
|
43
43
|
cwd: str | None = None,
|
44
|
-
|
44
|
+
render_cwd: bool = True,
|
45
45
|
max_output_line: int = 1000,
|
46
46
|
max_error_line: int = 1000,
|
47
47
|
execute_condition: BoolAttr = True,
|
@@ -77,22 +77,22 @@ class CmdTask(BaseTask):
|
|
77
77
|
fallback=fallback,
|
78
78
|
)
|
79
79
|
self._shell = shell
|
80
|
-
self.
|
80
|
+
self._render_shell = render_shell
|
81
81
|
self._shell_flag = shell_flag
|
82
|
-
self.
|
82
|
+
self._render_shell_flag = render_shell_flag
|
83
83
|
self._remote_host = remote_host
|
84
|
-
self.
|
84
|
+
self._render_remote_host = render_remote_host
|
85
85
|
self._remote_port = remote_port
|
86
86
|
self._remote_user = remote_user
|
87
|
-
self.
|
87
|
+
self._render_remote_user = render_remote_user
|
88
88
|
self._remote_password = remote_password
|
89
|
-
self.
|
89
|
+
self._render_remote_password = render_remote_password
|
90
90
|
self._remote_ssh_key = remote_ssh_key
|
91
|
-
self.
|
91
|
+
self._render_remote_ssh_key = render_remote_ssh_key
|
92
92
|
self._cmd = cmd
|
93
|
-
self.
|
93
|
+
self._render_cmd = render_cmd
|
94
94
|
self._cwd = cwd
|
95
|
-
self.
|
95
|
+
self._render_cwd = render_cwd
|
96
96
|
self._max_output_line = max_output_line
|
97
97
|
self._max_error_line = max_error_line
|
98
98
|
|
@@ -115,36 +115,45 @@ class CmdTask(BaseTask):
|
|
115
115
|
ctx.log_debug(f"Working directory: {cwd}")
|
116
116
|
env_map = self.__get_env_map(ctx)
|
117
117
|
ctx.log_debug(f"Environment map: {env_map}")
|
118
|
-
cmd_process =
|
119
|
-
|
120
|
-
|
121
|
-
|
122
|
-
|
123
|
-
|
124
|
-
|
125
|
-
|
126
|
-
|
127
|
-
|
128
|
-
|
129
|
-
|
130
|
-
|
131
|
-
|
132
|
-
|
133
|
-
|
134
|
-
|
135
|
-
|
136
|
-
|
137
|
-
|
138
|
-
|
139
|
-
|
140
|
-
|
141
|
-
|
142
|
-
|
143
|
-
|
118
|
+
cmd_process = None
|
119
|
+
try:
|
120
|
+
cmd_process = await asyncio.create_subprocess_exec(
|
121
|
+
shell,
|
122
|
+
shell_flag,
|
123
|
+
cmd_script,
|
124
|
+
cwd=cwd,
|
125
|
+
stdin=sys.stdin if sys.stdin.isatty() else None,
|
126
|
+
stdout=asyncio.subprocess.PIPE,
|
127
|
+
stderr=asyncio.subprocess.PIPE,
|
128
|
+
env=env_map,
|
129
|
+
bufsize=0,
|
130
|
+
)
|
131
|
+
stdout_task = asyncio.create_task(
|
132
|
+
self.__read_stream(cmd_process.stdout, ctx.print, self._max_output_line)
|
133
|
+
)
|
134
|
+
stderr_task = asyncio.create_task(
|
135
|
+
self.__read_stream(cmd_process.stderr, ctx.print, self._max_error_line)
|
136
|
+
)
|
137
|
+
# Wait for process to complete and gather stdout/stderr
|
138
|
+
return_code = await cmd_process.wait()
|
139
|
+
stdout = await stdout_task
|
140
|
+
stderr = await stderr_task
|
141
|
+
# Check for errors
|
142
|
+
if return_code != 0:
|
143
|
+
ctx.log_error(f"Exit status: {return_code}")
|
144
|
+
raise Exception(
|
145
|
+
f"Process {self._name} exited ({return_code}): {stderr}"
|
146
|
+
)
|
147
|
+
ctx.log_info(f"Exit status: {return_code}")
|
148
|
+
return CmdResult(stdout, stderr)
|
149
|
+
finally:
|
150
|
+
if cmd_process is not None and cmd_process.returncode is None:
|
151
|
+
cmd_process.terminate()
|
144
152
|
|
145
153
|
def __get_env_map(self, ctx: AnyContext) -> dict[str, str]:
|
146
154
|
envs = {key: val for key, val in ctx.env.items()}
|
147
155
|
envs["_ZRB_SSH_PASSWORD"] = self._get_remote_password(ctx)
|
156
|
+
envs["PYTHONBUFFERED"] = "1"
|
148
157
|
return envs
|
149
158
|
|
150
159
|
async def __read_stream(self, stream, log_method, max_lines):
|
@@ -162,7 +171,7 @@ class CmdTask(BaseTask):
|
|
162
171
|
|
163
172
|
def _get_shell(self, ctx: AnyContext) -> str:
|
164
173
|
return get_str_attr(
|
165
|
-
ctx, self._shell, DEFAULT_SHELL, auto_render=self.
|
174
|
+
ctx, self._shell, DEFAULT_SHELL, auto_render=self._render_shell
|
166
175
|
)
|
167
176
|
|
168
177
|
def _get_shell_flag(self, ctx: AnyContext) -> str:
|
@@ -177,12 +186,12 @@ class CmdTask(BaseTask):
|
|
177
186
|
ctx,
|
178
187
|
self._shell_flag,
|
179
188
|
default_shell_flag,
|
180
|
-
auto_render=self.
|
189
|
+
auto_render=self._render_shell_flag,
|
181
190
|
)
|
182
191
|
|
183
192
|
def _get_remote_host(self, ctx: AnyContext) -> str:
|
184
193
|
return get_str_attr(
|
185
|
-
ctx, self._remote_host, "", auto_render=self.
|
194
|
+
ctx, self._remote_host, "", auto_render=self._render_remote_host
|
186
195
|
)
|
187
196
|
|
188
197
|
def _get_remote_port(self, ctx: AnyContext) -> int:
|
@@ -190,7 +199,7 @@ class CmdTask(BaseTask):
|
|
190
199
|
|
191
200
|
def _get_remote_user(self, ctx: AnyContext) -> str:
|
192
201
|
return get_str_attr(
|
193
|
-
ctx, self._remote_user, "", auto_render=self.
|
202
|
+
ctx, self._remote_user, "", auto_render=self._render_remote_user
|
194
203
|
)
|
195
204
|
|
196
205
|
def _get_remote_password(self, ctx: AnyContext) -> str:
|
@@ -198,18 +207,16 @@ class CmdTask(BaseTask):
|
|
198
207
|
ctx,
|
199
208
|
self._remote_password,
|
200
209
|
"",
|
201
|
-
auto_render=self.
|
210
|
+
auto_render=self._render_remote_password,
|
202
211
|
)
|
203
212
|
|
204
213
|
def _get_remote_ssh_key(self, ctx: AnyContext) -> str:
|
205
214
|
return get_str_attr(
|
206
|
-
ctx, self._remote_ssh_key, "", auto_render=self.
|
215
|
+
ctx, self._remote_ssh_key, "", auto_render=self._render_remote_ssh_key
|
207
216
|
)
|
208
217
|
|
209
218
|
def _get_cwd(self, ctx: AnyContext) -> str:
|
210
|
-
cwd = get_str_attr(
|
211
|
-
ctx, self._cwd, os.getcwd(), auto_render=self._auto_render_cwd
|
212
|
-
)
|
219
|
+
cwd = get_str_attr(ctx, self._cwd, os.getcwd(), auto_render=self._render_cwd)
|
213
220
|
if cwd is None:
|
214
221
|
cwd = os.getcwd()
|
215
222
|
return os.path.abspath(cwd)
|
@@ -249,7 +256,7 @@ class CmdTask(BaseTask):
|
|
249
256
|
if callable(single_cmd_val):
|
250
257
|
return single_cmd_val(ctx)
|
251
258
|
if isinstance(single_cmd_val, str):
|
252
|
-
if self.
|
259
|
+
if self._render_cmd:
|
253
260
|
return ctx.render(single_cmd_val)
|
254
261
|
return single_cmd_val
|
255
262
|
if isinstance(single_cmd_val, AnyCmdVal):
|
zrb/task/http_check.py
CHANGED
@@ -1,16 +1,14 @@
|
|
1
1
|
import asyncio
|
2
2
|
from collections.abc import Callable
|
3
3
|
|
4
|
-
import
|
5
|
-
|
6
|
-
from
|
7
|
-
from
|
8
|
-
from
|
9
|
-
from
|
10
|
-
from
|
11
|
-
from
|
12
|
-
from .any_task import AnyTask
|
13
|
-
from .base_task import BaseTask
|
4
|
+
from zrb.attr.type import StrAttr
|
5
|
+
from zrb.context.any_context import AnyContext
|
6
|
+
from zrb.context.context import Context
|
7
|
+
from zrb.env.any_env import AnyEnv
|
8
|
+
from zrb.input.any_input import AnyInput
|
9
|
+
from zrb.task.any_task import AnyTask
|
10
|
+
from zrb.task.base_task import BaseTask
|
11
|
+
from zrb.util.attr import get_str_attr
|
14
12
|
|
15
13
|
|
16
14
|
class HttpCheck(BaseTask):
|
@@ -24,7 +22,7 @@ class HttpCheck(BaseTask):
|
|
24
22
|
input: list[AnyInput] | AnyInput | None = None,
|
25
23
|
env: list[AnyEnv] | AnyEnv | None = None,
|
26
24
|
url: StrAttr = "http://localhost",
|
27
|
-
|
25
|
+
render_url: bool = True,
|
28
26
|
http_method: StrAttr = "GET",
|
29
27
|
interval: int = 5,
|
30
28
|
execute_condition: bool | str | Callable[[Context], bool] = True,
|
@@ -45,19 +43,21 @@ class HttpCheck(BaseTask):
|
|
45
43
|
fallback=fallback,
|
46
44
|
)
|
47
45
|
self._url = url
|
48
|
-
self.
|
46
|
+
self._render_url = render_url
|
49
47
|
self._http_method = http_method
|
50
48
|
self._interval = interval
|
51
49
|
|
52
50
|
def _get_url(self, ctx: AnyContext) -> str:
|
53
51
|
return get_str_attr(
|
54
|
-
ctx, self._url, "http://localhost", auto_render=self.
|
52
|
+
ctx, self._url, "http://localhost", auto_render=self._render_url
|
55
53
|
)
|
56
54
|
|
57
55
|
def _get_http_method(self, ctx: AnyContext) -> str:
|
58
56
|
return get_str_attr(ctx, self._http_method, "GET", auto_render=True).upper()
|
59
57
|
|
60
58
|
async def _exec_action(self, ctx: AnyContext) -> bool:
|
59
|
+
import requests
|
60
|
+
|
61
61
|
url = self._get_url(ctx)
|
62
62
|
http_method = self._get_http_method(ctx)
|
63
63
|
while True:
|
zrb/task/llm_task.py
ADDED
@@ -0,0 +1,215 @@
|
|
1
|
+
import json
|
2
|
+
import os
|
3
|
+
from collections.abc import Callable
|
4
|
+
from typing import Any
|
5
|
+
|
6
|
+
from pydantic import BaseModel
|
7
|
+
|
8
|
+
from zrb.attr.type import StrAttr
|
9
|
+
from zrb.config import LLM_MODEL, LLM_SYSTEM_PROMPT
|
10
|
+
from zrb.context.any_context import AnyContext
|
11
|
+
from zrb.context.any_shared_context import AnySharedContext
|
12
|
+
from zrb.env.any_env import AnyEnv
|
13
|
+
from zrb.input.any_input import AnyInput
|
14
|
+
from zrb.task.any_task import AnyTask
|
15
|
+
from zrb.task.base_task import BaseTask
|
16
|
+
from zrb.util.attr import get_str_attr
|
17
|
+
from zrb.util.cli.style import stylize_faint
|
18
|
+
from zrb.util.llm.tool import callable_to_tool_schema
|
19
|
+
|
20
|
+
ListOfDict = list[dict[str, Any]]
|
21
|
+
|
22
|
+
|
23
|
+
class AdditionalTool(BaseModel):
|
24
|
+
fn: Callable
|
25
|
+
name: str | None
|
26
|
+
description: str | None
|
27
|
+
|
28
|
+
|
29
|
+
def scratchpad(thought: str) -> str:
|
30
|
+
"""Use this tool to note your thought and planning"""
|
31
|
+
return thought
|
32
|
+
|
33
|
+
|
34
|
+
class LLMTask(BaseTask):
|
35
|
+
|
36
|
+
def __init__(
|
37
|
+
self,
|
38
|
+
name: str,
|
39
|
+
color: int | None = None,
|
40
|
+
icon: str | None = None,
|
41
|
+
description: str | None = None,
|
42
|
+
cli_only: bool = False,
|
43
|
+
input: list[AnyInput | None] | AnyInput | None = None,
|
44
|
+
env: list[AnyEnv | None] | AnyEnv | None = None,
|
45
|
+
model: StrAttr | None = LLM_MODEL,
|
46
|
+
render_model: bool = True,
|
47
|
+
system_prompt: StrAttr | None = LLM_SYSTEM_PROMPT,
|
48
|
+
render_system_prompt: bool = True,
|
49
|
+
message: StrAttr | None = None,
|
50
|
+
tools: (
|
51
|
+
dict[str, Callable] | Callable[[AnySharedContext], dict[str, Callable]]
|
52
|
+
) = {},
|
53
|
+
history: ListOfDict | Callable[[AnySharedContext], ListOfDict] = [],
|
54
|
+
history_file: StrAttr | None = None,
|
55
|
+
render_history_file: bool = True,
|
56
|
+
model_kwargs: (
|
57
|
+
dict[str, Any] | Callable[[AnySharedContext], dict[str, Any]]
|
58
|
+
) = {},
|
59
|
+
execute_condition: bool | str | Callable[[AnySharedContext], bool] = True,
|
60
|
+
retries: int = 2,
|
61
|
+
retry_period: float = 0,
|
62
|
+
readiness_check: list[AnyTask] | AnyTask | None = None,
|
63
|
+
readiness_check_delay: float = 0.5,
|
64
|
+
readiness_check_period: float = 5,
|
65
|
+
readiness_failure_threshold: int = 1,
|
66
|
+
readiness_timeout: int = 60,
|
67
|
+
monitor_readiness: bool = False,
|
68
|
+
upstream: list[AnyTask] | AnyTask | None = None,
|
69
|
+
fallback: list[AnyTask] | AnyTask | None = None,
|
70
|
+
):
|
71
|
+
super().__init__(
|
72
|
+
name=name,
|
73
|
+
color=color,
|
74
|
+
icon=icon,
|
75
|
+
description=description,
|
76
|
+
cli_only=cli_only,
|
77
|
+
input=input,
|
78
|
+
env=env,
|
79
|
+
execute_condition=execute_condition,
|
80
|
+
retries=retries,
|
81
|
+
retry_period=retry_period,
|
82
|
+
readiness_check=readiness_check,
|
83
|
+
readiness_check_delay=readiness_check_delay,
|
84
|
+
readiness_check_period=readiness_check_period,
|
85
|
+
readiness_failure_threshold=readiness_failure_threshold,
|
86
|
+
readiness_timeout=readiness_timeout,
|
87
|
+
monitor_readiness=monitor_readiness,
|
88
|
+
upstream=upstream,
|
89
|
+
fallback=fallback,
|
90
|
+
)
|
91
|
+
self._model = model
|
92
|
+
self._render_model = render_model
|
93
|
+
self._model_kwargs = model_kwargs
|
94
|
+
self._system_prompt = system_prompt
|
95
|
+
self._render_system_prompt = render_system_prompt
|
96
|
+
self._message = message
|
97
|
+
self._tools = tools
|
98
|
+
self._history = history
|
99
|
+
self._history_file = history_file
|
100
|
+
self._render_history_file = render_history_file
|
101
|
+
self._additional_tools: list[AdditionalTool] = []
|
102
|
+
|
103
|
+
def add_tool(
|
104
|
+
self, tool: Callable, name: str | None = None, description: str | None = None
|
105
|
+
):
|
106
|
+
self._additional_tools.append(
|
107
|
+
AdditionalTool(fn=tool, name=name, description=description)
|
108
|
+
)
|
109
|
+
|
110
|
+
async def _exec_action(self, ctx: AnyContext) -> Any:
|
111
|
+
from litellm import acompletion
|
112
|
+
|
113
|
+
model_kwargs = self._get_model_kwargs(ctx)
|
114
|
+
ctx.log_debug("MODEL KWARGS", model_kwargs)
|
115
|
+
system_prompt = self._get_system_prompt(ctx)
|
116
|
+
ctx.log_debug("SYSTEM PROMPT", system_prompt)
|
117
|
+
history = self._get_history(ctx)
|
118
|
+
ctx.log_debug("HISTORY PROMPT", history)
|
119
|
+
user_message = {"role": "user", "content": self._get_message(ctx)}
|
120
|
+
ctx.print(stylize_faint(f"{user_message}"))
|
121
|
+
messages = history + [user_message]
|
122
|
+
available_tools = self._get_tools(ctx)
|
123
|
+
available_tools["scratchpad"] = scratchpad
|
124
|
+
tool_schema = [
|
125
|
+
callable_to_tool_schema(tool, name)
|
126
|
+
for name, tool in available_tools.items()
|
127
|
+
]
|
128
|
+
for additional_tool in self._additional_tools:
|
129
|
+
fn = additional_tool.fn
|
130
|
+
tool_name = additional_tool.name or fn.__name__
|
131
|
+
tool_description = additional_tool.description
|
132
|
+
available_tools[tool_name] = additional_tool.fn
|
133
|
+
tool_schema.append(
|
134
|
+
callable_to_tool_schema(
|
135
|
+
fn, name=tool_name, description=tool_description
|
136
|
+
)
|
137
|
+
)
|
138
|
+
ctx.log_debug("TOOL SCHEMA", tool_schema)
|
139
|
+
history_file = self._get_history_file(ctx)
|
140
|
+
while True:
|
141
|
+
response = await acompletion(
|
142
|
+
model=self._get_model(ctx),
|
143
|
+
messages=[{"role": "system", "content": system_prompt}] + messages,
|
144
|
+
tools=tool_schema,
|
145
|
+
**model_kwargs,
|
146
|
+
)
|
147
|
+
response_message = response.choices[0].message
|
148
|
+
ctx.print(stylize_faint(f"{response_message.to_dict()}"))
|
149
|
+
messages.append(response_message.to_dict())
|
150
|
+
tool_calls = response_message.tool_calls
|
151
|
+
if tool_calls:
|
152
|
+
# noqa Reference: https://docs.litellm.ai/docs/completion/function_call#full-code---parallel-function-calling-with-gpt-35-turbo-1106
|
153
|
+
for tool_call in tool_calls:
|
154
|
+
function_name = tool_call.function.name
|
155
|
+
function_to_call = available_tools[function_name]
|
156
|
+
function_kwargs = json.loads(tool_call.function.arguments)
|
157
|
+
function_response = function_to_call(**function_kwargs)
|
158
|
+
tool_call_message = {
|
159
|
+
"tool_call_id": tool_call.id,
|
160
|
+
"role": "tool",
|
161
|
+
"name": function_name,
|
162
|
+
"content": function_response,
|
163
|
+
}
|
164
|
+
ctx.print(stylize_faint(f"{tool_call_message}"))
|
165
|
+
messages.append(tool_call_message)
|
166
|
+
continue
|
167
|
+
if history_file != "":
|
168
|
+
os.makedirs(os.path.dirname(history_file), exist_ok=True)
|
169
|
+
with open(history_file, "w") as f:
|
170
|
+
f.write(json.dumps(messages, indent=2))
|
171
|
+
return response_message.content
|
172
|
+
|
173
|
+
def _get_model(self, ctx: AnyContext) -> str:
|
174
|
+
return get_str_attr(
|
175
|
+
ctx, self._model, "ollama_chat/llama3.1", auto_render=self._render_model
|
176
|
+
)
|
177
|
+
|
178
|
+
def _get_system_prompt(self, ctx: AnyContext) -> str:
|
179
|
+
return get_str_attr(
|
180
|
+
ctx,
|
181
|
+
self._system_prompt,
|
182
|
+
"You are a helpful assistant",
|
183
|
+
auto_render=self._render_system_prompt,
|
184
|
+
)
|
185
|
+
|
186
|
+
def _get_message(self, ctx: AnyContext) -> str:
|
187
|
+
return get_str_attr(ctx, self._message, "How are you?", auto_render=True)
|
188
|
+
|
189
|
+
def _get_model_kwargs(self, ctx: AnyContext) -> dict[str, Callable]:
|
190
|
+
if callable(self._model_kwargs):
|
191
|
+
return self._model_kwargs(ctx)
|
192
|
+
return self._model_kwargs
|
193
|
+
|
194
|
+
def _get_tools(self, ctx: AnyContext) -> dict[str, Callable]:
|
195
|
+
if callable(self._tools):
|
196
|
+
return self._tools(ctx)
|
197
|
+
return self._tools
|
198
|
+
|
199
|
+
def _get_history(self, ctx: AnyContext) -> ListOfDict:
|
200
|
+
if callable(self._history):
|
201
|
+
return self._history(ctx)
|
202
|
+
history_file = self._get_history_file(ctx)
|
203
|
+
if (
|
204
|
+
len(self._history) == 0
|
205
|
+
and history_file != ""
|
206
|
+
and os.path.isfile(history_file)
|
207
|
+
):
|
208
|
+
with open(history_file, "r") as f:
|
209
|
+
return json.loads(f.read())
|
210
|
+
return self._history
|
211
|
+
|
212
|
+
def _get_history_file(self, ctx: AnyContext) -> str:
|
213
|
+
return get_str_attr(
|
214
|
+
ctx, self._history_file, "", auto_render=self._render_history_file
|
215
|
+
)
|
zrb/task/make_task.py
CHANGED
@@ -1,13 +1,13 @@
|
|
1
1
|
from collections.abc import Callable
|
2
2
|
from typing import Any
|
3
3
|
|
4
|
-
from
|
5
|
-
from
|
6
|
-
from
|
7
|
-
from
|
8
|
-
from
|
9
|
-
from .any_task import AnyTask
|
10
|
-
from .base_task import BaseTask
|
4
|
+
from zrb.context.any_context import AnyContext
|
5
|
+
from zrb.context.any_shared_context import AnySharedContext
|
6
|
+
from zrb.env.any_env import AnyEnv
|
7
|
+
from zrb.group.any_group import AnyGroup
|
8
|
+
from zrb.input.any_input import AnyInput
|
9
|
+
from zrb.task.any_task import AnyTask
|
10
|
+
from zrb.task.base_task import BaseTask
|
11
11
|
|
12
12
|
|
13
13
|
def make_task(
|
@@ -16,8 +16,8 @@ def make_task(
|
|
16
16
|
icon: str | None = None,
|
17
17
|
description: str | None = None,
|
18
18
|
cli_only: bool = False,
|
19
|
-
input: list[AnyInput] | AnyInput | None = None,
|
20
|
-
env: list[AnyEnv] | AnyEnv | None = None,
|
19
|
+
input: list[AnyInput | None] | AnyInput | None = None,
|
20
|
+
env: list[AnyEnv | None] | AnyEnv | None = None,
|
21
21
|
execute_condition: bool | str | Callable[[AnySharedContext], bool] = True,
|
22
22
|
retries: int = 2,
|
23
23
|
retry_period: float = 0,
|
zrb/task/rsync_task.py
CHANGED
@@ -1,12 +1,12 @@
|
|
1
1
|
from collections.abc import Callable
|
2
2
|
|
3
|
-
from
|
4
|
-
from
|
5
|
-
from
|
6
|
-
from
|
7
|
-
from
|
8
|
-
from .
|
9
|
-
from .
|
3
|
+
from zrb.attr.type import IntAttr, StrAttr
|
4
|
+
from zrb.context.any_context import AnyContext
|
5
|
+
from zrb.env.any_env import AnyEnv
|
6
|
+
from zrb.input.any_input import AnyInput
|
7
|
+
from zrb.task.any_task import AnyTask
|
8
|
+
from zrb.task.cmd_task import CmdTask
|
9
|
+
from zrb.util.attr import get_str_attr
|
10
10
|
|
11
11
|
|
12
12
|
class RsyncTask(CmdTask):
|
@@ -32,13 +32,13 @@ class RsyncTask(CmdTask):
|
|
32
32
|
remote_ssh_key: StrAttr | None = None,
|
33
33
|
auto_render_remote_ssh_key: bool = True,
|
34
34
|
remote_source_path: StrAttr | None = None,
|
35
|
-
|
35
|
+
render_remote_source_path: bool = True,
|
36
36
|
remote_destination_path: StrAttr | None = None,
|
37
|
-
|
37
|
+
render_remote_destination_path: bool = True,
|
38
38
|
local_source_path: StrAttr | None = None,
|
39
|
-
|
39
|
+
render_local_source_path: bool = True,
|
40
40
|
local_destination_path: StrAttr | None = None,
|
41
|
-
|
41
|
+
render_local_destination_path: bool = True,
|
42
42
|
cwd: str | None = None,
|
43
43
|
auto_render_cwd: bool = True,
|
44
44
|
max_output_line: int = 1000,
|
@@ -59,19 +59,19 @@ class RsyncTask(CmdTask):
|
|
59
59
|
input=input,
|
60
60
|
env=env,
|
61
61
|
shell=shell,
|
62
|
-
|
62
|
+
render_shell=auto_render_shell,
|
63
63
|
remote_host=remote_host,
|
64
|
-
|
64
|
+
render_remote_host=auto_render_remote_host,
|
65
65
|
remote_port=remote_port,
|
66
66
|
auto_render_remote_port=auto_render_remote_port,
|
67
67
|
remote_user=remote_user,
|
68
|
-
|
68
|
+
render_remote_user=auto_render_remote_user,
|
69
69
|
remote_password=remote_password,
|
70
|
-
|
70
|
+
render_remote_password=auto_render_remote_password,
|
71
71
|
remote_ssh_key=remote_ssh_key,
|
72
|
-
|
72
|
+
render_remote_ssh_key=auto_render_remote_ssh_key,
|
73
73
|
cwd=cwd,
|
74
|
-
|
74
|
+
render_cwd=auto_render_cwd,
|
75
75
|
max_output_line=max_output_line,
|
76
76
|
max_error_line=max_error_line,
|
77
77
|
execute_condition=execute_condition,
|
@@ -82,13 +82,13 @@ class RsyncTask(CmdTask):
|
|
82
82
|
fallback=fallback,
|
83
83
|
)
|
84
84
|
self._remote_source_path = remote_source_path
|
85
|
-
self.
|
85
|
+
self._render_remote_source_path = render_remote_source_path
|
86
86
|
self._remote_destination_path = remote_destination_path
|
87
|
-
self.
|
87
|
+
self._render_remote_destination_path = render_remote_destination_path
|
88
88
|
self._local_source_path = local_source_path
|
89
|
-
self.
|
89
|
+
self._render_local_source_path = render_local_source_path
|
90
90
|
self._local_destination_path = local_destination_path
|
91
|
-
self.
|
91
|
+
self._render_local_destination_path = render_local_destination_path
|
92
92
|
|
93
93
|
def _get_source_path(self, ctx: AnyContext) -> str:
|
94
94
|
local_source_path = self._get_local_source_path(ctx)
|
@@ -113,7 +113,7 @@ class RsyncTask(CmdTask):
|
|
113
113
|
ctx,
|
114
114
|
self._remote_source_path,
|
115
115
|
"",
|
116
|
-
auto_render=self.
|
116
|
+
auto_render=self._render_remote_source_path,
|
117
117
|
)
|
118
118
|
|
119
119
|
def _get_remote_destination_path(self, ctx: AnyContext) -> str:
|
@@ -121,7 +121,7 @@ class RsyncTask(CmdTask):
|
|
121
121
|
ctx,
|
122
122
|
self._remote_destination_path,
|
123
123
|
"",
|
124
|
-
auto_render=self.
|
124
|
+
auto_render=self._render_remote_destination_path,
|
125
125
|
)
|
126
126
|
|
127
127
|
def _get_local_source_path(self, ctx: AnyContext) -> str:
|
@@ -129,7 +129,7 @@ class RsyncTask(CmdTask):
|
|
129
129
|
ctx,
|
130
130
|
self._local_source_path,
|
131
131
|
"",
|
132
|
-
auto_render=self.
|
132
|
+
auto_render=self._render_local_source_path,
|
133
133
|
)
|
134
134
|
|
135
135
|
def _get_local_destination_path(self, ctx: AnyContext) -> str:
|
@@ -137,7 +137,7 @@ class RsyncTask(CmdTask):
|
|
137
137
|
ctx,
|
138
138
|
self._local_destination_path,
|
139
139
|
"",
|
140
|
-
auto_render=self.
|
140
|
+
auto_render=self._render_local_destination_path,
|
141
141
|
)
|
142
142
|
|
143
143
|
def _get_cmd_script(self, ctx: AnyContext) -> str:
|