apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-17 09:11:03 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

All tests working

Browse source
This commit is contained in:
Cyber MacGeddon 2024-10-26 16:20:59 +01:00
parent 5e86a3c3ae
commit 629ab0d3ea
3 changed files with 27 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

29
trustgraph-base/trustgraph/clients/prompt_client.py
View file

@ -19,6 +19,13 @@ class Definition:
name: str
definition: str
@dataclasses.dataclass
class Relationship:
s: str
p: str
o: str
o_entity: str
class PromptClient(BaseClient):
def __init__(
@ -77,7 +84,7 @@ class PromptClient(BaseClient):
def request_relationships(self, chunk, timeout=300):
return self.call(
rels = self.request(
id="extract-relationships",
terms={
"text": chunk
@ -85,9 +92,19 @@ class PromptClient(BaseClient):
timeout=timeout
)
return [
Relationship(
s=d["subject"],
p=d["predicate"],
o=d["object"],
o_entity=d["object-entity"]
)
for d in rels
]
def request_topics(self, chunk, timeout=300):
return self.call(
return self.request(
id="extract-topics",
terms={
"text": chunk
@ -97,7 +114,7 @@ class PromptClient(BaseClient):
def request_rows(self, schema, chunk, timeout=300):
return self.call(
return self.request(
id="extract-rows",
terms={
"chunk": chunk,
@ -119,11 +136,11 @@ class PromptClient(BaseClient):
def request_kg_prompt(self, query, kg, timeout=300):
return self.call(
return self.request(
id="kg-prompt",
terms={
"query": query,
"kg": [
"knowledge": [
{ "s": v[0], "p": v[1], "o": v[2] }
for v in kg
]
@ -133,7 +150,7 @@ class PromptClient(BaseClient):
def request_document_prompt(self, query, documents, timeout=300):
return self.call(
return self.request(
id="document-prompt",
terms={
"query": query,