zrb 1.2.0__py3-none-any.whl → 1.2.2__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/builtin/llm/llm_chat.py +26 -3
- zrb/builtin/llm/tool/api.py +4 -2
- zrb/builtin/llm/tool/file.py +39 -0
- zrb/builtin/llm/tool/rag.py +37 -22
- zrb/builtin/llm/tool/web.py +46 -20
- zrb/builtin/project/add/fastapp/fastapp_task.py +1 -1
- zrb/builtin/project/add/fastapp/fastapp_template/my_app_name/_zrb/column/add_column_task.py +40 -40
- zrb/builtin/project/add/fastapp/fastapp_template/my_app_name/_zrb/column/add_column_util.py +6 -10
- zrb/builtin/project/add/fastapp/fastapp_template/my_app_name/_zrb/entity/add_entity_util.py +76 -56
- zrb/builtin/project/add/fastapp/fastapp_template/my_app_name/_zrb/module/add_module_util.py +16 -10
- zrb/builtin/project/add/fastapp/fastapp_template/my_app_name/common/logger_factory.py +1 -1
- zrb/builtin/project/add/fastapp/fastapp_util.py +6 -2
- zrb/config.py +3 -1
- zrb/content_transformer/content_transformer.py +7 -1
- zrb/context/context.py +8 -2
- zrb/input/text_input.py +9 -5
- zrb/task/llm_task.py +103 -16
- {zrb-1.2.0.dist-info → zrb-1.2.2.dist-info}/METADATA +3 -4
- {zrb-1.2.0.dist-info → zrb-1.2.2.dist-info}/RECORD +21 -20
- {zrb-1.2.0.dist-info → zrb-1.2.2.dist-info}/WHEEL +0 -0
- {zrb-1.2.0.dist-info → zrb-1.2.2.dist-info}/entry_points.txt +0 -0
@@ -17,8 +17,8 @@ from zrb.util.string.conversion import (
|
|
17
17
|
|
18
18
|
|
19
19
|
def is_gateway_navigation_config_file(ctx: AnyContext, file_path: str) -> bool:
|
20
|
-
return file_path == os.path.
|
21
|
-
APP_DIR, "module", "gateway", "config", "navigation.py"
|
20
|
+
return file_path == os.path.abspath(
|
21
|
+
os.path.join(APP_DIR, "module", "gateway", "config", "navigation.py")
|
22
22
|
)
|
23
23
|
|
24
24
|
|
@@ -76,106 +76,126 @@ def get_existing_auth_migration_xcom_key(ctx: AnyContext) -> str:
|
|
76
76
|
|
77
77
|
def is_in_app_schema_dir(ctx: AnyContext, file_path: str) -> bool:
|
78
78
|
return file_path.startswith(
|
79
|
-
os.path.
|
79
|
+
os.path.abspath(
|
80
|
+
os.path.join(APP_DIR, "schema", to_snake_case(ctx.input.entity))
|
81
|
+
)
|
80
82
|
)
|
81
83
|
|
82
84
|
|
83
85
|
def is_in_module_entity_test_dir(ctx: AnyContext, file_path: str) -> bool:
|
84
86
|
return file_path.startswith(
|
85
|
-
os.path.
|
86
|
-
|
87
|
-
|
88
|
-
|
89
|
-
|
87
|
+
os.path.abspath(
|
88
|
+
os.path.join(
|
89
|
+
APP_DIR,
|
90
|
+
"test",
|
91
|
+
to_snake_case(ctx.input.module),
|
92
|
+
to_snake_case(ctx.input.entity),
|
93
|
+
)
|
90
94
|
)
|
91
95
|
)
|
92
96
|
|
93
97
|
|
94
98
|
def is_in_module_entity_dir(ctx: AnyContext, file_path: str) -> bool:
|
95
99
|
return file_path.startswith(
|
96
|
-
os.path.
|
97
|
-
|
98
|
-
|
99
|
-
|
100
|
-
|
101
|
-
|
100
|
+
os.path.abspath(
|
101
|
+
os.path.join(
|
102
|
+
APP_DIR,
|
103
|
+
"module",
|
104
|
+
to_snake_case(ctx.input.module),
|
105
|
+
"service",
|
106
|
+
to_snake_case(ctx.input.entity),
|
107
|
+
)
|
102
108
|
)
|
103
109
|
)
|
104
110
|
|
105
111
|
|
106
112
|
def is_module_route_file(ctx: AnyContext, file_path: str) -> bool:
|
107
|
-
module_route_file = os.path.
|
108
|
-
|
109
|
-
|
110
|
-
|
111
|
-
|
113
|
+
module_route_file = os.path.abspath(
|
114
|
+
os.path.join(
|
115
|
+
APP_DIR,
|
116
|
+
"module",
|
117
|
+
to_snake_case(ctx.input.module),
|
118
|
+
"route.py",
|
119
|
+
)
|
112
120
|
)
|
113
121
|
return file_path == module_route_file
|
114
122
|
|
115
123
|
|
116
124
|
def is_module_migration_metadata_file(ctx: AnyContext, file_path: str) -> bool:
|
117
|
-
module_migration_metadata_file = os.path.
|
118
|
-
|
119
|
-
|
120
|
-
|
121
|
-
|
125
|
+
module_migration_metadata_file = os.path.abspath(
|
126
|
+
os.path.join(
|
127
|
+
APP_DIR,
|
128
|
+
"module",
|
129
|
+
to_snake_case(ctx.input.module),
|
130
|
+
"migration_metadata.py",
|
131
|
+
)
|
122
132
|
)
|
123
133
|
return file_path == module_migration_metadata_file
|
124
134
|
|
125
135
|
|
126
136
|
def is_module_client_file(ctx: AnyContext, file_path: str) -> bool:
|
127
|
-
module_any_client_file = os.path.
|
128
|
-
|
129
|
-
|
130
|
-
|
131
|
-
|
132
|
-
|
137
|
+
module_any_client_file = os.path.abspath(
|
138
|
+
os.path.join(
|
139
|
+
APP_DIR,
|
140
|
+
"module",
|
141
|
+
to_snake_case(ctx.input.module),
|
142
|
+
"client",
|
143
|
+
f"{to_snake_case(ctx.input.module)}_client.py",
|
144
|
+
)
|
133
145
|
)
|
134
146
|
return file_path == module_any_client_file
|
135
147
|
|
136
148
|
|
137
149
|
def is_module_api_client_file(ctx: AnyContext, file_path: str) -> bool:
|
138
|
-
module_api_client_file = os.path.
|
139
|
-
|
140
|
-
|
141
|
-
|
142
|
-
|
143
|
-
|
150
|
+
module_api_client_file = os.path.abspath(
|
151
|
+
os.path.join(
|
152
|
+
APP_DIR,
|
153
|
+
"module",
|
154
|
+
to_snake_case(ctx.input.module),
|
155
|
+
"client",
|
156
|
+
f"{to_snake_case(ctx.input.module)}_api_client.py",
|
157
|
+
)
|
144
158
|
)
|
145
159
|
return file_path == module_api_client_file
|
146
160
|
|
147
161
|
|
148
162
|
def is_module_direct_client_file(ctx: AnyContext, file_path: str) -> bool:
|
149
|
-
module_direct_client_file = os.path.
|
150
|
-
|
151
|
-
|
152
|
-
|
153
|
-
|
154
|
-
|
163
|
+
module_direct_client_file = os.path.abspath(
|
164
|
+
os.path.join(
|
165
|
+
APP_DIR,
|
166
|
+
"module",
|
167
|
+
to_snake_case(ctx.input.module),
|
168
|
+
"client",
|
169
|
+
f"{to_snake_case(ctx.input.module)}_direct_client.py",
|
170
|
+
)
|
155
171
|
)
|
156
172
|
return file_path == module_direct_client_file
|
157
173
|
|
158
174
|
|
159
175
|
def is_module_gateway_subroute_file(ctx: AnyContext, file_path: str) -> bool:
|
160
|
-
module_gateway_subroute_file = os.path.
|
161
|
-
|
162
|
-
|
163
|
-
|
164
|
-
|
165
|
-
|
176
|
+
module_gateway_subroute_file = os.path.abspath(
|
177
|
+
os.path.join(
|
178
|
+
APP_DIR,
|
179
|
+
"module",
|
180
|
+
"gateway",
|
181
|
+
"subroute",
|
182
|
+
f"{to_snake_case(ctx.input.module)}.py",
|
183
|
+
)
|
166
184
|
)
|
167
185
|
return file_path == module_gateway_subroute_file
|
168
186
|
|
169
187
|
|
170
188
|
def is_module_gateway_subroute_view_file(ctx: AnyContext, file_path: str) -> bool:
|
171
|
-
module_gateway_subroute_file = os.path.
|
172
|
-
|
173
|
-
|
174
|
-
|
175
|
-
|
176
|
-
|
177
|
-
|
178
|
-
|
189
|
+
module_gateway_subroute_file = os.path.abspath(
|
190
|
+
os.path.join(
|
191
|
+
APP_DIR,
|
192
|
+
"module",
|
193
|
+
"gateway",
|
194
|
+
"view",
|
195
|
+
"content",
|
196
|
+
f"{to_kebab_case(ctx.input.module)}",
|
197
|
+
f"{to_kebab_case(ctx.input.entity)}.html",
|
198
|
+
)
|
179
199
|
)
|
180
200
|
return file_path == module_gateway_subroute_file
|
181
201
|
|
@@ -16,41 +16,47 @@ from zrb.util.string.conversion import (
|
|
16
16
|
|
17
17
|
|
18
18
|
def is_app_config_file(ctx: AnyContext, file_path: str) -> bool:
|
19
|
-
return file_path == os.path.join(APP_DIR, "config.py")
|
19
|
+
return file_path == os.path.abspath(os.path.join(APP_DIR, "config.py"))
|
20
20
|
|
21
21
|
|
22
22
|
def is_app_main_file(ctx: AnyContext, file_path: str) -> bool:
|
23
|
-
return file_path == os.path.join(APP_DIR, "main.py")
|
23
|
+
return file_path == os.path.abspath(os.path.join(APP_DIR, "main.py"))
|
24
24
|
|
25
25
|
|
26
26
|
def is_gateway_route_file(ctx: AnyContext, file_path: str) -> bool:
|
27
|
-
return file_path == os.path.
|
27
|
+
return file_path == os.path.abspath(
|
28
|
+
os.path.join(APP_DIR, "module", "gateway", "route.py")
|
29
|
+
)
|
28
30
|
|
29
31
|
|
30
32
|
def is_gateway_navigation_config_file(ctx: AnyContext, file_path: str) -> bool:
|
31
|
-
return file_path == os.path.
|
32
|
-
APP_DIR, "module", "gateway", "config", "navigation.py"
|
33
|
+
return file_path == os.path.abspath(
|
34
|
+
os.path.join(APP_DIR, "module", "gateway", "config", "navigation.py")
|
33
35
|
)
|
34
36
|
|
35
37
|
|
36
38
|
def is_app_zrb_task_file(ctx: AnyContext, file_path: str) -> bool:
|
37
|
-
return file_path == os.path.join(APP_DIR, "_zrb", "task.py")
|
39
|
+
return file_path == os.path.abspath(os.path.join(APP_DIR, "_zrb", "task.py"))
|
38
40
|
|
39
41
|
|
40
42
|
def is_app_zrb_config_file(ctx: AnyContext, file_path: str) -> bool:
|
41
|
-
return file_path == os.path.join(APP_DIR, "_zrb", "config.py")
|
43
|
+
return file_path == os.path.abspath(os.path.join(APP_DIR, "_zrb", "config.py"))
|
42
44
|
|
43
45
|
|
44
46
|
def is_in_module_dir(ctx: AnyContext, file_path: str) -> bool:
|
45
47
|
return file_path.startswith(
|
46
|
-
os.path.
|
48
|
+
os.path.abspath(
|
49
|
+
os.path.join(APP_DIR, "module", to_snake_case(ctx.input.module))
|
50
|
+
)
|
47
51
|
)
|
48
52
|
|
49
53
|
|
50
54
|
def is_gateway_module_subroute_file(ctx: AnyContext, file_path: str) -> bool:
|
51
55
|
module_subroute_file_name = f"{to_snake_case(ctx.input.module)}.py"
|
52
|
-
return file_path == os.path.
|
53
|
-
|
56
|
+
return file_path == os.path.abspath(
|
57
|
+
os.path.join(
|
58
|
+
APP_DIR, "module", "gateway", "subroute", module_subroute_file_name
|
59
|
+
)
|
54
60
|
)
|
55
61
|
|
56
62
|
|
@@ -7,12 +7,16 @@ from zrb.util.string.conversion import double_quote, to_snake_case
|
|
7
7
|
|
8
8
|
def is_in_project_app_dir(ctx: AnyContext, file_path: str) -> bool:
|
9
9
|
return file_path.startswith(
|
10
|
-
os.path.
|
10
|
+
os.path.abspath(
|
11
|
+
os.path.join(ctx.input.project_dir, to_snake_case(ctx.input.app))
|
12
|
+
)
|
11
13
|
)
|
12
14
|
|
13
15
|
|
14
16
|
def is_project_zrb_init_file(ctx: AnyContext, file_path: str) -> bool:
|
15
|
-
return file_path == os.path.
|
17
|
+
return file_path == os.path.abspath(
|
18
|
+
os.path.join(ctx.input.project_dir, "zrb_init.py")
|
19
|
+
)
|
16
20
|
|
17
21
|
|
18
22
|
def update_project_zrb_init_file(ctx: AnyContext, zrb_init_path: str):
|
zrb/config.py
CHANGED
@@ -92,7 +92,8 @@ LLM_HISTORY_DIR = os.getenv(
|
|
92
92
|
LLM_HISTORY_FILE = os.getenv(
|
93
93
|
"ZRB_LLM_HISTORY_FILE", os.path.join(LLM_HISTORY_DIR, "history.json")
|
94
94
|
)
|
95
|
-
|
95
|
+
LLM_ALLOW_ACCESS_LOCAL_FILE = to_boolean(os.getenv("ZRB_LLM_ACCESS_LOCAL_FILE", "1"))
|
96
|
+
LLM_ALLOW_ACCESS_SHELL = to_boolean(os.getenv("ZRB_LLM_ACCESS_SHELL", "1"))
|
96
97
|
LLM_ALLOW_ACCESS_INTERNET = to_boolean(os.getenv("ZRB_LLM_ACCESS_INTERNET", "1"))
|
97
98
|
# noqa See: https://qdrant.github.io/fastembed/examples/Supported_Models/#supported-text-embedding-models
|
98
99
|
RAG_EMBEDDING_MODEL = os.getenv(
|
@@ -101,6 +102,7 @@ RAG_EMBEDDING_MODEL = os.getenv(
|
|
101
102
|
RAG_CHUNK_SIZE = int(os.getenv("ZRB_RAG_CHUNK_SIZE", "1024"))
|
102
103
|
RAG_OVERLAP = int(os.getenv("ZRB_RAG_OVERLAP", "128"))
|
103
104
|
RAG_MAX_RESULT_COUNT = int(os.getenv("ZRB_RAG_MAX_RESULT_COUNT", "5"))
|
105
|
+
SERP_API_KEY = os.getenv("SERP_API_KEY", "")
|
104
106
|
|
105
107
|
|
106
108
|
BANNER = f"""
|
@@ -1,4 +1,5 @@
|
|
1
1
|
import fnmatch
|
2
|
+
import os
|
2
3
|
import re
|
3
4
|
from collections.abc import Callable
|
4
5
|
|
@@ -40,7 +41,12 @@ class ContentTransformer(AnyContentTransformer):
|
|
40
41
|
return True
|
41
42
|
except re.error:
|
42
43
|
pass
|
43
|
-
|
44
|
+
if os.sep not in pattern and (
|
45
|
+
os.altsep is None or os.altsep not in pattern
|
46
|
+
):
|
47
|
+
# Pattern like "*.txt" – match only the basename.
|
48
|
+
return fnmatch.fnmatch(file_path, os.path.basename(file_path))
|
49
|
+
return fnmatch.fnmatch(file_path, file_path)
|
44
50
|
|
45
51
|
def transform_file(self, ctx: AnyContext, file_path: str):
|
46
52
|
if callable(self._transform_file):
|
zrb/context/context.py
CHANGED
@@ -88,7 +88,7 @@ class Context(AnyContext):
|
|
88
88
|
return template
|
89
89
|
return int(self.render(template))
|
90
90
|
|
91
|
-
def render_float(self, template: str) -> float:
|
91
|
+
def render_float(self, template: str | float) -> float:
|
92
92
|
if isinstance(template, float):
|
93
93
|
return template
|
94
94
|
return float(self.render(template))
|
@@ -102,9 +102,10 @@ class Context(AnyContext):
|
|
102
102
|
flush: bool = True,
|
103
103
|
plain: bool = False,
|
104
104
|
):
|
105
|
+
sep = " " if sep is None else sep
|
105
106
|
message = sep.join([f"{value}" for value in values])
|
106
107
|
if plain:
|
107
|
-
self.append_to_shared_log(remove_style(message))
|
108
|
+
# self.append_to_shared_log(remove_style(message))
|
108
109
|
print(message, sep=sep, end=end, file=file, flush=flush)
|
109
110
|
return
|
110
111
|
color = self._color
|
@@ -132,6 +133,7 @@ class Context(AnyContext):
|
|
132
133
|
flush: bool = True,
|
133
134
|
):
|
134
135
|
if self._shared_ctx.get_logging_level() <= logging.DEBUG:
|
136
|
+
sep = " " if sep is None else sep
|
135
137
|
message = sep.join([f"{value}" for value in values])
|
136
138
|
stylized_message = stylize_log(f"[DEBUG] {message}")
|
137
139
|
self.print(stylized_message, sep=sep, end=end, file=file, flush=flush)
|
@@ -145,6 +147,7 @@ class Context(AnyContext):
|
|
145
147
|
flush: bool = True,
|
146
148
|
):
|
147
149
|
if self._shared_ctx.get_logging_level() <= logging.INFO:
|
150
|
+
sep = " " if sep is None else sep
|
148
151
|
message = sep.join([f"{value}" for value in values])
|
149
152
|
stylized_message = stylize_log(f"[INFO] {message}")
|
150
153
|
self.print(stylized_message, sep=sep, end=end, file=file, flush=flush)
|
@@ -158,6 +161,7 @@ class Context(AnyContext):
|
|
158
161
|
flush: bool = True,
|
159
162
|
):
|
160
163
|
if self._shared_ctx.get_logging_level() <= logging.INFO:
|
164
|
+
sep = " " if sep is None else sep
|
161
165
|
message = sep.join([f"{value}" for value in values])
|
162
166
|
stylized_message = stylize_warning(f"[WARNING] {message}")
|
163
167
|
self.print(stylized_message, sep=sep, end=end, file=file, flush=flush)
|
@@ -171,6 +175,7 @@ class Context(AnyContext):
|
|
171
175
|
flush: bool = True,
|
172
176
|
):
|
173
177
|
if self._shared_ctx.get_logging_level() <= logging.ERROR:
|
178
|
+
sep = " " if sep is None else sep
|
174
179
|
message = sep.join([f"{value}" for value in values])
|
175
180
|
stylized_message = stylize_error(f"[ERROR] {message}")
|
176
181
|
self.print(stylized_message, sep=sep, end=end, file=file, flush=flush)
|
@@ -184,6 +189,7 @@ class Context(AnyContext):
|
|
184
189
|
flush: bool = True,
|
185
190
|
):
|
186
191
|
if self._shared_ctx.get_logging_level() <= logging.CRITICAL:
|
192
|
+
sep = " " if sep is None else sep
|
187
193
|
message = sep.join([f"{value}" for value in values])
|
188
194
|
stylized_message = stylize_error(f"[CRITICAL] {message}")
|
189
195
|
self.print(stylized_message, sep=sep, end=end, file=file, flush=flush)
|
zrb/input/text_input.py
CHANGED
@@ -69,16 +69,17 @@ class TextInput(BaseInput):
|
|
69
69
|
)
|
70
70
|
|
71
71
|
def _prompt_cli_str(self, shared_ctx: AnySharedContext) -> str:
|
72
|
-
prompt_message = (
|
73
|
-
|
72
|
+
prompt_message = super().prompt_message
|
73
|
+
comment_prompt_message = (
|
74
|
+
f"{self.comment_start}{prompt_message}{self.comment_end}"
|
74
75
|
)
|
75
|
-
|
76
|
+
comment_prompt_message_eol = f"{comment_prompt_message}\n"
|
76
77
|
default_value = self.get_default_str(shared_ctx)
|
77
78
|
with tempfile.NamedTemporaryFile(
|
78
79
|
delete=False, suffix=self._extension
|
79
80
|
) as temp_file:
|
80
81
|
temp_file_name = temp_file.name
|
81
|
-
temp_file.write(
|
82
|
+
temp_file.write(comment_prompt_message_eol.encode())
|
82
83
|
# Pre-fill with default content
|
83
84
|
if default_value:
|
84
85
|
temp_file.write(default_value.encode())
|
@@ -87,7 +88,10 @@ class TextInput(BaseInput):
|
|
87
88
|
subprocess.call([self._editor, temp_file_name])
|
88
89
|
# Read the edited content
|
89
90
|
edited_content = read_file(temp_file_name)
|
90
|
-
parts = [
|
91
|
+
parts = [
|
92
|
+
text.strip() for text in edited_content.split(comment_prompt_message, 1)
|
93
|
+
]
|
91
94
|
edited_content = "\n".join(parts).lstrip()
|
92
95
|
os.remove(temp_file_name)
|
96
|
+
print(f"{prompt_message}: {edited_content}")
|
93
97
|
return edited_content
|
zrb/task/llm_task.py
CHANGED
@@ -4,10 +4,20 @@ from collections.abc import Callable
|
|
4
4
|
from typing import Any
|
5
5
|
|
6
6
|
from pydantic_ai import Agent, Tool
|
7
|
-
from pydantic_ai.messages import
|
7
|
+
from pydantic_ai.messages import (
|
8
|
+
FinalResultEvent,
|
9
|
+
FunctionToolCallEvent,
|
10
|
+
FunctionToolResultEvent,
|
11
|
+
ModelMessagesTypeAdapter,
|
12
|
+
PartDeltaEvent,
|
13
|
+
PartStartEvent,
|
14
|
+
TextPartDelta,
|
15
|
+
ToolCallPartDelta,
|
16
|
+
)
|
17
|
+
from pydantic_ai.models import Model
|
8
18
|
from pydantic_ai.settings import ModelSettings
|
9
19
|
|
10
|
-
from zrb.attr.type import StrAttr
|
20
|
+
from zrb.attr.type import StrAttr, fstring
|
11
21
|
from zrb.config import LLM_MODEL, LLM_SYSTEM_PROMPT
|
12
22
|
from zrb.context.any_context import AnyContext
|
13
23
|
from zrb.context.any_shared_context import AnySharedContext
|
@@ -15,7 +25,7 @@ from zrb.env.any_env import AnyEnv
|
|
15
25
|
from zrb.input.any_input import AnyInput
|
16
26
|
from zrb.task.any_task import AnyTask
|
17
27
|
from zrb.task.base_task import BaseTask
|
18
|
-
from zrb.util.attr import get_str_attr
|
28
|
+
from zrb.util.attr import get_attr, get_str_attr
|
19
29
|
from zrb.util.cli.style import stylize_faint
|
20
30
|
from zrb.util.file import read_file, write_file
|
21
31
|
from zrb.util.run import run_async
|
@@ -34,7 +44,9 @@ class LLMTask(BaseTask):
|
|
34
44
|
cli_only: bool = False,
|
35
45
|
input: list[AnyInput | None] | AnyInput | None = None,
|
36
46
|
env: list[AnyEnv | None] | AnyEnv | None = None,
|
37
|
-
model:
|
47
|
+
model: (
|
48
|
+
Callable[[AnySharedContext], Model | str | fstring] | Model | None
|
49
|
+
) = LLM_MODEL,
|
38
50
|
model_settings: (
|
39
51
|
ModelSettings | Callable[[AnySharedContext], ModelSettings] | None
|
40
52
|
) = None,
|
@@ -93,7 +105,7 @@ class LLMTask(BaseTask):
|
|
93
105
|
successor=successor,
|
94
106
|
)
|
95
107
|
self._model = model
|
96
|
-
self._model_settings =
|
108
|
+
self._model_settings = model_settings
|
97
109
|
self._agent = agent
|
98
110
|
self._render_model = render_model
|
99
111
|
self._system_prompt = system_prompt
|
@@ -108,6 +120,9 @@ class LLMTask(BaseTask):
|
|
108
120
|
self._render_history_file = render_history_file
|
109
121
|
self._max_call_iteration = max_call_iteration
|
110
122
|
|
123
|
+
def set_model(self, model: Model | str):
|
124
|
+
self._model = model
|
125
|
+
|
111
126
|
def add_tool(self, tool: ToolOrCallable):
|
112
127
|
self._additional_tools.append(tool)
|
113
128
|
|
@@ -115,15 +130,85 @@ class LLMTask(BaseTask):
|
|
115
130
|
history = await self._read_conversation_history(ctx)
|
116
131
|
user_prompt = self._get_message(ctx)
|
117
132
|
agent = self._get_agent(ctx)
|
118
|
-
|
133
|
+
async with agent.iter(
|
119
134
|
user_prompt=user_prompt,
|
120
135
|
message_history=ModelMessagesTypeAdapter.validate_python(history),
|
121
|
-
)
|
122
|
-
|
123
|
-
|
124
|
-
|
136
|
+
) as agent_run:
|
137
|
+
async for node in agent_run:
|
138
|
+
# Each node represents a step in the agent's execution
|
139
|
+
await self._print_node(ctx, agent_run, node)
|
140
|
+
new_history = json.loads(agent_run.result.all_messages_json())
|
125
141
|
await self._write_conversation_history(ctx, new_history)
|
126
|
-
return result.data
|
142
|
+
return agent_run.result.data
|
143
|
+
|
144
|
+
async def _print_node(self, ctx: AnyContext, agent_run: Any, node: Any):
|
145
|
+
if Agent.is_user_prompt_node(node):
|
146
|
+
# A user prompt node => The user has provided input
|
147
|
+
ctx.print(stylize_faint(f">> UserPromptNode: {node.user_prompt}"))
|
148
|
+
elif Agent.is_model_request_node(node):
|
149
|
+
# A model request node => We can stream tokens from the model"s request
|
150
|
+
ctx.print(
|
151
|
+
stylize_faint(">> ModelRequestNode: streaming partial request tokens")
|
152
|
+
)
|
153
|
+
async with node.stream(agent_run.ctx) as request_stream:
|
154
|
+
is_streaming = False
|
155
|
+
async for event in request_stream:
|
156
|
+
if isinstance(event, PartStartEvent):
|
157
|
+
if is_streaming:
|
158
|
+
ctx.print("", plain=True)
|
159
|
+
ctx.print(
|
160
|
+
stylize_faint(
|
161
|
+
f"[Request] Starting part {event.index}: {event.part!r}"
|
162
|
+
),
|
163
|
+
)
|
164
|
+
is_streaming = False
|
165
|
+
elif isinstance(event, PartDeltaEvent):
|
166
|
+
if isinstance(event.delta, TextPartDelta):
|
167
|
+
ctx.print(
|
168
|
+
stylize_faint(f"{event.delta.content_delta}"),
|
169
|
+
end="",
|
170
|
+
plain=is_streaming,
|
171
|
+
)
|
172
|
+
elif isinstance(event.delta, ToolCallPartDelta):
|
173
|
+
ctx.print(
|
174
|
+
stylize_faint(f"{event.delta.args_delta}"),
|
175
|
+
end="",
|
176
|
+
plain=is_streaming,
|
177
|
+
)
|
178
|
+
is_streaming = True
|
179
|
+
elif isinstance(event, FinalResultEvent):
|
180
|
+
if is_streaming:
|
181
|
+
ctx.print("", plain=True)
|
182
|
+
ctx.print(
|
183
|
+
stylize_faint(f"[Result] tool_name={event.tool_name}"),
|
184
|
+
)
|
185
|
+
is_streaming = False
|
186
|
+
if is_streaming:
|
187
|
+
ctx.print("", plain=True)
|
188
|
+
elif Agent.is_call_tools_node(node):
|
189
|
+
# A handle-response node => The model returned some data, potentially calls a tool
|
190
|
+
ctx.print(
|
191
|
+
stylize_faint(
|
192
|
+
">> CallToolsNode: streaming partial response & tool usage"
|
193
|
+
)
|
194
|
+
)
|
195
|
+
async with node.stream(agent_run.ctx) as handle_stream:
|
196
|
+
async for event in handle_stream:
|
197
|
+
if isinstance(event, FunctionToolCallEvent):
|
198
|
+
ctx.print(
|
199
|
+
stylize_faint(
|
200
|
+
f"[Tools] The LLM calls tool={event.part.tool_name!r} with args={event.part.args} (tool_call_id={event.part.tool_call_id!r})" # noqa
|
201
|
+
)
|
202
|
+
)
|
203
|
+
elif isinstance(event, FunctionToolResultEvent):
|
204
|
+
ctx.print(
|
205
|
+
stylize_faint(
|
206
|
+
f"[Tools] Tool call {event.tool_call_id!r} returned => {event.result.content}" # noqa
|
207
|
+
)
|
208
|
+
)
|
209
|
+
elif Agent.is_end_node(node):
|
210
|
+
# Once an End node is reached, the agent run is complete
|
211
|
+
ctx.print(stylize_faint(f"{agent_run.result.data}"))
|
127
212
|
|
128
213
|
async def _write_conversation_history(
|
129
214
|
self, ctx: AnyContext, conversations: list[Any]
|
@@ -135,11 +220,9 @@ class LLMTask(BaseTask):
|
|
135
220
|
write_file(history_file, json.dumps(conversations, indent=2))
|
136
221
|
|
137
222
|
def _get_model_settings(self, ctx: AnyContext) -> ModelSettings | None:
|
138
|
-
if isinstance(self._model_settings, ModelSettings):
|
139
|
-
return self._model_settings
|
140
223
|
if callable(self._model_settings):
|
141
224
|
return self._model_settings(ctx)
|
142
|
-
return
|
225
|
+
return self._model_settings
|
143
226
|
|
144
227
|
def _get_agent(self, ctx: AnyContext) -> Agent:
|
145
228
|
if isinstance(self._agent, Agent):
|
@@ -158,12 +241,16 @@ class LLMTask(BaseTask):
|
|
158
241
|
self._get_model(ctx),
|
159
242
|
system_prompt=self._get_system_prompt(ctx),
|
160
243
|
tools=tools,
|
244
|
+
model_settings=self._get_model_settings(ctx),
|
161
245
|
)
|
162
246
|
|
163
|
-
def _get_model(self, ctx: AnyContext) -> str:
|
164
|
-
|
247
|
+
def _get_model(self, ctx: AnyContext) -> str | Model | None:
|
248
|
+
model = get_attr(
|
165
249
|
ctx, self._model, "ollama_chat/llama3.1", auto_render=self._render_model
|
166
250
|
)
|
251
|
+
if isinstance(model, (Model, str)) or model is None:
|
252
|
+
return model
|
253
|
+
raise ValueError("Invalid model")
|
167
254
|
|
168
255
|
def _get_system_prompt(self, ctx: AnyContext) -> str:
|
169
256
|
return get_str_attr(
|
@@ -1,6 +1,6 @@
|
|
1
1
|
Metadata-Version: 2.1
|
2
2
|
Name: zrb
|
3
|
-
Version: 1.2.
|
3
|
+
Version: 1.2.2
|
4
4
|
Summary: Your Automation Powerhouse
|
5
5
|
Home-page: https://github.com/state-alchemists/zrb
|
6
6
|
License: AGPL-3.0-or-later
|
@@ -24,7 +24,7 @@ Requires-Dist: isort (>=5.13.2,<5.14.0)
|
|
24
24
|
Requires-Dist: libcst (>=1.5.0,<2.0.0)
|
25
25
|
Requires-Dist: pdfplumber (>=0.11.4,<0.12.0) ; extra == "rag"
|
26
26
|
Requires-Dist: psutil (>=6.1.1,<7.0.0)
|
27
|
-
Requires-Dist: pydantic-ai (>=0.0.
|
27
|
+
Requires-Dist: pydantic-ai (>=0.0.31,<0.0.32)
|
28
28
|
Requires-Dist: python-dotenv (>=1.0.1,<2.0.0)
|
29
29
|
Requires-Dist: python-jose[cryptography] (>=3.4.0,<4.0.0)
|
30
30
|
Requires-Dist: requests (>=2.32.3,<3.0.0)
|
@@ -79,8 +79,7 @@ pip install --pre zrb
|
|
79
79
|
Alternatively, you can also use our installation script to install Zrb along with some prerequisites:
|
80
80
|
|
81
81
|
```bash
|
82
|
-
bash -c "$(curl -fsSL https://raw.githubusercontent.com/state-alchemists/zrb/
|
83
|
-
# bash -c "$(curl -fsSL https://raw.githubusercontent.com/state-alchemists/zrb/main/install.sh)"
|
82
|
+
bash -c "$(curl -fsSL https://raw.githubusercontent.com/state-alchemists/zrb/main/install.sh)"
|
84
83
|
```
|
85
84
|
|
86
85
|
# 🐞 Bug Report + Feature Request
|