diff --git a/trustgraph-base/trustgraph/clients/document_rag_client.py b/trustgraph-base/trustgraph/clients/document_rag_client.py index 1e8a706d..103cbb69 100644 --- a/trustgraph-base/trustgraph/clients/document_rag_client.py +++ b/trustgraph-base/trustgraph/clients/document_rag_client.py @@ -38,7 +38,7 @@ def __init__( output_schema=DocumentRagResponse, ) - def request(self, query, timeout=500): + def request(self, query, timeout=300): return self.call( query=query, timeout=timeout diff --git a/trustgraph-base/trustgraph/clients/triples_query_client.py b/trustgraph-base/trustgraph/clients/triples_query_client.py index 6c18ac3f..fc1e4b26 100644 --- a/trustgraph-base/trustgraph/clients/triples_query_client.py +++ b/trustgraph-base/trustgraph/clients/triples_query_client.py @@ -52,7 +52,7 @@ def request( self, s, p, o, user="trustgraph", collection="default", - limit=10, timeout=60, + limit=10, timeout=120, ): return self.call( s=self.create_value(s), diff --git a/trustgraph-parquet/scripts/load-graph-embeddings b/trustgraph-parquet/scripts/load-graph-embeddings index 2dc3c06f..0e6ecf93 100755 --- a/trustgraph-parquet/scripts/load-graph-embeddings +++ b/trustgraph-parquet/scripts/load-graph-embeddings @@ -6,7 +6,7 @@ Loads Graph embeddings into TrustGraph processing. import pulsar from pulsar.schema import JsonSchema -from trustgraph.schema import GraphEmbeddings, Value +from trustgraph.schema import GraphEmbeddings, Value, Metadata from trustgraph.schema import graph_embeddings_store_queue import argparse import os @@ -24,6 +24,8 @@ class Loader: output_queue, log_level, file, + user, + collection, ): self.client = pulsar.Client( @@ -38,6 +40,8 @@ class Loader: ) self.file = file + self.user = user + self.collection = collection def run(self): @@ -66,11 +70,16 @@ class Loader: n = ent.as_py() r = GraphEmbeddings( + metadata=Metadata( + metadata=[], + user=self.user, + collection=self.collection, + ), vectors=b, entity=Value( value=n, is_uri=n.startswith("https:") - ) + ), ) self.producer.send(r) @@ -90,6 +99,8 @@ def main(): default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650') default_output_queue = graph_embeddings_store_queue + default_user = 'trustgraph' + default_collection = 'default' parser.add_argument( '-p', '--pulsar-host', @@ -103,6 +114,18 @@ def main(): help=f'Output queue (default: {default_output_queue})' ) + parser.add_argument( + '-u', '--user', + default=default_user, + help=f'User ID (default: {default_user})' + ) + + parser.add_argument( + '-c', '--collection', + default=default_collection, + help=f'Collection ID (default: {default_collection})' + ) + parser.add_argument( '-l', '--log-level', type=LogLevel, @@ -127,6 +150,8 @@ def main(): output_queue=args.output_queue, log_level=args.log_level, file=args.file, + user=args.user, + collection=args.collection, ) p.run() diff --git a/trustgraph-parquet/scripts/load-triples b/trustgraph-parquet/scripts/load-triples index e03c065b..e6bb0ff7 100755 --- a/trustgraph-parquet/scripts/load-triples +++ b/trustgraph-parquet/scripts/load-triples @@ -6,7 +6,7 @@ Loads Graph embeddings into TrustGraph processing. import pulsar from pulsar.schema import JsonSchema -from trustgraph.schema import Triple, Value +from trustgraph.schema import Triples, Triple, Value, Metadata from trustgraph.schema import triples_store_queue import argparse import os @@ -24,6 +24,8 @@ class Loader: output_queue, log_level, file, + user, + collection, ): self.client = pulsar.Client( @@ -33,11 +35,13 @@ class Loader: self.producer = self.client.create_producer( topic=output_queue, - schema=JsonSchema(Triple), + schema=JsonSchema(Triples), chunking_enabled=True, ) self.file = file + self.user = user + self.collection = collection def run(self): @@ -66,10 +70,26 @@ class Loader: for s, p, o in zip(sc, pc, oc): - r = Triple( - s=Value(value=s.as_py(), is_uri=True), - p=Value(value=p.as_py(), is_uri=True), - o=Value(value=o.as_py(), is_uri=o.as_py().startswith("https:")) + r = Triples( + metadata=Metadata( + metadata=[], + user=self.user, + collection=self.collection, + ), + triples=[ + Triple( + s=Value( + value=s.as_py(), is_uri=True + ), + p=Value( + value=p.as_py(), is_uri=True + ), + o=Value( + value=o.as_py(), + is_uri=o.as_py().startswith("https:") + ) + ) + ] ) self.producer.send(r) @@ -89,6 +109,8 @@ def main(): default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650') default_output_queue = triples_store_queue + default_user = 'trustgraph' + default_collection = 'default' parser.add_argument( '-p', '--pulsar-host', @@ -102,6 +124,18 @@ def main(): help=f'Output queue (default: {default_output_queue})' ) + parser.add_argument( + '-u', '--user', + default=default_user, + help=f'User ID (default: {default_user})' + ) + + parser.add_argument( + '-c', '--collection', + default=default_collection, + help=f'Collection ID (default: {default_collection})' + ) + parser.add_argument( '-l', '--log-level', type=LogLevel, @@ -126,6 +160,8 @@ def main(): output_queue=args.output_queue, log_level=args.log_level, file=args.file, + user=args.user, + collection=args.collection, ) p.run()