Topics is working, added a test

This commit is contained in:
Cyber MacGeddon 2024-10-26 16:25:38 +01:00
parent 629ab0d3ea
commit 4984a205d9
2 changed files with 30 additions and 1 deletions

19
tests/test-lang-topics Executable file
View file

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

View file

@ -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):