waldiez 0.6.0__py3-none-any.whl → 0.6.1__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 (188) hide show
  1. waldiez/__init__.py +1 -1
  2. waldiez/_version.py +1 -1
  3. waldiez/cli.py +18 -7
  4. waldiez/cli_extras/jupyter.py +3 -0
  5. waldiez/cli_extras/runner.py +3 -1
  6. waldiez/cli_extras/studio.py +3 -1
  7. waldiez/exporter.py +9 -3
  8. waldiez/exporting/agent/exporter.py +9 -10
  9. waldiez/exporting/agent/extras/captain_agent_extras.py +6 -6
  10. waldiez/exporting/agent/extras/doc_agent_extras.py +6 -6
  11. waldiez/exporting/agent/extras/group_manager_agent_extas.py +34 -23
  12. waldiez/exporting/agent/extras/group_member_extras.py +6 -5
  13. waldiez/exporting/agent/extras/handoffs/after_work.py +1 -1
  14. waldiez/exporting/agent/extras/handoffs/available.py +1 -1
  15. waldiez/exporting/agent/extras/handoffs/condition.py +3 -2
  16. waldiez/exporting/agent/extras/handoffs/handoff.py +1 -1
  17. waldiez/exporting/agent/extras/handoffs/target.py +6 -4
  18. waldiez/exporting/agent/extras/rag/chroma_extras.py +27 -19
  19. waldiez/exporting/agent/extras/rag/mongo_extras.py +8 -8
  20. waldiez/exporting/agent/extras/rag/pgvector_extras.py +5 -5
  21. waldiez/exporting/agent/extras/rag/qdrant_extras.py +5 -4
  22. waldiez/exporting/agent/extras/rag/vector_db_extras.py +1 -1
  23. waldiez/exporting/agent/extras/rag_user_proxy_agent_extras.py +5 -7
  24. waldiez/exporting/agent/extras/reasoning_agent_extras.py +3 -5
  25. waldiez/exporting/chats/exporter.py +4 -4
  26. waldiez/exporting/chats/processor.py +1 -2
  27. waldiez/exporting/chats/utils/common.py +89 -48
  28. waldiez/exporting/chats/utils/group.py +9 -9
  29. waldiez/exporting/chats/utils/nested.py +7 -7
  30. waldiez/exporting/chats/utils/sequential.py +1 -1
  31. waldiez/exporting/chats/utils/single.py +2 -2
  32. waldiez/exporting/core/content.py +7 -7
  33. waldiez/exporting/core/context.py +5 -3
  34. waldiez/exporting/core/exporter.py +5 -3
  35. waldiez/exporting/core/exporters.py +2 -2
  36. waldiez/exporting/core/extras/agent_extras/captain_extras.py +2 -2
  37. waldiez/exporting/core/extras/agent_extras/group_manager_extras.py +2 -2
  38. waldiez/exporting/core/extras/agent_extras/rag_user_extras.py +2 -2
  39. waldiez/exporting/core/extras/agent_extras/standard_extras.py +3 -8
  40. waldiez/exporting/core/extras/base.py +7 -5
  41. waldiez/exporting/core/extras/flow_extras.py +4 -5
  42. waldiez/exporting/core/extras/model_extras.py +2 -2
  43. waldiez/exporting/core/extras/path_resolver.py +1 -2
  44. waldiez/exporting/core/extras/serializer.py +2 -2
  45. waldiez/exporting/core/protocols.py +6 -5
  46. waldiez/exporting/core/result.py +25 -28
  47. waldiez/exporting/core/types.py +10 -10
  48. waldiez/exporting/core/utils/llm_config.py +2 -2
  49. waldiez/exporting/core/validation.py +10 -11
  50. waldiez/exporting/flow/execution_generator.py +98 -10
  51. waldiez/exporting/flow/exporter.py +2 -2
  52. waldiez/exporting/flow/factory.py +2 -2
  53. waldiez/exporting/flow/file_generator.py +4 -2
  54. waldiez/exporting/flow/merger.py +5 -3
  55. waldiez/exporting/flow/orchestrator.py +72 -2
  56. waldiez/exporting/flow/utils/common.py +5 -5
  57. waldiez/exporting/flow/utils/importing.py +6 -7
  58. waldiez/exporting/flow/utils/linting.py +25 -9
  59. waldiez/exporting/flow/utils/logging.py +2 -2
  60. waldiez/exporting/models/exporter.py +8 -8
  61. waldiez/exporting/models/processor.py +5 -5
  62. waldiez/exporting/tools/exporter.py +2 -2
  63. waldiez/exporting/tools/processor.py +7 -4
  64. waldiez/io/__init__.py +8 -4
  65. waldiez/io/_ws.py +10 -6
  66. waldiez/io/models/constants.py +10 -10
  67. waldiez/io/models/content/audio.py +1 -0
  68. waldiez/io/models/content/base.py +20 -18
  69. waldiez/io/models/content/file.py +1 -0
  70. waldiez/io/models/content/image.py +1 -0
  71. waldiez/io/models/content/text.py +1 -0
  72. waldiez/io/models/content/video.py +1 -0
  73. waldiez/io/models/user_input.py +10 -5
  74. waldiez/io/models/user_response.py +17 -16
  75. waldiez/io/mqtt.py +18 -31
  76. waldiez/io/redis.py +18 -22
  77. waldiez/io/structured.py +52 -53
  78. waldiez/io/utils.py +3 -0
  79. waldiez/io/ws.py +5 -1
  80. waldiez/logger.py +16 -3
  81. waldiez/models/agents/__init__.py +3 -0
  82. waldiez/models/agents/agent/agent.py +23 -16
  83. waldiez/models/agents/agent/agent_data.py +25 -22
  84. waldiez/models/agents/agent/code_execution.py +9 -11
  85. waldiez/models/agents/agent/termination_message.py +10 -12
  86. waldiez/models/agents/agent/update_system_message.py +2 -4
  87. waldiez/models/agents/agents.py +8 -8
  88. waldiez/models/agents/assistant/assistant.py +6 -3
  89. waldiez/models/agents/assistant/assistant_data.py +2 -2
  90. waldiez/models/agents/captain/captain_agent.py +7 -4
  91. waldiez/models/agents/captain/captain_agent_data.py +5 -7
  92. waldiez/models/agents/doc_agent/doc_agent.py +7 -4
  93. waldiez/models/agents/doc_agent/doc_agent_data.py +9 -10
  94. waldiez/models/agents/doc_agent/rag_query_engine.py +10 -12
  95. waldiez/models/agents/extra_requirements.py +3 -3
  96. waldiez/models/agents/group_manager/group_manager.py +12 -7
  97. waldiez/models/agents/group_manager/group_manager_data.py +13 -12
  98. waldiez/models/agents/group_manager/speakers.py +17 -19
  99. waldiez/models/agents/rag_user_proxy/rag_user_proxy.py +7 -4
  100. waldiez/models/agents/rag_user_proxy/rag_user_proxy_data.py +4 -1
  101. waldiez/models/agents/rag_user_proxy/retrieve_config.py +69 -63
  102. waldiez/models/agents/rag_user_proxy/vector_db_config.py +19 -19
  103. waldiez/models/agents/reasoning/reasoning_agent.py +7 -4
  104. waldiez/models/agents/reasoning/reasoning_agent_data.py +3 -2
  105. waldiez/models/agents/reasoning/reasoning_agent_reason_config.py +8 -8
  106. waldiez/models/agents/user_proxy/user_proxy.py +6 -3
  107. waldiez/models/agents/user_proxy/user_proxy_data.py +1 -1
  108. waldiez/models/chat/chat.py +27 -20
  109. waldiez/models/chat/chat_data.py +22 -19
  110. waldiez/models/chat/chat_message.py +9 -9
  111. waldiez/models/chat/chat_nested.py +9 -9
  112. waldiez/models/chat/chat_summary.py +6 -6
  113. waldiez/models/common/__init__.py +2 -0
  114. waldiez/models/common/ag2_version.py +2 -0
  115. waldiez/models/common/dict_utils.py +8 -6
  116. waldiez/models/common/handoff.py +18 -17
  117. waldiez/models/common/method_utils.py +7 -7
  118. waldiez/models/common/naming.py +49 -0
  119. waldiez/models/flow/flow.py +11 -6
  120. waldiez/models/flow/flow_data.py +23 -17
  121. waldiez/models/flow/info.py +3 -3
  122. waldiez/models/flow/naming.py +2 -1
  123. waldiez/models/model/_aws.py +11 -13
  124. waldiez/models/model/_llm.py +5 -0
  125. waldiez/models/model/_price.py +2 -4
  126. waldiez/models/model/extra_requirements.py +1 -3
  127. waldiez/models/model/model.py +2 -2
  128. waldiez/models/model/model_data.py +21 -21
  129. waldiez/models/tool/extra_requirements.py +2 -4
  130. waldiez/models/tool/predefined/_duckduckgo.py +1 -0
  131. waldiez/models/tool/predefined/_email.py +1 -0
  132. waldiez/models/tool/predefined/_google.py +1 -0
  133. waldiez/models/tool/predefined/_perplexity.py +1 -0
  134. waldiez/models/tool/predefined/_searxng.py +1 -0
  135. waldiez/models/tool/predefined/_tavily.py +1 -0
  136. waldiez/models/tool/predefined/_wikipedia.py +1 -0
  137. waldiez/models/tool/predefined/_youtube.py +1 -0
  138. waldiez/models/tool/tool.py +8 -5
  139. waldiez/models/tool/tool_data.py +2 -2
  140. waldiez/models/waldiez.py +152 -4
  141. waldiez/runner.py +11 -5
  142. waldiez/running/async_utils.py +192 -0
  143. waldiez/running/base_runner.py +117 -264
  144. waldiez/running/dir_utils.py +52 -0
  145. waldiez/running/environment.py +10 -44
  146. waldiez/running/events_mixin.py +252 -0
  147. waldiez/running/exceptions.py +20 -0
  148. waldiez/running/gen_seq_diagram.py +18 -15
  149. waldiez/running/io_utils.py +216 -0
  150. waldiez/running/protocol.py +11 -5
  151. waldiez/running/requirements_mixin.py +65 -0
  152. waldiez/running/results_mixin.py +926 -0
  153. waldiez/running/standard_runner.py +22 -25
  154. waldiez/running/step_by_step/breakpoints_mixin.py +192 -60
  155. waldiez/running/step_by_step/command_handler.py +3 -0
  156. waldiez/running/step_by_step/events_processor.py +194 -14
  157. waldiez/running/step_by_step/step_by_step_models.py +110 -43
  158. waldiez/running/step_by_step/step_by_step_runner.py +107 -57
  159. waldiez/running/subprocess_runner/__base__.py +9 -1
  160. waldiez/running/subprocess_runner/_async_runner.py +5 -3
  161. waldiez/running/subprocess_runner/_sync_runner.py +6 -2
  162. waldiez/running/subprocess_runner/runner.py +39 -23
  163. waldiez/running/timeline_processor.py +1 -1
  164. waldiez/utils/__init__.py +2 -0
  165. waldiez/utils/conflict_checker.py +4 -4
  166. waldiez/utils/python_manager.py +415 -0
  167. waldiez/ws/_file_handler.py +18 -18
  168. waldiez/ws/_mock.py +2 -1
  169. waldiez/ws/cli.py +36 -12
  170. waldiez/ws/client_manager.py +35 -27
  171. waldiez/ws/errors.py +3 -0
  172. waldiez/ws/models.py +43 -52
  173. waldiez/ws/reloader.py +12 -4
  174. waldiez/ws/server.py +85 -55
  175. waldiez/ws/session_manager.py +8 -9
  176. waldiez/ws/session_stats.py +1 -1
  177. waldiez/ws/utils.py +4 -1
  178. {waldiez-0.6.0.dist-info → waldiez-0.6.1.dist-info}/METADATA +82 -93
  179. waldiez-0.6.1.dist-info/RECORD +254 -0
  180. waldiez/running/post_run.py +0 -186
  181. waldiez/running/pre_run.py +0 -281
  182. waldiez/running/run_results.py +0 -14
  183. waldiez/running/utils.py +0 -625
  184. waldiez-0.6.0.dist-info/RECORD +0 -251
  185. {waldiez-0.6.0.dist-info → waldiez-0.6.1.dist-info}/WHEEL +0 -0
  186. {waldiez-0.6.0.dist-info → waldiez-0.6.1.dist-info}/entry_points.txt +0 -0
  187. {waldiez-0.6.0.dist-info → waldiez-0.6.1.dist-info}/licenses/LICENSE +0 -0
  188. {waldiez-0.6.0.dist-info → waldiez-0.6.1.dist-info}/licenses/NOTICE.md +0 -0
@@ -3,7 +3,6 @@
3
3
  """Get chroma db related imports and content."""
4
4
 
5
5
  from pathlib import Path
6
- from typing import Set
7
6
 
8
7
  from waldiez.models import WaldiezRagUserProxy
9
8
 
@@ -33,7 +32,8 @@ def _get_chroma_client_string(agent: WaldiezRagUserProxy) -> tuple[str, str]:
33
32
  local_path = Path(agent.retrieve_config.db_config.local_storage_path)
34
33
  client_str += (
35
34
  "PersistentClient(\n"
36
- f' path=r"{local_path}",' + "\n"
35
+ f' path=r"{local_path}",'
36
+ "\n"
37
37
  " settings=Settings(anonymized_telemetry=False),\n"
38
38
  ")"
39
39
  )
@@ -80,7 +80,7 @@ def _get_chroma_embedding_function_string(
80
80
 
81
81
  def get_chroma_db_args(
82
82
  agent: WaldiezRagUserProxy, agent_name: str
83
- ) -> tuple[str, Set[str], str, str]:
83
+ ) -> tuple[str, set[str], str, str]:
84
84
  """Get the 'kwargs to use for ChromaVectorDB.
85
85
 
86
86
  Parameters
@@ -107,8 +107,10 @@ def get_chroma_db_args(
107
107
  if to_import_embedding:
108
108
  to_import.add(to_import_embedding)
109
109
  kwarg_string = (
110
- f" client={agent_name}_client," + "\n"
111
- f" embedding_function={embedding_function_arg}," + "\n"
110
+ f" client={agent_name}_client,"
111
+ "\n"
112
+ f" embedding_function={embedding_function_arg},"
113
+ "\n"
112
114
  )
113
115
  # The RAG example:
114
116
  # https://ag2ai.github.io/ag2/docs/notebooks/agentchat_groupchat_RAG/
@@ -123,30 +125,36 @@ def get_chroma_db_args(
123
125
  content_before += (
124
126
  f"{agent_name}_embedding_function = "
125
127
  "SentenceTransformerEmbeddingFunction(\n"
126
- f' model_name="{vector_db_model}",' + "\n"
127
- ")\n"
128
+ f' model_name="{vector_db_model}",'
129
+ "\n)\n"
128
130
  )
129
131
  collection_name = agent.retrieve_config.collection_name
130
132
  get_or_create = agent.retrieve_config.get_or_create
131
133
  if collection_name:
132
134
  if get_or_create:
133
135
  content_before += (
134
- f"{agent_name}_client.get_or_create_collection" + "(\n"
135
- f' "{collection_name}",' + "\n"
136
- f" embedding_function={embedding_function_arg}," + "\n"
137
- ")" + "\n"
136
+ f"{agent_name}_client.get_or_create_collection"
137
+ "(\n"
138
+ f' "{collection_name}",'
139
+ "\n"
140
+ f" embedding_function={embedding_function_arg},"
141
+ "\n)\n"
138
142
  )
139
143
  else:
140
144
  content_before += (
141
145
  "try:\n"
142
- f" {agent_name}_client.get_collection(" + "\n"
143
- f' "{collection_name}",' + "\n"
144
- f" embedding_function={embedding_function_arg}," + "\n"
145
- " )\n"
146
+ f" {agent_name}_client.get_collection("
147
+ "\n"
148
+ f' "{collection_name}",'
149
+ "\n"
150
+ f" embedding_function={embedding_function_arg},"
151
+ "\n )\n"
146
152
  "except ValueError:\n"
147
- f" {agent_name}_client.create_collection(" + "\n"
148
- f' "{collection_name}",' + "\n"
149
- f" embedding_function={embedding_function_arg}," + "\n"
150
- " )\n"
153
+ f" {agent_name}_client.create_collection("
154
+ "\n"
155
+ f' "{collection_name}",'
156
+ "\n"
157
+ f" embedding_function={embedding_function_arg},"
158
+ "\n )\n"
151
159
  )
152
160
  return kwarg_string, to_import, embedding_function_body, content_before
@@ -2,8 +2,6 @@
2
2
  # Copyright (c) 2024 - 2025 Waldiez and contributors.
3
3
  """Get mongodb related content and imports."""
4
4
 
5
- from typing import Set
6
-
7
5
  from waldiez.models import WaldiezRagUserProxy
8
6
 
9
7
 
@@ -45,7 +43,7 @@ def _get_mongodb_embedding_function_string(
45
43
 
46
44
  def get_mongodb_db_args(
47
45
  agent: WaldiezRagUserProxy, agent_name: str
48
- ) -> tuple[str, Set[str], str]:
46
+ ) -> tuple[str, set[str], str]:
49
47
  """Get the kwargs to use for MongoDBAtlasVectorDB.
50
48
 
51
49
  Parameters
@@ -63,14 +61,16 @@ def get_mongodb_db_args(
63
61
  embedding_function_arg, to_import_embedding, embedding_function_body = (
64
62
  _get_mongodb_embedding_function_string(agent, agent_name)
65
63
  )
66
- to_import: Set[str] = (
64
+ to_import: set[str] = (
67
65
  set() if not to_import_embedding else {to_import_embedding}
68
66
  )
69
- tab = " " * 12
67
+ tab: str = " " * 12
70
68
  db_config = agent.retrieve_config.db_config
71
- kwarg_string = (
72
- f'{tab}connection_string="{db_config.connection_url}",' + "\n"
73
- f"{tab}embedding_function={embedding_function_arg}," + "\n"
69
+ kwarg_string: str = (
70
+ f'{tab}connection_string="{db_config.connection_url}",'
71
+ "\n"
72
+ f"{tab}embedding_function={embedding_function_arg},"
73
+ "\n"
74
74
  )
75
75
  wait_until_document_ready = db_config.wait_until_document_ready
76
76
  wait_until_index_ready = db_config.wait_until_index_ready
@@ -2,8 +2,6 @@
2
2
  # Copyright (c) 2024 - 2025 Waldiez and contributors.
3
3
  """Get pgvector related content and imports."""
4
4
 
5
- from typing import Set
6
-
7
5
  from waldiez.models import WaldiezRagUserProxy
8
6
 
9
7
 
@@ -64,7 +62,7 @@ def _get_pgvector_embedding_function_string(
64
62
 
65
63
  def get_pgvector_db_args(
66
64
  agent: WaldiezRagUserProxy, agent_name: str
67
- ) -> tuple[str, Set[str], str]:
65
+ ) -> tuple[str, set[str], str]:
68
66
  """Get the kwargs to use for PGVectorDB.
69
67
 
70
68
  Parameters
@@ -89,7 +87,9 @@ def get_pgvector_db_args(
89
87
  else {to_import_client}
90
88
  )
91
89
  kwarg_str = (
92
- f" client={client_str}," + "\n"
93
- f" embedding_function={embedding_function_arg}," + "\n"
90
+ f" client={client_str},"
91
+ "\n"
92
+ f" embedding_function={embedding_function_arg},"
93
+ "\n"
94
94
  )
95
95
  return kwarg_str, to_import, embedding_function_body
@@ -3,7 +3,6 @@
3
3
  """Get qdrant db related imports and content."""
4
4
 
5
5
  from pathlib import Path
6
- from typing import Set
7
6
 
8
7
  from waldiez.models import WaldiezRagUserProxy
9
8
 
@@ -81,7 +80,7 @@ def _get_qdrant_embedding_function_string(
81
80
 
82
81
  def get_qdrant_db_args(
83
82
  agent: WaldiezRagUserProxy, agent_name: str
84
- ) -> tuple[str, Set[str], str]:
83
+ ) -> tuple[str, set[str], str]:
85
84
  """Get the kwargs to use for QdrantVectorDB.
86
85
 
87
86
  Parameters
@@ -106,7 +105,9 @@ def get_qdrant_db_args(
106
105
  else {to_import_client}
107
106
  )
108
107
  kwarg_string = (
109
- f" client={client_str}," + "\n"
110
- f" embedding_function={embedding_function_arg}," + "\n"
108
+ f" client={client_str},"
109
+ "\n"
110
+ f" embedding_function={embedding_function_arg},"
111
+ "\n"
111
112
  )
112
113
  return kwarg_string, to_import, embedding_function_body
@@ -21,7 +21,7 @@ class VectorDBExtras:
21
21
 
22
22
  before_arg: str
23
23
  vector_db_arg: str
24
- imports: set[str] = field(default_factory=set[str])
24
+ imports: set[str] = field(default_factory=set)
25
25
 
26
26
 
27
27
  def get_vector_db_extras(
@@ -2,8 +2,6 @@
2
2
  # Copyright (c) 2024 - 2025 Waldiez and contributors.
3
3
  """RAG user proxy agent configuration processor."""
4
4
 
5
- from typing import Union
6
-
7
5
  from waldiez.models import (
8
6
  WaldiezAgent,
9
7
  WaldiezRagUserProxy,
@@ -115,7 +113,7 @@ class RagUserProxyAgentProcessor:
115
113
  self.agent.retrieve_config
116
114
  )
117
115
  if not retrieve_config:
118
- return "", ""
116
+ return "", "" # pyright: ignore[reportUnreachable]
119
117
  args_dict = self._get_args_dict()
120
118
  if not args_dict:
121
119
  return "", ""
@@ -146,7 +144,7 @@ class RagUserProxyAgentProcessor:
146
144
 
147
145
  def _get_args_string(
148
146
  self,
149
- args_dict: dict[str, Union[str, list[str]]],
147
+ args_dict: dict[str, str | list[str]],
150
148
  vector_db_extras: VectorDBExtras,
151
149
  before_agent: str,
152
150
  ) -> str:
@@ -193,11 +191,11 @@ class RagUserProxyAgentProcessor:
193
191
  return f"{new_model_name}"
194
192
  return WaldiezRagUserProxyModels[self.agent.retrieve_config.vector_db]
195
193
 
196
- def _get_args_dict(self) -> dict[str, Union[str, list[str]]]:
194
+ def _get_args_dict(self) -> dict[str, str | list[str]]:
197
195
  if not isinstance(self.agent, WaldiezRagUserProxy):
198
196
  return {}
199
197
  model_arg = self._get_model_arg()
200
- args_dict: dict[str, Union[str, list[str]]] = {
198
+ args_dict: dict[str, str | list[str]] = {
201
199
  "task": self.agent.retrieve_config.task,
202
200
  "model": model_arg,
203
201
  }
@@ -212,7 +210,7 @@ class RagUserProxyAgentProcessor:
212
210
  if arg_value is not None:
213
211
  args_dict[arg] = arg_value
214
212
  args_dict[arg] = getattr(self.agent.retrieve_config, arg)
215
- docs_path: Union[str, list[str]] = []
213
+ docs_path: str | list[str] = []
216
214
  if self.agent.retrieve_config.docs_path:
217
215
  doc_paths = (
218
216
  self.agent.retrieve_config.docs_path
@@ -3,8 +3,6 @@
3
3
  # pylint: disable=too-few-public-methods,no-self-use
4
4
  """Reasoning agent configuration processor."""
5
5
 
6
- from typing import Optional
7
-
8
6
  from waldiez.models import WaldiezAgent, WaldiezReasoningAgent
9
7
 
10
8
  from ...core import (
@@ -37,9 +35,9 @@ class ReasoningAgentProcessor:
37
35
 
38
36
  def process(
39
37
  self,
40
- code_execution_config: Optional[CodeExecutionConfig] = None,
41
- termination_config: Optional[TerminationConfig] = None,
42
- system_message_config: Optional[SystemMessageConfig] = None,
38
+ code_execution_config: CodeExecutionConfig | None = None,
39
+ termination_config: TerminationConfig | None = None,
40
+ system_message_config: SystemMessageConfig | None = None,
43
41
  ) -> ReasoningExtras:
44
42
  """Process reasoning agent configuration.
45
43
 
@@ -2,7 +2,7 @@
2
2
  # Copyright (c) 2024 - 2025 Waldiez and contributors.
3
3
  """Chats exporter."""
4
4
 
5
- from typing import Any, Optional
5
+ from typing import Any
6
6
 
7
7
  from waldiez.models import (
8
8
  WaldiezAgent,
@@ -25,8 +25,8 @@ class ChatsExporter(Exporter[ChatExtras]):
25
25
  all_chats: list[WaldiezChat],
26
26
  chat_names: dict[str, str],
27
27
  main_chats: list[WaldiezAgentConnection],
28
- root_group_manager: Optional[WaldiezGroupManager],
29
- context: Optional[ExporterContext] = None,
28
+ root_group_manager: WaldiezGroupManager | None,
29
+ context: ExporterContext | None = None,
30
30
  **kwargs: Any,
31
31
  ) -> None:
32
32
  """Initialize the chats exporter.
@@ -92,6 +92,6 @@ class ChatsExporter(Exporter[ChatExtras]):
92
92
  extras.append_after_all_agents(extras.chat_registration)
93
93
  return extras
94
94
 
95
- def generate_main_content(self) -> Optional[str]:
95
+ def generate_main_content(self) -> str | None:
96
96
  """Generate the main content of the export."""
97
97
  return None
@@ -3,7 +3,6 @@
3
3
  """Chats processor."""
4
4
 
5
5
  from dataclasses import dataclass
6
- from typing import Optional
7
6
 
8
7
  from waldiez.models import (
9
8
  WaldiezAgent,
@@ -53,7 +52,7 @@ class ChatsProcessor:
53
52
  root_group_manager: WaldiezGroupManager | None,
54
53
  for_notebook: bool,
55
54
  is_async: bool,
56
- cache_seed: Optional[int],
55
+ cache_seed: int | None,
57
56
  serializer: Serializer,
58
57
  extras: ChatExtras,
59
58
  ) -> None:
@@ -5,7 +5,6 @@
5
5
  """Common utilities for exporting chats."""
6
6
 
7
7
  import json
8
- from typing import Optional
9
8
 
10
9
  from waldiez.models import WaldiezAgent, WaldiezChat
11
10
 
@@ -14,7 +13,7 @@ def get_chat_message_string(
14
13
  sender: WaldiezAgent,
15
14
  chat: WaldiezChat,
16
15
  chat_names: dict[str, str],
17
- ) -> tuple[str, Optional[str]]:
16
+ ) -> tuple[str, str | None]:
18
17
  """Get the agent's message as a string.
19
18
 
20
19
  Parameters
@@ -37,7 +36,7 @@ def get_chat_message_string(
37
36
  return "None", None
38
37
  if chat.message.type == "string":
39
38
  if chat.message.content is None: # pragma: no cover
40
- # should be coverred previousliy on pydantic validation
39
+ # should be covered previously on pydantic validation
41
40
  return "None", None
42
41
  if not chat.message.content:
43
42
  return "", None
@@ -52,16 +51,23 @@ def get_chat_message_string(
52
51
  return function_name, function_content
53
52
 
54
53
 
54
+ def _stop_logging(space: str, is_async: bool) -> str:
55
+ tab = " "
56
+ if is_async:
57
+ return f"{space}{tab}{tab}{tab}{tab}await stop_logging()\n"
58
+ return f"{space}{tab}{tab}{tab}{tab}stop_logging()\n"
59
+
60
+
55
61
  def get_event_handler_string(
56
- tab: str,
62
+ space: str,
57
63
  is_async: bool,
58
64
  ) -> str:
59
65
  """Get the event handler string.
60
66
 
61
67
  Parameters
62
68
  ----------
63
- tab : str
64
- The tab string.
69
+ space : str
70
+ The space before the content.
65
71
  is_async : bool
66
72
  Whether the handler is asynchronous.
67
73
 
@@ -70,56 +76,82 @@ def get_event_handler_string(
70
76
  str
71
77
  The event handler string.
72
78
  """
79
+ tab = " "
73
80
  content = (
74
- f"{tab}if on_event:\n"
75
- f"{tab} if not isinstance(results, list):\n"
76
- f"{tab} results = [results] # pylint: disable=redefined-variable-type\n"
77
- f"{tab} for index, result in enumerate(results):\n"
81
+ f"{space}if not isinstance(results, list):\n"
82
+ f"{space}{tab}results = [results] # pylint: disable=redefined-variable-type\n"
83
+ f"{space}got_agents = False\n"
84
+ f"{space}known_agents: list[ConversableAgent] = []\n"
85
+ f"{space}result_events: list[dict[str, Any]] = []\n"
86
+ f"{space}if on_event:\n"
87
+ f"{space}{tab}for index, result in enumerate(results):\n"
88
+ f"{space}{tab}{tab}result_events = []\n"
78
89
  )
90
+ await_str = "await " if is_async else ""
79
91
  if is_async:
80
92
  content += (
81
- f"{tab} async for event in result.events:\n"
82
- f"{tab} try:\n"
83
- f"{tab} should_continue = await on_event(event)\n"
84
- f"{tab} except BaseException as e:\n"
85
- f'{tab} print(f"Error in event handler: {{e}}")\n'
86
- f"{tab} raise SystemExit(\n"
87
- f'{tab} "Error in event handler: " + str(e)\n'
88
- f"{tab} ) from e\n"
93
+ f"{space}{tab}{tab}async for event in result.events:\n{_add_event_dump(space, tab)}"
94
+ f"{space}{tab}{tab}{tab}if not got_agents:\n"
95
+ f"{space}{tab}{tab}{tab}{tab}known_agents = _get_known_agents()\n"
96
+ f"{space}{tab}{tab}{tab}{tab}got_agents = True\n"
97
+ f"{space}{tab}{tab}{tab}try:\n"
98
+ f"{space}{tab}{tab}{tab}{tab}should_continue = await on_event(event, known_agents)\n"
99
+ f"{space}{tab}{tab}{tab}except BaseException as e:\n{_stop_logging(space, is_async)}"
100
+ f"{space}{tab}{tab}{tab}{tab}await store_error(e)\n"
101
+ f"{space}{tab}{tab}{tab}{tab}raise SystemExit(\n"
102
+ f'{space}{tab}{tab}{tab}{tab}{tab}"Error in event handler: " + str(e)\n'
103
+ f"{space}{tab}{tab}{tab}{tab}) from e\n"
89
104
  )
90
105
  else:
91
106
  content += (
92
- f"{tab} for event in result.events:\n"
93
- f"{tab} try:\n"
94
- f"{tab} should_continue = on_event(event)\n"
95
- f"{tab} except BaseException as e:\n"
96
- f'{tab} print(f"Error in event handler: {{e}}")\n'
97
- f"{tab} raise SystemExit(\n"
98
- f'{tab} "Error in event handler: " + str(e)\n'
99
- f"{tab} ) from e\n"
107
+ f"{space}{tab}{tab}for event in result.events:\n{_add_event_dump(space, tab)}"
108
+ f"{space}{tab}{tab}{tab}if not got_agents:\n"
109
+ f"{space}{tab}{tab}{tab}{tab}known_agents = _get_known_agents()\n"
110
+ f"{space}{tab}{tab}{tab}{tab}got_agents = True\n"
111
+ f"{space}{tab}{tab}{tab}try:\n"
112
+ f"{space}{tab}{tab}{tab}{tab}should_continue = on_event(event, known_agents)\n"
113
+ f"{space}{tab}{tab}{tab}except BaseException as e:\n{_stop_logging(space, is_async)}"
114
+ f"{space}{tab}{tab}{tab}{tab}store_error(e)\n"
115
+ f"{space}{tab}{tab}{tab}{tab}raise SystemExit(\n"
116
+ f'{space}{tab}{tab}{tab}{tab}{tab}"Error in event handler: " + str(e)\n'
117
+ f"{space}{tab}{tab}{tab}{tab}) from e\n"
100
118
  )
101
119
  content += (
102
- f'{tab} if event.type == "run_completion":\n'
103
- f"{tab} break\n"
104
- f"{tab} if not should_continue:\n"
105
- f'{tab} raise SystemExit("Event handler stopped processing")\n'
120
+ f'{space}{tab}{tab}{tab}if getattr(event, "type") == "run_completion":\n'
121
+ f"{space}{tab}{tab}{tab}{tab}break\n"
122
+ f"{space}{tab}{tab}{tab}if not should_continue:\n{_stop_logging(space, is_async)}"
123
+ f"{space}{tab}{tab}{tab}{tab}{await_str}store_error()\n"
124
+ f'{space}{tab}{tab}{tab}{tab}raise SystemExit("Event handler stopped processing")\n'
106
125
  )
107
- content += get_result_dicts_string(tab, is_async)
126
+ content += get_result_dicts_string(space, is_async)
108
127
  content += (
109
- f"{tab}else:\n"
110
- f"{tab} if not isinstance(results, list):\n"
111
- f"{tab} results = [results] # pylint: disable=redefined-variable-type\n"
112
- f"{tab} for index, result in enumerate(results):\n"
128
+ f"{space}else:\n{space}{tab}for index, result in enumerate(results):\n"
129
+ f"{space}{tab}{tab}result_events = []\n"
113
130
  )
114
131
  if is_async:
115
- content += f"{tab} await result.process()\n"
132
+ content += (
133
+ f"{space}{tab}{tab}await result.process()\n"
134
+ f"{space}{tab}{tab}async for event in result.events:\n"
135
+ )
116
136
  else:
117
- content += f"{tab} result.process()\n"
118
- content += get_result_dicts_string(tab, is_async)
119
-
137
+ content += (
138
+ f"{space}{tab}{tab}result.process()\n"
139
+ f"{space}{tab}{tab}for event in result.events:\n"
140
+ )
141
+ content += _add_event_dump(space, tab)
142
+ content += get_result_dicts_string(space, is_async)
120
143
  return content
121
144
 
122
145
 
146
+ def _add_event_dump(space: str, tab: str) -> str:
147
+ return (
148
+ f"{space}{tab}{tab}{tab}try:\n"
149
+ f'{space}{tab}{tab}{tab}{tab}result_events.append(event.model_dump(mode="json", fallback=str))\n'
150
+ f"{space}{tab}{tab}{tab}except BaseException: # pylint: disable=broad-exception-caught\n"
151
+ f"{space}{tab}{tab}{tab}{tab}pass\n"
152
+ )
153
+
154
+
123
155
  def get_result_dicts_string(tab: str, is_async: bool) -> str:
124
156
  """Get the result dicts string.
125
157
 
@@ -135,24 +167,33 @@ def get_result_dicts_string(tab: str, is_async: bool) -> str:
135
167
  str
136
168
  The result dicts string.
137
169
  """
170
+ # fmt: off
138
171
  space = f"{tab} "
139
- flow_content = f"{space}result_dict = {{\n"
172
+ flow_content = f"{space}result_cost = "
173
+ if is_async:
174
+ flow_content += "await "
175
+ flow_content += "result.cost\n"
176
+ flow_content += f"{space}result_context_variables = "
177
+ if is_async:
178
+ flow_content += "await "
179
+ flow_content += "result.context_variables\n"
180
+ flow_content += f"{space}result_dict = {{\n"
140
181
  flow_content += f'{space} "index": index,\n'
182
+ flow_content += f'{space} "uuid": str(result.uuid),\n'
183
+ flow_content += f'{space} "events": result_events,\n'
141
184
  if is_async:
142
185
  flow_content += f'{space} "messages": await result.messages,\n'
143
186
  flow_content += f'{space} "summary": await result.summary,\n'
144
- flow_content += f'{space} "cost": (await result.cost).model_dump(mode="json", fallback=str) if await result.cost else None,\n'
145
- flow_content += f'{space} "context_variables": (await result.context_variables).model_dump(mode="json", fallback=str) if await result.context_variables else None,\n'
146
- flow_content += (
147
- f'{space} "last_speaker": await result.last_speaker,\n'
148
- )
187
+ flow_content += f'{space} "cost": result_cost.model_dump(mode="json", fallback=str) if result_cost else None,\n'
188
+ flow_content += f'{space} "context_variables": result_context_variables.model_dump(mode="json", fallback=str) if result_context_variables else None,\n'
189
+ flow_content += f'{space} "last_speaker": await result.last_speaker,\n'
149
190
  else:
150
191
  flow_content += f'{space} "messages": result.messages,\n'
151
192
  flow_content += f'{space} "summary": result.summary,\n'
152
- flow_content += f'{space} "cost": result.cost.model_dump(mode="json", fallback=str) if result.cost else None,\n'
153
- flow_content += f'{space} "context_variables": result.context_variables.model_dump(mode="json", fallback=str) if result.context_variables else None,\n'
193
+ flow_content += f'{space} "cost": result_cost.model_dump(mode="json", fallback=str) if result_cost else None,\n'
194
+ flow_content += f'{space} "context_variables": result_context_variables.model_dump(mode="json", fallback=str) if result_context_variables else None,\n'
154
195
  flow_content += f'{space} "last_speaker": result.last_speaker,\n'
155
- flow_content += f'{space} "uuid": str(result.uuid),\n'
156
196
  flow_content += f"{space}}}\n"
157
197
  flow_content += f"{space}result_dicts.append(result_dict)\n"
198
+ # fmt: on
158
199
  return flow_content
@@ -40,19 +40,19 @@ def export_group_chats(
40
40
  tuple[str, str]
41
41
  The group chat string and the import string.
42
42
  """
43
- tab = " " * tabs
44
- run_group_chat = "run_group_chat"
43
+ space = " " * tabs
44
+ run_group_chat: str = "run_group_chat"
45
45
  if is_async:
46
46
  run_group_chat = "await a_run_group_chat"
47
47
  manager_name = agent_names[manager.id]
48
48
  pattern_name = f"{manager_name}_pattern"
49
- content = f"{tab}results = {run_group_chat}(" + "\n"
50
- content += f"{tab} pattern={pattern_name}," + "\n"
49
+ content = f"{space}results = {run_group_chat}(" + "\n"
50
+ content += f"{space} pattern={pattern_name}," + "\n"
51
51
  if initial_chat:
52
- content += f"{tab} messages={json.dumps(initial_chat)}," + "\n"
52
+ content += f"{space} messages={json.dumps(initial_chat)}," + "\n"
53
53
  else:
54
- content += f'{tab} messages="",\n'
55
- content += f"{tab} max_rounds={manager.data.max_round},\n"
56
- content += f"{tab})\n"
57
- content += get_event_handler_string(tab=tab, is_async=is_async)
54
+ content += f'{space} messages="",\n'
55
+ content += f"{space} max_rounds={manager.data.max_round},\n"
56
+ content += f"{space})\n"
57
+ content += get_event_handler_string(space=space, is_async=is_async)
58
58
  return content
@@ -3,7 +3,7 @@
3
3
  # pylint: disable=too-many-locals
4
4
  """Nested chats exporting."""
5
5
 
6
- from typing import Any, Callable, Optional
6
+ from typing import Any, Callable
7
7
 
8
8
  from waldiez.models import (
9
9
  WaldiezAgent,
@@ -73,7 +73,7 @@ def export_nested_chat_registration(
73
73
  )
74
74
  content += f"{var_name}: list[dict[str, Any]] = {chat_queue}" + "\n"
75
75
  content += f"""
76
- {agent_name}.register_nested_chats( # pyright: ignore
76
+ {agent_name}.register_nested_chats(
77
77
  trigger={trigger_names},
78
78
  chat_queue={var_name},
79
79
  use_async={is_async},
@@ -125,7 +125,7 @@ def get_nested_chat_message_string(
125
125
  agent_names: dict[str, str],
126
126
  chat_names: dict[str, str],
127
127
  serializer: Callable[..., str],
128
- ) -> tuple[str, Optional[str]]:
128
+ ) -> tuple[str, str | None]:
129
129
  """Get the nested chat message string.
130
130
 
131
131
  Parameters
@@ -145,10 +145,10 @@ def get_nested_chat_message_string(
145
145
 
146
146
  Returns
147
147
  -------
148
- tuple[str, Optional[str]]
148
+ tuple[str, str | None]
149
149
  The message string and the method name if the message is a method.
150
150
  """
151
- sender_name: Optional[str] = None
151
+ sender_name: str | None = None
152
152
  sender_id = waldiez_chat.target if message.is_reply else waldiez_chat.source
153
153
  recipient_id = (
154
154
  waldiez_chat.source if message.is_reply else waldiez_chat.target
@@ -240,7 +240,7 @@ def get_chat_nested_string(
240
240
  chat: WaldiezChat,
241
241
  is_reply: bool,
242
242
  chat_names: dict[str, str],
243
- ) -> tuple[str, Optional[str]]:
243
+ ) -> tuple[str, str | None]:
244
244
  """Get the nested chat message.
245
245
 
246
246
  Parameters
@@ -254,7 +254,7 @@ def get_chat_nested_string(
254
254
 
255
255
  Returns
256
256
  -------
257
- tuple[str, Optional[str]]
257
+ tuple[str, str | None]
258
258
  If the message is a string, the message content and None.
259
259
  If the message is a method, the method name and the method content.
260
260
  If the message is None, 'None' and None.
@@ -68,7 +68,7 @@ def export_sequential_chat(
68
68
  additional_methods_string += additional_methods
69
69
  content += "\n" + f"{tab} {chat_string}"
70
70
  content += "\n" + " " * tabs + "])\n"
71
- content += get_event_handler_string(tab=tab, is_async=is_async)
71
+ content += get_event_handler_string(space=tab, is_async=is_async)
72
72
  return content, additional_methods_string
73
73
 
74
74
 
@@ -157,7 +157,7 @@ def get_simple_chat_string(
157
157
  )
158
158
  chat_string += message_arg
159
159
  chat_string += "\n" + f"{tab})" + "\n"
160
- chat_string += get_event_handler_string(tab=tab, is_async=is_async)
160
+ chat_string += get_event_handler_string(space=tab, is_async=is_async)
161
161
  return chat_string, additional_methods_string
162
162
 
163
163
 
@@ -212,7 +212,7 @@ def get_empty_simple_chat_string(
212
212
  )
213
213
  content += message_arg
214
214
  content += f"{tab})" + "\n"
215
- content += get_event_handler_string(tab=tab, is_async=is_async)
215
+ content += get_event_handler_string(space=tab, is_async=is_async)
216
216
  return content, ""
217
217
 
218
218