- Fix topics extract (#126)

- Add schemas to extraction
trustgraph-ai · Oct 27, 2024 · b07d539 · b07d539
1 parent 1e13776
commit b07d539
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git 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
@@ -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
         ]