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

token errors

Browse source
This commit is contained in:
Cyber MacGeddon 2026-04-12 16:42:39 +01:00
parent fc9c0cf0b4
commit 05c29de5bc
11 changed files with 92 additions and 61 deletions
Download patch file Download diff file Expand all files Collapse all files

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: