Changed output to be only background context (#375)

* Changed output to be only background context

Signed-off-by: Mandana Vaziri <[email protected]>

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":

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:

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)