waldiez 0.1.17__py3-none-any.whl → 0.1.19__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.

waldiez/__init__.py CHANGED
@@ -1,17 +1,50 @@
1
1
  """Waldiez package."""
2
2
 
3
+ import logging
4
+
3
5
  from ._version import __version__
4
6
  from .conflict_checker import check_conflicts
5
7
  from .exporter import WaldiezExporter
6
8
  from .models import Waldiez
7
9
  from .runner import WaldiezRunner
8
10
 
11
+
12
+ # pylint: disable=too-few-public-methods
13
+ class FlamlFilter(logging.Filter):
14
+ """Filter out flaml.automl is not available message."""
15
+
16
+ def filter(self, record: logging.LogRecord) -> bool:
17
+ """Filter out flaml.automl is not available message.
18
+
19
+ this is just annoying:
20
+
21
+ ```
22
+ flaml.automl is not available.
23
+ Please install flaml[automl] to enable AutoML functionalities.
24
+ ```
25
+
26
+ Parameters
27
+ ----------
28
+ record : logging.LogRecord
29
+ Log record to filter.
30
+
31
+ Returns
32
+ -------
33
+ bool
34
+ Whether to filter out the log record.
35
+ """
36
+ return "flaml.automl is not available" not in record.getMessage()
37
+
38
+
9
39
  # flag to check if ag2 and autogen-agentchat
10
40
  # are installed at the same time
11
41
  __WALDIEZ_CHECKED_FOR_CONFLICTS = False
42
+ # flag to handle flaml logging
43
+ # suppress the annoying message about flaml.automl
44
+ __WALDIEZ_HANDLED_FLAML_LOGGING = False
12
45
 
13
46
 
14
- def _check_conflicts_once() -> None:
47
+ def _check_conflicts() -> None:
15
48
  """Check for conflicts once."""
16
49
  # pylint: disable=global-statement
17
50
  global __WALDIEZ_CHECKED_FOR_CONFLICTS
@@ -20,8 +53,18 @@ def _check_conflicts_once() -> None:
20
53
  __WALDIEZ_CHECKED_FOR_CONFLICTS = True
21
54
 
22
55
 
23
- _check_conflicts_once()
56
+ def _handle_flaml_logging() -> None:
57
+ """Handle flaml logging once."""
58
+ # pylint: disable=global-statement
59
+ global __WALDIEZ_HANDLED_FLAML_LOGGING
60
+ if __WALDIEZ_HANDLED_FLAML_LOGGING is False:
61
+ __WALDIEZ_HANDLED_FLAML_LOGGING = True
62
+ flam_logger = logging.getLogger("flaml")
63
+ flam_logger.addFilter(FlamlFilter())
64
+
24
65
 
66
+ _check_conflicts()
67
+ _handle_flaml_logging()
25
68
 
26
69
  __all__ = [
27
70
  "Waldiez",
waldiez/_version.py CHANGED
@@ -1,3 +1,3 @@
1
1
  """Version information for Waldiez."""
2
2
 
3
- __version__ = "0.1.17"
3
+ __version__ = "0.1.19"
waldiez/cli.py CHANGED
@@ -90,7 +90,7 @@ def run(
90
90
  raise typer.Exit(code=1) from error
91
91
  waldiez = Waldiez.from_dict(data)
92
92
  runner = WaldiezRunner(waldiez)
93
- results = runner.run(stream=None, output_path=output_path)
93
+ results = runner.run(output_path=output_path)
94
94
  logger = _get_logger()
95
95
  if isinstance(results, list):
96
96
  logger.info("Results:")
waldiez/runner.py CHANGED
@@ -7,16 +7,18 @@ return the results. Before running the flow, any additional environment
7
7
  variables specified in the waldiez file are set.
8
8
  """
9
9
 
10
+ # pylint: disable=import-outside-toplevel,reimported
11
+
10
12
  import datetime
11
13
  import importlib.util
12
14
  import io
13
15
  import os
14
16
  import shutil
17
+ import site
15
18
  import subprocess # nosemgrep # nosec
16
19
  import sys
17
20
  import tempfile
18
21
  from contextlib import contextmanager
19
- from contextvars import ContextVar
20
22
  from pathlib import Path
21
23
  from types import TracebackType
22
24
  from typing import (
@@ -36,8 +38,6 @@ from .models.waldiez import Waldiez
36
38
  if TYPE_CHECKING:
37
39
  from autogen import ChatResult # type: ignore
38
40
 
39
- from .io import WaldiezIOStream
40
-
41
41
 
42
42
  @contextmanager
43
43
  def _chdir(to: Union[str, Path]) -> Iterator[None]:
@@ -71,10 +71,8 @@ class WaldiezRunner:
71
71
  self._waldiez = waldiez
72
72
  self._running = False
73
73
  self._file_path = file_path
74
- self._stream: ContextVar[Optional["WaldiezIOStream"]] = ContextVar(
75
- "waldiez_stream", default=None
76
- )
77
74
  self._exporter = WaldiezExporter(waldiez)
75
+ self._called_install_requirements = False
78
76
 
79
77
  @classmethod
80
78
  def load(
@@ -136,10 +134,6 @@ class WaldiezRunner:
136
134
  """Exit the context manager."""
137
135
  if self._running:
138
136
  self._running = False
139
- token = self._stream.get()
140
- if token is not None:
141
- self._stream.reset(token)
142
- del token
143
137
 
144
138
  @property
145
139
  def waldiez(self) -> Waldiez:
@@ -151,24 +145,15 @@ class WaldiezRunner:
151
145
  """Get the running status."""
152
146
  return self._running
153
147
 
154
- def _get_print_function(self) -> Callable[..., None]:
155
- """Get the print function."""
156
- token = self._stream.get()
157
- if token is not None:
158
- return token.print
159
- return print
160
-
161
- def _install_requirements(self) -> None:
148
+ def install_requirements(self) -> None:
162
149
  """Install the requirements for the flow."""
150
+ self._called_install_requirements = True
151
+ printer = get_printer()
163
152
  extra_requirements = set(
164
153
  req for req in self.waldiez.requirements if req not in sys.modules
165
154
  )
166
155
  if extra_requirements:
167
- print_function = self._get_print_function()
168
- # pylint: disable=inconsistent-quotes
169
- print_function(
170
- f"Installing requirements: {', '.join(extra_requirements)}"
171
- )
156
+ printer(f"Installing requirements: {', '.join(extra_requirements)}")
172
157
  pip_install = [sys.executable, "-m", "pip", "install"]
173
158
  if not in_virtualenv():
174
159
  pip_install.append("--user")
@@ -180,18 +165,17 @@ class WaldiezRunner:
180
165
  ) as proc:
181
166
  if proc.stdout:
182
167
  for line in io.TextIOWrapper(proc.stdout, encoding="utf-8"):
183
- print_function(line.strip())
168
+ printer(line.strip())
184
169
  if proc.stderr:
185
170
  for line in io.TextIOWrapper(proc.stderr, encoding="utf-8"):
186
- print_function(line.strip())
187
- print_function(
188
- "Requirements installed.\n"
189
- "NOTE: If new packages were added and you are using Jupyter, "
190
- "you might need to restart the kernel."
191
- )
171
+ printer(line.strip())
172
+ refresh_environment()
192
173
 
174
+ @staticmethod
193
175
  def _after_run(
194
- self, temp_dir: Path, output_path: Optional[Union[str, Path]]
176
+ temp_dir: Path,
177
+ output_path: Optional[Union[str, Path]],
178
+ printer: Callable[..., None],
195
179
  ) -> None:
196
180
  if output_path:
197
181
  destination_dir = Path(output_path).parent
@@ -202,9 +186,7 @@ class WaldiezRunner:
202
186
  )
203
187
  destination_dir.mkdir(parents=True, exist_ok=True)
204
188
  # copy the contents of the temp dir to the destination dir
205
- self._get_print_function()(
206
- f"Copying the results to {destination_dir}"
207
- )
189
+ printer(f"Copying the results to {destination_dir}")
208
190
  for item in temp_dir.iterdir():
209
191
  # skip cache files
210
192
  if (
@@ -257,6 +239,16 @@ class WaldiezRunner:
257
239
  Union[ChatResult, List[ChatResult]]
258
240
  The result(s) of the chat(s).
259
241
  """
242
+ if not self._called_install_requirements:
243
+ self.install_requirements()
244
+ else:
245
+ refresh_environment()
246
+ printer = get_printer()
247
+ printer(
248
+ "Requirements installed.\n"
249
+ "NOTE: If new packages were added and you are using Jupyter, "
250
+ "you might need to restart the kernel."
251
+ )
260
252
  results: Union["ChatResult", List["ChatResult"]] = []
261
253
  if not uploads_root:
262
254
  uploads_root = Path(tempfile.mkdtemp())
@@ -284,32 +276,15 @@ class WaldiezRunner:
284
276
  old_vars = self._set_env_vars()
285
277
  module = importlib.util.module_from_spec(spec)
286
278
  spec.loader.exec_module(module)
287
- print_function = self._get_print_function()
288
- print_function("Starting workflow...")
279
+ printer("<Waldiez> - Starting workflow...")
289
280
  results = module.main()
290
281
  sys.path.pop(0)
291
282
  self._reset_env_vars(old_vars)
292
- self._after_run(temp_dir, output_path)
283
+ self._after_run(temp_dir, output_path, printer)
293
284
  return results
294
285
 
295
- def _run(
296
- self,
297
- output_path: Optional[Union[str, Path]],
298
- uploads_root: Optional[Union[str, Path]],
299
- ) -> Union["ChatResult", List["ChatResult"]]:
300
- self._install_requirements()
301
- token = self._stream.get()
302
- if token is not None:
303
- # pylint: disable=import-outside-toplevel
304
- from .io import WaldiezIOStream
305
-
306
- with WaldiezIOStream.set_default(token):
307
- return self._do_run(output_path, uploads_root)
308
- return self._do_run(output_path, uploads_root)
309
-
310
286
  def run(
311
287
  self,
312
- stream: Optional["WaldiezIOStream"] = None,
313
288
  output_path: Optional[Union[str, Path]] = None,
314
289
  uploads_root: Optional[Union[str, Path]] = None,
315
290
  ) -> Union["ChatResult", List["ChatResult"]]:
@@ -317,8 +292,6 @@ class WaldiezRunner:
317
292
 
318
293
  Parameters
319
294
  ----------
320
- stream : Optional[WaldiezIOStream], optional
321
- The stream to use, by default None.
322
295
  output_path : Optional[Union[str, Path]], optional
323
296
  The output path, by default None.
324
297
  uploads_root : Optional[Union[str, Path]], optional
@@ -337,14 +310,11 @@ class WaldiezRunner:
337
310
  if self._running is True:
338
311
  raise RuntimeError("Workflow already running")
339
312
  self._running = True
340
- token = self._stream.set(stream)
341
313
  file_path = output_path or self._file_path
342
314
  try:
343
- return self._run(file_path, uploads_root)
315
+ return self._do_run(file_path, uploads_root)
344
316
  finally:
345
317
  self._running = False
346
- self._stream.reset(token)
347
- del token
348
318
 
349
319
 
350
320
  def in_virtualenv() -> bool:
@@ -358,3 +328,37 @@ def in_virtualenv() -> bool:
358
328
  return hasattr(sys, "real_prefix") or (
359
329
  hasattr(sys, "base_prefix") and sys.base_prefix != sys.prefix
360
330
  )
331
+
332
+
333
+ def refresh_environment() -> None:
334
+ """Refresh the environment."""
335
+ # backup the default IOStream
336
+ from autogen.io import IOStream # type: ignore
337
+
338
+ default_io_stream = IOStream.get_default()
339
+ site.main()
340
+ # pylint: disable=import-outside-toplevel
341
+ modules_to_reload = [
342
+ mod for mod in sys.modules if mod.startswith("autogen")
343
+ ]
344
+ for mod in modules_to_reload:
345
+ del sys.modules[mod]
346
+ import autogen
347
+ from autogen.io import IOStream
348
+
349
+ importlib.reload(autogen)
350
+ # restore the default IOStream
351
+ IOStream.set_global_default(default_io_stream)
352
+
353
+
354
+ def get_printer() -> Callable[..., None]:
355
+ """Get the printer function.
356
+
357
+ Returns
358
+ -------
359
+ Callable[..., None]
360
+ The printer function.
361
+ """
362
+ from autogen.io import IOStream
363
+
364
+ return IOStream.get_default().print
@@ -1,10 +1,11 @@
1
- Metadata-Version: 2.3
1
+ Metadata-Version: 2.4
2
2
  Name: waldiez
3
- Version: 0.1.17
3
+ Version: 0.1.19
4
4
  Summary: waldiez
5
5
  Project-URL: homepage, https://waldiez.github.io/waldiez/
6
6
  Project-URL: repository, https://github.com/waldiez/waldiez.git
7
7
  Author-email: Panagiotis Kasnesis <pkasnesis@thingenious.io>, Lazaros Toumanidis <laztoum@protonmail.com>, Stella Ioannidou <stella@humancentered.gr>
8
+ License-File: LICENSE
8
9
  Classifier: Development Status :: 3 - Alpha
9
10
  Classifier: Intended Audience :: Developers
10
11
  Classifier: Intended Audience :: Science/Research
@@ -60,7 +61,7 @@ Requires-Dist: mdx-include==1.4.2; extra == 'docs'
60
61
  Requires-Dist: mdx-truly-sane-lists==1.3; extra == 'docs'
61
62
  Requires-Dist: mkdocs-jupyter==0.25.1; extra == 'docs'
62
63
  Requires-Dist: mkdocs-macros-plugin==1.3.7; extra == 'docs'
63
- Requires-Dist: mkdocs-material==9.5.48; extra == 'docs'
64
+ Requires-Dist: mkdocs-material==9.5.49; extra == 'docs'
64
65
  Requires-Dist: mkdocs-minify-html-plugin==0.2.3; extra == 'docs'
65
66
  Requires-Dist: mkdocs==1.6.1; extra == 'docs'
66
67
  Requires-Dist: mkdocstrings-python==1.12.2; extra == 'docs'
@@ -88,7 +89,7 @@ To a python script or a jupyter notebook with the corresponding [ag2](https://gi
88
89
 
89
90
  - Convert .waldiez flows to .py or .ipynb
90
91
  - Run a .waldiez flow
91
- - Provide a custom [IOSStream](https://ag2ai.github.io/ag2/docs/reference/io/base#iostream) to handle input and output.
92
+ - Store the runtime logs of a flow to csv for further analysis
92
93
 
93
94
  ## Installation
94
95
 
@@ -178,73 +179,6 @@ runner = WaldiezRunner.load(flow_path)
178
179
  runner.run(output_path=output_path)
179
180
  ```
180
181
 
181
- #### Run a flow with a custom IOStream
182
-
183
- ```python
184
- # Run the flow with a custom IOStream
185
- # In case the standard 'input' and 'print' functions cannot be used
186
- import time
187
- import threading
188
-
189
- from typing import Any
190
-
191
- from waldiez import WaldiezRunner
192
- from waldiez.io import WaldiezIOStream
193
-
194
- flow_path = "/path/to/a/flow.waldiez"
195
- output_path = "/path/to/an/output.py"
196
-
197
-
198
- def custom_print_function(*args: Any, sep: str = " ", **kwargs: Any) -> None:
199
- """Custom print function."""
200
- print(*args, sep=sep, **kwargs)
201
-
202
-
203
- # Custom input handler
204
- class InputProcessorWrapper:
205
- """Wrapper input processor.
206
-
207
- To manage the interaction between the custom input processor and IOStream.
208
- """
209
-
210
- def __init__(self):
211
- self.stream = None # Placeholder for the WaldiezIOStream instance
212
- self.lock = threading.Lock() # Ensure thread-safe operations
213
-
214
- def custom_input_processor(self, prompt: str) -> None:
215
- """Simulate external input and send it back to the IOStream."""
216
- def external_input_simulation():
217
- with self.lock: # Ensure thread-safe access
218
- time.sleep(2) # Simulate delay for network input
219
- if self.stream:
220
- self.stream.set_input("Simulated external input")
221
- else:
222
- raise RuntimeError("Stream reference not set!")
223
-
224
- threading.Thread(target=external_input_simulation, daemon=True).start()
225
-
226
- def set_stream(self, stream: "WaldiezIOStream"):
227
- """Set the WaldiezIOStream instance."""
228
- with self.lock: # Ensure thread-safe setting of the stream
229
- self.stream = stream
230
-
231
- processor_wrapper = InputProcessorWrapper()
232
-
233
- stream = WaldiezIOStream(
234
- input_timeout=30,
235
- print_function=
236
- on_prompt_input=processor_wrapper.custom_input_processor,
237
- )
238
-
239
- # Link the processor wrapper to the WaldiezIOStream instance
240
- processor_wrapper.set_stream(custom_stream)
241
-
242
- with WaldiezIOStream.set_default(io_stream):
243
- runner = WaldiezRunner.load(flow_path)
244
- runner.run(stream=io_stream, output_path=output_path)
245
-
246
- ```
247
-
248
182
  ### Tools
249
183
 
250
184
  - [ag2 (formerly AutoGen)](https://github.com/ag2ai/ag2)
@@ -1,11 +1,11 @@
1
- waldiez/__init__.py,sha256=ASGmxP3JcASIv89_tMYhYxvW0FlHK1jM0Nfb1hMN-z0,722
1
+ waldiez/__init__.py,sha256=9BLMLUURndkGEwnSNF1RXg1XKm4WTYM3cUeuYo-M0qY,1936
2
2
  waldiez/__main__.py,sha256=mUQWu4CF03Jbbscvcfb_9EFTnIMliJJJmAuWf0sRRZU,110
3
- waldiez/_version.py,sha256=JZiuUXn7Efmws1BQsY7xqwPgrly1krkrpN-yh6tNz8s,63
4
- waldiez/cli.py,sha256=Ayp6Xu2HSBBFIouYh8cArjOjgXv5kbtZd_PYewcDklU,6347
3
+ waldiez/_version.py,sha256=EX_WkICFRi_lmrwAjPX_8gLn_pS_uIsuK3aso-wK5Hk,63
4
+ waldiez/cli.py,sha256=19LFqep8urzcLm6TzVSyJEOKI2Q8XXZbuuFsgV0k9LQ,6334
5
5
  waldiez/conflict_checker.py,sha256=E-w0TfTivDAVpNvjd_NSBeaaFsWtWAyxKoSz0-g6x2U,880
6
6
  waldiez/exporter.py,sha256=iKe-l_Me8NRWsXHIdBcrOrnLT9XIyp4iYi4HLuuj2jA,9342
7
7
  waldiez/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
8
- waldiez/runner.py,sha256=v4A4NHdSnh96TQ5Vx1vs8UrTsshc2oAwDZrEHmq4RkU,11537
8
+ waldiez/runner.py,sha256=10GoT27cxEJdrapmQ4ny9aN7ZcnmsNHApzwIft5ZxP0,11220
9
9
  waldiez/exporting/__init__.py,sha256=GMY7qTRpNmc7tpaCFNGLT5wX9eu26NvoNPuYX8MzP50,344
10
10
  waldiez/exporting/agents/__init__.py,sha256=v5KA112W_EFYwXE2TSBKYyO8rRKUOUAOpFS5CMSnfRs,110
11
11
  waldiez/exporting/agents/agent.py,sha256=d2ut2B90-aG4uBRhPWS7B4MKIvSJqQva6DjlXiqleVQ,7513
@@ -39,7 +39,6 @@ waldiez/exporting/utils/method_utils.py,sha256=7-RUMTylNM2W0iM1bPX2_Gn3553XZSl2s
39
39
  waldiez/exporting/utils/naming.py,sha256=VdoVODQduhXIs9hQFWUVEVqTaSyNDt7rkECsuIgXYwI,3196
40
40
  waldiez/exporting/utils/object_string.py,sha256=2kdIu4in3iUV92a2KbLWwp9POhvY-fzF_r2AGVnCKls,2166
41
41
  waldiez/exporting/utils/path_check.py,sha256=aO49sbk6hUHEr65UjNNUSKO5WCGnQitiT733W-kGVtI,1062
42
- waldiez/io/__init__.py,sha256=-w6ywIjG8ByxGtpSeRt1bjhnP4CZzpHf8WRWckjs4ng,5172
43
42
  waldiez/models/__init__.py,sha256=IMq8vzuAgmv92kHSSuZQLF38vVd31ojgOHonuHqWYj0,2888
44
43
  waldiez/models/waldiez.py,sha256=yI38WEXPJa-FmRF6Try-MabWoPCeqlxq8VUt0sD-fJg,9428
45
44
  waldiez/models/agents/__init__.py,sha256=3ZyVYBHMFzZjRMIdPrBF6HLg82LPAlEubL6utL6KhfU,1856
@@ -85,8 +84,8 @@ waldiez/models/model/model_data.py,sha256=VLPb60rJeZEgVZCjjkQGiwTrKz7OddVLrXCrEu
85
84
  waldiez/models/skill/__init__.py,sha256=rU88bajKOGMYoHFcE8MP0jW9H0MswbQmvz5wxS35BYE,169
86
85
  waldiez/models/skill/skill.py,sha256=fhsAI413an2_d4DBIkf7dzEuWk6rGs2t4sl97a4dj20,3473
87
86
  waldiez/models/skill/skill_data.py,sha256=RTWn8Od6w7g-nRIpsS29sqZ8sPm5dCPiK7-qXmU-KD4,815
88
- waldiez-0.1.17.dist-info/METADATA,sha256=xfStnh_3pvXf9OU87UM3VMr3l8mchce7g2-OnfCI3mg,10262
89
- waldiez-0.1.17.dist-info/WHEEL,sha256=C2FUgwZgiLbznR-k0b_5k3Ai_1aASOXDss3lzCUsUug,87
90
- waldiez-0.1.17.dist-info/entry_points.txt,sha256=9MQ8Y1rD19CU7UwjNPwoyTRpQsPs2QimjrtwTD0bD6k,44
91
- waldiez-0.1.17.dist-info/licenses/LICENSE,sha256=VQEHM6WMQLRu1qaGl3GWsoOknDwro-69eGo4NLIJPIM,1064
92
- waldiez-0.1.17.dist-info/RECORD,,
87
+ waldiez-0.1.19.dist-info/METADATA,sha256=w6B5mTtRqZ80Y42gnIU1MbjKtCOWK94WdXKTbhgKpw0,8161
88
+ waldiez-0.1.19.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
89
+ waldiez-0.1.19.dist-info/entry_points.txt,sha256=9MQ8Y1rD19CU7UwjNPwoyTRpQsPs2QimjrtwTD0bD6k,44
90
+ waldiez-0.1.19.dist-info/licenses/LICENSE,sha256=VQEHM6WMQLRu1qaGl3GWsoOknDwro-69eGo4NLIJPIM,1064
91
+ waldiez-0.1.19.dist-info/RECORD,,
@@ -1,4 +1,4 @@
1
1
  Wheel-Version: 1.0
2
- Generator: hatchling 1.26.3
2
+ Generator: hatchling 1.27.0
3
3
  Root-Is-Purelib: true
4
4
  Tag: py3-none-any
waldiez/io/__init__.py DELETED
@@ -1,142 +0,0 @@
1
- """Custom IOStream class to use with autogen.
2
-
3
- It is meant to be used when we want to use custom
4
- `print` and `input`. For example, when a websocket
5
- is used to trigger a UI element that requires user input.
6
- and sends back the user's input to the websocket. In the same
7
- way, we can use it to forward what is meant to be printed.
8
- """
9
-
10
- import threading
11
- from typing import Any, Callable, Optional
12
-
13
- from autogen.io import IOStream # type: ignore[import-untyped]
14
-
15
-
16
- class WaldiezIOStream(IOStream):
17
- """Custom IOStream class to handle `print` and `input` functions."""
18
-
19
- def __init__(
20
- self,
21
- input_timeout: float = 60.0,
22
- print_function: Optional[Callable[..., None]] = None,
23
- on_prompt_input: Optional[Callable[[str], None]] = None,
24
- ) -> None:
25
- """
26
- Initialize the IOStream.
27
-
28
- Parameters
29
- ----------
30
- input_timeout : float, optional
31
- The input timeout in seconds, by default 60.0.
32
- print_function : Optional[Callable[..., None]], optional
33
- The function to handle print operations, by default None.
34
- on_prompt_input : Optional[Callable[[str], None]], optional
35
- The function to call for processing input prompts, by default None.
36
-
37
- Notes
38
- -----
39
- - on_prompt_input: It does not return a string (like 'input' does).
40
- Instead, it is meant to be used to forward the prompt somewhere else
41
- (e.g., a websocket). When we get the input, we can call
42
- `waldiez_io_stream.set_input(input_data)` with the input data.
43
- """
44
- self.input_timeout = input_timeout # Timeout for input
45
- self.print_function = print_function # Custom print handler
46
- self._on_prompt_input = on_prompt_input # Custom input prompt handler
47
- self.current_input: Optional[str] = None # Store the current input
48
- self._input_event = threading.Event() # Event to signal input readiness
49
- self.allow_input = True # Flag to allow or block input setting
50
-
51
- def print(
52
- self,
53
- *objects: Any,
54
- sep: str = " ",
55
- end: str = "\n",
56
- flush: bool = False,
57
- ) -> None:
58
- """
59
- Mock the `print` function.
60
-
61
- Parameters
62
- ----------
63
- objects : Any
64
- The objects to print.
65
- sep : str, optional
66
- The separator, by default " ".
67
- end : str, optional
68
- The ending string, by default a new line.
69
- flush : bool, optional
70
- Whether to flush the output, by default False.
71
- """
72
- print_function: Callable[..., None] = self.print_function or print
73
- print_function(*objects, sep=sep, end=end, flush=flush)
74
-
75
- def input(self, prompt: str = "", *, password: bool = False) -> str:
76
- """
77
- Mock the `input` function with optional timeout handling.
78
-
79
- Parameters
80
- ----------
81
- prompt : str, optional
82
- The prompt to show, by default "".
83
- password : bool, optional
84
- Whether to hide the input as password (not used), by default False.
85
-
86
- Returns
87
- -------
88
- str
89
- The user's input or '\n' if timeout occurs.
90
- """
91
- _prompt = prompt or "Your input:"
92
- if _prompt in (">", "> "): # pragma: no cover
93
- _prompt = "Your input:"
94
- if prompt:
95
- if self._on_prompt_input:
96
- self._on_prompt_input(_prompt)
97
- self.print(_prompt, end="")
98
-
99
- # Only reset if no input is currently set
100
- # e.g. handle the case when we call set_input(..) before input(..)
101
- already_set = self._input_event.is_set()
102
- if not self._input_event.is_set():
103
- self.current_input = None # Reset previous input
104
- self._input_event.clear() # Clear the event before waiting
105
-
106
- # Wait for input or timeout
107
- if not self._input_event.wait(self.input_timeout):
108
- # Timeout occurred, return what we have so far
109
- to_return = (
110
- self.current_input if self.current_input is not None else "\n"
111
- )
112
- self.current_input = None
113
- self._input_event.clear()
114
- # if we had already set the input, return it
115
- return to_return if already_set else "\n"
116
-
117
- # Input is ready, return it
118
- to_return = (
119
- self.current_input if self.current_input is not None else "\n"
120
- )
121
- self.current_input = None
122
- self._input_event.clear() # Clear the event after waiting
123
- return to_return
124
-
125
- def set_input(self, value: str) -> None:
126
- """
127
- Set the input value and signal that input is ready.
128
-
129
- Parameters
130
- ----------
131
- value : str
132
- The value to set as input.
133
- """
134
- if self.allow_input: # Respect the allow_input flag
135
- self.current_input = value
136
- self._input_event.set() # Signal that input is ready
137
-
138
- def reset_state(self) -> None:
139
- """Reset the IOStream state for testing."""
140
- self.current_input = None
141
- self._input_event.clear()
142
- self.allow_input = True # Re-enable input setting