- Fix topics extract (#126)

- Add schemas to extraction
This commit is contained in:
cybermaggedon 2024-10-27 08:22:12 +00:00 committed by GitHub
parent 1e137768ca
commit b07d539802
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 12 additions and 2 deletions

View file

@ -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");

View file

@ -28,7 +28,7 @@ class Relationship:
@dataclasses.dataclass
class Topic:
topic: str
name: str
definition: str
class PromptClient(BaseClient):
@ -118,7 +118,7 @@ class PromptClient(BaseClient):
)
return [
Topic(topic=d["topic"], definition=d["definition"])
Topic(name=d["topic"], definition=d["definition"])
for d in topics
]