waldiez 0.3.11__py3-none-any.whl → 0.4.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of waldiez might be problematic. Click here for more details.

Files changed (66) hide show
  1. waldiez/_version.py +1 -1
  2. waldiez/cli.py +1 -3
  3. waldiez/exporting/agent/agent_exporter.py +26 -15
  4. waldiez/exporting/agent/utils/__init__.py +2 -4
  5. waldiez/exporting/agent/utils/captain_agent.py +250 -0
  6. waldiez/exporting/agent/utils/swarm_agent.py +12 -7
  7. waldiez/exporting/base/utils/comments.py +1 -0
  8. waldiez/exporting/chats/utils/swarm.py +1 -1
  9. waldiez/exporting/flow/flow_exporter.py +5 -6
  10. waldiez/exporting/flow/utils/__init__.py +3 -6
  11. waldiez/exporting/flow/utils/def_main.py +5 -4
  12. waldiez/exporting/flow/utils/flow_content.py +38 -0
  13. waldiez/exporting/flow/utils/importing_utils.py +64 -29
  14. waldiez/exporting/skills/skills_exporter.py +13 -6
  15. waldiez/exporting/skills/utils.py +92 -6
  16. waldiez/models/__init__.py +6 -0
  17. waldiez/models/agents/__init__.py +14 -0
  18. waldiez/models/agents/agent/__init__.py +2 -1
  19. waldiez/models/agents/agent/agent.py +71 -11
  20. waldiez/models/agents/agent/agent_type.py +11 -0
  21. waldiez/models/agents/agents.py +11 -1
  22. waldiez/models/agents/captain_agent/__init__.py +15 -0
  23. waldiez/models/agents/captain_agent/captain_agent.py +45 -0
  24. waldiez/models/agents/captain_agent/captain_agent_data.py +62 -0
  25. waldiez/models/agents/captain_agent/captain_agent_lib_entry.py +38 -0
  26. waldiez/models/agents/extra_requirements.py +88 -0
  27. waldiez/models/agents/group_manager/speakers.py +3 -0
  28. waldiez/models/agents/rag_user/retrieve_config.py +3 -0
  29. waldiez/models/agents/reasoning/reasoning_agent_reason_config.py +1 -0
  30. waldiez/models/agents/swarm_agent/after_work.py +13 -11
  31. waldiez/models/agents/swarm_agent/on_condition.py +3 -2
  32. waldiez/models/agents/swarm_agent/on_condition_available.py +1 -0
  33. waldiez/models/agents/swarm_agent/swarm_agent_data.py +3 -3
  34. waldiez/models/agents/swarm_agent/update_system_message.py +1 -0
  35. waldiez/models/chat/chat_message.py +1 -0
  36. waldiez/models/chat/chat_summary.py +1 -0
  37. waldiez/models/common/__init__.py +4 -0
  38. waldiez/models/common/ag2_version.py +30 -0
  39. waldiez/models/common/base.py +1 -1
  40. waldiez/models/common/date_utils.py +2 -0
  41. waldiez/models/common/dict_utils.py +2 -0
  42. waldiez/models/common/method_utils.py +98 -0
  43. waldiez/models/flow/__init__.py +2 -0
  44. waldiez/models/flow/utils.py +61 -1
  45. waldiez/models/model/__init__.py +2 -0
  46. waldiez/models/model/extra_requirements.py +57 -0
  47. waldiez/models/model/model.py +5 -2
  48. waldiez/models/model/model_data.py +3 -1
  49. waldiez/models/skill/__init__.py +4 -0
  50. waldiez/models/skill/extra_requirements.py +39 -0
  51. waldiez/models/skill/skill.py +157 -13
  52. waldiez/models/skill/skill_data.py +14 -0
  53. waldiez/models/skill/skill_type.py +8 -0
  54. waldiez/models/waldiez.py +47 -76
  55. waldiez/runner.py +19 -7
  56. waldiez/running/environment.py +30 -1
  57. waldiez/running/running.py +0 -6
  58. waldiez/utils/pysqlite3_checker.py +18 -5
  59. {waldiez-0.3.11.dist-info → waldiez-0.4.0.dist-info}/METADATA +42 -30
  60. {waldiez-0.3.11.dist-info → waldiez-0.4.0.dist-info}/RECORD +64 -55
  61. waldiez/exporting/agent/utils/agent_class_name.py +0 -36
  62. waldiez/exporting/agent/utils/agent_imports.py +0 -55
  63. {waldiez-0.3.11.dist-info → waldiez-0.4.0.dist-info}/WHEEL +0 -0
  64. {waldiez-0.3.11.dist-info → waldiez-0.4.0.dist-info}/entry_points.txt +0 -0
  65. {waldiez-0.3.11.dist-info → waldiez-0.4.0.dist-info}/licenses/LICENSE +0 -0
  66. {waldiez-0.3.11.dist-info → waldiez-0.4.0.dist-info}/licenses/NOTICE.md +0 -0
