mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-21 19:21:03 +02:00
Fetch and list mostly there
This commit is contained in:
Cyber MacGeddon
parent
06dadf971e
commit
cb529db523
7 changed files with 300 additions and 218 deletions
|
|
@ -31,14 +31,11 @@ class KnowledgeRequest(Record):
|
|||
user = String()
|
||||
|
||||
# fetch-kg-core, list-kg-cores, delete-kg-core
|
||||
collection = String()
|
||||
|
||||
# fetch-kg-core, list-kg-cores, delete-kg-core
|
||||
document_id = String()
|
||||
id = String()
|
||||
|
||||
class KnowledgeResponse(Record):
|
||||
error = Error()
|
||||
document_ids = Array(String())
|
||||
ids = Array(String())
|
||||
eos = Boolean() # Indicates end of knowledge core stream
|
||||
triples = Triples()
|
||||
graph_embeddings = GraphEmbeddings()
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
|
||||
from .. schema import KnowledgeResponse, Error, Triple
|
||||
from .. schema import KnowledgeResponse, Error
|
||||
from .. knowledge import hash
|
||||
from .. exceptions import RequestError
|
||||
from .. tables.knowledge import KnowledgeTableStore
|
||||
|
|
@ -18,82 +18,7 @@ class KnowledgeManager:
|
|||
cassandra_host, cassandra_user, cassandra_password, keyspace
|
||||
)
|
||||
|
||||
async def ASDlist_cores(self, request):
|
||||
|
||||
if request.document_metadata.kind not in (
|
||||
"text/plain", "application/pdf"
|
||||
):
|
||||
raise RequestError(
|
||||
"Invalid document kind: " + request.document_metadata.kind
|
||||
)
|
||||
|
||||
if await self.table_store.document_exists(
|
||||
request.document_metadata.user,
|
||||
request.document_metadata.id
|
||||
):
|
||||
raise RuntimeError("Document already exists")
|
||||
|
||||
# Create object ID for blob
|
||||
object_id = uuid.uuid4()
|
||||
|
||||
print("Add blob...")
|
||||
|
||||
await self.blob_store.add(
|
||||
object_id, base64.b64decode(request.content),
|
||||
request.document_metadata.kind
|
||||
)
|
||||
|
||||
print("Add table...")
|
||||
|
||||
await self.table_store.add_document(
|
||||
request.document_metadata, object_id
|
||||
)
|
||||
|
||||
print("Add complete", flush=True)
|
||||
|
||||
return LibrarianResponse(
|
||||
error = None,
|
||||
document_metadata = None,
|
||||
content = None,
|
||||
document_metadatas = None,
|
||||
processing_metadatas = None,
|
||||
)
|
||||
|
||||
async def fetch_core(self, request):
|
||||
|
||||
print("Removing doc...")
|
||||
|
||||
if not await self.table_store.document_exists(
|
||||
request.user,
|
||||
request.document_id,
|
||||
):
|
||||
raise RuntimeError("Document does not exist")
|
||||
|
||||
object_id = await self.table_store.get_document_object_id(
|
||||
request.user,
|
||||
request.document_id
|
||||
)
|
||||
|
||||
# Remove blob...
|
||||
await self.blob_store.remove(object_id)
|
||||
|
||||
# Remove doc table row
|
||||
await self.table_store.remove_document(
|
||||
request.user,
|
||||
request.document_id
|
||||
)
|
||||
|
||||
print("Remove complete", flush=True)
|
||||
|
||||
return LibrarianResponse(
|
||||
error = None,
|
||||
document_metadata = None,
|
||||
content = None,
|
||||
document_metadatas = None,
|
||||
processing_metadatas = None,
|
||||
)
|
||||
|
||||
async def update_document(self, request):
|
||||
async def delete_kg_core(self, request):
|
||||
|
||||
print("Updating doc...")
|
||||
|
||||
|
|
@ -117,129 +42,48 @@ class KnowledgeManager:
|
|||
processing_metadatas = None,
|
||||
)
|
||||
|
||||
async def get_document_metadata(self, request):
|
||||
async def fetch_kg_core(self, request, flow):
|
||||
|
||||
print("Get doc...")
|
||||
print("Fetch core...")
|
||||
|
||||
doc = await self.table_store.get_document(
|
||||
request.user,
|
||||
request.document_id
|
||||
)
|
||||
|
||||
print("Get complete", flush=True)
|
||||
|
||||
return LibrarianResponse(
|
||||
error = None,
|
||||
document_metadata = doc,
|
||||
content = None,
|
||||
document_metadatas = None,
|
||||
processing_metadatas = None,
|
||||
)
|
||||
|
||||
async def get_document_content(self, request):
|
||||
|
||||
print("Get doc content...")
|
||||
|
||||
object_id = await self.table_store.get_document_object_id(
|
||||
request.user,
|
||||
request.document_id
|
||||
)
|
||||
|
||||
content = await self.blob_store.get(
|
||||
object_id
|
||||
)
|
||||
|
||||
print("Get complete", flush=True)
|
||||
|
||||
return LibrarianResponse(
|
||||
error = None,
|
||||
document_metadata = None,
|
||||
content = base64.b64encode(content),
|
||||
document_metadatas = None,
|
||||
processing_metadatas = None,
|
||||
)
|
||||
|
||||
async def add_processing(self, request):
|
||||
|
||||
print("Add processing")
|
||||
|
||||
if await self.table_store.processing_exists(
|
||||
request.processing_metadata.user,
|
||||
request.processing_metadata.id
|
||||
):
|
||||
raise RuntimeError("Processing already exists")
|
||||
|
||||
doc = await self.table_store.get_document(
|
||||
request.processing_metadata.user,
|
||||
request.processing_metadata.document_id
|
||||
)
|
||||
|
||||
object_id = await self.table_store.get_document_object_id(
|
||||
request.processing_metadata.user,
|
||||
request.processing_metadata.document_id
|
||||
)
|
||||
|
||||
content = await self.blob_store.get(
|
||||
object_id
|
||||
)
|
||||
|
||||
print("Got content")
|
||||
|
||||
print("Add processing...")
|
||||
|
||||
await self.table_store.add_processing(request.processing_metadata)
|
||||
|
||||
print("Invoke document processing...")
|
||||
|
||||
await self.load_document(
|
||||
document = doc,
|
||||
processing = request.processing_metadata,
|
||||
content = content,
|
||||
)
|
||||
|
||||
print("Add complete", flush=True)
|
||||
|
||||
return LibrarianResponse(
|
||||
error = None,
|
||||
document_metadata = None,
|
||||
content = None,
|
||||
document_metadatas = None,
|
||||
processing_metadatas = None,
|
||||
)
|
||||
|
||||
async def remove_processing(self, request):
|
||||
|
||||
print("Removing processing...")
|
||||
|
||||
if not await self.table_store.processing_exists(
|
||||
request.user,
|
||||
request.processing_id,
|
||||
):
|
||||
raise RuntimeError("Processing object does not exist")
|
||||
async def publish(obj):
|
||||
print(obj)
|
||||
# await publisher.send(triples)
|
||||
|
||||
# Remove doc table row
|
||||
await self.table_store.remove_processing(
|
||||
await self.table_store.get_triples(
|
||||
request.user,
|
||||
request.processing_id
|
||||
request.id,
|
||||
publish
|
||||
)
|
||||
|
||||
print("Remove complete", flush=True)
|
||||
async def publish_ge(obj):
|
||||
print(obj)
|
||||
|
||||
return LibrarianResponse(
|
||||
error = None,
|
||||
document_metadata = None,
|
||||
content = None,
|
||||
document_metadatas = None,
|
||||
processing_metadatas = None,
|
||||
# Remove doc table row
|
||||
await self.table_store.get_graph_embeddings(
|
||||
request.user,
|
||||
request.id,
|
||||
publish
|
||||
)
|
||||
|
||||
async def list_cores(self, request):
|
||||
|
||||
ids = await self.table_store.list_cores(request.user)
|
||||
print("Fetch complete", flush=True)
|
||||
|
||||
return KnowledgeResponse(
|
||||
error = None,
|
||||
document_ids = ids,
|
||||
ids = None,
|
||||
eos = True,
|
||||
triples = None,
|
||||
graph_embeddings = None,
|
||||
)
|
||||
|
||||
async def list_kg_cores(self, request, flow):
|
||||
|
||||
ids = await self.table_store.list_kg_cores(request.user)
|
||||
|
||||
return KnowledgeResponse(
|
||||
error = None,
|
||||
ids = ids,
|
||||
eos = False,
|
||||
triples = None,
|
||||
graph_embeddings = None
|
||||
|
|
|
|||
|
|
@ -181,7 +181,7 @@ class Processor(AsyncProcessor):
|
|||
|
||||
print("Document submitted")
|
||||
|
||||
async def process_request(self, v):
|
||||
async def process_request(self, v, flow):
|
||||
|
||||
if v.operation is None:
|
||||
raise RequestError("Null operation")
|
||||
|
|
@ -189,15 +189,15 @@ class Processor(AsyncProcessor):
|
|||
print("request", v.operation)
|
||||
|
||||
impls = {
|
||||
"list-cores": self.knowledge.list_cores,
|
||||
"fetch-core": self.knowledge.fetch_core,
|
||||
"delet-core": self.knowledge.delete_core,
|
||||
"list-kg-cores": self.knowledge.list_kg_cores,
|
||||
"fetch-kg-core": self.knowledge.fetch_kg_core,
|
||||
"delete-kg-core": self.knowledge.delete_kg_core,
|
||||
}
|
||||
|
||||
if v.operation not in impls:
|
||||
raise RequestError(f"Invalid operation: {v.operation}")
|
||||
|
||||
return await impls[v.operation](v)
|
||||
return await impls[v.operation](v, flow)
|
||||
|
||||
async def on_knowledge_request(self, msg, consumer, flow):
|
||||
|
||||
|
|
@ -211,7 +211,7 @@ class Processor(AsyncProcessor):
|
|||
|
||||
try:
|
||||
|
||||
resp = await self.process_request(v)
|
||||
resp = await self.process_request(v, flow)
|
||||
|
||||
await self.knowledge_response_producer.send(
|
||||
resp, properties={"id": id}
|
||||
|
|
|
|||
55
trustgraph-flow/trustgraph/gateway/dispatch/knowledge.py
Normal file
55
trustgraph-flow/trustgraph/gateway/dispatch/knowledge.py
Normal file
|
|
@ -0,0 +1,55 @@
|
|||
|
||||
import base64
|
||||
|
||||
from ... schema import KnowledgeRequest, KnowledgeResponse
|
||||
from ... schema import knowledge_request_queue
|
||||
from ... schema import knowledge_response_queue
|
||||
|
||||
from . requestor import ServiceRequestor
|
||||
from . serialize import serialize_graph_embeddings
|
||||
from . serialize import serialize_triples
|
||||
from . serialize import to_document_metadata, to_processing_metadata
|
||||
|
||||
class KnowledgeRequestor(ServiceRequestor):
|
||||
def __init__(self, pulsar_client, consumer, subscriber, timeout=120):
|
||||
|
||||
super(KnowledgeRequestor, self).__init__(
|
||||
pulsar_client=pulsar_client,
|
||||
consumer_name = consumer,
|
||||
subscription = subscriber,
|
||||
request_queue=knowledge_request_queue,
|
||||
response_queue=knowledge_response_queue,
|
||||
request_schema=KnowledgeRequest,
|
||||
response_schema=KnowledgeResponse,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
def to_request(self, body):
|
||||
|
||||
return KnowledgeRequest(
|
||||
operation = body.get("operation", None),
|
||||
user = body.get("user", None),
|
||||
id = body.get("id", None),
|
||||
)
|
||||
|
||||
def from_response(self, message):
|
||||
|
||||
print(message)
|
||||
|
||||
response = {
|
||||
"eos": message.eos
|
||||
}
|
||||
|
||||
if message.ids is not None:
|
||||
response["ids"] = message.ids
|
||||
|
||||
if message.triples:
|
||||
response["triples"] = serialize_triples(message.triples)
|
||||
|
||||
if message.graph_embeddings:
|
||||
response["graph-embeddings"] = serialize_graph_embeddings(
|
||||
message.graph_embeddings
|
||||
)
|
||||
|
||||
return response, True
|
||||
|
||||
|
|
@ -5,6 +5,7 @@ import uuid
|
|||
from . config import ConfigRequestor
|
||||
from . flow import FlowRequestor
|
||||
from . librarian import LibrarianRequestor
|
||||
from . knowledge import KnowledgeRequestor
|
||||
|
||||
from . embeddings import EmbeddingsRequestor
|
||||
from . agent import AgentRequestor
|
||||
|
|
@ -44,6 +45,7 @@ global_dispatchers = {
|
|||
"config": ConfigRequestor,
|
||||
"flow": FlowRequestor,
|
||||
"librarian": LibrarianRequestor,
|
||||
"knowledge": KnowledgeRequestor,
|
||||
}
|
||||
|
||||
sender_dispatchers = {
|
||||
|
|
|
|||
|
|
@ -21,7 +21,7 @@ class Librarian:
|
|||
minio_host, minio_access_key, minio_secret_key, bucket_name
|
||||
)
|
||||
|
||||
self.table_store = TableStore(
|
||||
self.table_store = LibraryTableStore(
|
||||
cassandra_host, cassandra_user, cassandra_password, keyspace
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,7 @@
|
|||
|
||||
from .. schema import KnowledgeResponse, Triple, Triples, EntityEmbeddings
|
||||
from .. schema import Metadata, Value, GraphEmbeddings
|
||||
|
||||
from cassandra.cluster import Cluster
|
||||
from cassandra.auth import PlainTextAuthProvider
|
||||
from ssl import SSLContext, PROTOCOL_TLSv1_2
|
||||
|
|
@ -61,7 +64,6 @@ class KnowledgeTableStore:
|
|||
self.cassandra.execute("""
|
||||
CREATE TABLE IF NOT EXISTS triples (
|
||||
user text,
|
||||
collection text,
|
||||
document_id text,
|
||||
id uuid,
|
||||
time timestamp,
|
||||
|
|
@ -71,7 +73,7 @@ class KnowledgeTableStore:
|
|||
triples list<tuple<
|
||||
text, boolean, text, boolean, text, boolean
|
||||
>>,
|
||||
PRIMARY KEY (user, collection, document_id, id)
|
||||
PRIMARY KEY ((user, document_id), id)
|
||||
);
|
||||
""");
|
||||
|
||||
|
|
@ -80,7 +82,6 @@ class KnowledgeTableStore:
|
|||
self.cassandra.execute("""
|
||||
create table if not exists graph_embeddings (
|
||||
user text,
|
||||
collection text,
|
||||
document_id text,
|
||||
id uuid,
|
||||
time timestamp,
|
||||
|
|
@ -93,16 +94,20 @@ class KnowledgeTableStore:
|
|||
list<list<double>>
|
||||
>
|
||||
>,
|
||||
PRIMARY KEY (user, collection, document_id, id)
|
||||
PRIMARY KEY ((user, document_id), id)
|
||||
);
|
||||
""");
|
||||
|
||||
self.cassandra.execute("""
|
||||
CREATE INDEX IF NOT EXISTS graph_embeddings_user ON
|
||||
graph_embeddings ( user );
|
||||
""");
|
||||
|
||||
print("document_embeddings table...", flush=True)
|
||||
|
||||
self.cassandra.execute("""
|
||||
create table if not exists document_embeddings (
|
||||
user text,
|
||||
collection text,
|
||||
document_id text,
|
||||
id uuid,
|
||||
time timestamp,
|
||||
|
|
@ -115,10 +120,15 @@ class KnowledgeTableStore:
|
|||
list<list<double>>
|
||||
>
|
||||
>,
|
||||
PRIMARY KEY (user, collection, document_id, id)
|
||||
PRIMARY KEY ((user, document_id), id)
|
||||
);
|
||||
""");
|
||||
|
||||
self.cassandra.execute("""
|
||||
CREATE INDEX IF NOT EXISTS document_embeddings_user ON
|
||||
document_embeddings ( user );
|
||||
""");
|
||||
|
||||
print("Cassandra schema OK.", flush=True)
|
||||
|
||||
def prepare_statements(self):
|
||||
|
|
@ -126,28 +136,49 @@ class KnowledgeTableStore:
|
|||
self.insert_triples_stmt = self.cassandra.prepare("""
|
||||
INSERT INTO triples
|
||||
(
|
||||
id, user, collection, document_id, time,
|
||||
metadata, triples
|
||||
id, user, document_id,
|
||||
time, metadata, triples
|
||||
)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?)
|
||||
VALUES (?, ?, ?, ?, ?, ?)
|
||||
""")
|
||||
|
||||
self.insert_graph_embeddings_stmt = self.cassandra.prepare("""
|
||||
INSERT INTO graph_embeddings
|
||||
(
|
||||
id, user, collection, document_id, time,
|
||||
metadata, entity_embeddings
|
||||
id, user, document_id, time, metadata, entity_embeddings
|
||||
)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?)
|
||||
VALUES (?, ?, ?, ?, ?, ?)
|
||||
""")
|
||||
|
||||
self.insert_document_embeddings_stmt = self.cassandra.prepare("""
|
||||
INSERT INTO document_embeddings
|
||||
(
|
||||
id, user, collection, document_id, time,
|
||||
metadata, chunks
|
||||
id, user, document_id, time, metadata, chunks
|
||||
)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?)
|
||||
VALUES (?, ?, ?, ?, ?, ?)
|
||||
""")
|
||||
|
||||
self.list_cores_stmt = self.cassandra.prepare("""
|
||||
SELECT DISTINCT user, document_id FROM graph_embeddings
|
||||
WHERE user = ?
|
||||
""")
|
||||
|
||||
self.get_triples_stmt = self.cassandra.prepare("""
|
||||
SELECT id, time, metadata, triples
|
||||
FROM triples
|
||||
WHERE user = ? AND document_id = ?
|
||||
""")
|
||||
|
||||
self.get_graph_embeddings_stmt = self.cassandra.prepare("""
|
||||
SELECT id, time, metadata, entity_embeddings
|
||||
FROM graph_embeddings
|
||||
WHERE user = ? AND document_id = ?
|
||||
""")
|
||||
|
||||
self.get_document_embeddings_stmt = self.cassandra.prepare("""
|
||||
SELECT id, time, metadata, chunks
|
||||
FROM document_embeddings
|
||||
WHERE user = ? AND document_id = ?
|
||||
""")
|
||||
|
||||
async def add_triples(self, m):
|
||||
|
|
@ -181,7 +212,7 @@ class KnowledgeTableStore:
|
|||
self.insert_triples_stmt,
|
||||
(
|
||||
uuid.uuid4(), m.metadata.user,
|
||||
m.metadata.collection, m.metadata.id, when,
|
||||
m.metadata.id, when,
|
||||
metadata, triples,
|
||||
)
|
||||
)
|
||||
|
|
@ -225,7 +256,7 @@ class KnowledgeTableStore:
|
|||
self.insert_graph_embeddings_stmt,
|
||||
(
|
||||
uuid.uuid4(), m.metadata.user,
|
||||
m.metadata.collection, m.metadata.id, when,
|
||||
m.metadata.id, when,
|
||||
metadata, entities,
|
||||
)
|
||||
)
|
||||
|
|
@ -269,7 +300,7 @@ class KnowledgeTableStore:
|
|||
self.insert_document_embeddings_stmt,
|
||||
(
|
||||
uuid.uuid4(), m.metadata.user,
|
||||
m.metadata.collection, m.metadata.id, when,
|
||||
m.metadata.id, when,
|
||||
metadata, chunks,
|
||||
)
|
||||
)
|
||||
|
|
@ -282,4 +313,157 @@ class KnowledgeTableStore:
|
|||
print(f"{e}, retry...", flush=True)
|
||||
await asyncio.sleep(1)
|
||||
|
||||
|
||||
async def list_kg_cores(self, user):
|
||||
|
||||
print("List kg cores...")
|
||||
|
||||
while True:
|
||||
|
||||
try:
|
||||
|
||||
resp = self.cassandra.execute(
|
||||
self.list_cores_stmt,
|
||||
(user,)
|
||||
)
|
||||
|
||||
break
|
||||
|
||||
except Exception as e:
|
||||
print("Exception:", type(e))
|
||||
print(f"{e}, retry...", flush=True)
|
||||
await asyncio.sleep(1)
|
||||
|
||||
|
||||
lst = [
|
||||
row[1]
|
||||
for row in resp
|
||||
]
|
||||
|
||||
print("Done")
|
||||
|
||||
return lst
|
||||
|
||||
async def get_triples(self, user, document_id, receiver):
|
||||
|
||||
print("Get triples...")
|
||||
|
||||
while True:
|
||||
|
||||
try:
|
||||
|
||||
print("Executing!")
|
||||
print(self.get_triples_stmt)
|
||||
print(user, document_id)
|
||||
|
||||
resp = self.cassandra.execute(
|
||||
self.get_triples_stmt,
|
||||
(user, document_id)
|
||||
)
|
||||
print("HERE")
|
||||
|
||||
print("Executed")
|
||||
|
||||
break
|
||||
|
||||
except Exception as e:
|
||||
print("Exception:", type(e))
|
||||
print(f"{e}, retry...", flush=True)
|
||||
await asyncio.sleep(1)
|
||||
|
||||
print("Unpack...")
|
||||
for row in resp:
|
||||
|
||||
print("ROW")
|
||||
print(row)
|
||||
|
||||
if row[2]:
|
||||
metadata = [
|
||||
Triple(
|
||||
s = Value(value = elt[0], is_uri = elt[1]),
|
||||
p = Value(value = elt[2], is_uri = elt[3]),
|
||||
o = Value(value = elt[4], is_uri = elt[5]),
|
||||
)
|
||||
for elt in row[2]
|
||||
]
|
||||
else:
|
||||
metadata = []
|
||||
|
||||
triples = [
|
||||
Triple(
|
||||
s = Value(value = elt[0], is_uri = elt[1]),
|
||||
p = Value(value = elt[2], is_uri = elt[3]),
|
||||
o = Value(value = elt[4], is_uri = elt[5]),
|
||||
)
|
||||
for elt in row[3]
|
||||
]
|
||||
|
||||
await receiver(
|
||||
Triples(
|
||||
metadata = Metadata(
|
||||
id = document_id,
|
||||
user = user,
|
||||
metadata = metadata,
|
||||
),
|
||||
triples = triples
|
||||
)
|
||||
)
|
||||
|
||||
print("Done")
|
||||
|
||||
async def get_graph_embeddings(self, user, document_id, receiver):
|
||||
|
||||
print("Get GE...")
|
||||
|
||||
while True:
|
||||
|
||||
try:
|
||||
|
||||
resp = self.cassandra.execute(
|
||||
self.get_graph_embeddings_stmt,
|
||||
(user, document_id)
|
||||
)
|
||||
|
||||
break
|
||||
|
||||
except Exception as e:
|
||||
print("Exception:", type(e))
|
||||
print(f"{e}, retry...", flush=True)
|
||||
await asyncio.sleep(1)
|
||||
|
||||
for row in resp:
|
||||
|
||||
print("MD...")
|
||||
if row[2]:
|
||||
metadata = [
|
||||
Triple(
|
||||
s = Value(value = elt[0], is_uri = elt[1]),
|
||||
p = Value(value = elt[2], is_uri = elt[3]),
|
||||
o = Value(value = elt[4], is_uri = elt[5]),
|
||||
)
|
||||
for elt in row[2]
|
||||
]
|
||||
else:
|
||||
metadata = []
|
||||
|
||||
print("EE...")
|
||||
entities = [
|
||||
EntityEmbeddings(
|
||||
entity = Value(value = ent[0][0], is_uri = ent[0][1]),
|
||||
vectors = ent[1]
|
||||
)
|
||||
for ent in row[3]
|
||||
]
|
||||
|
||||
await receiver(
|
||||
GraphEmbeddings(
|
||||
metadata = Metadata(
|
||||
id = document_id,
|
||||
user = user,
|
||||
metadata = metadata,
|
||||
),
|
||||
entities = entities
|
||||
)
|
||||
)
|
||||
|
||||
print("Done")
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue