mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-24 04:31:02 +02:00
Phase 2
This commit is contained in:
Cyber MacGeddon
parent
d87f244872
commit
561f7f1dda
5 changed files with 274 additions and 421 deletions
|
|
@ -5,12 +5,10 @@ Uses the agent service to answer a question
|
|||
import argparse
|
||||
import os
|
||||
import textwrap
|
||||
import uuid
|
||||
import asyncio
|
||||
import json
|
||||
from websockets.asyncio.client import connect
|
||||
from trustgraph.api import Api
|
||||
|
||||
default_url = os.getenv("TRUSTGRAPH_URL", 'ws://localhost:8088/')
|
||||
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
|
||||
default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
|
||||
default_user = 'trustgraph'
|
||||
default_collection = 'default'
|
||||
|
||||
|
|
@ -99,79 +97,47 @@ def output(text, prefix="> ", width=78):
|
|||
)
|
||||
print(out)
|
||||
|
||||
async def question(
|
||||
def question(
|
||||
url, question, flow_id, user, collection,
|
||||
plan=None, state=None, group=None, verbose=False, streaming=True
|
||||
plan=None, state=None, group=None, verbose=False, streaming=True,
|
||||
token=None
|
||||
):
|
||||
|
||||
if not url.endswith("/"):
|
||||
url += "/"
|
||||
|
||||
url = url + "api/v1/socket"
|
||||
|
||||
if verbose:
|
||||
output(wrap(question), "\U00002753 ")
|
||||
print()
|
||||
|
||||
# Track last chunk type and current outputter for streaming
|
||||
last_chunk_type = None
|
||||
current_outputter = None
|
||||
# Create API client
|
||||
api = Api(url=url, token=token)
|
||||
socket = api.socket()
|
||||
flow = socket.flow(flow_id)
|
||||
|
||||
def think(x):
|
||||
if verbose:
|
||||
output(wrap(x), "\U0001f914 ")
|
||||
print()
|
||||
# Prepare request parameters
|
||||
request_params = {
|
||||
"question": question,
|
||||
"user": user,
|
||||
"streaming": streaming,
|
||||
}
|
||||
|
||||
def observe(x):
|
||||
if verbose:
|
||||
output(wrap(x), "\U0001f4a1 ")
|
||||
print()
|
||||
# Only add optional fields if they have values
|
||||
if state is not None:
|
||||
request_params["state"] = state
|
||||
if group is not None:
|
||||
request_params["group"] = group
|
||||
|
||||
mid = str(uuid.uuid4())
|
||||
try:
|
||||
# Call agent
|
||||
response = flow.agent(**request_params)
|
||||
|
||||
async with connect(url) as ws:
|
||||
# Handle streaming response
|
||||
if streaming:
|
||||
# Track last chunk type and current outputter for streaming
|
||||
last_chunk_type = None
|
||||
current_outputter = None
|
||||
|
||||
req = {
|
||||
"id": mid,
|
||||
"service": "agent",
|
||||
"flow": flow_id,
|
||||
"request": {
|
||||
"question": question,
|
||||
"user": user,
|
||||
"history": [],
|
||||
"streaming": streaming
|
||||
}
|
||||
}
|
||||
|
||||
# Only add optional fields if they have values
|
||||
if state is not None:
|
||||
req["request"]["state"] = state
|
||||
if group is not None:
|
||||
req["request"]["group"] = group
|
||||
|
||||
req = json.dumps(req)
|
||||
|
||||
await ws.send(req)
|
||||
|
||||
while True:
|
||||
|
||||
msg = await ws.recv()
|
||||
|
||||
obj = json.loads(msg)
|
||||
|
||||
if "error" in obj:
|
||||
raise RuntimeError(obj["error"])
|
||||
|
||||
if obj["id"] != mid:
|
||||
print("Ignore message")
|
||||
continue
|
||||
|
||||
response = obj["response"]
|
||||
|
||||
# Handle streaming format (new format with chunk_type)
|
||||
if "chunk_type" in response:
|
||||
chunk_type = response["chunk_type"]
|
||||
content = response.get("content", "")
|
||||
for chunk in response:
|
||||
chunk_type = chunk.chunk_type
|
||||
content = chunk.content
|
||||
|
||||
# Check if we're switching to a new message type
|
||||
if last_chunk_type != chunk_type:
|
||||
|
|
@ -195,33 +161,27 @@ async def question(
|
|||
# Output the chunk
|
||||
if current_outputter:
|
||||
current_outputter.output(content)
|
||||
elif chunk_type == "answer":
|
||||
elif chunk_type == "final-answer":
|
||||
print(content, end="", flush=True)
|
||||
else:
|
||||
# Handle legacy format (backward compatibility)
|
||||
if "thought" in response:
|
||||
think(response["thought"])
|
||||
|
||||
if "observation" in response:
|
||||
observe(response["observation"])
|
||||
# Close any remaining outputter
|
||||
if current_outputter:
|
||||
current_outputter.__exit__(None, None, None)
|
||||
current_outputter = None
|
||||
# Add final newline if we were outputting answer
|
||||
elif last_chunk_type == "final-answer":
|
||||
print()
|
||||
|
||||
if "answer" in response:
|
||||
print(response["answer"])
|
||||
else:
|
||||
# Non-streaming response
|
||||
if "answer" in response:
|
||||
print(response["answer"])
|
||||
if "error" in response:
|
||||
raise RuntimeError(response["error"])
|
||||
|
||||
if "error" in response:
|
||||
raise RuntimeError(response["error"])
|
||||
|
||||
if obj["complete"]:
|
||||
# Close any remaining outputter
|
||||
if current_outputter:
|
||||
current_outputter.__exit__(None, None, None)
|
||||
current_outputter = None
|
||||
# Add final newline if we were outputting answer
|
||||
elif last_chunk_type == "answer":
|
||||
print()
|
||||
break
|
||||
|
||||
await ws.close()
|
||||
finally:
|
||||
# Clean up socket connection
|
||||
socket.close()
|
||||
|
||||
def main():
|
||||
|
||||
|
|
@ -236,6 +196,12 @@ def main():
|
|||
help=f'API URL (default: {default_url})',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'-t', '--token',
|
||||
default=default_token,
|
||||
help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'-f', '--flow-id',
|
||||
default="default",
|
||||
|
|
@ -292,19 +258,18 @@ def main():
|
|||
|
||||
try:
|
||||
|
||||
asyncio.run(
|
||||
question(
|
||||
url = args.url,
|
||||
flow_id = args.flow_id,
|
||||
question = args.question,
|
||||
user = args.user,
|
||||
collection = args.collection,
|
||||
plan = args.plan,
|
||||
state = args.state,
|
||||
group = args.group,
|
||||
verbose = args.verbose,
|
||||
streaming = not args.no_streaming,
|
||||
)
|
||||
question(
|
||||
url = args.url,
|
||||
flow_id = args.flow_id,
|
||||
question = args.question,
|
||||
user = args.user,
|
||||
collection = args.collection,
|
||||
plan = args.plan,
|
||||
state = args.state,
|
||||
group = args.group,
|
||||
verbose = args.verbose,
|
||||
streaming = not args.no_streaming,
|
||||
token = args.token,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -4,89 +4,50 @@ Uses the DocumentRAG service to answer a question
|
|||
|
||||
import argparse
|
||||
import os
|
||||
import asyncio
|
||||
import json
|
||||
import uuid
|
||||
from websockets.asyncio.client import connect
|
||||
from trustgraph.api import Api
|
||||
|
||||
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
|
||||
default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
|
||||
default_user = 'trustgraph'
|
||||
default_collection = 'default'
|
||||
default_doc_limit = 10
|
||||
|
||||
async def question_streaming(url, flow_id, question, user, collection, doc_limit):
|
||||
"""Streaming version using websockets"""
|
||||
def question(url, flow_id, question, user, collection, doc_limit, streaming=True, token=None):
|
||||
|
||||
# Convert http:// to ws://
|
||||
if url.startswith('http://'):
|
||||
url = 'ws://' + url[7:]
|
||||
elif url.startswith('https://'):
|
||||
url = 'wss://' + url[8:]
|
||||
# Create API client
|
||||
api = Api(url=url, token=token)
|
||||
|
||||
if not url.endswith("/"):
|
||||
url += "/"
|
||||
if streaming:
|
||||
# Use socket client for streaming
|
||||
socket = api.socket()
|
||||
flow = socket.flow(flow_id)
|
||||
|
||||
url = url + "api/v1/socket"
|
||||
try:
|
||||
response = flow.document_rag(
|
||||
question=question,
|
||||
user=user,
|
||||
collection=collection,
|
||||
doc_limit=doc_limit,
|
||||
streaming=True
|
||||
)
|
||||
|
||||
mid = str(uuid.uuid4())
|
||||
# Stream output
|
||||
for chunk in response:
|
||||
print(chunk.content, end="", flush=True)
|
||||
print() # Final newline
|
||||
|
||||
async with connect(url) as ws:
|
||||
req = {
|
||||
"id": mid,
|
||||
"service": "document-rag",
|
||||
"flow": flow_id,
|
||||
"request": {
|
||||
"query": question,
|
||||
"user": user,
|
||||
"collection": collection,
|
||||
"doc-limit": doc_limit,
|
||||
"streaming": True
|
||||
}
|
||||
}
|
||||
|
||||
req = json.dumps(req)
|
||||
await ws.send(req)
|
||||
|
||||
while True:
|
||||
msg = await ws.recv()
|
||||
obj = json.loads(msg)
|
||||
|
||||
if "error" in obj:
|
||||
raise RuntimeError(obj["error"])
|
||||
|
||||
if obj["id"] != mid:
|
||||
print("Ignore message")
|
||||
continue
|
||||
|
||||
response = obj["response"]
|
||||
|
||||
# Handle streaming format (chunk)
|
||||
if "chunk" in response:
|
||||
chunk = response["chunk"]
|
||||
print(chunk, end="", flush=True)
|
||||
elif "response" in response:
|
||||
# Final response with complete text
|
||||
# Already printed via chunks, just add newline
|
||||
pass
|
||||
|
||||
if obj["complete"]:
|
||||
print() # Final newline
|
||||
break
|
||||
|
||||
await ws.close()
|
||||
|
||||
def question_non_streaming(url, flow_id, question, user, collection, doc_limit):
|
||||
"""Non-streaming version using HTTP API"""
|
||||
|
||||
api = Api(url).flow().id(flow_id)
|
||||
|
||||
resp = api.document_rag(
|
||||
question=question, user=user, collection=collection,
|
||||
doc_limit=doc_limit,
|
||||
)
|
||||
|
||||
print(resp)
|
||||
finally:
|
||||
socket.close()
|
||||
else:
|
||||
# Use REST API for non-streaming
|
||||
flow = api.flow().id(flow_id)
|
||||
resp = flow.document_rag(
|
||||
question=question,
|
||||
user=user,
|
||||
collection=collection,
|
||||
doc_limit=doc_limit,
|
||||
)
|
||||
print(resp)
|
||||
|
||||
def main():
|
||||
|
||||
|
|
@ -101,6 +62,12 @@ def main():
|
|||
help=f'API URL (default: {default_url})',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'-t', '--token',
|
||||
default=default_token,
|
||||
help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'-f', '--flow-id',
|
||||
default="default",
|
||||
|
|
@ -127,6 +94,7 @@ def main():
|
|||
|
||||
parser.add_argument(
|
||||
'-d', '--doc-limit',
|
||||
type=int,
|
||||
default=default_doc_limit,
|
||||
help=f'Document limit (default: {default_doc_limit})'
|
||||
)
|
||||
|
|
@ -141,30 +109,20 @@ def main():
|
|||
|
||||
try:
|
||||
|
||||
if not args.no_streaming:
|
||||
asyncio.run(
|
||||
question_streaming(
|
||||
url=args.url,
|
||||
flow_id=args.flow_id,
|
||||
question=args.question,
|
||||
user=args.user,
|
||||
collection=args.collection,
|
||||
doc_limit=args.doc_limit,
|
||||
)
|
||||
)
|
||||
else:
|
||||
question_non_streaming(
|
||||
url=args.url,
|
||||
flow_id=args.flow_id,
|
||||
question=args.question,
|
||||
user=args.user,
|
||||
collection=args.collection,
|
||||
doc_limit=args.doc_limit,
|
||||
)
|
||||
question(
|
||||
url=args.url,
|
||||
flow_id=args.flow_id,
|
||||
question=args.question,
|
||||
user=args.user,
|
||||
collection=args.collection,
|
||||
doc_limit=args.doc_limit,
|
||||
streaming=not args.no_streaming,
|
||||
token=args.token,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
|
||||
print("Exception:", e, flush=True)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
main()
|
||||
|
|
|
|||
|
|
@ -4,13 +4,10 @@ Uses the GraphRAG service to answer a question
|
|||
|
||||
import argparse
|
||||
import os
|
||||
import asyncio
|
||||
import json
|
||||
import uuid
|
||||
from websockets.asyncio.client import connect
|
||||
from trustgraph.api import Api
|
||||
|
||||
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
|
||||
default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
|
||||
default_user = 'trustgraph'
|
||||
default_collection = 'default'
|
||||
default_entity_limit = 50
|
||||
|
|
@ -18,89 +15,51 @@ default_triple_limit = 30
|
|||
default_max_subgraph_size = 150
|
||||
default_max_path_length = 2
|
||||
|
||||
async def question_streaming(
|
||||
def question(
|
||||
url, flow_id, question, user, collection, entity_limit, triple_limit,
|
||||
max_subgraph_size, max_path_length
|
||||
max_subgraph_size, max_path_length, streaming=True, token=None
|
||||
):
|
||||
"""Streaming version using websockets"""
|
||||
|
||||
# Convert http:// to ws://
|
||||
if url.startswith('http://'):
|
||||
url = 'ws://' + url[7:]
|
||||
elif url.startswith('https://'):
|
||||
url = 'wss://' + url[8:]
|
||||
# Create API client
|
||||
api = Api(url=url, token=token)
|
||||
|
||||
if not url.endswith("/"):
|
||||
url += "/"
|
||||
if streaming:
|
||||
# Use socket client for streaming
|
||||
socket = api.socket()
|
||||
flow = socket.flow(flow_id)
|
||||
|
||||
url = url + "api/v1/socket"
|
||||
try:
|
||||
response = flow.graph_rag(
|
||||
question=question,
|
||||
user=user,
|
||||
collection=collection,
|
||||
entity_limit=entity_limit,
|
||||
triple_limit=triple_limit,
|
||||
max_subgraph_size=max_subgraph_size,
|
||||
max_path_length=max_path_length,
|
||||
streaming=True
|
||||
)
|
||||
|
||||
mid = str(uuid.uuid4())
|
||||
# Stream output
|
||||
for chunk in response:
|
||||
print(chunk.content, end="", flush=True)
|
||||
print() # Final newline
|
||||
|
||||
async with connect(url) as ws:
|
||||
req = {
|
||||
"id": mid,
|
||||
"service": "graph-rag",
|
||||
"flow": flow_id,
|
||||
"request": {
|
||||
"query": question,
|
||||
"user": user,
|
||||
"collection": collection,
|
||||
"entity-limit": entity_limit,
|
||||
"triple-limit": triple_limit,
|
||||
"max-subgraph-size": max_subgraph_size,
|
||||
"max-path-length": max_path_length,
|
||||
"streaming": True
|
||||
}
|
||||
}
|
||||
|
||||
req = json.dumps(req)
|
||||
await ws.send(req)
|
||||
|
||||
while True:
|
||||
msg = await ws.recv()
|
||||
obj = json.loads(msg)
|
||||
|
||||
if "error" in obj:
|
||||
raise RuntimeError(obj["error"])
|
||||
|
||||
if obj["id"] != mid:
|
||||
print("Ignore message")
|
||||
continue
|
||||
|
||||
response = obj["response"]
|
||||
|
||||
# Handle streaming format (chunk)
|
||||
if "chunk" in response:
|
||||
chunk = response["chunk"]
|
||||
print(chunk, end="", flush=True)
|
||||
elif "response" in response:
|
||||
# Final response with complete text
|
||||
# Already printed via chunks, just add newline
|
||||
pass
|
||||
|
||||
if obj["complete"]:
|
||||
print() # Final newline
|
||||
break
|
||||
|
||||
await ws.close()
|
||||
|
||||
def question_non_streaming(
|
||||
url, flow_id, question, user, collection, entity_limit, triple_limit,
|
||||
max_subgraph_size, max_path_length
|
||||
):
|
||||
"""Non-streaming version using HTTP API"""
|
||||
|
||||
api = Api(url).flow().id(flow_id)
|
||||
|
||||
resp = api.graph_rag(
|
||||
question=question, user=user, collection=collection,
|
||||
entity_limit=entity_limit, triple_limit=triple_limit,
|
||||
max_subgraph_size=max_subgraph_size,
|
||||
max_path_length=max_path_length
|
||||
)
|
||||
|
||||
print(resp)
|
||||
finally:
|
||||
socket.close()
|
||||
else:
|
||||
# Use REST API for non-streaming
|
||||
flow = api.flow().id(flow_id)
|
||||
resp = flow.graph_rag(
|
||||
question=question,
|
||||
user=user,
|
||||
collection=collection,
|
||||
entity_limit=entity_limit,
|
||||
triple_limit=triple_limit,
|
||||
max_subgraph_size=max_subgraph_size,
|
||||
max_path_length=max_path_length
|
||||
)
|
||||
print(resp)
|
||||
|
||||
def main():
|
||||
|
||||
|
|
@ -115,6 +74,12 @@ def main():
|
|||
help=f'API URL (default: {default_url})',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'-t', '--token',
|
||||
default=default_token,
|
||||
help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'-f', '--flow-id',
|
||||
default="default",
|
||||
|
|
@ -141,24 +106,28 @@ def main():
|
|||
|
||||
parser.add_argument(
|
||||
'-e', '--entity-limit',
|
||||
type=int,
|
||||
default=default_entity_limit,
|
||||
help=f'Entity limit (default: {default_entity_limit})'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'-t', '--triple-limit',
|
||||
'--triple-limit',
|
||||
type=int,
|
||||
default=default_triple_limit,
|
||||
help=f'Triple limit (default: {default_triple_limit})'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'-s', '--max-subgraph-size',
|
||||
type=int,
|
||||
default=default_max_subgraph_size,
|
||||
help=f'Max subgraph size (default: {default_max_subgraph_size})'
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'-p', '--max-path-length',
|
||||
type=int,
|
||||
default=default_max_path_length,
|
||||
help=f'Max path length (default: {default_max_path_length})'
|
||||
)
|
||||
|
|
@ -173,36 +142,23 @@ def main():
|
|||
|
||||
try:
|
||||
|
||||
if not args.no_streaming:
|
||||
asyncio.run(
|
||||
question_streaming(
|
||||
url=args.url,
|
||||
flow_id=args.flow_id,
|
||||
question=args.question,
|
||||
user=args.user,
|
||||
collection=args.collection,
|
||||
entity_limit=args.entity_limit,
|
||||
triple_limit=args.triple_limit,
|
||||
max_subgraph_size=args.max_subgraph_size,
|
||||
max_path_length=args.max_path_length,
|
||||
)
|
||||
)
|
||||
else:
|
||||
question_non_streaming(
|
||||
url=args.url,
|
||||
flow_id=args.flow_id,
|
||||
question=args.question,
|
||||
user=args.user,
|
||||
collection=args.collection,
|
||||
entity_limit=args.entity_limit,
|
||||
triple_limit=args.triple_limit,
|
||||
max_subgraph_size=args.max_subgraph_size,
|
||||
max_path_length=args.max_path_length,
|
||||
)
|
||||
question(
|
||||
url=args.url,
|
||||
flow_id=args.flow_id,
|
||||
question=args.question,
|
||||
user=args.user,
|
||||
collection=args.collection,
|
||||
entity_limit=args.entity_limit,
|
||||
triple_limit=args.triple_limit,
|
||||
max_subgraph_size=args.max_subgraph_size,
|
||||
max_path_length=args.max_path_length,
|
||||
streaming=not args.no_streaming,
|
||||
token=args.token,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
|
||||
print("Exception:", e, flush=True)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
main()
|
||||
|
|
|
|||
|
|
@ -5,64 +5,39 @@ and user prompt. Both arguments are required.
|
|||
|
||||
import argparse
|
||||
import os
|
||||
import json
|
||||
import uuid
|
||||
import asyncio
|
||||
from websockets.asyncio.client import connect
|
||||
from trustgraph.api import Api
|
||||
|
||||
default_url = os.getenv("TRUSTGRAPH_URL", 'ws://localhost:8088/')
|
||||
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
|
||||
default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
|
||||
|
||||
async def query(url, flow_id, system, prompt, streaming=True):
|
||||
def query(url, flow_id, system, prompt, streaming=True, token=None):
|
||||
|
||||
if not url.endswith("/"):
|
||||
url += "/"
|
||||
# Create API client
|
||||
api = Api(url=url, token=token)
|
||||
socket = api.socket()
|
||||
flow = socket.flow(flow_id)
|
||||
|
||||
url = url + "api/v1/socket"
|
||||
try:
|
||||
# Call text completion
|
||||
response = flow.text_completion(
|
||||
system=system,
|
||||
prompt=prompt,
|
||||
streaming=streaming
|
||||
)
|
||||
|
||||
mid = str(uuid.uuid4())
|
||||
if streaming:
|
||||
# Stream output to stdout without newline
|
||||
for chunk in response:
|
||||
print(chunk.content, end="", flush=True)
|
||||
# Add final newline after streaming
|
||||
print()
|
||||
else:
|
||||
# Non-streaming: print complete response
|
||||
print(response)
|
||||
|
||||
async with connect(url) as ws:
|
||||
|
||||
req = {
|
||||
"id": mid,
|
||||
"service": "text-completion",
|
||||
"flow": flow_id,
|
||||
"request": {
|
||||
"system": system,
|
||||
"prompt": prompt,
|
||||
"streaming": streaming
|
||||
}
|
||||
}
|
||||
|
||||
await ws.send(json.dumps(req))
|
||||
|
||||
while True:
|
||||
|
||||
msg = await ws.recv()
|
||||
|
||||
obj = json.loads(msg)
|
||||
|
||||
if "error" in obj:
|
||||
raise RuntimeError(obj["error"])
|
||||
|
||||
if obj["id"] != mid:
|
||||
continue
|
||||
|
||||
if "response" in obj["response"]:
|
||||
if streaming:
|
||||
# Stream output to stdout without newline
|
||||
print(obj["response"]["response"], end="", flush=True)
|
||||
else:
|
||||
# Non-streaming: print complete response
|
||||
print(obj["response"]["response"])
|
||||
|
||||
if obj["complete"]:
|
||||
if streaming:
|
||||
# Add final newline after streaming
|
||||
print()
|
||||
break
|
||||
|
||||
await ws.close()
|
||||
finally:
|
||||
# Clean up socket connection
|
||||
socket.close()
|
||||
|
||||
def main():
|
||||
|
||||
|
|
@ -77,6 +52,12 @@ def main():
|
|||
help=f'API URL (default: {default_url})',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'-t', '--token',
|
||||
default=default_token,
|
||||
help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'system',
|
||||
nargs=1,
|
||||
|
|
@ -105,17 +86,18 @@ def main():
|
|||
|
||||
try:
|
||||
|
||||
asyncio.run(query(
|
||||
query(
|
||||
url=args.url,
|
||||
flow_id=args.flow_id,
|
||||
system=args.system[0],
|
||||
prompt=args.prompt[0],
|
||||
streaming=not args.no_streaming
|
||||
))
|
||||
streaming=not args.no_streaming,
|
||||
token=args.token,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
|
||||
print("Exception:", e, flush=True)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
main()
|
||||
|
|
|
|||
|
|
@ -10,76 +10,61 @@ using key=value arguments on the command line, and these replace
|
|||
import argparse
|
||||
import os
|
||||
import json
|
||||
import uuid
|
||||
import asyncio
|
||||
from websockets.asyncio.client import connect
|
||||
from trustgraph.api import Api
|
||||
|
||||
default_url = os.getenv("TRUSTGRAPH_URL", 'ws://localhost:8088/')
|
||||
default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/')
|
||||
default_token = os.getenv("TRUSTGRAPH_TOKEN", None)
|
||||
|
||||
async def query(url, flow_id, template_id, variables, streaming=True):
|
||||
def query(url, flow_id, template_id, variables, streaming=True, token=None):
|
||||
|
||||
if not url.endswith("/"):
|
||||
url += "/"
|
||||
# Create API client
|
||||
api = Api(url=url, token=token)
|
||||
socket = api.socket()
|
||||
flow = socket.flow(flow_id)
|
||||
|
||||
url = url + "api/v1/socket"
|
||||
try:
|
||||
# Call prompt
|
||||
response = flow.prompt(
|
||||
id=template_id,
|
||||
variables=variables,
|
||||
streaming=streaming
|
||||
)
|
||||
|
||||
mid = str(uuid.uuid4())
|
||||
if streaming:
|
||||
full_response = {"text": "", "object": ""}
|
||||
|
||||
async with connect(url) as ws:
|
||||
# Stream output
|
||||
for chunk in response:
|
||||
content = chunk.content
|
||||
if content:
|
||||
print(content, end="", flush=True)
|
||||
full_response["text"] += content
|
||||
|
||||
req = {
|
||||
"id": mid,
|
||||
"service": "prompt",
|
||||
"flow": flow_id,
|
||||
"request": {
|
||||
"id": template_id,
|
||||
"variables": variables,
|
||||
"streaming": streaming
|
||||
}
|
||||
}
|
||||
# Check if this is an object response (JSON)
|
||||
if hasattr(chunk, 'object') and chunk.object:
|
||||
full_response["object"] = chunk.object
|
||||
|
||||
await ws.send(json.dumps(req))
|
||||
# Handle final output
|
||||
if full_response["text"]:
|
||||
# Add final newline after streaming text
|
||||
print()
|
||||
elif full_response["object"]:
|
||||
# Print JSON object (pretty-printed)
|
||||
print(json.dumps(json.loads(full_response["object"]), indent=4))
|
||||
|
||||
full_response = {"text": "", "object": ""}
|
||||
|
||||
while True:
|
||||
|
||||
msg = await ws.recv()
|
||||
|
||||
obj = json.loads(msg)
|
||||
|
||||
if "error" in obj:
|
||||
raise RuntimeError(obj["error"])
|
||||
|
||||
if obj["id"] != mid:
|
||||
continue
|
||||
|
||||
response = obj["response"]
|
||||
|
||||
# Handle text responses (streaming)
|
||||
if "text" in response and response["text"]:
|
||||
if streaming:
|
||||
# Stream output to stdout without newline
|
||||
print(response["text"], end="", flush=True)
|
||||
full_response["text"] += response["text"]
|
||||
else:
|
||||
# Non-streaming: print complete response
|
||||
else:
|
||||
# Non-streaming: handle response
|
||||
if isinstance(response, str):
|
||||
print(response)
|
||||
elif isinstance(response, dict):
|
||||
if "text" in response:
|
||||
print(response["text"])
|
||||
elif "object" in response:
|
||||
print(json.dumps(json.loads(response["object"]), indent=4))
|
||||
|
||||
# Handle object responses (JSON, never streamed)
|
||||
if "object" in response and response["object"]:
|
||||
full_response["object"] = response["object"]
|
||||
|
||||
if obj["complete"]:
|
||||
if streaming and full_response["text"]:
|
||||
# Add final newline after streaming text
|
||||
print()
|
||||
elif full_response["object"]:
|
||||
# Print JSON object (pretty-printed)
|
||||
print(json.dumps(json.loads(full_response["object"]), indent=4))
|
||||
break
|
||||
|
||||
await ws.close()
|
||||
finally:
|
||||
# Clean up socket connection
|
||||
socket.close()
|
||||
|
||||
def main():
|
||||
|
||||
|
|
@ -94,6 +79,12 @@ def main():
|
|||
help=f'API URL (default: {default_url})',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'-t', '--token',
|
||||
default=default_token,
|
||||
help='Authentication token (default: $TRUSTGRAPH_TOKEN)',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'-f', '--flow-id',
|
||||
default="default",
|
||||
|
|
@ -135,17 +126,18 @@ specified multiple times''',
|
|||
|
||||
try:
|
||||
|
||||
asyncio.run(query(
|
||||
query(
|
||||
url=args.url,
|
||||
flow_id=args.flow_id,
|
||||
template_id=args.id[0],
|
||||
variables=variables,
|
||||
streaming=not args.no_streaming
|
||||
))
|
||||
streaming=not args.no_streaming,
|
||||
token=args.token,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
|
||||
print("Exception:", e, flush=True)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
main()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue