apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-04-25 00:16:23 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

Expose LLM token usage across all service layers (#782)

Browse source 
Expose LLM token usage (in_token, out_token, model) across all
service layers

Propagate token counts from LLM services through the prompt,
text-completion, graph-RAG, document-RAG, and agent orchestrator
pipelines to the API gateway and Python SDK. All fields are Optional
— None means "not available", distinguishing from a real zero count.

Key changes:

- Schema: Add in_token/out_token/model to TextCompletionResponse,
  PromptResponse, GraphRagResponse, DocumentRagResponse,
  AgentResponse

- TextCompletionClient: New TextCompletionResult return type. Split
  into text_completion() (non-streaming) and
  text_completion_stream() (streaming with per-chunk handler
  callback)

- PromptClient: New PromptResult with response_type
  (text/json/jsonl), typed fields (text/object/objects), and token
  usage. All callers updated.

- RAG services: Accumulate token usage across all prompt calls
  (extract-concepts, edge-scoring, edge-reasoning,
  synthesis). Non-streaming path sends single combined response
  instead of chunk + end_of_session.

- Agent orchestrator: UsageTracker accumulates tokens across
  meta-router, pattern prompt calls, and react reasoning. Attached
  to end_of_dialog.

- Translators: Encode token fields when not None (is not None, not truthy)

- Python SDK: RAG and text-completion methods return
  TextCompletionResult (non-streaming) or RAGChunk/AgentAnswer with
  token fields (streaming)

- CLI: --show-usage flag on tg-invoke-llm, tg-invoke-prompt,
  tg-invoke-graph-rag, tg-invoke-document-rag, tg-invoke-agent
This commit is contained in:
cybermaggedon 2026-04-13 14:38:34 +01:00 • committed by GitHub
parent 67cfa80836
commit 14e49d83c7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
60 changed files with 1252 additions and 577 deletions
Download patch file Download diff file Expand all files Collapse all files

20
trustgraph-cli/trustgraph/cli/invoke_agent.py
View file

@ -272,7 +272,8 @@ def question(
url, question, flow_id, user, collection,
plan=None, state=None, group=None, pattern=None,
verbose=False, streaming=True,
token=None, explainable=False, debug=False
token=None, explainable=False, debug=False,
show_usage=False
):
# Explainable mode uses the API to capture and process provenance events
if explainable:
@ -323,6 +324,7 @@ def question(
# Track last chunk type and current outputter for streaming
last_chunk_type = None
current_outputter = None
last_answer_chunk = None
for chunk in response:
chunk_type = chunk.chunk_type
@ -357,6 +359,7 @@ def question(
current_outputter.word_buffer = ""
elif chunk_type == "final-answer":
print(content, end="", flush=True)
last_answer_chunk = chunk
# Close any remaining outputter
if current_outputter:
@ -366,6 +369,14 @@ def question(
elif last_chunk_type == "final-answer":
print()
if show_usage and last_answer_chunk:
print(
f"Input tokens: {last_answer_chunk.in_token} "
f"Output tokens: {last_answer_chunk.out_token} "
f"Model: {last_answer_chunk.model}",
file=sys.stderr,
)
else:
# Non-streaming response - but agents use multipart messaging
# so we iterate through the chunks (which are complete messages, not text chunks)
@ -477,6 +488,12 @@ def main():
help='Show debug output for troubleshooting'
)
parser.add_argument(
'--show-usage',
action='store_true',
help='Show token usage and model on stderr'
)
args = parser.parse_args()
try:
@ -496,6 +513,7 @@ def main():
token = args.token,
explainable = args.explainable,
debug = args.debug,
show_usage = args.show_usage,
)
except Exception as e:

34
trustgraph-cli/trustgraph/cli/invoke_document_rag.py
View file

@ -99,7 +99,8 @@ def question_explainable(
def question(
url, flow_id, question_text, user, collection, doc_limit,
streaming=True, token=None, explainable=False, debug=False
streaming=True, token=None, explainable=False, debug=False,
show_usage=False
):
# Explainable mode uses the API to capture and process provenance events
if explainable:
@ -133,22 +134,40 @@ def question(
)
# Stream output
last_chunk = None
for chunk in response:
print(chunk, end="", flush=True)
print(chunk.content, end="", flush=True)
last_chunk = chunk
print() # Final newline
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=sys.stderr,
)
finally:
socket.close()
else:
# Use REST API for non-streaming
flow = api.flow().id(flow_id)
resp = flow.document_rag(
result = flow.document_rag(
query=question_text,
user=user,
collection=collection,
doc_limit=doc_limit,
)
print(resp)
print(result.text)
if show_usage:
print(
f"Input tokens: {result.in_token} "
f"Output tokens: {result.out_token} "
f"Model: {result.model}",
file=sys.stderr,
)
def main():
@ -219,6 +238,12 @@ def main():
help='Show debug output for troubleshooting'
)
parser.add_argument(
'--show-usage',
action='store_true',
help='Show token usage and model on stderr'
)
args = parser.parse_args()
try:
@ -234,6 +259,7 @@ def main():
token=args.token,
explainable=args.explainable,
debug=args.debug,
show_usage=args.show_usage,
)
except Exception as e:

33
trustgraph-cli/trustgraph/cli/invoke_graph_rag.py
View file

@ -753,7 +753,7 @@ def question(
url, flow_id, question, user, collection, entity_limit, triple_limit,
max_subgraph_size, max_path_length, edge_score_limit=50,
edge_limit=25, streaming=True, token=None,
explainable=False, debug=False
explainable=False, debug=False, show_usage=False
):
# Explainable mode uses the API to capture and process provenance events
@ -798,16 +798,26 @@ def question(
)
# Stream output
last_chunk = None
for chunk in response:
print(chunk, end="", flush=True)
print(chunk.content, end="", flush=True)
last_chunk = chunk
print() # Final newline
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=sys.stderr,
)
finally:
socket.close()
else:
# Use REST API for non-streaming
flow = api.flow().id(flow_id)
resp = flow.graph_rag(
result = flow.graph_rag(
query=question,
user=user,
collection=collection,
@ -818,7 +828,15 @@ def question(
edge_score_limit=edge_score_limit,
edge_limit=edge_limit,
)
print(resp)
print(result.text)
if show_usage:
print(
f"Input tokens: {result.in_token} "
f"Output tokens: {result.out_token} "
f"Model: {result.model}",
file=sys.stderr,
)
def main():
@ -923,6 +941,12 @@ def main():
help='Show debug output for troubleshooting'
)
parser.add_argument(
'--show-usage',
action='store_true',
help='Show token usage and model on stderr'
)
args = parser.parse_args()
try:
@ -943,6 +967,7 @@ def main():
token=args.token,
explainable=args.explainable,
debug=args.debug,
show_usage=args.show_usage,
)
except Exception as e:

35
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_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,14 +27,29 @@ 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, end="", flush=True)
# Add final newline after streaming
print(chunk.content, end="", flush=True)
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)
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
@ -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:

36
trustgraph-cli/trustgraph/cli/invoke_prompt.py
View file

@ -15,7 +15,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, template_id, variables, streaming=True, token=None):
def query(url, flow_id, template_id, variables, streaming=True, token=None,
show_usage=False):
# Create API client
api = Api(url=url, token=token)
@ -31,16 +32,30 @@ def query(url, flow_id, template_id, variables, streaming=True, token=None):
)
if streaming:
# Stream output (prompt yields strings directly)
last_chunk = None
for chunk in response:
if chunk:
print(chunk, end="", flush=True)
# Add final newline after streaming
if chunk.content:
print(chunk.content, end="", flush=True)
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)
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
@ -92,6 +107,12 @@ specified multiple times''',
help='Disable streaming (default: streaming enabled for text responses)'
)
parser.add_argument(
'--show-usage',
action='store_true',
help='Show token usage and model on stderr'
)
args = parser.parse_args()
variables = {}
@ -113,6 +134,7 @@ specified multiple times''',
variables=variables,
streaming=not args.no_streaming,
token=args.token,
show_usage=args.show_usage,
)
except Exception as e: