diff --git a/src/pdl/pdl.py b/src/pdl/pdl.py index 7745ca38..845741a5 100644 --- a/src/pdl/pdl.py +++ b/src/pdl/pdl.py @@ -174,9 +174,9 @@ def main(): ) parser.add_argument( "--stream", - choices=["result", "background", "none"], - default="result", - help="stream the result, the background messages, or nothing on the standard output", + choices=["context", "none"], + default="context", + help="stream the background context, or nothing on the standard output", ) parser.add_argument( "-t", @@ -244,10 +244,7 @@ def main(): validate_scope(initial_scope) match args.stream: - case "result": - stream_result = True - stream_background = False - case "background": + case "context": stream_result = False stream_background = True case "none": diff --git a/src/pdl/pdl_interpreter.py b/src/pdl/pdl_interpreter.py index 47a08e51..651c3fec 100644 --- a/src/pdl/pdl_interpreter.py +++ b/src/pdl/pdl_interpreter.py @@ -179,16 +179,8 @@ def generate( prog, loc = parse_file(pdl_file) if state is None: state = InterpreterState(cwd=Path(pdl_file).parent) - result, _, _, trace = process_prog(state, initial_scope, prog, loc) - if not state.yield_result: - if state.yield_background: - print("\n----------------") - if result is None: - print() - else: - print(stringify(result)) - else: - print() + _, _, _, trace = process_prog(state, initial_scope, prog, loc) + print() if trace_file: write_trace(trace_file, trace) except PDLParseError as exc: diff --git a/src/pdl/pdl_scheduler.py b/src/pdl/pdl_scheduler.py index af59a2f6..9f47a48d 100644 --- a/src/pdl/pdl_scheduler.py +++ b/src/pdl/pdl_scheduler.py @@ -53,6 +53,20 @@ def color_of(kind: BlockKind): return color +def color_of_role(role: str): + color: Optional[Color] = None + match role: + case "assistant": + color = "green" + case "user": + color = None + case "system": + color = "cyan" + case "available_tools": + color = "magenta" + return color + + def yield_result(result: Any, kind: BlockKind) -> None: if color_of(kind) is None: text = stringify(result) @@ -62,6 +76,7 @@ def yield_result(result: Any, kind: BlockKind) -> None: _LAST_ROLE = None +ROLE_COLOR = "blue" def yield_background(background) -> None: @@ -72,5 +87,10 @@ def yield_background(background) -> None: background = background[1:] else: s = "\n" - s += "\n".join([f"{msg['role']}: {msg['content']}" for msg in background]) + s += "\n".join( + [ + f"{colored(msg['role'], ROLE_COLOR)}: {colored(msg['content'], color_of_role(msg['role']))}" + for msg in background + ] + ) print(s, end="", flush=True)