From d1ac03f0ce953ccb7ba74d32353d02adf13e1079 Mon Sep 17 00:00:00 2001
From: Cyber MacGeddon <cybermaggedon@gmail.com>
Date: Sun, 12 Apr 2026 16:02:59 +0100
Subject: [PATCH] Output token counts

---
 trustgraph-cli/trustgraph/cli/invoke_llm.py | 31 ++++++++++++++++++---
 1 file changed, 27 insertions(+), 4 deletions(-)

diff --git a/trustgraph-cli/trustgraph/cli/invoke_llm.py b/trustgraph-cli/trustgraph/cli/invoke_llm.py
index 47603442..3bf521f6 100644
--- a/trustgraph-cli/trustgraph/cli/invoke_llm.py
+++ b/trustgraph-cli/trustgraph/cli/invoke_llm.py
@@ -10,7 +10,8 @@ from trustgraph.api import Api
 default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
 default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
 
-def query(url, flow_id, system, prompt, streaming=True, token=None):
+def query(url, flow_id, system, prompt, streaming=True, token=None,
+          show_usage=False):
 
     # Create API client
     api = Api(url=url, token=token)
@@ -26,15 +27,30 @@ def query(url, flow_id, system, prompt, streaming=True, token=None):
         )
 
         if streaming:
-            # Stream output to stdout without newline
+            last_chunk = None
             for chunk in response:
                 print(chunk.content, end="", flush=True)
-            # Add final newline after streaming
+                last_chunk = chunk
             print()
+
+            if show_usage and last_chunk:
+                print(
+                    f"Input tokens: {last_chunk.in_token}  "
+                    f"Output tokens: {last_chunk.out_token}  "
+                    f"Model: {last_chunk.model}",
+                    file=__import__('sys').stderr,
+                )
         else:
-            # Non-streaming: print complete response
             print(response.text)
 
+            if show_usage:
+                print(
+                    f"Input tokens: {response.in_token}  "
+                    f"Output tokens: {response.out_token}  "
+                    f"Model: {response.model}",
+                    file=__import__('sys').stderr,
+                )
+
     finally:
         # Clean up socket connection
         socket.close()
@@ -82,6 +98,12 @@ def main():
         help='Disable streaming (default: streaming enabled)'
     )
 
+    parser.add_argument(
+        '--show-usage',
+        action='store_true',
+        help='Show token usage and model on stderr'
+    )
+
     args = parser.parse_args()
 
     try:
@@ -93,6 +115,7 @@ def main():
             prompt=args.prompt[0],
             streaming=not args.no_streaming,
             token=args.token,
+            show_usage=args.show_usage,
         )
 
     except Exception as e: