diff --git a/trustgraph-flow/trustgraph/extract/kg/agent/extract.py b/trustgraph-flow/trustgraph/extract/kg/agent/extract.py index d1b5ca85..8c2b1aee 100644 --- a/trustgraph-flow/trustgraph/extract/kg/agent/extract.py +++ b/trustgraph-flow/trustgraph/extract/kg/agent/extract.py @@ -13,7 +13,7 @@ from .... base import AgentClientSpec default_ident = "kg-extract-relationships" default_concurrency = 1 default_template_id = "agent-kg-extract" -default_config_key = "prompt" +default_config_type = "prompt" class Processor(FlowProcessor): @@ -22,12 +22,12 @@ class Processor(FlowProcessor): id = params.get("id") concurrency = params.get("concurrency", 1) template_id = params.get("template-id", default_template_id) - config_key = params.get("config-key", default_config_key) + config_key = params.get("config-type", default_config_type) super().__init__(**params | { "id": id, "template-id": template_id, - "config-key": config_key, + "config-type": config_key, "concurrency": concurrency, }) diff --git a/trustgraph-flow/trustgraph/model/prompt/template/service.py b/trustgraph-flow/trustgraph/model/prompt/template/service.py index 641f764a..981c76cc 100755 --- a/trustgraph-flow/trustgraph/model/prompt/template/service.py +++ b/trustgraph-flow/trustgraph/model/prompt/template/service.py @@ -33,6 +33,7 @@ class Processor(FlowProcessor): super(Processor, self).__init__( **params | { "id": id, + "config-type": config_key, "concurrency": concurrency, } ) diff --git a/trustgraph-flow/trustgraph/templates/__init__.py b/trustgraph-flow/trustgraph/templates/__init__.py new file mode 100644 index 00000000..cabd9e97 --- /dev/null +++ b/trustgraph-flow/trustgraph/templates/__init__.py @@ -0,0 +1,3 @@ + +from .prompt_manager import * + diff --git a/trustgraph-flow/trustgraph/model/prompt/template/prompt_manager.py b/trustgraph-flow/trustgraph/templates/prompt_manager.py similarity index 100% rename from trustgraph-flow/trustgraph/model/prompt/template/prompt_manager.py rename to trustgraph-flow/trustgraph/templates/prompt_manager.py