Structuring the service

trustgraph-base/trustgraph/base/__init__.py

@@ -30,4 +30,5 @@ from . agent_service import AgentService
 from . graph_rag_client import GraphRagClientSpec
 from . tool_service import ToolService
 from . tool_client import ToolClientSpec
+from . agent_client import AgentClientSpec
 

trustgraph-base/trustgraph/base/agent_client.py

@@ -25,15 +25,15 @@ class AgentClient(RequestResponse):
 
         return resp
 
-class GraphEmbeddingsClientSpec(RequestResponseSpec):
+class AgentClientSpec(RequestResponseSpec):
     def __init__(
             self, request_name, response_name,
     ):
-        super(GraphEmbeddingsClientSpec, self).__init__(
+        super(AgentClientSpec, self).__init__(
             request_name = request_name,
             request_schema = GraphEmbeddingsRequest,
             response_name = response_name,
             response_schema = GraphEmbeddingsResponse,
-            impl = GraphEmbeddingsClient,
+            impl = AgentClient,
         )
 

trustgraph-flow/trustgraph/extract/kg/agent/extract.py

@@ -8,7 +8,7 @@ from .... schema import ConfigRequest, ConfigResponse
 from .... rdf import TRUSTGRAPH_ENTITIES, RDF_LABEL, SUBJECT_OF, DEFINITION
 
 from .... base import FlowProcessor, ConsumerSpec, ProducerSpec
-from .... base import AgentClientSpec, ConfigClientSpec
+from .... base import AgentClientSpec
 
 default_ident = "kg-extract-relationships"
 default_concurrency = 1
@@ -53,13 +53,6 @@ class Processor(FlowProcessor):
             )
         )
 
-        self.register_specification(
-            ConfigClientSpec(
-                request_name="config-request",
-                response_name="config-response",
-            )
-        )
-
         self.register_specification(
             ProducerSpec(
                 name="triples",
@@ -277,6 +270,14 @@ class Processor(FlowProcessor):
 
     @staticmethod
     def add_args(parser):
+
+        parser.add_argument(
+            '-c', '--concurrency',
+            type=int,
+            default=default_concurrency,
+            help=f'Concurrent processing threads (default: {default_concurrency})'
+        )
+
         parser.add_argument(
             "--template-id",
             type=str,
@@ -284,19 +285,14 @@ class Processor(FlowProcessor):
             help="Template ID to use for agent extraction"
         )
 
-    @staticmethod
-    def run():
-        import argparse
-        from .... base import run_processor
-        
-        parser = argparse.ArgumentParser(
-            description="Agent-based knowledge extractor"
+        parser.add_argument(
+            '--config-type',
+            default="prompt",
+            help=f'Configuration key for prompts (default: prompt)',
         )
-        Processor.add_args(parser)
-        args = parser.parse_args()
-        
-        processor = Processor(
-            template_id=args.template_id
-        )
-        
-        run_processor(processor)
+
+        FlowProcessor.add_args(parser)
+
+def run():
+
+    Processor.launch(default_ident, __doc__)

trustgraph-flow/trustgraph/model/prompt/template/service.py

@@ -230,14 +230,14 @@ class Processor(FlowProcessor):
             help=f'Concurrent processing threads (default: {default_concurrency})'
         )
 
-        FlowProcessor.add_args(parser)
-
         parser.add_argument(
             '--config-type',
             default="prompt",
             help=f'Configuration key for prompts (default: prompt)',
         )
 
+        FlowProcessor.add_args(parser)
+
 def run():
 
     Processor.launch(default_ident, __doc__)