mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-16 08:41:03 +02:00
c23e28aa66
The Metadata dataclass dropped its `metadata: list[Triple]` field
and EntityEmbeddings/ChunkEmbeddings settled on a singular
`vector: list[float]` field, but several call sites kept passing
`Metadata(metadata=...)` and `EntityEmbeddings(vectors=...)`. The
bugs were latent until a websocket client first hit
`/api/v1/flow/default/import/entity-contexts`, at which point the
dispatcher TypeError'd on construction.
Production fixes (5 call sites on the same migration tail):
* trustgraph-flow gateway dispatchers entity_contexts_import.py
and graph_embeddings_import.py — drop the stale
Metadata(metadata=...) kwarg; switch graph_embeddings_import
to the singular `vector` wire key.
* trustgraph-base messaging translators knowledge.py and
document_loading.py — fix decode side to read the singular
`"vector"` key, matching what their own encode sides have
always written.
* trustgraph-flow tables/knowledge.py — fix Cassandra row
deserialiser to construct EntityEmbeddings(vector=...)
instead of vectors=.
* trustgraph-flow gateway core_import/core_export — switch the
kg-core msgpack wire format to the singular `"v"`/`"vector"`
key and drop the dead `m["m"]` envelope field that referenced
the removed Metadata.metadata triples list (it was a
guaranteed KeyError on the export side).
Defense-in-depth regression coverage (32 new tests across 7 files):
* tests/contract/test_schema_field_contracts.py — pin the field
set of Metadata, EntityEmbeddings, ChunkEmbeddings,
EntityContext so any future schema rename fails CI loudly
with a clear diff.
* tests/unit/test_translators/test_knowledge_translator_roundtrip.py
and test_document_embeddings_translator_roundtrip.py -
encode→decode round-trip the affected translators end to end,
locking in the singular `"vector"` wire key.
* tests/unit/test_gateway/test_entity_contexts_import_dispatcher.py
and test_graph_embeddings_import_dispatcher.py — exercise the
websocket dispatchers' receive() path with realistic
payloads, the direct regression test for the original
production crash.
* tests/unit/test_gateway/test_core_import_export_roundtrip.py
— pack/unpack the kg-core msgpack format through the real
dispatcher classes (with KnowledgeRequestor mocked),
including a full export→import round-trip.
* tests/unit/test_tables/test_knowledge_table_store.py —
exercise the Cassandra row → schema conversion via __new__ to
bypass the live cluster connection.
Also fixes an unrelated leaked-coroutine RuntimeWarning in
test_gateway/test_service.py::test_run_method_calls_web_run_app: the
mocked aiohttp.web.run_app now closes the coroutine that Api.run() hands
it, mirroring what the real run_app would do, instead of leaving it for
the GC to complain about.
465 lines
12 KiB
Python
465 lines
12 KiB
Python
|
|
from .. schema import KnowledgeResponse, Triple, Triples, EntityEmbeddings
|
|
from .. schema import Metadata, Term, IRI, LITERAL, GraphEmbeddings
|
|
|
|
from cassandra.cluster import Cluster
|
|
|
|
|
|
def term_to_tuple(term):
|
|
"""Convert Term to (value, is_uri) tuple for database storage."""
|
|
if term.type == IRI:
|
|
return (term.iri, True)
|
|
else: # LITERAL
|
|
return (term.value, False)
|
|
|
|
|
|
def tuple_to_term(value, is_uri):
|
|
"""Convert (value, is_uri) tuple from database to Term."""
|
|
if is_uri:
|
|
return Term(type=IRI, iri=value)
|
|
else:
|
|
return Term(type=LITERAL, value=value)
|
|
from cassandra.auth import PlainTextAuthProvider
|
|
from ssl import SSLContext, PROTOCOL_TLSv1_2
|
|
|
|
import uuid
|
|
import time
|
|
import asyncio
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class KnowledgeTableStore:
|
|
|
|
def __init__(
|
|
self,
|
|
cassandra_host, cassandra_username, cassandra_password, keyspace,
|
|
):
|
|
|
|
self.keyspace = keyspace
|
|
|
|
logger.info("Connecting to Cassandra...")
|
|
|
|
# Ensure cassandra_host is a list
|
|
if isinstance(cassandra_host, str):
|
|
cassandra_host = [h.strip() for h in cassandra_host.split(',')]
|
|
|
|
if cassandra_username and cassandra_password:
|
|
ssl_context = SSLContext(PROTOCOL_TLSv1_2)
|
|
auth_provider = PlainTextAuthProvider(
|
|
username=cassandra_username, password=cassandra_password
|
|
)
|
|
self.cluster = Cluster(
|
|
cassandra_host,
|
|
auth_provider=auth_provider,
|
|
ssl_context=ssl_context
|
|
)
|
|
else:
|
|
self.cluster = Cluster(cassandra_host)
|
|
|
|
self.cassandra = self.cluster.connect()
|
|
|
|
logger.info("Connected.")
|
|
|
|
self.ensure_cassandra_schema()
|
|
|
|
self.prepare_statements()
|
|
|
|
def ensure_cassandra_schema(self):
|
|
|
|
logger.debug("Ensure Cassandra schema...")
|
|
|
|
logger.debug("Keyspace...")
|
|
|
|
# FIXME: Replication factor should be configurable
|
|
self.cassandra.execute(f"""
|
|
create keyspace if not exists {self.keyspace}
|
|
with replication = {{
|
|
'class' : 'SimpleStrategy',
|
|
'replication_factor' : 1
|
|
}};
|
|
""");
|
|
|
|
self.cassandra.set_keyspace(self.keyspace)
|
|
|
|
logger.debug("triples table...")
|
|
|
|
self.cassandra.execute("""
|
|
CREATE TABLE IF NOT EXISTS triples (
|
|
user text,
|
|
document_id text,
|
|
id uuid,
|
|
time timestamp,
|
|
metadata list<tuple<
|
|
text, boolean, text, boolean, text, boolean
|
|
>>,
|
|
triples list<tuple<
|
|
text, boolean, text, boolean, text, boolean
|
|
>>,
|
|
PRIMARY KEY ((user, document_id), id)
|
|
);
|
|
""");
|
|
|
|
logger.debug("graph_embeddings table...")
|
|
|
|
self.cassandra.execute("""
|
|
create table if not exists graph_embeddings (
|
|
user text,
|
|
document_id text,
|
|
id uuid,
|
|
time timestamp,
|
|
metadata list<tuple<
|
|
text, boolean, text, boolean, text, boolean
|
|
>>,
|
|
entity_embeddings list<
|
|
tuple<
|
|
tuple<text, boolean>,
|
|
list<double>
|
|
>
|
|
>,
|
|
PRIMARY KEY ((user, document_id), id)
|
|
);
|
|
""");
|
|
|
|
self.cassandra.execute("""
|
|
CREATE INDEX IF NOT EXISTS graph_embeddings_user ON
|
|
graph_embeddings ( user );
|
|
""");
|
|
|
|
logger.debug("document_embeddings table...")
|
|
|
|
self.cassandra.execute("""
|
|
create table if not exists document_embeddings (
|
|
user text,
|
|
document_id text,
|
|
id uuid,
|
|
time timestamp,
|
|
metadata list<tuple<
|
|
text, boolean, text, boolean, text, boolean
|
|
>>,
|
|
chunks list<
|
|
tuple<
|
|
blob,
|
|
list<double>
|
|
>
|
|
>,
|
|
PRIMARY KEY ((user, document_id), id)
|
|
);
|
|
""");
|
|
|
|
self.cassandra.execute("""
|
|
CREATE INDEX IF NOT EXISTS document_embeddings_user ON
|
|
document_embeddings ( user );
|
|
""");
|
|
|
|
logger.info("Cassandra schema OK.")
|
|
|
|
def prepare_statements(self):
|
|
|
|
self.insert_triples_stmt = self.cassandra.prepare("""
|
|
INSERT INTO triples
|
|
(
|
|
id, user, document_id,
|
|
time, metadata, triples
|
|
)
|
|
VALUES (?, ?, ?, ?, ?, ?)
|
|
""")
|
|
|
|
self.insert_graph_embeddings_stmt = self.cassandra.prepare("""
|
|
INSERT INTO graph_embeddings
|
|
(
|
|
id, user, document_id, time, metadata, entity_embeddings
|
|
)
|
|
VALUES (?, ?, ?, ?, ?, ?)
|
|
""")
|
|
|
|
self.insert_document_embeddings_stmt = self.cassandra.prepare("""
|
|
INSERT INTO document_embeddings
|
|
(
|
|
id, user, document_id, time, metadata, chunks
|
|
)
|
|
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 = ?
|
|
""")
|
|
|
|
self.delete_triples_stmt = self.cassandra.prepare("""
|
|
DELETE FROM triples
|
|
WHERE user = ? AND document_id = ?
|
|
""")
|
|
|
|
self.delete_graph_embeddings_stmt = self.cassandra.prepare("""
|
|
DELETE FROM graph_embeddings
|
|
WHERE user = ? AND document_id = ?
|
|
""")
|
|
|
|
async def add_triples(self, m):
|
|
|
|
when = int(time.time() * 1000)
|
|
|
|
triples = [
|
|
(
|
|
*term_to_tuple(v.s), *term_to_tuple(v.p), *term_to_tuple(v.o)
|
|
)
|
|
for v in m.triples
|
|
]
|
|
|
|
while True:
|
|
|
|
try:
|
|
|
|
resp = self.cassandra.execute(
|
|
self.insert_triples_stmt,
|
|
(
|
|
uuid.uuid4(), m.metadata.user,
|
|
m.metadata.root or m.metadata.id, when,
|
|
[], triples,
|
|
)
|
|
)
|
|
|
|
break
|
|
|
|
except Exception as e:
|
|
|
|
logger.error("Exception occurred", exc_info=True)
|
|
raise e
|
|
|
|
async def add_graph_embeddings(self, m):
|
|
|
|
when = int(time.time() * 1000)
|
|
|
|
entities = [
|
|
(
|
|
term_to_tuple(v.entity),
|
|
v.vector
|
|
)
|
|
for v in m.entities
|
|
]
|
|
|
|
while True:
|
|
|
|
try:
|
|
|
|
resp = self.cassandra.execute(
|
|
self.insert_graph_embeddings_stmt,
|
|
(
|
|
uuid.uuid4(), m.metadata.user,
|
|
m.metadata.root or m.metadata.id, when,
|
|
[], entities,
|
|
)
|
|
)
|
|
|
|
break
|
|
|
|
except Exception as e:
|
|
|
|
logger.error("Exception occurred", exc_info=True)
|
|
raise e
|
|
|
|
async def add_document_embeddings(self, m):
|
|
|
|
when = int(time.time() * 1000)
|
|
|
|
chunks = [
|
|
(
|
|
v.chunk_id,
|
|
v.vector,
|
|
)
|
|
for v in m.chunks
|
|
]
|
|
|
|
while True:
|
|
|
|
try:
|
|
|
|
resp = self.cassandra.execute(
|
|
self.insert_document_embeddings_stmt,
|
|
(
|
|
uuid.uuid4(), m.metadata.user,
|
|
m.metadata.root or m.metadata.id, when,
|
|
[], chunks,
|
|
)
|
|
)
|
|
|
|
break
|
|
|
|
except Exception as e:
|
|
|
|
logger.error("Exception occurred", exc_info=True)
|
|
raise e
|
|
|
|
async def list_kg_cores(self, user):
|
|
|
|
logger.debug("List kg cores...")
|
|
|
|
while True:
|
|
|
|
try:
|
|
|
|
resp = self.cassandra.execute(
|
|
self.list_cores_stmt,
|
|
(user,)
|
|
)
|
|
|
|
break
|
|
|
|
except Exception as e:
|
|
logger.error("Exception occurred", exc_info=True)
|
|
raise e
|
|
|
|
|
|
lst = [
|
|
row[1]
|
|
for row in resp
|
|
]
|
|
|
|
logger.debug("Done")
|
|
|
|
return lst
|
|
|
|
async def delete_kg_core(self, user, document_id):
|
|
|
|
logger.debug("Delete kg cores...")
|
|
|
|
while True:
|
|
|
|
try:
|
|
|
|
resp = self.cassandra.execute(
|
|
self.delete_triples_stmt,
|
|
(user, document_id)
|
|
)
|
|
|
|
break
|
|
|
|
except Exception as e:
|
|
logger.error("Exception occurred", exc_info=True)
|
|
raise e
|
|
|
|
while True:
|
|
|
|
try:
|
|
|
|
resp = self.cassandra.execute(
|
|
self.delete_graph_embeddings_stmt,
|
|
(user, document_id)
|
|
)
|
|
|
|
break
|
|
|
|
except Exception as e:
|
|
logger.error("Exception occurred", exc_info=True)
|
|
raise e
|
|
|
|
async def get_triples(self, user, document_id, receiver):
|
|
|
|
logger.debug("Get triples...")
|
|
|
|
while True:
|
|
|
|
try:
|
|
|
|
resp = self.cassandra.execute(
|
|
self.get_triples_stmt,
|
|
(user, document_id)
|
|
)
|
|
|
|
break
|
|
|
|
except Exception as e:
|
|
logger.error("Exception occurred", exc_info=True)
|
|
raise e
|
|
|
|
for row in resp:
|
|
|
|
if row[3]:
|
|
triples = [
|
|
Triple(
|
|
s = tuple_to_term(elt[0], elt[1]),
|
|
p = tuple_to_term(elt[2], elt[3]),
|
|
o = tuple_to_term(elt[4], elt[5]),
|
|
)
|
|
for elt in row[3]
|
|
]
|
|
else:
|
|
triples = []
|
|
|
|
await receiver(
|
|
Triples(
|
|
metadata = Metadata(
|
|
id = document_id,
|
|
user = user,
|
|
collection = "default", # FIXME: What to put here?
|
|
),
|
|
triples = triples
|
|
)
|
|
)
|
|
|
|
logger.debug("Done")
|
|
|
|
async def get_graph_embeddings(self, user, document_id, receiver):
|
|
|
|
logger.debug("Get GE...")
|
|
|
|
while True:
|
|
|
|
try:
|
|
|
|
resp = self.cassandra.execute(
|
|
self.get_graph_embeddings_stmt,
|
|
(user, document_id)
|
|
)
|
|
|
|
break
|
|
|
|
except Exception as e:
|
|
logger.error("Exception occurred", exc_info=True)
|
|
raise e
|
|
|
|
for row in resp:
|
|
|
|
if row[3]:
|
|
entities = [
|
|
EntityEmbeddings(
|
|
entity = tuple_to_term(ent[0][0], ent[0][1]),
|
|
vector = ent[1]
|
|
)
|
|
for ent in row[3]
|
|
]
|
|
else:
|
|
entities = []
|
|
|
|
await receiver(
|
|
GraphEmbeddings(
|
|
metadata = Metadata(
|
|
id = document_id,
|
|
user = user,
|
|
collection = "default", # FIXME: What to put here?
|
|
),
|
|
entities = entities
|
|
)
|
|
)
|
|
|
|
logger.debug("Done")
|
|
|