waldiez/_version.py CHANGED
@@ -2,4 +2,4 @@
2
2
  # Copyright (c) 2024 - 2025 Waldiez and contributors.
3
3
  """Version information for Waldiez."""
4
4
 
5
- __version__ = "0.3.11"
5
+ __version__ = "0.4.0"
waldiez/cli.py CHANGED
@@ -86,10 +86,8 @@ def run(
86
86
  ),
87
87
  ) -> None:
88
88
  """Run a Waldiez flow."""
89
- # a swarm chat without a user agent
90
- # creates a new user (this has a default code execution with docker)
91
- # temp (until we handle/detect docker setup)
92
89
  os.environ["AUTOGEN_USE_DOCKER"] = "0"
90
+ os.environ["NEP50_DISABLE_WARNING"] = "1"
93
91
  output_path = _get_output_path(output, force)
94
92
  with file.open("r", encoding="utf-8") as _file:
95
93
  try:
@@ -6,7 +6,7 @@
6
6
  from pathlib import Path
7
7
  from typing import Callable, Dict, List, Optional, Tuple, Union
8
8
 
9
- from waldiez.models import WaldiezAgent, WaldiezChat
9
+ from waldiez.models import WaldiezAgent, WaldiezChat, WaldiezModel
10
10
 
11
11
  from ..base import (
12
12
  AgentPosition,
@@ -18,9 +18,8 @@ from ..base import (
18
18
  ImportPosition,
19
19
  )
20
20
  from .utils import (
21
- get_agent_class_name,
22
21
  get_agent_code_execution_config,
23
- get_agent_imports,
22
+ get_captain_agent_extras,
24
23
  get_group_manager_extras,
25
24
  get_is_termination_message,
26
25
  get_rag_user_extras,
@@ -36,9 +35,9 @@ class AgentExporter(BaseExporter, ExporterMixin):
36
35
  self,
37
36
  agent: WaldiezAgent,
38
37
  agent_names: Dict[str, str],
39
- model_names: Dict[str, str],
40
- skill_names: Dict[str, str],
38
+ models: Tuple[List[WaldiezModel], Dict[str, str]],
41
39
  chats: Tuple[List[WaldiezChat], Dict[str, str]],
40
+ skill_names: Dict[str, str],
42
41
  is_async: bool,
43
42
  group_chat_members: List[WaldiezAgent],
44
43
  for_notebook: bool,
@@ -53,12 +52,12 @@ class AgentExporter(BaseExporter, ExporterMixin):
53
52
  The agent to export.
54
53
  agent_names : Dict[str, str]
55
54
  The agent ids to names mapping.
56
- model_names : Dict[str, str]
57
- The model ids to names mapping.
58
- skill_names : Dict[str, str]
59
- The skill ids to names mapping.
55
+ models : Tuple[List[WaldiezModel], Dict[str, str]]
56
+ All the models and the model ids to names mapping.
60
57
  chats : Tuple[List[WaldiezChat], Dict[str, str]]
61
58
  All the chats and the chat ids to names mapping.
59
+ skill_names : Dict[str, str]
60
+ The skill ids to names mapping.
62
61
  is_async : bool
63
62
  Whether the whole flow is async.
64
63
  for_notebook : bool
@@ -72,14 +71,14 @@ class AgentExporter(BaseExporter, ExporterMixin):
72
71
  if output_dir is not None and not isinstance(output_dir, Path):
73
72
  output_dir = Path(output_dir)
74
73
  self.output_dir = output_dir
75
- self.model_names = model_names
74
+ self.models = models[0]
75
+ self.model_names = models[1]
76
76
  self.skill_names = skill_names
77
77
  self.arguments_resolver = arguments_resolver
78
78
  self.group_chat_members = group_chat_members
79
79
  self.chats = chats
80
80
  self.is_async = is_async
81
81
  self._agent_name = agent_names[agent.id]
82
- self._agent_class = get_agent_class_name(self.agent)
83
82
  # content, argument, import
84
83
  self._code_execution = get_agent_code_execution_config(
85
84
  agent=self.agent,
@@ -119,6 +118,13 @@ class AgentExporter(BaseExporter, ExporterMixin):
119
118
  agent=self.agent,
120
119
  serializer=self.serializer,
121
120
  )
121
+ self._captain = get_captain_agent_extras(
122
+ agent=self.agent,
123
+ agent_names=self.agent_names,
124
+ all_models=self.models,
125
+ serializer=self.serializer,
126
+ output_dir=self.output_dir,
127
+ )
122
128
 
123
129
  def get_imports(self) -> Optional[List[Tuple[str, ImportPosition]]]:
124
130
  """Get the imports.
@@ -130,7 +136,7 @@ class AgentExporter(BaseExporter, ExporterMixin):
130
136
  """
131
137
  position = ImportPosition.THIRD_PARTY
132
138
  # default imports based on the agent class.
133
- agent_imports = get_agent_imports(self._agent_class)
139
+ agent_imports = self.agent.ag2_imports
134
140
  # if code execution is enabled, update the imports.
135
141
  if self._code_execution[2]:
136
142
  agent_imports.add(self._code_execution[2])
@@ -220,7 +226,6 @@ class AgentExporter(BaseExporter, ExporterMixin):
220
226
  """
221
227
  agent = self.agent
222
228
  agent_name = self._agent_name
223
- agent_class = self._agent_class
224
229
  retrieve_arg = self._rag[1]
225
230
  group_chat_arg = self._group_chat[1]
226
231
  is_termination = self._termination[0]
@@ -231,8 +236,14 @@ class AgentExporter(BaseExporter, ExporterMixin):
231
236
  default_auto_reply = (
232
237
  f'"{self.string_escape(agent.data.agent_default_auto_reply)}"'
233
238
  )
234
- extras = f"{group_chat_arg}{retrieve_arg}{self._reasoning}"
235
- agent_str = f"""{agent_name} = {agent_class}(
239
+ extras = (
240
+ f"{group_chat_arg}{retrieve_arg}{self._reasoning}{self._captain}"
241
+ )
242
+ ag2_class = self.agent.ag2_class
243
+ if agent.agent_type == "swarm":
244
+ # SwarmAgent is deprecated.
245
+ ag2_class = "ConversableAgent"
246
+ agent_str = f"""{agent_name} = {ag2_class}(
236
247
  name="{agent_name}",
237
248
  description="{agent.description}"{system_message_arg},
238
249
  human_input_mode="{agent.data.human_input_mode}",
@@ -2,8 +2,7 @@
2
2
  # Copyright (c) 2024 - 2025 Waldiez and contributors.
3
3
  """Utility functions for generating agent related strings."""
4
4
 
5
- from .agent_class_name import get_agent_class_name
6
- from .agent_imports import get_agent_imports
5
+ from .captain_agent import get_captain_agent_extras
7
6
  from .code_execution import get_agent_code_execution_config
8
7
  from .group_manager import get_group_manager_extras
9
8
  from .rag_user import get_rag_user_extras
@@ -13,10 +12,9 @@ from .teachability import get_agent_teachability_string
13
12
  from .termination_message import get_is_termination_message
14
13
 
15
14
  __all__ = [
16
- "get_agent_class_name",
17
- "get_agent_imports",
18
15
  "get_agent_code_execution_config",
19
16
  "get_agent_teachability_string",
17
+ "get_captain_agent_extras",
20
18
  "get_group_manager_extras",
21
19
  "get_is_termination_message",
22
20
  "get_rag_user_extras",
@@ -0,0 +1,250 @@
1
+ # SPDX-License-Identifier: Apache-2.0.
2
+ # Copyright (c) 2024 - 2025 Waldiez and contributors.
3
+ """ "Extras for exporting a captain agent."""
4
+
5
+ import json
6
+ import os
7
+ from datetime import datetime, timezone
8
+ from pathlib import Path
9
+ from typing import Any, Callable, Dict, List, Optional, Union
10
+
11
+ from waldiez.models import (
12
+ WaldiezAgent,
13
+ WaldiezCaptainAgent,
14
+ WaldiezModel,
15
+ WaldiezModelData,
16
+ )
17
+
18
+
19
+ def get_captain_agent_extras(
20
+ agent: WaldiezAgent,
21
+ agent_names: Dict[str, str],
22
+ all_models: List[WaldiezModel],
23
+ serializer: Callable[..., str],
24
+ output_dir: Optional[Union[str, Path]],
25
+ ) -> str:
26
+ """Get the extra args for the captain agent.
27
+
28
+ Parameters
29
+ ----------
30
+ agent : WaldiezAgent
31
+ The agent.
32
+ agent_names : Dict[str, str]
33
+ A mapping of agent ids to agent names.
34
+ all_models : List[WaldiezModel]
35
+ All the models in the flow.
36
+ serializer : Callable[..., str]
37
+ The serializer to use.
38
+ output_dir : Optional[Union[str, Path]]
39
+ The output directory to save the agent lib and nested config.
40
+ Returns
41
+ -------
42
+ str
43
+ The extra args to use in the captain agent.
44
+ """
45
+ # extra args: nested_config, agent_lib, tool_lib
46
+ if not isinstance(agent, WaldiezCaptainAgent):
47
+ return ""
48
+ agent_name = agent_names[agent.id]
49
+ save_path = str(output_dir) if output_dir else "."
50
+ extra_args_content = "\n" + " agent_config_save_path=os.getcwd(),"
51
+ if agent.data.agent_lib:
52
+ lib_dict = [
53
+ lib.model_dump(by_alias=False) for lib in agent.data.agent_lib
54
+ ]
55
+ lib_json_name = f"{agent_name}_agent_lib.json"
56
+ agent_lib_path = os.path.join(save_path, lib_json_name)
57
+ with open(agent_lib_path, "w", encoding="utf-8", newline="\n") as f:
58
+ json.dump(lib_dict, f, ensure_ascii=False, indent=4)
59
+ extra_args_content += "\n" + f' agent_lib="{lib_json_name}",'
60
+ if agent.data.tool_lib:
61
+ extra_args_content += "\n" + f' tool_lib="{agent.data.tool_lib}",'
62
+ nested_config = generate_nested_config(
63
+ agent,
64
+ agent_name,
65
+ all_models,
66
+ save_path,
67
+ )
68
+ serialized_nested_config = serializer(nested_config)
69
+ extra_args_content += (
70
+ "\n" + f" nested_config={serialized_nested_config},"
71
+ )
72
+ return extra_args_content
73
+
74
+
75
+ def generate_nested_config(
76
+ agent: WaldiezCaptainAgent,
77
+ agent_name: str,
78
+ all_models: List[WaldiezModel],
79
+ save_path: str,
80
+ ) -> Dict[str, Any]:
81
+ """Generate the nested config for the captain agent.
82
+
83
+ Parameters
84
+ ----------
85
+ agent : WaldiezCaptainAgent
86
+ The captain agent.
87
+ agent_name : str
88
+ The agent name.
89
+ all_models : List[WaldiezModel]
90
+ All the models in the flow.
91
+ save_path : str
92
+ The path to save the nested config.
93
+ Returns
94
+ -------
95
+ Dict[str, Any]
96
+ The nested config.
97
+ """
98
+ config_file_or_env_name = f"{agent_name}_llm_config.json"
99
+ llm_config = get_llm_config(agent, all_models)
100
+ os.makedirs(save_path, exist_ok=True)
101
+ config_file_or_env_path = os.path.join(save_path, config_file_or_env_name)
102
+ with open(
103
+ config_file_or_env_path, "w", encoding="utf-8", newline="\n"
104
+ ) as f:
105
+ json.dump([llm_config], f, ensure_ascii=False, indent=4)
106
+ nested_config = {
107
+ "autobuild_init_config": {
108
+ "config_file_or_env": config_file_or_env_name,
109
+ "builder_model": llm_config["model"],
110
+ "agent_model": llm_config["model"],
111
+ },
112
+ "autobuild_build_config": get_auto_build_build_config(
113
+ agent, llm_config
114
+ ),
115
+ "group_chat_config": {"max_round": agent.data.max_round},
116
+ "group_chat_llm_config": None,
117
+ "max_turns": agent.data.max_turns,
118
+ }
119
+ return nested_config
120
+
121
+
122
+ def get_llm_config(
123
+ agent: WaldiezAgent,
124
+ all_models: List[WaldiezModel],
125
+ ) -> Dict[str, Any]:
126
+ """Get the config list environment variable name and its dict value.
127
+
128
+ Parameters
129
+ ----------
130
+ agent : WaldiezAgent
131
+ The agent.
132
+ all_models : List[WaldiezModel]
133
+ All the models in the flow.
134
+ Returns
135
+ -------
136
+ Dict[str, str]
137
+ The llm config dict.
138
+ """
139
+ model_name = "gpt-4o"
140
+ temperature: Optional[float] = 1
141
+ top_p: Optional[float] = 0.95
142
+ max_tokens: Optional[int] = 2048
143
+ if agent.data.model_ids:
144
+ waldiez_model = get_waldiez_model(agent.data.model_ids[0], all_models)
145
+ model_name = waldiez_model.name
146
+ temperature = waldiez_model.data.temperature
147
+ top_p = waldiez_model.data.top_p
148
+ max_tokens = waldiez_model.data.max_tokens
149
+ config_dict = {
150
+ "model": model_name,
151
+ "temperature": temperature,
152
+ "top_p": top_p,
153
+ "max_tokens": max_tokens,
154
+ }
155
+ return config_dict
156
+
157
+
158
+ def get_auto_build_build_config(
159
+ agent: WaldiezAgent,
160
+ llm_config: Dict[str, Any],
161
+ ) -> Dict[str, Any]:
162
+ """Get the auto build build config.
163
+
164
+ Parameters
165
+ ----------
166
+ agent : WaldiezAgent
167
+ The agent.
168
+ llm_config : Dict[str, Any]
169
+ The llm config.
170
+
171
+ Returns
172
+ -------
173
+ Dict[str, Any]
174
+ The auto build build config.
175
+ """
176
+ coding = False
177
+ code_execution_config = {
178
+ "timeout": 300,
179
+ "work_dir": "groupchat",
180
+ "last_n_messages": 1,
181
+ "use_docker": False,
182
+ }
183
+ if agent.data.code_execution_config is not False:
184
+ coding = True
185
+ code_execution_config["work_dir"] = (
186
+ agent.data.code_execution_config.work_dir or "groupchat"
187
+ )
188
+ code_execution_config["last_n_messages"] = (
189
+ agent.data.code_execution_config.last_n_messages or 1
190
+ )
191
+ code_execution_config["timeout"] = (
192
+ agent.data.code_execution_config.timeout or 300
193
+ )
194
+ return {
195
+ "default_llm_config": {
196
+ "temperature": llm_config["temperature"],
197
+ "top_p": llm_config["top_p"],
198
+ "max_tokens": llm_config["max_tokens"],
199
+ },
200
+ "code_execution_config": code_execution_config,
201
+ "coding": coding,
202
+ }
203
+
204
+
205
+ def get_waldiez_model(
206
+ model_id: str, all_models: List[WaldiezModel]
207
+ ) -> WaldiezModel:
208
+ """Get the model name from the model id.
209
+
210
+ Parameters
211
+ ----------
212
+ model_id : str
213
+ The model id.
214
+ all_models : List[WaldiezModel]
215
+ All the models in the flow.
216
+
217
+ Returns
218
+ -------
219
+ str
220
+ The model name.
221
+ """
222
+ for model in all_models:
223
+ if model.id == model_id:
224
+ return model
225
+ now = (
226
+ datetime.now(tz=timezone.utc)
227
+ .isoformat(timespec="milliseconds")
228
+ .replace("+00:00", "Z")
229
+ )
230
+ return WaldiezModel(
231
+ id=model_id,
232
+ type="model",
233
+ name="gpt-4o",
234
+ description="The GPT-4o model.",
235
+ tags=["gpt-4o"],
236
+ requirements=[],
237
+ created_at=now,
238
+ updated_at=now,
239
+ data=WaldiezModelData(
240
+ api_type="openai",
241
+ temperature=1,
242
+ top_p=0.95,
243
+ max_tokens=2048,
244
+ base_url=None,
245
+ api_key=os.environ.get("OPENAI_API_KEY", "REPLACE_ME"),
246
+ api_version=None,
247
+ default_headers={},
248
+ price=None,
249
+ ),
250
+ )
@@ -21,11 +21,11 @@ from waldiez.models import (
21
21
  # functions (List[Callable]):
22
22
  # -A list of functions to register with the agent.
23
23
  # update_agent_state_before_reply (List[Callable]):
24
- # - A list of functions, including UPDATE_SYSTEM_MESSAGEs,
24
+ # - A list of functions, including UpdateSystemMessage,
25
25
  # called to update the agent before it replies.
26
26
 
27
27
  # Additional methods:
28
- # register_hand_off(hand_offs: List[AfterWork|OnCondition]):
28
+ # register_hand_off(agent, hand_offs: List[AfterWork|OnCondition]):
29
29
 
30
30
 
31
31
  def get_swarm_extras(
@@ -173,14 +173,14 @@ def get_update_agent_state_before_reply_arg(
173
173
  name_suffix=agent_names[agent.id],
174
174
  )
175
175
  arg_string += (
176
- "\n" + f"{tab}{tab}UPDATE_SYSTEM_MESSAGE({function_name}),"
176
+ "\n" + f"{tab}{tab}UpdateSystemMessage({function_name}),"
177
177
  )
178
178
  before_agent += "\n" + function_content + "\n"
179
179
  else:
180
180
  escaped_function = string_escape(function.update_function)
181
181
  arg_string += (
182
182
  "\n"
183
- + f'{tab}{tab}UPDATE_SYSTEM_MESSAGE("{escaped_function}"),'
183
+ + f'{tab}{tab}UpdateSystemMessage("{escaped_function}"),'
184
184
  )
185
185
  else:
186
186
  skill_name = skill_names.get(function, "")
@@ -233,7 +233,12 @@ def get_agent_handoff_registrations(
233
233
  if not agent.handoffs:
234
234
  return before_agent, after_agent
235
235
  tab = " "
236
- after_agent = f"{agent_name}.register_hand_off(" + "\n" + f"{tab}[" + "\n"
236
+ # change {agent}.register_hand_off([...
237
+ # to register_hand_off({agent}, [...
238
+ # after_agent = f"{agent_name}.register_hand_off(" + "\n" + f"{tab}[" + "\n"
239
+ after_agent = (
240
+ "register_hand_off(\n" + f"{tab}{agent_name}," + "\n" + f"{tab}[" + "\n"
241
+ )
237
242
  for hand_off in agent.handoffs:
238
243
  if isinstance(hand_off, WaldiezSwarmOnCondition):
239
244
  registration, before_handoff = get_agent_on_condition_handoff(
@@ -385,7 +390,7 @@ def _get_agent_on_condition_handoff_to_agent(
385
390
  before_agent = ""
386
391
  tab = " "
387
392
  on_condition = (
388
- f"{tab}{tab}ON_CONDITION(" + "\n"
393
+ f"{tab}{tab}OnCondition(" + "\n"
389
394
  f"{tab}{tab}{tab}target={recipient}," + "\n"
390
395
  f'{tab}{tab}{tab}condition="{condition}",' + "\n"
391
396
  )
@@ -431,7 +436,7 @@ def _get_agent_on_condition_handoff_to_nested_chat(
431
436
  before_agent += f"{chat_queue_var_name} = {chat_queue} " + "\n"
432
437
  condition_string = string_escape(condition)
433
438
  on_condition = (
434
- f"{tab}{tab}ON_CONDITION(" + "\n"
439
+ f"{tab}{tab}OnCondition(" + "\n"
435
440
  f"{tab}{tab}{tab}target=" + "{\n"
436
441
  f'{tab}{tab}{tab}{tab}"chat_queue": {chat_queue_var_name},' + "\n"
437
442
  f'{tab}{tab}{tab}{tab}"config": None,' + "\n"
@@ -13,6 +13,7 @@ from typing_extensions import Literal
13
13
  CommentKey = Literal[
14
14
  "agents", "imports", "skills", "models", "nested", "run", "logging"
15
15
  ]
16
+ """Possible keys for comments."""
16
17
 
17
18
 
18
19
  def comment(for_notebook: bool, hashtags: int = 1) -> str:
@@ -196,7 +196,7 @@ def get_swarm_after_work_string(
196
196
  The after work string and the additional methods string.
197
197
  """
198
198
  if not chat.after_work:
199
- return "AFTER_WORK(AfterWorkOption.TERMINATE)", ""
199
+ return "AfterWork(AfterWorkOption.TERMINATE)", ""
200
200
  additional_methods = ""
201
201
  after_work_string, function_content = chat.after_work.get_recipient(
202
202
  agent_names=agent_names,
@@ -52,6 +52,7 @@ from .utils import (
52
52
  get_after_run_content,
53
53
  get_def_main,
54
54
  get_ipynb_content_start,
55
+ get_np_no_nep50_handle,
55
56
  get_py_content_start,
56
57
  get_sqlite_out,
57
58
  get_start_logging,
@@ -210,6 +211,7 @@ class FlowExporter(BaseExporter, ExporterMixin):
210
211
  )
211
212
  content += self.get_comment("imports", self.for_notebook) + "\n"
212
213
  content += imports[0] + "\n"
214
+ content += get_np_no_nep50_handle() + "\n"
213
215
  content += self.get_comment("logging", self.for_notebook) + "\n"
214
216
  content += get_start_logging(tabs=0) + "\n"
215
217
  content += "start_logging()\n\n"
@@ -240,13 +242,10 @@ class FlowExporter(BaseExporter, ExporterMixin):
240
242
  cache_seed=cache_seed,
241
243
  )
242
244
  else:
243
- # content += f" with Cache.disk(cache_seed={cache_seed}" + "):\n"
244
- # content += f"{flow_chats}" + "\n"
245
- # content += "\n" + chats_content + "\n"
246
245
  if chats_content.startswith("\n"):
247
246
  chats_content = chats_content[1:]
248
247
  content += (
249
- "\n" + f"with Cache.disk(cache_seed={cache_seed}):"
248
+ "\n" + f"with Cache.disk(cache_seed={cache_seed}) as cache:"
250
249
  "\n" + chats_content + "\n"
251
250
  )
252
251
  if is_async:
@@ -431,9 +430,9 @@ class FlowExporter(BaseExporter, ExporterMixin):
431
430
  exporter = AgentExporter(
432
431
  agent=agent,
433
432
  agent_names=self.agent_names,
434
- model_names=self.model_names,
435
- skill_names=self.skill_names,
433
+ models=(self.models, self.model_names),
436
434
  chats=(self.chats, self.chat_names),
435
+ skill_names=self.skill_names,
437
436
  is_async=self.waldiez.is_async,
438
437
  for_notebook=self.for_notebook,
439
438
  output_dir=self.output_dir,
@@ -14,14 +14,11 @@ from .def_main import get_def_main
14
14
  from .flow_content import (
15
15
  get_after_run_content,
16
16
  get_ipynb_content_start,
17
+ get_np_no_nep50_handle,
17
18
  get_py_content_start,
18
19
  )
19
20
  from .flow_names import ensure_unique_names
20
- from .importing_utils import (
21
- gather_imports,
22
- get_standard_imports,
23
- get_the_imports_string,
24
- )
21
+ from .importing_utils import gather_imports, get_the_imports_string
25
22
  from .logging_utils import (
26
23
  get_sqlite_out,
27
24
  get_sqlite_out_call,
@@ -41,12 +38,12 @@ __all__ = [
41
38
  "gather_imports",
42
39
  "get_after_run_content",
43
40
  "get_def_main",
41
+ "get_np_no_nep50_handle",
44
42
  "get_py_content_start",
45
43
  "get_ipynb_content_start",
46
44
  "get_start_logging",
47
45
  "get_stop_logging",
48
46
  "get_sqlite_out",
49
47
  "get_sqlite_out_call",
50
- "get_standard_imports",
51
48
  "get_the_imports_string",
52
49
  ]
@@ -53,7 +53,7 @@ def get_def_main(
53
53
  else:
54
54
  content += "def call_main() -> None:\n"
55
55
  content += ' """Run the main function and print the results."""\n'
56
- content += " results = "
56
+ content += " results: Union[ChatResult, List[ChatResult], Dict[int, ChatResult]] = "
57
57
  if is_async:
58
58
  content += "await "
59
59
  content += "main()\n"
@@ -62,11 +62,12 @@ def get_def_main(
62
62
  content += " ordered_results = dict(sorted(results.items()))\n"
63
63
  content += " for _, result in ordered_results.items():\n"
64
64
  content += " pprint(asdict(result))\n"
65
- content += " elif isinstance(results, list):\n"
65
+ content += " else:\n"
66
+ content += " if not isinstance(results, list):\n"
67
+ content += " results = [results]\n"
66
68
  content += " for result in results:\n"
67
69
  content += " pprint(asdict(result))\n"
68
- content += " else:\n"
69
- content += " pprint(asdict(results))\n\n\n"
70
+ content += "\n\n"
70
71
  content += 'if __name__ == "__main__":\n'
71
72
  if is_async:
72
73
  content += " anyio.run(call_main)\n"
@@ -71,12 +71,14 @@ PYLINT_RULES = [
71
71
  "unknown-option-value",
72
72
  "unused-argument",
73
73
  "unused-import",
74
+ "unused-variable",
74
75
  "invalid-name",
75
76
  "import-error",
76
77
  "inconsistent-quotes",
77
78
  "missing-function-docstring",
78
79
  "missing-param-doc",
79
80
  "missing-return-doc",
81
+ "ungrouped-imports",
80
82
  ]
81
83
 
82
84
 
@@ -159,3 +161,39 @@ def get_after_run_content(
159
161
  {space}{tab}pass
160
162
  """
161
163
  return content
164
+
165
+
166
+ def get_np_no_nep50_handle() -> str:
167
+ """Handle catching the "module numpy has no attribute _no_pep50_warning" error.
168
+
169
+ Returns
170
+ -------
171
+ str
172
+ The content to handle the error.
173
+ """
174
+ # https://github.com/numpy/numpy/blob/v2.2.2/\
175
+ # doc/source/release/2.2.0-notes.rst#nep-50-promotion-state-option-removed
176
+ content = '''
177
+ # try to make sure we don't get:
178
+ # module 'numpy' has no attribute '_no_nep50_warning'"
179
+ os.environ["NEP50_DEPRECATION_WARNING"] = "0"
180
+ os.environ["NEP50_DISABLE_WARNING"] = "1"
181
+ os.environ["NPY_PROMOTION_STATE"] = "weak"
182
+ if not hasattr(np, "_no_pep50_warning"):
183
+
184
+ import contextlib
185
+ from typing import Generator
186
+
187
+ @contextlib.contextmanager
188
+ def _np_no_nep50_warning() -> Generator[None, None, None]:
189
+ """Dummy function to avoid the warning.
190
+
191
+ Yields
192
+ ------
193
+ None
194
+ Nothing.
195
+ """
196
+ yield
197
+ setattr(np, "_no_pep50_warning", _np_no_nep50_warning) # noqa
198
+ '''
199
+ return content