diff --git a/trustgraph-base/trustgraph/messaging/translators/knowledge.py b/trustgraph-base/trustgraph/messaging/translators/knowledge.py index 3f09b41b..281474b7 100644 --- a/trustgraph-base/trustgraph/messaging/translators/knowledge.py +++ b/trustgraph-base/trustgraph/messaging/translators/knowledge.py @@ -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 \ No newline at end of file diff --git a/trustgraph-base/trustgraph/schema/knowledge/knowledge.py b/trustgraph-base/trustgraph/schema/knowledge/knowledge.py index 4353065b..4b209b2a 100644 --- a/trustgraph-base/trustgraph/schema/knowledge/knowledge.py +++ b/trustgraph-base/trustgraph/schema/knowledge/knowledge.py @@ -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 diff --git a/trustgraph-cli/trustgraph/cli/get_kg_core.py b/trustgraph-cli/trustgraph/cli/get_kg_core.py index 2ff1a3cc..2a55b0d0 100644 --- a/trustgraph-cli/trustgraph/cli/get_kg_core.py +++ b/trustgraph-cli/trustgraph/cli/get_kg_core.py @@ -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"]) diff --git a/trustgraph-flow/trustgraph/cores/knowledge.py b/trustgraph-flow/trustgraph/cores/knowledge.py index 6f017c43..afa87300 100644 --- a/trustgraph-flow/trustgraph/cores/knowledge.py +++ b/trustgraph-flow/trustgraph/cores/knowledge.py @@ -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, diff --git a/trustgraph-flow/trustgraph/cores/service.py b/trustgraph-flow/trustgraph/cores/service.py index 5c50c207..9d656c6d 100755 --- a/trustgraph-flow/trustgraph/cores/service.py +++ b/trustgraph-flow/trustgraph/cores/service.py @@ -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): diff --git a/trustgraph-flow/trustgraph/tables/knowledge.py b/trustgraph-flow/trustgraph/tables/knowledge.py index 53a12b35..04705507 100644 --- a/trustgraph-flow/trustgraph/tables/knowledge.py +++ b/trustgraph-flow/trustgraph/tables/knowledge.py @@ -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 = "<>" +BLANK_MARKER = "<>" + + +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