From 561f7f1dda9afe9b1bda6e47e3cdf7e883221f6b Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Thu, 4 Dec 2025 09:36:07 +0000 Subject: [PATCH] Phase 2 --- trustgraph-cli/trustgraph/cli/invoke_agent.py | 169 +++++++---------- .../trustgraph/cli/invoke_document_rag.py | 142 +++++---------- .../trustgraph/cli/invoke_graph_rag.py | 172 +++++++----------- trustgraph-cli/trustgraph/cli/invoke_llm.py | 94 ++++------ .../trustgraph/cli/invoke_prompt.py | 118 ++++++------ 5 files changed, 274 insertions(+), 421 deletions(-) diff --git a/trustgraph-cli/trustgraph/cli/invoke_agent.py b/trustgraph-cli/trustgraph/cli/invoke_agent.py index e6e82edd..6af12cd5 100644 --- a/trustgraph-cli/trustgraph/cli/invoke_agent.py +++ b/trustgraph-cli/trustgraph/cli/invoke_agent.py @@ -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: diff --git a/trustgraph-cli/trustgraph/cli/invoke_document_rag.py b/trustgraph-cli/trustgraph/cli/invoke_document_rag.py index e6a040ac..6b3a44bc 100644 --- a/trustgraph-cli/trustgraph/cli/invoke_document_rag.py +++ b/trustgraph-cli/trustgraph/cli/invoke_document_rag.py @@ -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() \ No newline at end of file + main() diff --git a/trustgraph-cli/trustgraph/cli/invoke_graph_rag.py b/trustgraph-cli/trustgraph/cli/invoke_graph_rag.py index 45d02b6d..725206c8 100644 --- a/trustgraph-cli/trustgraph/cli/invoke_graph_rag.py +++ b/trustgraph-cli/trustgraph/cli/invoke_graph_rag.py @@ -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() \ No newline at end of file + main() diff --git a/trustgraph-cli/trustgraph/cli/invoke_llm.py b/trustgraph-cli/trustgraph/cli/invoke_llm.py index da69dcd6..261993d9 100644 --- a/trustgraph-cli/trustgraph/cli/invoke_llm.py +++ b/trustgraph-cli/trustgraph/cli/invoke_llm.py @@ -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() \ No newline at end of file + main() diff --git a/trustgraph-cli/trustgraph/cli/invoke_prompt.py b/trustgraph-cli/trustgraph/cli/invoke_prompt.py index c996c57d..b1eba5aa 100644 --- a/trustgraph-cli/trustgraph/cli/invoke_prompt.py +++ b/trustgraph-cli/trustgraph/cli/invoke_prompt.py @@ -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() \ No newline at end of file + main()