diff --git a/tests/test-lang-topics b/tests/test-lang-topics new file mode 100755 index 00000000..2b668524 --- /dev/null +++ b/tests/test-lang-topics @@ -0,0 +1,19 @@ +#!/usr/bin/env python3 + +import pulsar +from trustgraph.clients.prompt_client import PromptClient + +p = PromptClient(pulsar_host="pulsar://localhost:6650") + +chunk = """I noticed a cat in my garden. It is a four-legged animal +which is a mammal and can be tame or wild. I wonder if it will be friends +with me. I think the cat's name is Fred and it has 4 legs""" + +resp = p.request_topics( + chunk=chunk, +) + +for d in resp: + print(d.topic) + print(" ", d.definition) + diff --git a/trustgraph-base/trustgraph/clients/prompt_client.py b/trustgraph-base/trustgraph/clients/prompt_client.py index 9a41cbf9..b43a854d 100644 --- a/trustgraph-base/trustgraph/clients/prompt_client.py +++ b/trustgraph-base/trustgraph/clients/prompt_client.py @@ -26,6 +26,11 @@ class Relationship: o: str o_entity: str +@dataclasses.dataclass +class Topic: + topic: str + definition: str + class PromptClient(BaseClient): def __init__( @@ -104,13 +109,18 @@ class PromptClient(BaseClient): def request_topics(self, chunk, timeout=300): - return self.request( + topics = self.request( id="extract-topics", terms={ "text": chunk }, timeout=timeout ) + + return [ + Topic(topic=d["topic"], definition=d["definition"]) + for d in topics + ] def request_rows(self, schema, chunk, timeout=300):