mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-16 16:51:02 +02:00
Topics is working, added a test
This commit is contained in:
parent
629ab0d3ea
commit
4984a205d9
2 changed files with 30 additions and 1 deletions
19
tests/test-lang-topics
Executable file
19
tests/test-lang-topics
Executable 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)
|
||||
|
||||
|
|
@ -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):
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue