diff --git a/templates/components/prompt-template.jsonnet b/templates/components/prompt-template.jsonnet index cc4195aa..dbf2514f 100644 --- a/templates/components/prompt-template.jsonnet +++ b/templates/components/prompt-template.jsonnet @@ -49,6 +49,11 @@ local default_prompts = import "prompts/default-prompts.jsonnet"; "document-prompt=text", "extract-rows=json", + "--prompt-schema", + 'extract-definitions={ "type": "array", "items": { "type": "object", "properties": { "entity": { "type": "string" }, "definition": { "type": "string" } }, "required": [ "entity", "definition" ] } }', + 'extract-relationships={ "type": "array", "items": { "type": "object", "properties": { "subject": { "type": "string" }, "predicate": { "type": "string" }, "object": { "type": "string" }, "object-entity": { "type": "boolean" } }, "required": [ "subject", "predicate", "object", "object-entity" ] } }', + 'extract-topics={ "type": "array", "items": { "type": "object", "properties": { "topic": { "type": "string" }, "definition": { "type": "string" } }, "required": [ "topic", "definition" ] } }', + ]) .with_limits("0.5", "128M") .with_reservations("0.1", "128M"); @@ -114,6 +119,11 @@ local default_prompts = import "prompts/default-prompts.jsonnet"; "document-prompt=text", "extract-rows=json", + "--prompt-schema", + 'extract-definitions={ "type": "array", "items": { "type": "object", "properties": { "entity": { "type": "string" }, "definition": { "type": "string" } }, "required": [ "entity", "definition" ] } }', + 'extract-relationships={ "type": "array", "items": { "type": "object", "properties": { "subject": { "type": "string" }, "predicate": { "type": "string" }, "object": { "type": "string" }, "object-entity": { "type": "boolean" } }, "required": [ "subject", "predicate", "object", "object-entity" ] } }', + 'extract-topics={ "type": "array", "items": { "type": "object", "properties": { "topic": { "type": "string" }, "definition": { "type": "string" } }, "required": [ "topic", "definition" ] } }', + ]) .with_limits("0.5", "128M") .with_reservations("0.1", "128M"); diff --git a/trustgraph-base/trustgraph/clients/prompt_client.py b/trustgraph-base/trustgraph/clients/prompt_client.py index b43a854d..51a9c520 100644 --- a/trustgraph-base/trustgraph/clients/prompt_client.py +++ b/trustgraph-base/trustgraph/clients/prompt_client.py @@ -28,7 +28,7 @@ class Relationship: @dataclasses.dataclass class Topic: - topic: str + name: str definition: str class PromptClient(BaseClient): @@ -118,7 +118,7 @@ def request_topics(self, chunk, timeout=300): ) return [ - Topic(topic=d["topic"], definition=d["definition"]) + Topic(name=d["topic"], definition=d["definition"]) for d in topics ]