diff --git a/trustgraph-cli/trustgraph/cli/delete_config_item.py b/trustgraph-cli/trustgraph/cli/delete_config_item.py index 1de02890..cf4cba93 100644 --- a/trustgraph-cli/trustgraph/cli/delete_config_item.py +++ b/trustgraph-cli/trustgraph/cli/delete_config_item.py @@ -8,10 +8,11 @@ from trustgraph.api import Api from trustgraph.api.types import ConfigKey default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') +default_token = os.getenv("TRUSTGRAPH_TOKEN", None) -def delete_config_item(url, config_type, key): +def delete_config_item(url, config_type, key, token=None): - api = Api(url).config() + api = Api(url, token=token).config() config_key = ConfigKey(type=config_type, key=key) api.delete([config_key]) @@ -43,6 +44,12 @@ def main(): help=f'API URL (default: {default_url})', ) + parser.add_argument( + '-t', '--token', + default=default_token, + help='Authentication token (default: $TRUSTGRAPH_TOKEN)', + ) + args = parser.parse_args() try: @@ -51,6 +58,7 @@ def main(): url=args.api_url, config_type=args.type, key=args.key, + token=args.token, ) except Exception as e: diff --git a/trustgraph-cli/trustgraph/cli/get_config_item.py b/trustgraph-cli/trustgraph/cli/get_config_item.py index 832d2711..c2421e94 100644 --- a/trustgraph-cli/trustgraph/cli/get_config_item.py +++ b/trustgraph-cli/trustgraph/cli/get_config_item.py @@ -9,10 +9,11 @@ from trustgraph.api import Api from trustgraph.api.types import ConfigKey default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') +default_token = os.getenv("TRUSTGRAPH_TOKEN", None) -def get_config_item(url, config_type, key, format_type): +def get_config_item(url, config_type, key, format_type, token=None): - api = Api(url).config() + api = Api(url, token=token).config() config_key = ConfigKey(type=config_type, key=key) values = api.get([config_key]) @@ -59,6 +60,12 @@ def main(): help=f'API URL (default: {default_url})', ) + parser.add_argument( + '-t', '--token', + default=default_token, + help='Authentication token (default: $TRUSTGRAPH_TOKEN)', + ) + args = parser.parse_args() try: @@ -68,6 +75,7 @@ def main(): config_type=args.type, key=args.key, format_type=args.format, + token=args.token, ) except Exception as e: diff --git a/trustgraph-cli/trustgraph/cli/get_kg_core.py b/trustgraph-cli/trustgraph/cli/get_kg_core.py index 6e0a8bc0..b75f7155 100644 --- a/trustgraph-cli/trustgraph/cli/get_kg_core.py +++ b/trustgraph-cli/trustgraph/cli/get_kg_core.py @@ -14,6 +14,7 @@ import msgpack default_url = os.getenv("TRUSTGRAPH_URL", 'ws://localhost:8088/') default_user = 'trustgraph' +default_token = os.getenv("TRUSTGRAPH_TOKEN", None) def write_triple(f, data): msg = ( @@ -51,13 +52,16 @@ def write_ge(f, data): ) f.write(msgpack.packb(msg, use_bin_type=True)) -async def fetch(url, user, id, output): +async def fetch(url, user, id, output, token=None): if not url.endswith("/"): url += "/" url = url + "api/v1/socket" + if token: + url = f"{url}?token={token}" + mid = str(uuid.uuid4()) async with connect(url) as ws: @@ -138,6 +142,12 @@ def main(): help=f'Output file' ) + parser.add_argument( + '-t', '--token', + default=default_token, + help='Authentication token (default: $TRUSTGRAPH_TOKEN)', + ) + args = parser.parse_args() try: @@ -148,6 +158,7 @@ def main(): user = args.user, id = args.id, output = args.output, + token = args.token, ) ) diff --git a/trustgraph-cli/trustgraph/cli/list_config_items.py b/trustgraph-cli/trustgraph/cli/list_config_items.py index 33e8f7ba..5cd0f233 100644 --- a/trustgraph-cli/trustgraph/cli/list_config_items.py +++ b/trustgraph-cli/trustgraph/cli/list_config_items.py @@ -8,10 +8,11 @@ import json from trustgraph.api import Api default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') +default_token = os.getenv("TRUSTGRAPH_TOKEN", None) -def list_config_items(url, config_type, format_type): +def list_config_items(url, config_type, format_type, token=None): - api = Api(url).config() + api = Api(url, token=token).config() keys = api.list(config_type) @@ -47,6 +48,12 @@ def main(): help=f'API URL (default: {default_url})', ) + parser.add_argument( + '-t', '--token', + default=default_token, + help='Authentication token (default: $TRUSTGRAPH_TOKEN)', + ) + args = parser.parse_args() try: @@ -55,6 +62,7 @@ def main(): url=args.api_url, config_type=args.type, format_type=args.format, + token=args.token, ) except Exception as e: diff --git a/trustgraph-cli/trustgraph/cli/load_sample_documents.py b/trustgraph-cli/trustgraph/cli/load_sample_documents.py index fd6751be..186006a8 100644 --- a/trustgraph-cli/trustgraph/cli/load_sample_documents.py +++ b/trustgraph-cli/trustgraph/cli/load_sample_documents.py @@ -13,6 +13,7 @@ from trustgraph.api.types import hash, Uri, Literal, Triple default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') default_user = 'trustgraph' +default_token = os.getenv("TRUSTGRAPH_TOKEN", None) from requests.adapters import HTTPAdapter @@ -655,10 +656,10 @@ documents = [ class Loader: def __init__( - self, url, user + self, url, user, token=None ): - self.api = Api(url).library() + self.api = Api(url, token=token).library() self.user = user def load(self, documents): @@ -719,6 +720,12 @@ def main(): help=f'User ID (default: {default_user})' ) + parser.add_argument( + '-t', '--token', + default=default_token, + help='Authentication token (default: $TRUSTGRAPH_TOKEN)', + ) + args = parser.parse_args() try: @@ -726,6 +733,7 @@ def main(): p = Loader( url=args.url, user=args.user, + token=args.token, ) p.load(documents) diff --git a/trustgraph-cli/trustgraph/cli/load_structured_data.py b/trustgraph-cli/trustgraph/cli/load_structured_data.py index 9bb9f78c..bf112417 100644 --- a/trustgraph-cli/trustgraph/cli/load_structured_data.py +++ b/trustgraph-cli/trustgraph/cli/load_structured_data.py @@ -22,6 +22,7 @@ import logging logger = logging.getLogger(__name__) default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') +default_token = os.getenv("TRUSTGRAPH_TOKEN", None) def load_structured_data( @@ -41,7 +42,8 @@ def load_structured_data( user: str = 'trustgraph', collection: str = 'default', dry_run: bool = False, - verbose: bool = False + verbose: bool = False, + token: str = None ): """ Load structured data using a descriptor configuration. @@ -133,9 +135,9 @@ def load_structured_data( # Get batch size from descriptor batch_size = descriptor.get('output', {}).get('options', {}).get('batch_size', 1000) - + # Send to TrustGraph using shared function - imported_count = _send_to_trustgraph(output_objects, api_url, flow, batch_size) + imported_count = _send_to_trustgraph(output_objects, api_url, flow, batch_size, token=token) # Summary format_info = descriptor.get('format', {}) @@ -288,10 +290,10 @@ def load_structured_data( # Get batch size from descriptor or use default batch_size = descriptor.get('output', {}).get('options', {}).get('batch_size', 1000) - + # Send to TrustGraph print(f"šŸš€ Importing {len(output_records)} records to TrustGraph...") - imported_count = _send_to_trustgraph(output_records, api_url, flow, batch_size) + imported_count = _send_to_trustgraph(output_records, api_url, flow, batch_size, token=token) # Get summary info from descriptor format_info = descriptor.get('format', {}) @@ -571,66 +573,30 @@ def _process_data_pipeline(input_file, descriptor_file, user, collection, sample return output_records, descriptor -def _send_to_trustgraph(objects, api_url, flow, batch_size=1000): - """Send ExtractedObject records to TrustGraph using WebSocket""" - import json - import asyncio - from websockets.asyncio.client import connect - +def _send_to_trustgraph(objects, api_url, flow, batch_size=1000, token=None): + """Send ExtractedObject records to TrustGraph using Python API""" + from trustgraph.api import Api + try: - # Construct objects import URL similar to load_knowledge pattern - if not api_url.endswith("/"): - api_url += "/" - - # Convert HTTP URL to WebSocket URL if needed - ws_url = api_url.replace("http://", "ws://").replace("https://", "wss://") - objects_url = ws_url + f"api/v1/flow/{flow}/import/objects" - - logger.info(f"Connecting to objects import endpoint: {objects_url}") - - async def import_objects(): - async with connect(objects_url) as ws: - imported_count = 0 - - for record in objects: - try: - # Send individual ExtractedObject records - await ws.send(json.dumps(record)) - imported_count += 1 - - if imported_count % 100 == 0: - logger.debug(f"Imported {imported_count}/{len(objects)} records...") - - except Exception as e: - logger.error(f"Failed to send record {imported_count + 1}: {e}") - print(f"āŒ Failed to send record {imported_count + 1}: {e}") - - logger.info(f"Successfully imported {imported_count} records to TrustGraph") - return imported_count - - # Run the async import - imported_count = asyncio.run(import_objects()) - - # Summary total_records = len(objects) - failed_count = total_records - imported_count - + logger.info(f"Importing {total_records} records to TrustGraph...") + + # Use Python API bulk import + api = Api(api_url, token=token) + bulk = api.bulk() + + bulk.import_objects(flow=flow, objects=iter(objects)) + + logger.info(f"Successfully imported {total_records} records to TrustGraph") + + # Summary print(f"\nšŸ“Š Import Summary:") print(f"- Total records: {total_records}") - print(f"- Successfully imported: {imported_count}") - print(f"- Failed: {failed_count}") - - if failed_count > 0: - print(f"āš ļø {failed_count} records failed to import. Check logs for details.") - else: - print("āœ… All records imported successfully!") - - return imported_count - - except ImportError as e: - logger.error(f"Failed to import required modules: {e}") - print(f"Error: Required modules not available - {e}") - raise + print(f"- Successfully imported: {total_records}") + print("āœ… All records imported successfully!") + + return total_records + except Exception as e: logger.error(f"Failed to import data to TrustGraph: {e}") print(f"Import failed: {e}") @@ -1024,7 +990,13 @@ For more information on the descriptor format, see: '--error-file', help='Path to write error records (optional)' ) - + + parser.add_argument( + '-t', '--token', + default=default_token, + help='Authentication token (default: $TRUSTGRAPH_TOKEN)', + ) + args = parser.parse_args() # Input validation @@ -1077,7 +1049,8 @@ For more information on the descriptor format, see: user=args.user, collection=args.collection, dry_run=args.dry_run, - verbose=args.verbose + verbose=args.verbose, + token=args.token ) except FileNotFoundError as e: print(f"Error: File not found - {e}", file=sys.stderr) diff --git a/trustgraph-cli/trustgraph/cli/put_flow_class.py b/trustgraph-cli/trustgraph/cli/put_flow_class.py index 5b4bc44b..6a88421d 100644 --- a/trustgraph-cli/trustgraph/cli/put_flow_class.py +++ b/trustgraph-cli/trustgraph/cli/put_flow_class.py @@ -9,10 +9,11 @@ from trustgraph.api import Api import json default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') +default_token = os.getenv("TRUSTGRAPH_TOKEN", None) -def put_flow_class(url, class_name, config): +def put_flow_class(url, class_name, config, token=None): - api = Api(url) + api = Api(url, token=token) class_names = api.flow().put_class(class_name, config) @@ -29,6 +30,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( '-n', '--class-name', help=f'Flow class name', @@ -47,6 +54,7 @@ def main(): url=args.api_url, class_name=args.class_name, config=json.loads(args.config), + token=args.token, ) except Exception as e: diff --git a/trustgraph-cli/trustgraph/cli/put_kg_core.py b/trustgraph-cli/trustgraph/cli/put_kg_core.py index 6374e2f6..cd0738fe 100644 --- a/trustgraph-cli/trustgraph/cli/put_kg_core.py +++ b/trustgraph-cli/trustgraph/cli/put_kg_core.py @@ -13,6 +13,7 @@ import msgpack default_url = os.getenv("TRUSTGRAPH_URL", 'ws://localhost:8088/') default_user = 'trustgraph' +default_token = os.getenv("TRUSTGRAPH_TOKEN", None) def read_message(unpacked, id, user): @@ -47,13 +48,16 @@ def read_message(unpacked, id, user): else: raise RuntimeError("Unpacked unexpected messsage type", unpacked[0]) -async def put(url, user, id, input): +async def put(url, user, id, input, token=None): if not url.endswith("/"): url += "/" url = url + "api/v1/socket" + if token: + url = f"{url}?token={token}" + async with connect(url) as ws: @@ -160,6 +164,12 @@ def main(): help=f'Input file' ) + parser.add_argument( + '-t', '--token', + default=default_token, + help='Authentication token (default: $TRUSTGRAPH_TOKEN)', + ) + args = parser.parse_args() try: @@ -170,6 +180,7 @@ def main(): user = args.user, id = args.id, input = args.input, + token = args.token, ) ) diff --git a/trustgraph-cli/trustgraph/cli/remove_library_document.py b/trustgraph-cli/trustgraph/cli/remove_library_document.py index f6e6813c..07a1fd59 100644 --- a/trustgraph-cli/trustgraph/cli/remove_library_document.py +++ b/trustgraph-cli/trustgraph/cli/remove_library_document.py @@ -10,11 +10,12 @@ from trustgraph.api import Api default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') default_user = 'trustgraph' +default_token = os.getenv("TRUSTGRAPH_TOKEN", None) -def remove_doc(url, user, id): +def remove_doc(url, user, id, token=None): - api = Api(url).library() + api = Api(url, token=token).library() api.remove_document(user=user, id=id) @@ -43,11 +44,17 @@ def main(): help=f'Document ID' ) + parser.add_argument( + '-t', '--token', + default=default_token, + help='Authentication token (default: $TRUSTGRAPH_TOKEN)', + ) + args = parser.parse_args() try: - remove_doc(args.url, args.user, args.identifier) + remove_doc(args.url, args.user, args.identifier, token=args.token) except Exception as e: diff --git a/trustgraph-cli/trustgraph/cli/set_collection.py b/trustgraph-cli/trustgraph/cli/set_collection.py index e987c4c8..1a39eb08 100644 --- a/trustgraph-cli/trustgraph/cli/set_collection.py +++ b/trustgraph-cli/trustgraph/cli/set_collection.py @@ -9,10 +9,11 @@ from trustgraph.api import Api default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') default_user = "trustgraph" +default_token = os.getenv("TRUSTGRAPH_TOKEN", None) -def set_collection(url, user, collection, name, description, tags): +def set_collection(url, user, collection, name, description, tags, token=None): - api = Api(url).collection() + api = Api(url, token=token).collection() result = api.update_collection( user=user, @@ -82,6 +83,12 @@ def main(): help='Collection tags (can be specified multiple times)' ) + parser.add_argument( + '--token', + default=default_token, + help='Authentication token (default: $TRUSTGRAPH_TOKEN)', + ) + args = parser.parse_args() try: @@ -92,7 +99,8 @@ def main(): collection = args.collection, name = args.name, description = args.description, - tags = args.tags + tags = args.tags, + token = args.token ) except Exception as e: diff --git a/trustgraph-cli/trustgraph/cli/set_prompt.py b/trustgraph-cli/trustgraph/cli/set_prompt.py index f287a9cc..bffc2cf2 100644 --- a/trustgraph-cli/trustgraph/cli/set_prompt.py +++ b/trustgraph-cli/trustgraph/cli/set_prompt.py @@ -10,10 +10,11 @@ import tabulate import textwrap default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') +default_token = os.getenv("TRUSTGRAPH_TOKEN", None) -def set_system(url, system): +def set_system(url, system, token=None): - api = Api(url).config() + api = Api(url, token=token).config() api.put([ ConfigValue(type="prompt", key="system", value=json.dumps(system)) @@ -21,9 +22,9 @@ def set_system(url, system): print("System prompt set.") -def set_prompt(url, id, prompt, response, schema): +def set_prompt(url, id, prompt, response, schema, token=None): - api = Api(url).config() + api = Api(url, token=token).config() values = api.get([ ConfigKey(type="prompt", key="template-index") @@ -71,6 +72,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( '--id', help=f'Prompt ID', @@ -103,9 +110,9 @@ def main(): if args.system: if args.id or args.prompt or args.schema or args.response: raise RuntimeError("Can't use --system with other args") - + set_system( - url=args.api_url, system=args.system + url=args.api_url, system=args.system, token=args.token ) else: @@ -130,7 +137,7 @@ def main(): set_prompt( url=args.api_url, id=args.id, prompt=args.prompt, - response=args.response, schema=schobj + response=args.response, schema=schobj, token=args.token ) except Exception as e: diff --git a/trustgraph-cli/trustgraph/cli/set_token_costs.py b/trustgraph-cli/trustgraph/cli/set_token_costs.py index 87a4e264..19b8c703 100644 --- a/trustgraph-cli/trustgraph/cli/set_token_costs.py +++ b/trustgraph-cli/trustgraph/cli/set_token_costs.py @@ -10,10 +10,11 @@ import tabulate import textwrap default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') +default_token = os.getenv("TRUSTGRAPH_TOKEN", None) -def set_costs(api_url, model, input_costs, output_costs): +def set_costs(api_url, model, input_costs, output_costs, token=None): - api = Api(api_url).config() + api = Api(api_url, token=token).config() api.put([ ConfigValue( @@ -95,6 +96,12 @@ def main(): help=f'Input costs in $ per 1M tokens', ) + parser.add_argument( + '-t', '--token', + default=default_token, + help='Authentication token (default: $TRUSTGRAPH_TOKEN)', + ) + args = parser.parse_args() try: diff --git a/trustgraph-cli/trustgraph/cli/show_flow_state.py b/trustgraph-cli/trustgraph/cli/show_flow_state.py index ca6d2b1d..6ca4df8f 100644 --- a/trustgraph-cli/trustgraph/cli/show_flow_state.py +++ b/trustgraph-cli/trustgraph/cli/show_flow_state.py @@ -9,10 +9,11 @@ import os default_metrics_url = "http://localhost:8088/api/metrics" default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') +default_token = os.getenv("TRUSTGRAPH_TOKEN", None) -def dump_status(metrics_url, api_url, flow_id): +def dump_status(metrics_url, api_url, flow_id, token=None): - api = Api(api_url).flow() + api = Api(api_url, token=token).flow() flow = api.get(flow_id) class_name = flow["class-name"] @@ -77,11 +78,17 @@ def main(): help=f'Metrics URL (default: {default_metrics_url})', ) + parser.add_argument( + '-t', '--token', + default=default_token, + help='Authentication token (default: $TRUSTGRAPH_TOKEN)', + ) + args = parser.parse_args() try: - dump_status(args.metrics_url, args.api_url, args.flow_id) + dump_status(args.metrics_url, args.api_url, args.flow_id, token=args.token) except Exception as e: diff --git a/trustgraph-cli/trustgraph/cli/show_graph.py b/trustgraph-cli/trustgraph/cli/show_graph.py index 232ebb34..b5b15e3c 100644 --- a/trustgraph-cli/trustgraph/cli/show_graph.py +++ b/trustgraph-cli/trustgraph/cli/show_graph.py @@ -9,10 +9,11 @@ from trustgraph.api import Api default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') default_user = 'trustgraph' default_collection = 'default' +default_token = os.getenv("TRUSTGRAPH_TOKEN", None) -def show_graph(url, flow_id, user, collection): +def show_graph(url, flow_id, user, collection, token=None): - api = Api(url).flow().id(flow_id) + api = Api(url, token=token).flow().id(flow_id) rows = api.triples_query( user=user, collection=collection, @@ -53,6 +54,12 @@ def main(): help=f'Collection ID (default: {default_collection})' ) + parser.add_argument( + '-t', '--token', + default=default_token, + help='Authentication token (default: $TRUSTGRAPH_TOKEN)', + ) + args = parser.parse_args() try: @@ -62,6 +69,7 @@ def main(): flow_id = args.flow_id, user = args.user, collection = args.collection, + token = args.token, ) except Exception as e: diff --git a/trustgraph-cli/trustgraph/cli/show_library_processing.py b/trustgraph-cli/trustgraph/cli/show_library_processing.py index 51dbe865..9ab69355 100644 --- a/trustgraph-cli/trustgraph/cli/show_library_processing.py +++ b/trustgraph-cli/trustgraph/cli/show_library_processing.py @@ -9,10 +9,11 @@ import json default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') default_user = "trustgraph" +default_token = os.getenv("TRUSTGRAPH_TOKEN", None) -def show_procs(url, user): +def show_procs(url, user, token=None): - api = Api(url).library() + api = Api(url, token=token).library() procs = api.get_processings(user = user) @@ -57,12 +58,18 @@ def main(): help=f'User ID (default: {default_user})' ) + parser.add_argument( + '-t', '--token', + default=default_token, + help='Authentication token (default: $TRUSTGRAPH_TOKEN)', + ) + args = parser.parse_args() try: show_procs( - url = args.api_url, user = args.user + url = args.api_url, user = args.user, token = args.token ) except Exception as e: diff --git a/trustgraph-cli/trustgraph/cli/show_token_costs.py b/trustgraph-cli/trustgraph/cli/show_token_costs.py index 2f889eef..9e7c352a 100644 --- a/trustgraph-cli/trustgraph/cli/show_token_costs.py +++ b/trustgraph-cli/trustgraph/cli/show_token_costs.py @@ -12,10 +12,11 @@ import textwrap tabulate.PRESERVE_WHITESPACE = True default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') +default_token = os.getenv("TRUSTGRAPH_TOKEN", None) -def show_config(url): +def show_config(url, token=None): - api = Api(url).config() + api = Api(url, token=token).config() models = api.list("token-costs") @@ -61,12 +62,19 @@ def main(): help=f'API URL (default: {default_url})', ) + parser.add_argument( + '-t', '--token', + default=default_token, + help='Authentication token (default: $TRUSTGRAPH_TOKEN)', + ) + args = parser.parse_args() try: show_config( url=args.api_url, + token=args.token, ) except Exception as e: