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:
cybermaggedon 2026-07-16 00:31:47 +01:00 committed by GitHub
parent a5a302b0d0
commit 60529c3b3d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 194 additions and 34 deletions

View file

@ -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

View file

@ -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

View file

@ -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"])

View file

@ -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,

View file

@ -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):

View file

@ -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