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.
Files changed (174) hide show
  1. zrb/__init__.py +49 -40
  2. zrb/__main__.py +5 -3
  3. zrb/attr/type.py +2 -1
  4. zrb/builtin/__init__.py +42 -2
  5. zrb/builtin/base64.py +34 -0
  6. zrb/builtin/git.py +156 -0
  7. zrb/builtin/git_subtree.py +88 -0
  8. zrb/builtin/group.py +34 -0
  9. zrb/builtin/llm/llm_chat.py +47 -0
  10. zrb/builtin/llm/tool/cli.py +9 -0
  11. zrb/builtin/llm/tool/rag.py +189 -0
  12. zrb/builtin/llm/tool/web.py +74 -0
  13. zrb/builtin/md5.py +36 -0
  14. zrb/builtin/project/add/fastapp.py +72 -0
  15. zrb/builtin/project/add/fastapp_template/.gitignore +4 -0
  16. zrb/builtin/project/add/fastapp_template/README.md +7 -0
  17. zrb/builtin/project/add/fastapp_template/_zrb/config.py +17 -0
  18. zrb/builtin/project/add/fastapp_template/_zrb/group.py +16 -0
  19. zrb/builtin/project/add/fastapp_template/_zrb/helper.py +97 -0
  20. zrb/builtin/project/add/fastapp_template/_zrb/main.py +132 -0
  21. zrb/builtin/project/add/fastapp_template/_zrb/venv_task.py +22 -0
  22. zrb/builtin/project/add/fastapp_template/common/app.py +18 -0
  23. zrb/builtin/project/add/fastapp_template/common/db_engine.py +5 -0
  24. zrb/builtin/project/add/fastapp_template/common/db_repository.py +134 -0
  25. zrb/builtin/project/add/fastapp_template/common/error.py +8 -0
  26. zrb/builtin/project/add/fastapp_template/common/schema.py +5 -0
  27. zrb/builtin/project/add/fastapp_template/common/usecase.py +232 -0
  28. zrb/builtin/project/add/fastapp_template/config.py +29 -0
  29. zrb/builtin/project/add/fastapp_template/main.py +7 -0
  30. zrb/builtin/project/add/fastapp_template/migrate.py +3 -0
  31. zrb/builtin/project/add/fastapp_template/module/__init__.py +0 -0
  32. zrb/builtin/project/add/fastapp_template/module/auth/alembic.ini +117 -0
  33. zrb/builtin/project/add/fastapp_template/module/auth/client/api_client.py +7 -0
  34. zrb/builtin/project/add/fastapp_template/module/auth/client/base_client.py +27 -0
  35. zrb/builtin/project/add/fastapp_template/module/auth/client/direct_client.py +6 -0
  36. zrb/builtin/project/add/fastapp_template/module/auth/client/factory.py +9 -0
  37. zrb/builtin/project/add/fastapp_template/module/auth/migration/README +1 -0
  38. zrb/builtin/project/add/fastapp_template/module/auth/migration/env.py +108 -0
  39. zrb/builtin/project/add/fastapp_template/module/auth/migration/script.py.mako +26 -0
  40. zrb/builtin/project/add/fastapp_template/module/auth/migration/versions/3093c7336477_add_user_table.py +37 -0
  41. zrb/builtin/project/add/fastapp_template/module/auth/migration_metadata.py +6 -0
  42. zrb/builtin/project/add/fastapp_template/module/auth/route.py +22 -0
  43. zrb/builtin/project/add/fastapp_template/module/auth/service/__init__.py +0 -0
  44. zrb/builtin/project/add/fastapp_template/module/auth/service/user/__init__.py +0 -0
  45. zrb/builtin/project/add/fastapp_template/module/auth/service/user/repository/__init__.py +0 -0
  46. zrb/builtin/project/add/fastapp_template/module/auth/service/user/repository/db_repository.py +39 -0
  47. zrb/builtin/project/add/fastapp_template/module/auth/service/user/repository/factory.py +13 -0
  48. zrb/builtin/project/add/fastapp_template/module/auth/service/user/repository/repository.py +34 -0
  49. zrb/builtin/project/add/fastapp_template/module/auth/service/user/usecase.py +45 -0
  50. zrb/builtin/project/add/fastapp_template/module/gateway/alembic.ini +117 -0
  51. zrb/builtin/project/add/fastapp_template/module/gateway/migration/README +1 -0
  52. zrb/builtin/project/add/fastapp_template/module/gateway/migration/env.py +108 -0
  53. zrb/builtin/project/add/fastapp_template/module/gateway/migration/script.py.mako +26 -0
  54. zrb/builtin/project/add/fastapp_template/module/gateway/migration/versions/.gitkeep +0 -0
  55. zrb/builtin/project/add/fastapp_template/module/gateway/migration_metadata.py +3 -0
  56. zrb/builtin/project/add/fastapp_template/module/gateway/route.py +27 -0
  57. zrb/builtin/project/add/fastapp_template/requirements.txt +6 -0
  58. zrb/builtin/project/add/fastapp_template/schema/__init__.py +0 -0
  59. zrb/builtin/project/add/fastapp_template/schema/role.py +31 -0
  60. zrb/builtin/project/add/fastapp_template/schema/user.py +31 -0
  61. zrb/builtin/project/add/fastapp_template/template.env +2 -0
  62. zrb/builtin/project/create/__init__.py +0 -0
  63. zrb/builtin/project/create/create.py +41 -0
  64. zrb/builtin/project/create/project-template/README.md +3 -0
  65. zrb/builtin/project/create/project-template/zrb_init.py +7 -0
  66. zrb/builtin/python.py +11 -0
  67. zrb/builtin/shell/__init__.py +0 -5
  68. zrb/builtin/shell/autocomplete/__init__.py +0 -9
  69. zrb/builtin/shell/autocomplete/bash.py +5 -6
  70. zrb/builtin/shell/autocomplete/subcmd.py +7 -8
  71. zrb/builtin/shell/autocomplete/zsh.py +5 -6
  72. zrb/builtin/todo.py +219 -0
  73. zrb/callback/any_callback.py +1 -1
  74. zrb/callback/callback.py +5 -5
  75. zrb/cmd/cmd_val.py +2 -2
  76. zrb/config.py +16 -3
  77. zrb/content_transformer/any_content_transformer.py +1 -1
  78. zrb/content_transformer/content_transformer.py +2 -2
  79. zrb/context/any_context.py +1 -1
  80. zrb/context/any_shared_context.py +3 -3
  81. zrb/context/context.py +10 -8
  82. zrb/context/shared_context.py +9 -8
  83. zrb/env/__init__.py +0 -3
  84. zrb/env/any_env.py +1 -1
  85. zrb/env/env.py +3 -4
  86. zrb/env/env_file.py +4 -4
  87. zrb/env/env_map.py +2 -2
  88. zrb/group/__init__.py +0 -3
  89. zrb/group/any_group.py +3 -3
  90. zrb/group/group.py +7 -6
  91. zrb/input/any_input.py +1 -1
  92. zrb/input/base_input.py +4 -4
  93. zrb/input/bool_input.py +5 -5
  94. zrb/input/float_input.py +3 -3
  95. zrb/input/int_input.py +3 -3
  96. zrb/input/option_input.py +51 -0
  97. zrb/input/password_input.py +2 -2
  98. zrb/input/str_input.py +1 -1
  99. zrb/input/text_input.py +12 -10
  100. zrb/runner/cli.py +80 -45
  101. zrb/runner/web_app.py +150 -0
  102. zrb/runner/web_controller/__init__.py +0 -0
  103. zrb/runner/web_controller/group_info_ui/__init__.py +0 -0
  104. zrb/runner/{web_app → web_controller}/group_info_ui/controller.py +7 -8
  105. zrb/runner/{web_app → web_controller}/group_info_ui/view.html +2 -2
  106. zrb/runner/web_controller/home_page/__init__.py +0 -0
  107. zrb/runner/{web_app → web_controller}/home_page/controller.py +7 -6
  108. zrb/runner/{web_app → web_controller}/home_page/view.html +2 -2
  109. zrb/runner/web_controller/task_ui/__init__.py +0 -0
  110. zrb/runner/{web_app → web_controller}/task_ui/controller.py +8 -12
  111. zrb/runner/{web_app → web_controller}/task_ui/view.html +2 -2
  112. zrb/runner/web_util.py +5 -35
  113. zrb/session/any_session.py +13 -7
  114. zrb/session/session.py +78 -40
  115. zrb/session_state_log/session_state_log.py +7 -5
  116. zrb/session_state_logger/any_session_state_logger.py +1 -1
  117. zrb/session_state_logger/default_session_state_logger.py +2 -2
  118. zrb/session_state_logger/file_session_state_logger.py +19 -27
  119. zrb/task/any_task.py +4 -4
  120. zrb/task/base_task.py +33 -23
  121. zrb/task/base_trigger.py +11 -12
  122. zrb/task/cmd_task.py +72 -65
  123. zrb/task/http_check.py +13 -13
  124. zrb/task/llm_task.py +215 -0
  125. zrb/task/make_task.py +9 -9
  126. zrb/task/rsync_task.py +25 -25
  127. zrb/task/scaffolder.py +18 -15
  128. zrb/task/scheduler.py +6 -7
  129. zrb/task/task.py +1 -1
  130. zrb/task/tcp_check.py +11 -13
  131. zrb/util/attr.py +19 -3
  132. zrb/util/cli/style.py +71 -2
  133. zrb/util/cli/subcommand.py +2 -2
  134. zrb/util/codemod/__init__.py +0 -0
  135. zrb/util/codemod/add_code_to_class.py +35 -0
  136. zrb/util/codemod/add_code_to_function.py +36 -0
  137. zrb/util/codemod/add_code_to_method.py +55 -0
  138. zrb/util/codemod/add_key_to_dict.py +51 -0
  139. zrb/util/codemod/add_param_to_function_call.py +39 -0
  140. zrb/util/codemod/add_property_to_class.py +55 -0
  141. zrb/util/git.py +156 -0
  142. zrb/util/git_subtree.py +94 -0
  143. zrb/util/group.py +2 -2
  144. zrb/util/llm/tool.py +63 -0
  145. zrb/util/string/conversion.py +7 -0
  146. zrb/util/todo.py +259 -0
  147. {zrb-1.0.0a2.dist-info → zrb-1.0.0a4.dist-info}/METADATA +13 -5
  148. zrb-1.0.0a4.dist-info/RECORD +197 -0
  149. zrb/builtin/shell/_group.py +0 -9
  150. zrb/builtin/shell/autocomplete/_group.py +0 -6
  151. zrb/runner/web_app/any_request_handler.py +0 -24
  152. zrb/runner/web_server.py +0 -224
  153. zrb-1.0.0a2.dist-info/RECORD +0 -120
  154. /zrb/{runner/web_app → builtin/project}/__init__.py +0 -0
  155. /zrb/{runner/web_app/group_info_ui → builtin/project/add}/__init__.py +0 -0
  156. /zrb/{runner/web_app/home_page → builtin/project/add/fastapp_template}/__init__.py +0 -0
  157. /zrb/{runner/web_app/task_ui → builtin/project/add/fastapp_template/common}/__init__.py +0 -0
  158. /zrb/runner/{web_app → web_controller}/group_info_ui/partial/group_info.html +0 -0
  159. /zrb/runner/{web_app → web_controller}/group_info_ui/partial/group_li.html +0 -0
  160. /zrb/runner/{web_app → web_controller}/group_info_ui/partial/task_info.html +0 -0
  161. /zrb/runner/{web_app → web_controller}/group_info_ui/partial/task_li.html +0 -0
  162. /zrb/runner/{web_app → web_controller}/home_page/partial/group_info.html +0 -0
  163. /zrb/runner/{web_app → web_controller}/home_page/partial/group_li.html +0 -0
  164. /zrb/runner/{web_app → web_controller}/home_page/partial/task_info.html +0 -0
  165. /zrb/runner/{web_app → web_controller}/home_page/partial/task_li.html +0 -0
  166. /zrb/runner/{web_app → web_controller}/static/favicon-32x32.png +0 -0
  167. /zrb/runner/{web_app → web_controller}/static/pico.min.css +0 -0
  168. /zrb/runner/{web_app → web_controller}/task_ui/partial/common-util.js +0 -0
  169. /zrb/runner/{web_app → web_controller}/task_ui/partial/input.html +0 -0
  170. /zrb/runner/{web_app → web_controller}/task_ui/partial/main.js +0 -0
  171. /zrb/runner/{web_app → web_controller}/task_ui/partial/show-existing-session.js +0 -0
  172. /zrb/runner/{web_app → web_controller}/task_ui/partial/visualize-history.js +0 -0
  173. {zrb-1.0.0a2.dist-info → zrb-1.0.0a4.dist-info}/WHEEL +0 -0
  174. {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 ..attr.type import BoolAttr, IntAttr, StrAttr
6
- from ..cmd.cmd_result import CmdResult
7
- from ..cmd.cmd_val import AnyCmdVal, CmdVal, SingleCmdVal
8
- from ..config import DEFAULT_SHELL
9
- from ..context.any_context import AnyContext
10
- from ..env.any_env import AnyEnv
11
- from ..input.any_input import AnyInput
12
- from ..util.attr import get_int_attr, get_str_attr
13
- from ..util.cmd.remote import get_remote_cmd_script
14
- from .any_task import AnyTask
15
- from .base_task import BaseTask
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
- auto_render_shell: bool = True,
29
+ render_shell: bool = True,
30
30
  shell_flag: StrAttr | None = None,
31
- auto_render_shell_flag: bool = True,
31
+ render_shell_flag: bool = True,
32
32
  remote_host: StrAttr | None = None,
33
- auto_render_remote_host: bool = True,
33
+ render_remote_host: bool = True,
34
34
  remote_port: IntAttr | None = None,
35
35
  remote_user: StrAttr | None = None,
36
- auto_render_remote_user: bool = True,
36
+ render_remote_user: bool = True,
37
37
  remote_password: StrAttr | None = None,
38
- auto_render_remote_password: bool = True,
38
+ render_remote_password: bool = True,
39
39
  remote_ssh_key: StrAttr | None = None,
40
- auto_render_remote_ssh_key: bool = True,
40
+ render_remote_ssh_key: bool = True,
41
41
  cmd: CmdVal = "",
42
- auto_render_cmd: bool = True,
42
+ render_cmd: bool = True,
43
43
  cwd: str | None = None,
44
- auto_render_cwd: bool = True,
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._auto_render_shell = auto_render_shell
80
+ self._render_shell = render_shell
81
81
  self._shell_flag = shell_flag
82
- self._auto_render_shell_flag = auto_render_shell_flag
82
+ self._render_shell_flag = render_shell_flag
83
83
  self._remote_host = remote_host
84
- self._auto_render_remote_host = auto_render_remote_host
84
+ self._render_remote_host = render_remote_host
85
85
  self._remote_port = remote_port
86
86
  self._remote_user = remote_user
87
- self._auto_render_remote_user = auto_render_remote_user
87
+ self._render_remote_user = render_remote_user
88
88
  self._remote_password = remote_password
89
- self._auto_render_remote_password = auto_render_remote_password
89
+ self._render_remote_password = render_remote_password
90
90
  self._remote_ssh_key = remote_ssh_key
91
- self._auto_render_remote_ssh_key = auto_render_remote_ssh_key
91
+ self._render_remote_ssh_key = render_remote_ssh_key
92
92
  self._cmd = cmd
93
- self._auto_render_cmd = auto_render_cmd
93
+ self._render_cmd = render_cmd
94
94
  self._cwd = cwd
95
- self._auto_render_cwd = auto_render_cwd
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 = await asyncio.create_subprocess_exec(
119
- shell,
120
- shell_flag,
121
- cmd_script,
122
- cwd=cwd,
123
- stdin=sys.stdin if sys.stdin.isatty() else None,
124
- stdout=asyncio.subprocess.PIPE,
125
- stderr=asyncio.subprocess.PIPE,
126
- env=env_map,
127
- bufsize=0,
128
- )
129
- stdout_task = asyncio.create_task(
130
- self.__read_stream(cmd_process.stdout, ctx.print, self._max_output_line)
131
- )
132
- stderr_task = asyncio.create_task(
133
- self.__read_stream(cmd_process.stderr, ctx.print, self._max_error_line)
134
- )
135
- # Wait for process to complete and gather stdout/stderr
136
- return_code = await cmd_process.wait()
137
- stdout = await stdout_task
138
- stderr = await stderr_task
139
- # Check for errors
140
- if return_code != 0:
141
- ctx.log_error(f"Exit status: {return_code}")
142
- raise Exception(f"Process {self._name} exited ({return_code}): {stderr}")
143
- return CmdResult(stdout, stderr)
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._auto_render_shell
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._auto_render_shell_flag,
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._auto_render_remote_host
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._auto_render_remote_user
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._auto_render_remote_password,
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._auto_render_remote_ssh_key
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._auto_render_cmd:
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 requests
5
-
6
- from ..attr.type import StrAttr
7
- from ..context.any_context import AnyContext
8
- from ..context.context import Context
9
- from ..env.any_env import AnyEnv
10
- from ..input.any_input import AnyInput
11
- from ..util.attr import get_str_attr
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
- auto_render_url: bool = True,
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._auto_render_url = auto_render_url
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._auto_render_url
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 ..context.any_context import AnyContext
5
- from ..context.any_shared_context import AnySharedContext
6
- from ..env.any_env import AnyEnv
7
- from ..group.any_group import AnyGroup
8
- from ..input.any_input import AnyInput
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 ..attr.type import IntAttr, StrAttr
4
- from ..context.any_context import AnyContext
5
- from ..env.any_env import AnyEnv
6
- from ..input.any_input import AnyInput
7
- from ..util.attr import get_str_attr
8
- from .any_task import AnyTask
9
- from .cmd_task import CmdTask
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
- auto_render_remote_source_path: bool = True,
35
+ render_remote_source_path: bool = True,
36
36
  remote_destination_path: StrAttr | None = None,
37
- auto_render_remote_destination_path: bool = True,
37
+ render_remote_destination_path: bool = True,
38
38
  local_source_path: StrAttr | None = None,
39
- auto_render_local_source_path: bool = True,
39
+ render_local_source_path: bool = True,
40
40
  local_destination_path: StrAttr | None = None,
41
- auto_render_local_destination_path: bool = True,
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
- auto_render_shell=auto_render_shell,
62
+ render_shell=auto_render_shell,
63
63
  remote_host=remote_host,
64
- auto_render_remote_host=auto_render_remote_host,
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
- auto_render_remote_user=auto_render_remote_user,
68
+ render_remote_user=auto_render_remote_user,
69
69
  remote_password=remote_password,
70
- auto_render_remote_password=auto_render_remote_password,
70
+ render_remote_password=auto_render_remote_password,
71
71
  remote_ssh_key=remote_ssh_key,
72
- auto_render_remote_ssh_key=auto_render_remote_ssh_key,
72
+ render_remote_ssh_key=auto_render_remote_ssh_key,
73
73
  cwd=cwd,
74
- auto_render_cwd=auto_render_cwd,
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._auto_render_remote_source_path = auto_render_remote_source_path
85
+ self._render_remote_source_path = render_remote_source_path
86
86
  self._remote_destination_path = remote_destination_path
87
- self._auto_render_remote_destination_path = auto_render_remote_destination_path
87
+ self._render_remote_destination_path = render_remote_destination_path
88
88
  self._local_source_path = local_source_path
89
- self._auto_render_local_source_path = auto_render_local_source_path
89
+ self._render_local_source_path = render_local_source_path
90
90
  self._local_destination_path = local_destination_path
91
- self._auto_render_local_destination_path = auto_render_local_destination_path
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._auto_render_remote_source_path,
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._auto_render_remote_destination_path,
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._auto_render_local_source_path,
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._auto_render_local_destination_path,
140
+ auto_render=self._render_local_destination_path,
141
141
  )
142
142
 
143
143
  def _get_cmd_script(self, ctx: AnyContext) -> str: