mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-26 21:51:03 +02:00
fix: knowledge core round-trip for workspace-scoped librarian and all term types (#1049)
Change librarian from a single shared client to per-workspace clients, ensuring knowledge core get/put operations use workspace-scoped queues. LibrarianClient is now created in on_workspace_created and cleaned up in on_workspace_deleted. Fix term_to_tuple/tuple_to_term to handle TRIPLE and BLANK term types. Previously only IRI and LITERAL were supported, causing ~1000 edges (ns/contains with quoted triple objects) to be silently lost during Cassandra storage. TRIPLE terms are now JSON-serialized with a <<TRIPLE>> marker prefix, BLANK terms with <<BLANK>>. Make document tree export recursive via _stream_doc_tree so that grandchild documents (chunks under pages) are included in knowledge core export. Without this, explainability source links break after a get/put/load round-trip because chunk documents are not restored. Add error response encoding to KnowledgeResponseTranslator.encode() and fix encode_with_completion to treat error responses as non-final. Change KnowledgeResponse.ids from list[str] with default factory to Optional[list[str]] so list responses can be distinguished from empty responses. Add error display to get_kg_core CLI.
This commit is contained in:
parent
a5a302b0d0
commit
60529c3b3d
6 changed files with 194 additions and 34 deletions
|
|
@ -1,3 +1,4 @@
|
|||
import logging
|
||||
from typing import Dict, Any, Tuple, Optional
|
||||
from ...schema import (
|
||||
KnowledgeRequest, KnowledgeResponse, Triples, GraphEmbeddings,
|
||||
|
|
@ -8,6 +9,8 @@ from ...schema import (
|
|||
from .base import MessageTranslator
|
||||
from .primitives import ValueTranslator, SubgraphTranslator
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class KnowledgeRequestTranslator(MessageTranslator):
|
||||
"""Translator for KnowledgeRequest schema objects"""
|
||||
|
|
@ -264,6 +267,15 @@ class KnowledgeResponseTranslator(MessageTranslator):
|
|||
}
|
||||
}
|
||||
|
||||
# Error response
|
||||
if obj.error:
|
||||
return {
|
||||
"error": {
|
||||
"type": obj.error.type,
|
||||
"message": obj.error.message,
|
||||
}
|
||||
}
|
||||
|
||||
# End of stream marker
|
||||
if obj.eos is True:
|
||||
return {"eos": True}
|
||||
|
|
@ -274,14 +286,13 @@ class KnowledgeResponseTranslator(MessageTranslator):
|
|||
def encode_with_completion(self, obj: KnowledgeResponse) -> Tuple[Dict[str, Any], bool]:
|
||||
"""Returns (response_dict, is_final)"""
|
||||
response = self.encode(obj)
|
||||
|
||||
# Check if this is a final response
|
||||
|
||||
is_final = (
|
||||
obj.ids is not None or # List response
|
||||
obj.eos is True or # End of stream
|
||||
(not obj.triples and not obj.graph_embeddings
|
||||
obj.ids is not None or
|
||||
obj.eos is True or
|
||||
(not obj.error and not obj.triples and not obj.graph_embeddings
|
||||
and not obj.document_embeddings
|
||||
and not obj.library_metadata and not obj.library_blob) # Empty response
|
||||
and not obj.library_metadata and not obj.library_blob)
|
||||
)
|
||||
|
||||
|
||||
return response, is_final
|
||||
|
|
@ -66,7 +66,7 @@ class KnowledgeRequest:
|
|||
@dataclass
|
||||
class KnowledgeResponse:
|
||||
error: Error | None = None
|
||||
ids: list[str] = field(default_factory=list)
|
||||
ids: list[str] | None = None
|
||||
eos: bool = False # Indicates end of knowledge core stream
|
||||
triples: Triples | None = None
|
||||
graph_embeddings: GraphEmbeddings | None = None
|
||||
|
|
|
|||
|
|
@ -87,6 +87,13 @@ def fetch(url, workspace, id, output, token=None):
|
|||
|
||||
for response in socket.get_kg_core(id):
|
||||
|
||||
if response.get("error"):
|
||||
err = response["error"]
|
||||
print(
|
||||
f"Error: {err.get('type', 'unknown')}: "
|
||||
f"{err.get('message', 'Unknown error')}"
|
||||
)
|
||||
|
||||
if "triples" in response:
|
||||
t += 1
|
||||
write_triple(f, response["triples"])
|
||||
|
|
|
|||
|
|
@ -15,11 +15,24 @@ import logging
|
|||
# Module logger
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def _librarian_error(respond, message):
|
||||
await respond(
|
||||
KnowledgeResponse(
|
||||
error=Error(
|
||||
type="librarian-error",
|
||||
message=message,
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class KnowledgeManager:
|
||||
|
||||
def __init__(
|
||||
self, cassandra_host, cassandra_username, cassandra_password,
|
||||
keyspace, flow_config, librarian=None, replication_factor=1,
|
||||
keyspace, flow_config, librarian_clients=None,
|
||||
replication_factor=1,
|
||||
):
|
||||
|
||||
self.table_store = KnowledgeTableStore(
|
||||
|
|
@ -27,7 +40,7 @@ class KnowledgeManager:
|
|||
replication_factor
|
||||
)
|
||||
|
||||
self.librarian = librarian
|
||||
self.librarian_clients = librarian_clients if librarian_clients is not None else {}
|
||||
self._pending_library_metadata = {}
|
||||
|
||||
self.loader_queue = asyncio.Queue(maxsize=20)
|
||||
|
|
@ -90,8 +103,19 @@ class KnowledgeManager:
|
|||
publish_ge,
|
||||
)
|
||||
|
||||
if self.librarian:
|
||||
await self._stream_library_docs(request.id, respond)
|
||||
librarian = self.librarian_clients.get(workspace)
|
||||
if librarian is None:
|
||||
logger.error(
|
||||
f"No librarian client for workspace {workspace}"
|
||||
)
|
||||
await _librarian_error(
|
||||
respond,
|
||||
f"No librarian client for workspace {workspace}",
|
||||
)
|
||||
else:
|
||||
await self._stream_library_docs(
|
||||
librarian, request.id, respond,
|
||||
)
|
||||
|
||||
logger.debug("Knowledge core retrieval complete")
|
||||
|
||||
|
|
@ -129,11 +153,22 @@ class KnowledgeManager:
|
|||
workspace, request.graph_embeddings
|
||||
)
|
||||
|
||||
if request.library_metadata and self.librarian:
|
||||
await self._put_library_metadata(request.library_metadata, workspace)
|
||||
librarian = self.librarian_clients.get(workspace)
|
||||
|
||||
if request.library_blob and self.librarian:
|
||||
await self._put_library_blob(request.library_blob, workspace)
|
||||
if request.library_metadata or request.library_blob:
|
||||
if librarian is None:
|
||||
logger.error(
|
||||
f"No librarian client for workspace {workspace}"
|
||||
)
|
||||
else:
|
||||
if request.library_metadata:
|
||||
await self._put_library_metadata(
|
||||
request.library_metadata, workspace,
|
||||
)
|
||||
if request.library_blob:
|
||||
await self._put_library_blob(
|
||||
librarian, request.library_blob, workspace,
|
||||
)
|
||||
|
||||
await respond(
|
||||
KnowledgeResponse(
|
||||
|
|
@ -263,36 +298,42 @@ class KnowledgeManager:
|
|||
|
||||
await self.loader_queue.put((request, respond, workspace))
|
||||
|
||||
async def _stream_library_docs(self, document_id, respond):
|
||||
async def _stream_library_docs(self, librarian, document_id, respond):
|
||||
|
||||
try:
|
||||
root_meta = await self.librarian.fetch_document_metadata(
|
||||
root_meta = await librarian.fetch_document_metadata(
|
||||
document_id
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not fetch library metadata for {document_id}: {e}")
|
||||
await _librarian_error(respond, str(e))
|
||||
return
|
||||
|
||||
if root_meta is None:
|
||||
return
|
||||
|
||||
await self._stream_one_doc(root_meta, respond)
|
||||
await self._stream_doc_tree(librarian, root_meta, respond)
|
||||
|
||||
async def _stream_doc_tree(self, librarian, doc_meta, respond):
|
||||
|
||||
await self._stream_one_doc(librarian, doc_meta, respond)
|
||||
|
||||
try:
|
||||
resp = await self.librarian.request(
|
||||
resp = await librarian.request(
|
||||
LibrarianRequest(
|
||||
operation="list-children",
|
||||
document_id=document_id,
|
||||
document_id=doc_meta.id,
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not list children for {document_id}: {e}")
|
||||
logger.warning(f"Could not list children for {doc_meta.id}: {e}")
|
||||
await _librarian_error(respond, str(e))
|
||||
return
|
||||
|
||||
for child_meta in resp.document_metadatas:
|
||||
await self._stream_one_doc(child_meta, respond)
|
||||
await self._stream_doc_tree(librarian, child_meta, respond)
|
||||
|
||||
async def _stream_one_doc(self, doc_meta, respond):
|
||||
async def _stream_one_doc(self, librarian, doc_meta, respond):
|
||||
|
||||
lm = LibraryMetadata(
|
||||
id=doc_meta.id,
|
||||
|
|
@ -309,11 +350,12 @@ class KnowledgeManager:
|
|||
)
|
||||
|
||||
try:
|
||||
content = await self.librarian.fetch_document_content(
|
||||
content = await librarian.fetch_document_content(
|
||||
doc_meta.id
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not fetch content for {doc_meta.id}: {e}")
|
||||
await _librarian_error(respond, str(e))
|
||||
return
|
||||
|
||||
await respond(
|
||||
|
|
@ -328,7 +370,7 @@ class KnowledgeManager:
|
|||
async def _put_library_metadata(self, lm, workspace):
|
||||
self._pending_library_metadata[lm.id] = lm
|
||||
|
||||
async def _put_library_blob(self, lb, workspace):
|
||||
async def _put_library_blob(self, librarian, lb, workspace):
|
||||
|
||||
lm = self._pending_library_metadata.pop(lb.id, None)
|
||||
if lm is None:
|
||||
|
|
@ -353,7 +395,7 @@ class KnowledgeManager:
|
|||
operation = "add-document"
|
||||
|
||||
try:
|
||||
await self.librarian.request(
|
||||
await librarian.request(
|
||||
LibrarianRequest(
|
||||
operation=operation,
|
||||
document_id=lm.id,
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ from .. base import LibrarianClient
|
|||
|
||||
from .. schema import KnowledgeRequest, KnowledgeResponse, Error
|
||||
from .. schema import knowledge_request_queue, knowledge_response_queue
|
||||
from .. schema import librarian_request_queue, librarian_response_queue
|
||||
|
||||
from .. schema import Document, Metadata
|
||||
from .. schema import TextDocument, Metadata
|
||||
|
|
@ -79,9 +80,7 @@ class Processor(WorkspaceProcessor):
|
|||
}
|
||||
)
|
||||
|
||||
self.librarian_client = LibrarianClient(
|
||||
id=id, backend=self.pubsub, taskgroup=self.taskgroup,
|
||||
)
|
||||
self.librarian_clients = {}
|
||||
|
||||
self.knowledge = KnowledgeManager(
|
||||
cassandra_host = self.cassandra_host,
|
||||
|
|
@ -89,7 +88,7 @@ class Processor(WorkspaceProcessor):
|
|||
cassandra_password = self.cassandra_password,
|
||||
keyspace = keyspace,
|
||||
flow_config = self,
|
||||
librarian = self.librarian_client,
|
||||
librarian_clients = self.librarian_clients,
|
||||
replication_factor = replication_factor,
|
||||
)
|
||||
|
||||
|
|
@ -142,18 +141,39 @@ class Processor(WorkspaceProcessor):
|
|||
),
|
||||
)
|
||||
|
||||
librarian_client = LibrarianClient(
|
||||
id=self.id,
|
||||
backend=self.pubsub,
|
||||
taskgroup=self.taskgroup,
|
||||
librarian_request_queue=workspace_queue(
|
||||
librarian_request_queue, workspace,
|
||||
),
|
||||
librarian_response_queue=workspace_queue(
|
||||
librarian_response_queue, workspace,
|
||||
),
|
||||
librarian_subscriber=(
|
||||
f"{self.id}--{workspace}--librarian"
|
||||
),
|
||||
)
|
||||
|
||||
await response_producer.start()
|
||||
await consumer.start()
|
||||
await librarian_client.start()
|
||||
|
||||
self.librarian_clients[workspace] = librarian_client
|
||||
|
||||
self.workspace_consumers[workspace] = {
|
||||
"consumer": consumer,
|
||||
"response": response_producer,
|
||||
"librarian": librarian_client,
|
||||
}
|
||||
|
||||
logger.info(f"Subscribed to workspace queue: {workspace}")
|
||||
|
||||
async def on_workspace_deleted(self, workspace):
|
||||
|
||||
self.librarian_clients.pop(workspace, None)
|
||||
|
||||
clients = self.workspace_consumers.pop(workspace, None)
|
||||
if clients:
|
||||
for client in clients.values():
|
||||
|
|
@ -163,7 +183,6 @@ class Processor(WorkspaceProcessor):
|
|||
async def start(self):
|
||||
|
||||
await super(Processor, self).start()
|
||||
await self.librarian_client.start()
|
||||
|
||||
async def on_knowledge_config(self, workspace, config, version):
|
||||
|
||||
|
|
|
|||
|
|
@ -1,18 +1,92 @@
|
|||
|
||||
import json
|
||||
|
||||
from .. schema import KnowledgeResponse, Triple, Triples, EntityEmbeddings
|
||||
from .. schema import Metadata, Term, IRI, LITERAL, GraphEmbeddings
|
||||
from .. schema import Metadata, Term, IRI, BLANK, LITERAL, TRIPLE
|
||||
from .. schema import GraphEmbeddings
|
||||
from .. schema import DocumentEmbeddings, ChunkEmbeddings
|
||||
|
||||
from cassandra.cluster import Cluster
|
||||
|
||||
from . cassandra_async import async_execute, async_execute_paged
|
||||
|
||||
TRIPLE_MARKER = "<<TRIPLE>>"
|
||||
BLANK_MARKER = "<<BLANK>>"
|
||||
|
||||
|
||||
def _serialize_term(term):
|
||||
if term.type == IRI:
|
||||
return {"t": "i", "i": term.iri}
|
||||
elif term.type == LITERAL:
|
||||
r = {"t": "l", "v": term.value}
|
||||
if term.datatype:
|
||||
r["dt"] = term.datatype
|
||||
if term.language:
|
||||
r["ln"] = term.language
|
||||
return r
|
||||
elif term.type == BLANK:
|
||||
return {"t": "b", "d": term.id}
|
||||
elif term.type == TRIPLE and term.triple:
|
||||
return {"t": "t", "tr": _serialize_triple(term.triple)}
|
||||
return {"t": term.type}
|
||||
|
||||
|
||||
def _deserialize_term(data):
|
||||
t = data.get("t", "")
|
||||
if t == "i":
|
||||
return Term(type=IRI, iri=data.get("i", ""))
|
||||
elif t == "l":
|
||||
return Term(
|
||||
type=LITERAL, value=data.get("v", ""),
|
||||
datatype=data.get("dt", ""), language=data.get("ln", ""),
|
||||
)
|
||||
elif t == "b":
|
||||
return Term(type=BLANK, id=data.get("d", ""))
|
||||
elif t == "t":
|
||||
tr = data.get("tr")
|
||||
return Term(
|
||||
type=TRIPLE,
|
||||
triple=_deserialize_triple(tr) if tr else None,
|
||||
)
|
||||
return Term(type=t)
|
||||
|
||||
|
||||
def _serialize_triple(triple):
|
||||
r = {}
|
||||
if triple.s:
|
||||
r["s"] = _serialize_term(triple.s)
|
||||
if triple.p:
|
||||
r["p"] = _serialize_term(triple.p)
|
||||
if triple.o:
|
||||
r["o"] = _serialize_term(triple.o)
|
||||
if triple.g:
|
||||
r["g"] = triple.g
|
||||
return r
|
||||
|
||||
|
||||
def _deserialize_triple(data):
|
||||
return Triple(
|
||||
s=_deserialize_term(data["s"]) if "s" in data else None,
|
||||
p=_deserialize_term(data["p"]) if "p" in data else None,
|
||||
o=_deserialize_term(data["o"]) if "o" in data else None,
|
||||
g=data.get("g"),
|
||||
)
|
||||
|
||||
|
||||
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
|
||||
elif term.type == TRIPLE:
|
||||
return (
|
||||
TRIPLE_MARKER + json.dumps(
|
||||
_serialize_triple(term.triple), separators=(",", ":")
|
||||
),
|
||||
False,
|
||||
)
|
||||
elif term.type == BLANK:
|
||||
return (BLANK_MARKER + term.id, False)
|
||||
else:
|
||||
return (term.value, False)
|
||||
|
||||
|
||||
|
|
@ -20,6 +94,13 @@ 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)
|
||||
elif value.startswith(TRIPLE_MARKER):
|
||||
triple_data = json.loads(value[len(TRIPLE_MARKER):])
|
||||
return Term(
|
||||
type=TRIPLE, triple=_deserialize_triple(triple_data),
|
||||
)
|
||||
elif value.startswith(BLANK_MARKER):
|
||||
return Term(type=BLANK, id=value[len(BLANK_MARKER):])
|
||||
else:
|
||||
return Term(type=LITERAL, value=value)
|
||||
from cassandra.auth import PlainTextAuthProvider
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue