mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-17 09:11:03 +02:00
Relationships extraction working
This commit is contained in:
Cyber MacGeddon
parent
3883391f90
commit
e6b1bb4919
3 changed files with 70 additions and 74 deletions
|
|
@ -2,7 +2,7 @@
|
|||
from pulsar.schema import Record, Bytes, String, Boolean, Array, Map, Integer
|
||||
|
||||
from . topic import topic
|
||||
from . types import Error, RowSchema
|
||||
from . types import Error
|
||||
|
||||
############################################################################
|
||||
|
||||
|
|
|
|||
|
|
@ -7,8 +7,6 @@ entity/context definitions for embedding.
|
|||
|
||||
import json
|
||||
import urllib.parse
|
||||
from pulsar.schema import JsonSchema
|
||||
import uuid
|
||||
|
||||
from .... schema import Chunk, Triple, Triples, Metadata, Value
|
||||
from .... schema import EntityContext, EntityContexts
|
||||
|
|
@ -75,10 +73,6 @@ class Processor(FlowProcessor):
|
|||
|
||||
return uri
|
||||
|
||||
def get_definitions(self, chunk):
|
||||
|
||||
return self.prompt.request_definitions(chunk)
|
||||
|
||||
async def emit_triples(self, pub, metadata, triples):
|
||||
|
||||
t = Triples(
|
||||
|
|
@ -97,8 +91,6 @@ class Processor(FlowProcessor):
|
|||
|
||||
async def on_message(self, msg, consumer, flow):
|
||||
|
||||
id = str(uuid.uuid4())
|
||||
|
||||
v = msg.value()
|
||||
print(f"Indexing {v.metadata.id}...", flush=True)
|
||||
|
||||
|
|
|
|||
|
|
@ -5,59 +5,56 @@ relationship analysis to get entity relationship edges which are output as
|
|||
graph edges.
|
||||
"""
|
||||
|
||||
import json
|
||||
import urllib.parse
|
||||
|
||||
from .... schema import Chunk, Triple, Triples
|
||||
from .... schema import Metadata, Value
|
||||
from .... schema import chunk_ingest_queue, triples_store_queue
|
||||
from .... schema import prompt_request_queue
|
||||
from .... schema import prompt_response_queue
|
||||
from .... log_level import LogLevel
|
||||
from .... clients.prompt_client import PromptClient
|
||||
from .... schema import PromptRequest, PromptResponse
|
||||
from .... rdf import RDF_LABEL, TRUSTGRAPH_ENTITIES, SUBJECT_OF
|
||||
from .... base import ConsumerProducer
|
||||
|
||||
from .... base import FlowProcessor, RequestResponseSpec, ConsumerSpec
|
||||
from .... base import ProducerSpec
|
||||
|
||||
RDF_LABEL_VALUE = Value(value=RDF_LABEL, is_uri=True)
|
||||
SUBJECT_OF_VALUE = Value(value=SUBJECT_OF, is_uri=True)
|
||||
|
||||
module = "kg-extract-relationships"
|
||||
default_ident = "kg-extract-relationships"
|
||||
|
||||
default_input_queue = chunk_ingest_queue
|
||||
default_output_queue = triples_store_queue
|
||||
default_subscriber = module
|
||||
|
||||
class Processor(ConsumerProducer):
|
||||
class Processor(FlowProcessor):
|
||||
|
||||
def __init__(self, **params):
|
||||
|
||||
input_queue = params.get("input_queue", default_input_queue)
|
||||
output_queue = params.get("output_queue", default_output_queue)
|
||||
subscriber = params.get("subscriber", default_subscriber)
|
||||
pr_request_queue = params.get(
|
||||
"prompt_request_queue", prompt_request_queue
|
||||
)
|
||||
pr_response_queue = params.get(
|
||||
"prompt_response_queue", prompt_response_queue
|
||||
)
|
||||
id = params.get("id")
|
||||
|
||||
super(Processor, self).__init__(
|
||||
**params | {
|
||||
"input_queue": input_queue,
|
||||
"output_queue": output_queue,
|
||||
"subscriber": subscriber,
|
||||
"input_schema": Chunk,
|
||||
"output_schema": Triples,
|
||||
"prompt_request_queue": pr_request_queue,
|
||||
"prompt_response_queue": pr_response_queue,
|
||||
"id": id,
|
||||
}
|
||||
)
|
||||
|
||||
self.prompt = PromptClient(
|
||||
pulsar_host=self.pulsar_host,
|
||||
pulsar_api_key=self.pulsar_api_key,
|
||||
input_queue=pr_request_queue,
|
||||
output_queue=pr_response_queue,
|
||||
subscriber = module + "-prompt",
|
||||
self.register_specification(
|
||||
ConsumerSpec(
|
||||
name = "input",
|
||||
schema = Chunk,
|
||||
handler = self.on_message
|
||||
)
|
||||
)
|
||||
|
||||
self.register_specification(
|
||||
RequestResponseSpec(
|
||||
request_name = "prompt-request",
|
||||
request_schema = PromptRequest,
|
||||
response_name = "prompt-response",
|
||||
response_schema = PromptResponse,
|
||||
)
|
||||
)
|
||||
|
||||
self.register_specification(
|
||||
ProducerSpec(
|
||||
name = "triples",
|
||||
schema = Triples
|
||||
)
|
||||
)
|
||||
|
||||
def to_uri(self, text):
|
||||
|
|
@ -68,28 +65,49 @@ class Processor(ConsumerProducer):
|
|||
|
||||
return uri
|
||||
|
||||
def get_relationships(self, chunk):
|
||||
|
||||
return self.prompt.request_relationships(chunk)
|
||||
|
||||
async def emit_edges(self, metadata, triples):
|
||||
async def emit_triples(self, pub, metadata, triples):
|
||||
|
||||
t = Triples(
|
||||
metadata=metadata,
|
||||
triples=triples,
|
||||
)
|
||||
await self.send(t)
|
||||
await pub.send(t)
|
||||
|
||||
async def handle(self, msg):
|
||||
async def on_message(self, msg, consumer, flow):
|
||||
|
||||
v = msg.value()
|
||||
print(f"Indexing {v.metadata.id}...", flush=True)
|
||||
|
||||
chunk = v.chunk.decode("utf-8")
|
||||
|
||||
print(chunk, flush=True)
|
||||
|
||||
try:
|
||||
|
||||
rels = self.get_relationships(chunk)
|
||||
try:
|
||||
|
||||
resp = await flow("prompt-request").request(
|
||||
PromptRequest(
|
||||
id="extract-relationships",
|
||||
terms={
|
||||
"text": json.dumps(chunk)
|
||||
},
|
||||
)
|
||||
)
|
||||
print("Response", resp, flush=True)
|
||||
|
||||
if resp.error is not None:
|
||||
print("Error:", resp.error.message, flush=True)
|
||||
raise RuntimeError(resp.error.message)
|
||||
|
||||
if resp.object is None:
|
||||
raise RuntimeError("Expecting object in prompt response")
|
||||
|
||||
rels = json.loads(resp.object)
|
||||
|
||||
except Exception as e:
|
||||
print("Prompt exception:", e, flush=True)
|
||||
raise e
|
||||
|
||||
triples = []
|
||||
|
||||
|
|
@ -100,9 +118,9 @@ class Processor(ConsumerProducer):
|
|||
|
||||
for rel in rels:
|
||||
|
||||
s = rel.s
|
||||
p = rel.p
|
||||
o = rel.o
|
||||
s = rel["s"]
|
||||
p = rel["p"]
|
||||
o = rel["o"]
|
||||
|
||||
if s == "": continue
|
||||
if p == "": continue
|
||||
|
|
@ -118,7 +136,7 @@ class Processor(ConsumerProducer):
|
|||
p_uri = self.to_uri(p)
|
||||
p_value = Value(value=str(p_uri), is_uri=True)
|
||||
|
||||
if rel.o_entity:
|
||||
if rel["o_entity"]:
|
||||
o_uri = self.to_uri(o)
|
||||
o_value = Value(value=str(o_uri), is_uri=True)
|
||||
else:
|
||||
|
|
@ -144,7 +162,7 @@ class Processor(ConsumerProducer):
|
|||
o=Value(value=str(p), is_uri=False)
|
||||
))
|
||||
|
||||
if rel.o_entity:
|
||||
if rel["o_entity"]:
|
||||
# Label for o
|
||||
triples.append(Triple(
|
||||
s=o_value,
|
||||
|
|
@ -159,7 +177,7 @@ class Processor(ConsumerProducer):
|
|||
o=Value(value=v.metadata.id, is_uri=True)
|
||||
))
|
||||
|
||||
if rel.o_entity:
|
||||
if rel["o_entity"]:
|
||||
# 'Subject of' for o
|
||||
triples.append(Triple(
|
||||
s=o_value,
|
||||
|
|
@ -168,6 +186,7 @@ class Processor(ConsumerProducer):
|
|||
))
|
||||
|
||||
await self.emit_edges(
|
||||
flow("triples"),
|
||||
Metadata(
|
||||
id=v.metadata.id,
|
||||
metadata=[],
|
||||
|
|
@ -185,24 +204,9 @@ class Processor(ConsumerProducer):
|
|||
@staticmethod
|
||||
def add_args(parser):
|
||||
|
||||
ConsumerProducer.add_args(
|
||||
parser, default_input_queue, default_subscriber,
|
||||
default_output_queue,
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--prompt-request-queue',
|
||||
default=prompt_request_queue,
|
||||
help=f'Prompt request queue (default: {prompt_request_queue})',
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
'--prompt-response-queue',
|
||||
default=prompt_response_queue,
|
||||
help=f'Prompt response queue (default: {prompt_response_queue})',
|
||||
)
|
||||
FlowProcessor.add_args(parser)
|
||||
|
||||
def run():
|
||||
|
||||
Processor.launch(module, __doc__)
|
||||
Processor.launch(default_ident, __doc__)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue