Merge pull request #11 from tanevanwifferen/master

no-intermediates feature

Author: adhikasp

README.md

@@ -101,6 +101,15 @@ To bypass tool confirmation requirements, use the `--no-confirmations` flag:
 $ llm --no-confirmations "What is the top article on hackernews today?"
 ```
 
+To use in bash scripts, add the --no-intermediates, so it doesn't print intermediate messages, only the concluding end message.
+```bash
+$ llm --no-intermediates "What is the time in Tokyo right now?"
+```
+
+```bash
+$ llm --no-confirmations "What is the top article on hackernews today?"
+```
+
 ### Continuation
 
 Add a `c ` prefix to your message to continue the last conversation.

src/mcp_client_cli/cli.py

@@ -46,6 +46,7 @@ class AgentState(TypedDict):
     today_datetime: str
     # The user's memories.
     memories: str = "no memories"
+    remaining_steps: int = 5
 
 async def run() -> None:
     """Run the LLM agent."""
@@ -100,6 +101,8 @@ def setup_argument_parser() -> argparse.Namespace:
                        help='Print output as raw text instead of parsing markdown')
     parser.add_argument('--no-tools', action='store_true',
                        help='Do not add any tools')
+    parser.add_argument('--no-intermediates', action='store_true',
+                       help='Only print the final message')
     parser.add_argument('--show-memories', action='store_true',
                        help='Show user memories')
     return parser.parse_args()
@@ -234,9 +237,10 @@ async def handle_conversation(args: argparse.Namespace, query: HumanMessage,
             messages=[query], 
             today_datetime=datetime.now().isoformat(),
             memories=formatted_memories,
+            remaining_steps=3
         )
 
-        output = OutputHandler(text_only=args.text_only)
+        output = OutputHandler(text_only=args.text_only, only_last_message=args.no_intermediates)
         output.start()
         try:
             async for chunk in agent_executor.astream(

src/mcp_client_cli/output.py

@@ -1,13 +1,16 @@
+import json
 from langchain_core.messages import BaseMessage, AIMessage, AIMessageChunk, ToolMessage
 from rich.console import Console, ConsoleDimensions
 from rich.live import Live
 from rich.markdown import Markdown
 from rich.prompt import Confirm
 
 class OutputHandler:
-    def __init__(self, text_only: bool = False):
+    def __init__(self, text_only: bool = False, only_last_message: bool = False):
         self.console = Console()
         self.text_only = text_only
+        self.only_last_message = only_last_message
+        self.last_message = ""
         if self.text_only:
             self.md = ""
         else:
@@ -26,6 +29,9 @@ def start(self):
 
     def update(self, chunk: any):
         self.md = self._parse_chunk(chunk, self.md)
+        if(self.only_last_message and self.text_only):
+            # when only_last_message, we print in finish()
+            return
         if self.text_only:
             self.console.print(self._parse_chunk(chunk), end="")
         else:
@@ -36,9 +42,13 @@ def update(self, chunk: any):
 
     def update_error(self, error: Exception):
         import traceback
-        self.md += f"Error: {error}\n\nStack trace:\n```\n{traceback.format_exc()}```"
+        error = f"Error: {error}\n\nStack trace:\n```\n{traceback.format_exc()}```"
+        self.md += error;
+        if(self.only_last_message):
+            self.console.print(error)
+            return;
         if self.text_only:
-            self.console.print(self.md)
+            self.console.print_exception(self.md)
         else:
             partial_md = self._truncate_md_to_fit(self.md, self.console.size)
             self._live.update(Markdown(partial_md), refresh=True)
@@ -63,9 +73,12 @@ def confirm_tool_call(self, config: dict, chunk: any) -> bool:
 
     def finish(self):
         self.stop()
-        if not self.text_only:
+        to_print = self.last_message if self.only_last_message else Markdown(self.md)
+        if not self.text_only and not self.only_last_message:
             self.console.clear()
             self.console.print(Markdown(self.md))
+        if self.only_last_message:
+            self.console.print(to_print)
 
     def _parse_chunk(self, chunk: any, md: str = "") -> str:
         """
@@ -77,6 +90,7 @@ def _parse_chunk(self, chunk: any, md: str = "") -> str:
             message_chunk = chunk[1][0]  # Get the message content
             if isinstance(message_chunk, AIMessageChunk):
                 content = message_chunk.content
+                self.last_message += content
                 if isinstance(content, str):
                     md += content
                 elif isinstance(content, list) and len(content) > 0 and isinstance(content[0], dict) and "text" in content[0]:
@@ -85,6 +99,7 @@ def _parse_chunk(self, chunk: any, md: str = "") -> str:
         elif isinstance(chunk, dict) and "messages" in chunk:
             # Print a newline after the complete message
             md += "\n"
+            self.last_message = ""
         elif isinstance(chunk, tuple) and chunk[0] == "values":
             message: BaseMessage = chunk[1]['messages'][-1]
             if isinstance(message, AIMessage) and message.tool_calls:
@@ -108,6 +123,7 @@ def _parse_chunk(self, chunk: any, md: str = "") -> str:
                             lines.append(f"{arg}: {value}")
                     lines.append("```\n")
                     md += "\n".join(lines)
+                self.last_message = ""
             elif isinstance(message, ToolMessage) and message.status != "success":
                 md += "Failed call with error:"
                 md += f"\n\n{message.content}"