apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-21 11:11:03 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

Output token counts

Browse source
This commit is contained in:
Cyber MacGeddon 2026-04-12 16:02:59 +01:00
parent 1630346bd9
commit d1ac03f0ce
1 changed files with 27 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

31
trustgraph-cli/trustgraph/cli/invoke_llm.py
View file

@ -10,7 +10,8 @@ from trustgraph.api import Api
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
default_token = os.getenv("TRUSTGRAPH_TOKEN", None) 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 # Create API client
api = Api(url=url, token=token) api = Api(url=url, token=token)
@ -26,15 +27,30 @@ def query(url, flow_id, system, prompt, streaming=True, token=None):
) )
if streaming: if streaming:
# Stream output to stdout without newline last_chunk = None
for chunk in response: for chunk in response:
print(chunk.content, end="", flush=True) print(chunk.content, end="", flush=True)
# Add final newline after streaming last_chunk = chunk
print() 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: else:
# Non-streaming: print complete response
print(response.text) 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: finally:
# Clean up socket connection # Clean up socket connection
socket.close() socket.close()
@ -82,6 +98,12 @@ def main():
help='Disable streaming (default: streaming enabled)' 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() args = parser.parse_args()
try: try:
@ -93,6 +115,7 @@ def main():
prompt=args.prompt[0], prompt=args.prompt[0],
streaming=not args.no_streaming, streaming=not args.no_streaming,
token=args.token, token=args.token,
show_usage=args.show_usage,
) )
except Exception as e: except Exception as e: