apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-11 22:32:12 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

Use UUID-based URNs for page and chunk IDs

Browse source 
Page and chunk document IDs were deterministic ({doc_id}/p{num},
{doc_id}/p{num}/c{num}), causing "Document already exists" errors
when reprocessing documents through different flows. Content may
differ between runs due to different parameters or extractors, so
deterministic IDs are incorrect.

Pages now use urn:page:{uuid}, chunks use
urn:chunk:{uuid}. Parent- child relationships are tracked via
librarian metadata and provenance triples.

Also brings Mistral OCR and Tesseract OCR decoders up to parity
with the PDF decoder: librarian fetch/save support, per-page
output with unique IDs, and provenance triple emission. Fixes
Mistral OCR bug where only the first 5 pages were processed.
This commit is contained in:
Cyber MacGeddon 2026-03-21 20:18:43 +00:00
parent 1a7b654bd3
commit 8488a8dc4c
8 changed files with 565 additions and 109 deletions
Download patch file Download diff file Expand all files Collapse all files

38
tests/unit/test_provenance/test_uris.py
View file

@ -10,8 +10,7 @@ from trustgraph.provenance.uris import (
_encode_id,
document_uri,
page_uri,
chunk_uri_from_page,
chunk_uri_from_doc,
chunk_uri,
activity_uri,
subgraph_uri,
agent_uri,
@ -60,31 +59,22 @@ class TestDocumentUris:
assert document_uri(iri) == iri
def test_page_uri_format(self):
result = page_uri("https://example.com/doc/123", 5)
assert result == "https://example.com/doc/123/p5"
result = page_uri()
assert result.startswith("urn:page:")
def test_page_uri_page_zero(self):
result = page_uri("https://example.com/doc/123", 0)
assert result == "https://example.com/doc/123/p0"
def test_page_uri_unique(self):
r1 = page_uri()
r2 = page_uri()
assert r1 != r2
def test_chunk_uri_from_page_format(self):
result = chunk_uri_from_page("https://example.com/doc/123", 2, 3)
assert result == "https://example.com/doc/123/p2/c3"
def test_chunk_uri_format(self):
result = chunk_uri()
assert result.startswith("urn:chunk:")
def test_chunk_uri_from_doc_format(self):
result = chunk_uri_from_doc("https://example.com/doc/123", 7)
assert result == "https://example.com/doc/123/c7"
def test_page_uri_preserves_doc_iri(self):
doc = "urn:isbn:978-3-16-148410-0"
result = page_uri(doc, 1)
assert result.startswith(doc)
def test_chunk_from_page_hierarchy(self):
"""Chunk URI should contain both page and chunk identifiers."""
result = chunk_uri_from_page("https://example.com/doc", 3, 5)
assert "/p3/" in result
assert result.endswith("/c5")
def test_chunk_uri_unique(self):
r1 = chunk_uri()
r2 = chunk_uri()
assert r1 != r2
class TestActivityAndSubgraphUris:

10
trustgraph-base/trustgraph/provenance/__init__.py
View file

@ -9,14 +9,14 @@ Provides helpers for:
Usage example:
from trustgraph.provenance import (
document_uri, page_uri, chunk_uri_from_page,
document_uri, page_uri, chunk_uri,
document_triples, derived_entity_triples,
get_vocabulary_triples,
)
# Generate URIs
doc_uri = document_uri("my-doc-123")
page_uri = page_uri("my-doc-123", page_number=1)
pg_uri = page_uri()
# Build provenance triples
triples = document_triples(
@ -35,8 +35,7 @@ from . uris import (
TRUSTGRAPH_BASE,
document_uri,
page_uri,
chunk_uri_from_page,
chunk_uri_from_doc,
chunk_uri,
activity_uri,
subgraph_uri,
agent_uri,
@ -138,8 +137,7 @@ __all__ = [
"TRUSTGRAPH_BASE",
"document_uri",
"page_uri",
"chunk_uri_from_page",
"chunk_uri_from_doc",
"chunk_uri",
"activity_uri",
"subgraph_uri",
"agent_uri",

28
trustgraph-base/trustgraph/provenance/uris.py
View file

@ -1,12 +1,11 @@
"""
URI generation for provenance entities.
Document IDs are already IRIs (e.g., https://trustgraph.ai/doc/abc123).
Child entities (pages, chunks) append path segments to the parent IRI:
- Document: {doc_iri} (as provided)
- Page: {doc_iri}/p{page_number}
- Chunk: {page_iri}/c{chunk_index} (from page)
{doc_iri}/c{chunk_index} (from text doc)
Document IDs are externally provided (e.g., https://trustgraph.ai/doc/abc123).
Child entities (pages, chunks) use UUID-based URNs:
- Document: {doc_iri} (as provided, not generated here)
- Page: urn:page:{uuid}
- Chunk: urn:chunk:{uuid}
- Activity: https://trustgraph.ai/activity/{uuid}
- Subgraph: https://trustgraph.ai/subgraph/{uuid}
"""
@ -28,19 +27,14 @@ def document_uri(doc_iri: str) -> str:
return doc_iri
def page_uri(doc_iri: str, page_number: int) -> str:
"""Generate URI for a page by appending to document IRI."""
return f"{doc_iri}/p{page_number}"
def page_uri() -> str:
"""Generate a unique URI for a page."""
return f"urn:page:{uuid.uuid4()}"
def chunk_uri_from_page(doc_iri: str, page_number: int, chunk_index: int) -> str:
"""Generate URI for a chunk extracted from a page."""
return f"{doc_iri}/p{page_number}/c{chunk_index}"
def chunk_uri_from_doc(doc_iri: str, chunk_index: int) -> str:
"""Generate URI for a chunk extracted directly from a text document."""
return f"{doc_iri}/c{chunk_index}"
def chunk_uri() -> str:
"""Generate a unique URI for a chunk."""
return f"urn:chunk:{uuid.uuid4()}"
def activity_uri(activity_id: str = None) -> str:

15
trustgraph-flow/trustgraph/chunking/recursive/chunker.py
View file

@ -12,7 +12,7 @@ from ... schema import TextDocument, Chunk, Metadata, Triples
from ... base import ChunkingService, ConsumerSpec, ProducerSpec
from ... provenance import (
derived_entity_triples,
chunk_uri as make_chunk_uri, derived_entity_triples,
set_graph, GRAPH_SOURCE,
)
@ -124,10 +124,9 @@ class Processor(ChunkingService):
logger.debug(f"Created chunk of size {len(chunk.page_content)}")
# Generate chunk document ID by appending /c{index} to parent
# Works for both page URIs (doc/p3 -> doc/p3/c1) and doc URIs (doc -> doc/c1)
chunk_doc_id = f"{parent_doc_id}/c{chunk_index}"
chunk_uri = chunk_doc_id # URI is same as document ID
# Generate unique chunk ID
c_uri = make_chunk_uri()
chunk_doc_id = c_uri
parent_uri = parent_doc_id
chunk_content = chunk.page_content.encode("utf-8")
@ -145,7 +144,7 @@ class Processor(ChunkingService):
# Emit provenance triples (stored in source graph for separation from core knowledge)
prov_triples = derived_entity_triples(
entity_uri=chunk_uri,
entity_uri=c_uri,
parent_uri=parent_uri,
component_name=COMPONENT_NAME,
component_version=COMPONENT_VERSION,
@ -159,7 +158,7 @@ class Processor(ChunkingService):
await flow("triples").send(Triples(
metadata=Metadata(
id=chunk_uri,
id=c_uri,
root=v.metadata.root,
user=v.metadata.user,
collection=v.metadata.collection,
@ -170,7 +169,7 @@ class Processor(ChunkingService):
# Forward chunk ID + content (post-chunker optimization)
r = Chunk(
metadata=Metadata(
id=chunk_uri,
id=c_uri,
root=v.metadata.root,
user=v.metadata.user,
collection=v.metadata.collection,

15
trustgraph-flow/trustgraph/chunking/token/chunker.py
View file

@ -12,7 +12,7 @@ from ... schema import TextDocument, Chunk, Metadata, Triples
from ... base import ChunkingService, ConsumerSpec, ProducerSpec
from ... provenance import (
derived_entity_triples,
chunk_uri as make_chunk_uri, derived_entity_triples,
set_graph, GRAPH_SOURCE,
)
@ -122,10 +122,9 @@ class Processor(ChunkingService):
logger.debug(f"Created chunk of size {len(chunk.page_content)}")
# Generate chunk document ID by appending /c{index} to parent
# Works for both page URIs (doc/p3 -> doc/p3/c1) and doc URIs (doc -> doc/c1)
chunk_doc_id = f"{parent_doc_id}/c{chunk_index}"
chunk_uri = chunk_doc_id # URI is same as document ID
# Generate unique chunk ID
c_uri = make_chunk_uri()
chunk_doc_id = c_uri
parent_uri = parent_doc_id
chunk_content = chunk.page_content.encode("utf-8")
@ -143,7 +142,7 @@ class Processor(ChunkingService):
# Emit provenance triples (stored in source graph for separation from core knowledge)
prov_triples = derived_entity_triples(
entity_uri=chunk_uri,
entity_uri=c_uri,
parent_uri=parent_uri,
component_name=COMPONENT_NAME,
component_version=COMPONENT_VERSION,
@ -157,7 +156,7 @@ class Processor(ChunkingService):
await flow("triples").send(Triples(
metadata=Metadata(
id=chunk_uri,
id=c_uri,
root=v.metadata.root,
user=v.metadata.user,
collection=v.metadata.collection,
@ -168,7 +167,7 @@ class Processor(ChunkingService):
# Forward chunk ID + content (post-chunker optimization)
r = Chunk(
metadata=Metadata(
id=chunk_uri,
id=c_uri,
root=v.metadata.root,
user=v.metadata.user,
collection=v.metadata.collection,

310
trustgraph-flow/trustgraph/decoding/mistral_ocr/processor.py
View file

@ -1,29 +1,48 @@
"""
Simple decoder, accepts PDF documents on input, outputs pages from the
PDF document as text as separate output objects.
Mistral OCR decoder, accepts PDF documents on input, outputs pages from the
PDF document as markdown text as separate output objects.
Supports both inline document data and fetching from librarian via Pulsar
for large documents.
"""
from pypdf import PdfWriter, PdfReader
from io import BytesIO
import asyncio
import base64
import uuid
import os
from mistralai import Mistral
from mistralai import DocumentURLChunk, ImageURLChunk, TextChunk
from mistralai.models import OCRResponse
from ... schema import Document, TextDocument, Metadata
from ... schema import LibrarianRequest, LibrarianResponse, DocumentMetadata
from ... schema import librarian_request_queue, librarian_response_queue
from ... schema import Triples
from ... base import FlowProcessor, ConsumerSpec, ProducerSpec
from ... base import Consumer, Producer, ConsumerMetrics, ProducerMetrics
from ... provenance import (
document_uri, page_uri as make_page_uri, derived_entity_triples,
set_graph, GRAPH_SOURCE,
)
import logging
logger = logging.getLogger(__name__)
# Component identification for provenance
COMPONENT_NAME = "mistral-ocr-decoder"
COMPONENT_VERSION = "1.0.0"
default_ident = "pdf-decoder"
default_api_key = os.getenv("MISTRAL_TOKEN")
default_librarian_request_queue = librarian_request_queue
default_librarian_response_queue = librarian_response_queue
pages_per_chunk = 5
def chunks(lst, n):
@ -48,27 +67,6 @@ def replace_images_in_markdown(markdown_str: str, images_dict: dict) -> str:
)
return markdown_str
def get_combined_markdown(ocr_response: OCRResponse) -> str:
"""
Combine OCR text and images into a single markdown document.
Args:
ocr_response: Response from OCR processing containing text and images
Returns:
Combined markdown string with embedded images
"""
markdowns: list[str] = []
# Extract images from page
for page in ocr_response.pages:
image_data = {}
for img in page.images:
image_data[img.id] = img.image_base64
# Replace image placeholders with actual images
markdowns.append(replace_images_in_markdown(page.markdown, image_data))
return "\n\n".join(markdowns)
class Processor(FlowProcessor):
def __init__(self, **params):
@ -97,6 +95,50 @@ class Processor(FlowProcessor):
)
)
self.register_specification(
ProducerSpec(
name = "triples",
schema = Triples,
)
)
# Librarian client for fetching document content
librarian_request_q = params.get(
"librarian_request_queue", default_librarian_request_queue
)
librarian_response_q = params.get(
"librarian_response_queue", default_librarian_response_queue
)
librarian_request_metrics = ProducerMetrics(
processor = id, flow = None, name = "librarian-request"
)
self.librarian_request_producer = Producer(
backend = self.pubsub,
topic = librarian_request_q,
schema = LibrarianRequest,
metrics = librarian_request_metrics,
)
librarian_response_metrics = ConsumerMetrics(
processor = id, flow = None, name = "librarian-response"
)
self.librarian_response_consumer = Consumer(
taskgroup = self.taskgroup,
backend = self.pubsub,
flow = None,
topic = librarian_response_q,
subscriber = f"{id}-librarian",
schema = LibrarianResponse,
handler = self.on_librarian_response,
metrics = librarian_response_metrics,
)
# Pending librarian requests: request_id -> asyncio.Future
self.pending_requests = {}
if api_key is None:
raise RuntimeError("Mistral API key not specified")
@ -107,15 +149,125 @@ class Processor(FlowProcessor):
logger.info("Mistral OCR processor initialized")
async def start(self):
await super(Processor, self).start()
await self.librarian_request_producer.start()
await self.librarian_response_consumer.start()
async def on_librarian_response(self, msg, consumer, flow):
"""Handle responses from the librarian service."""
response = msg.value()
request_id = msg.properties().get("id")
if request_id and request_id in self.pending_requests:
future = self.pending_requests.pop(request_id)
future.set_result(response)
else:
logger.warning(f"Received unexpected librarian response: {request_id}")
async def fetch_document_content(self, document_id, user, timeout=120):
"""
Fetch document content from librarian via Pulsar.
"""
request_id = str(uuid.uuid4())
request = LibrarianRequest(
operation="get-document-content",
document_id=document_id,
user=user,
)
# Create future for response
future = asyncio.get_event_loop().create_future()
self.pending_requests[request_id] = future
try:
# Send request
await self.librarian_request_producer.send(
request, properties={"id": request_id}
)
# Wait for response
response = await asyncio.wait_for(future, timeout=timeout)
if response.error:
raise RuntimeError(
f"Librarian error: {response.error.type}: {response.error.message}"
)
return response.content
except asyncio.TimeoutError:
self.pending_requests.pop(request_id, None)
raise RuntimeError(f"Timeout fetching document {document_id}")
async def save_child_document(self, doc_id, parent_id, user, content,
document_type="page", title=None, timeout=120):
"""
Save a child document to the librarian.
"""
request_id = str(uuid.uuid4())
doc_metadata = DocumentMetadata(
id=doc_id,
user=user,
kind="text/plain",
title=title or doc_id,
parent_id=parent_id,
document_type=document_type,
)
request = LibrarianRequest(
operation="add-child-document",
document_metadata=doc_metadata,
content=base64.b64encode(content).decode("utf-8"),
)
# Create future for response
future = asyncio.get_event_loop().create_future()
self.pending_requests[request_id] = future
try:
# Send request
await self.librarian_request_producer.send(
request, properties={"id": request_id}
)
# Wait for response
response = await asyncio.wait_for(future, timeout=timeout)
if response.error:
raise RuntimeError(
f"Librarian error saving child document: {response.error.type}: {response.error.message}"
)
return doc_id
except asyncio.TimeoutError:
self.pending_requests.pop(request_id, None)
raise RuntimeError(f"Timeout saving child document {doc_id}")
def ocr(self, blob):
"""
Run Mistral OCR on a PDF blob, returning per-page markdown strings.
Args:
blob: Raw PDF bytes
Returns:
List of (page_markdown, page_number) tuples, 1-indexed
"""
logger.debug("Parse PDF...")
pdfbuf = BytesIO(blob)
pdf = PdfReader(pdfbuf)
pages = []
global_page_num = 0
for chunk in chunks(pdf.pages, pages_per_chunk):
logger.debug("Get next pages...")
part = PdfWriter()
@ -152,11 +304,19 @@ class Processor(FlowProcessor):
logger.debug("Extract markdown...")
markdown = get_combined_markdown(processed)
for page in processed.pages:
global_page_num += 1
image_data = {}
for img in page.images:
image_data[img.id] = img.image_base64
markdown = replace_images_in_markdown(
page.markdown, image_data
)
pages.append((markdown, global_page_num))
logger.info("OCR complete.")
logger.info(f"OCR complete, {len(pages)} pages.")
return markdown
return pages
async def on_message(self, msg, consumer, flow):
@ -166,16 +326,83 @@ class Processor(FlowProcessor):
logger.info(f"Decoding {v.metadata.id}...")
markdown = self.ocr(base64.b64decode(v.data))
# Get PDF content - fetch from librarian or use inline data
if v.document_id:
logger.info(f"Fetching document {v.document_id} from librarian...")
content = await self.fetch_document_content(
document_id=v.document_id,
user=v.metadata.user,
)
if isinstance(content, str):
content = content.encode('utf-8')
blob = base64.b64decode(content)
logger.info(f"Fetched {len(blob)} bytes from librarian")
else:
blob = base64.b64decode(v.data)
r = TextDocument(
metadata=v.metadata,
text=markdown.encode("utf-8"),
)
# Get the source document ID
source_doc_id = v.document_id or v.metadata.id
await flow("output").send(r)
# Run OCR, get per-page markdown
pages = self.ocr(blob)
logger.info("Done.")
for markdown, page_num in pages:
logger.debug(f"Processing page {page_num}")
# Generate unique page ID
pg_uri = make_page_uri()
page_doc_id = pg_uri
page_content = markdown.encode("utf-8")
# Save page as child document in librarian
await self.save_child_document(
doc_id=page_doc_id,
parent_id=source_doc_id,
user=v.metadata.user,
content=page_content,
document_type="page",
title=f"Page {page_num}",
)
# Emit provenance triples
doc_uri = document_uri(source_doc_id)
prov_triples = derived_entity_triples(
entity_uri=pg_uri,
parent_uri=doc_uri,
component_name=COMPONENT_NAME,
component_version=COMPONENT_VERSION,
label=f"Page {page_num}",
page_number=page_num,
)
await flow("triples").send(Triples(
metadata=Metadata(
id=pg_uri,
root=v.metadata.root,
user=v.metadata.user,
collection=v.metadata.collection,
),
triples=set_graph(prov_triples, GRAPH_SOURCE),
))
# Forward page document ID to chunker
# Chunker will fetch content from librarian
r = TextDocument(
metadata=Metadata(
id=pg_uri,
root=v.metadata.root,
user=v.metadata.user,
collection=v.metadata.collection,
),
document_id=page_doc_id,
text=b"", # Empty, chunker will fetch from librarian
)
await flow("output").send(r)
logger.debug("PDF decoding complete")
@staticmethod
def add_args(parser):
@ -188,7 +415,18 @@ class Processor(FlowProcessor):
help=f'Mistral API Key'
)
parser.add_argument(
'--librarian-request-queue',
default=default_librarian_request_queue,
help=f'Librarian request queue (default: {default_librarian_request_queue})',
)
parser.add_argument(
'--librarian-response-queue',
default=default_librarian_response_queue,
help=f'Librarian response queue (default: {default_librarian_response_queue})',
)
def run():
Processor.launch(default_ident, __doc__)

8
trustgraph-flow/trustgraph/decoding/pdf/pdf_decoder.py
View file

@ -23,7 +23,7 @@ from ... base import FlowProcessor, ConsumerSpec, ProducerSpec
from ... base import Consumer, Producer, ConsumerMetrics, ProducerMetrics
from ... provenance import (
document_uri, page_uri, derived_entity_triples,
document_uri, page_uri as make_page_uri, derived_entity_triples,
set_graph, GRAPH_SOURCE,
)
@ -272,8 +272,9 @@ class Processor(FlowProcessor):
logger.debug(f"Processing page {page_num}")
# Generate page document ID
page_doc_id = f"{source_doc_id}/p{page_num}"
# Generate unique page ID
pg_uri = make_page_uri()
page_doc_id = pg_uri
page_content = page.page_content.encode("utf-8")
# Save page as child document in librarian
@ -288,7 +289,6 @@ class Processor(FlowProcessor):
# Emit provenance triples (stored in source graph for separation from core knowledge)
doc_uri = document_uri(source_doc_id)
pg_uri = page_uri(source_doc_id, page_num)
prov_triples = derived_entity_triples(
entity_uri=pg_uri,

250
trustgraph-ocr/trustgraph/decoding/ocr/pdf_decoder.py
View file

@ -2,22 +2,42 @@
"""
Simple decoder, accepts PDF documents on input, outputs pages from the
PDF document as text as separate output objects.
Supports both inline document data and fetching from librarian via Pulsar
for large documents.
"""
import tempfile
import asyncio
import base64
import logging
import uuid
import pytesseract
from pdf2image import convert_from_bytes
from ... schema import Document, TextDocument, Metadata
from ... schema import LibrarianRequest, LibrarianResponse, DocumentMetadata
from ... schema import librarian_request_queue, librarian_response_queue
from ... schema import Triples
from ... base import FlowProcessor, ConsumerSpec, ProducerSpec
from ... base import Consumer, Producer, ConsumerMetrics, ProducerMetrics
from ... provenance import (
document_uri, page_uri as make_page_uri, derived_entity_triples,
set_graph, GRAPH_SOURCE,
)
# Component identification for provenance
COMPONENT_NAME = "tesseract-ocr-decoder"
COMPONENT_VERSION = "1.0.0"
# Module logger
logger = logging.getLogger(__name__)
default_ident = "pdf-decoder"
default_librarian_request_queue = librarian_request_queue
default_librarian_response_queue = librarian_response_queue
class Processor(FlowProcessor):
def __init__(self, **params):
@ -45,8 +65,150 @@ class Processor(FlowProcessor):
)
)
self.register_specification(
ProducerSpec(
name = "triples",
schema = Triples,
)
)
# Librarian client for fetching document content
librarian_request_q = params.get(
"librarian_request_queue", default_librarian_request_queue
)
librarian_response_q = params.get(
"librarian_response_queue", default_librarian_response_queue
)
librarian_request_metrics = ProducerMetrics(
processor = id, flow = None, name = "librarian-request"
)
self.librarian_request_producer = Producer(
backend = self.pubsub,
topic = librarian_request_q,
schema = LibrarianRequest,
metrics = librarian_request_metrics,
)
librarian_response_metrics = ConsumerMetrics(
processor = id, flow = None, name = "librarian-response"
)
self.librarian_response_consumer = Consumer(
taskgroup = self.taskgroup,
backend = self.pubsub,
flow = None,
topic = librarian_response_q,
subscriber = f"{id}-librarian",
schema = LibrarianResponse,
handler = self.on_librarian_response,
metrics = librarian_response_metrics,
)
# Pending librarian requests: request_id -> asyncio.Future
self.pending_requests = {}
logger.info("PDF OCR processor initialized")
async def start(self):
await super(Processor, self).start()
await self.librarian_request_producer.start()
await self.librarian_response_consumer.start()
async def on_librarian_response(self, msg, consumer, flow):
"""Handle responses from the librarian service."""
response = msg.value()
request_id = msg.properties().get("id")
if request_id and request_id in self.pending_requests:
future = self.pending_requests.pop(request_id)
future.set_result(response)
else:
logger.warning(f"Received unexpected librarian response: {request_id}")
async def fetch_document_content(self, document_id, user, timeout=120):
"""
Fetch document content from librarian via Pulsar.
"""
request_id = str(uuid.uuid4())
request = LibrarianRequest(
operation="get-document-content",
document_id=document_id,
user=user,
)
# Create future for response
future = asyncio.get_event_loop().create_future()
self.pending_requests[request_id] = future
try:
# Send request
await self.librarian_request_producer.send(
request, properties={"id": request_id}
)
# Wait for response
response = await asyncio.wait_for(future, timeout=timeout)
if response.error:
raise RuntimeError(
f"Librarian error: {response.error.type}: {response.error.message}"
)
return response.content
except asyncio.TimeoutError:
self.pending_requests.pop(request_id, None)
raise RuntimeError(f"Timeout fetching document {document_id}")
async def save_child_document(self, doc_id, parent_id, user, content,
document_type="page", title=None, timeout=120):
"""
Save a child document to the librarian.
"""
request_id = str(uuid.uuid4())
doc_metadata = DocumentMetadata(
id=doc_id,
user=user,
kind="text/plain",
title=title or doc_id,
parent_id=parent_id,
document_type=document_type,
)
request = LibrarianRequest(
operation="add-child-document",
document_metadata=doc_metadata,
content=base64.b64encode(content).decode("utf-8"),
)
# Create future for response
future = asyncio.get_event_loop().create_future()
self.pending_requests[request_id] = future
try:
# Send request
await self.librarian_request_producer.send(
request, properties={"id": request_id}
)
# Wait for response
response = await asyncio.wait_for(future, timeout=timeout)
if response.error:
raise RuntimeError(
f"Librarian error saving child document: {response.error.type}: {response.error.message}"
)
return doc_id
except asyncio.TimeoutError:
self.pending_requests.pop(request_id, None)
raise RuntimeError(f"Timeout saving child document {doc_id}")
async def on_message(self, msg, consumer, flow):
logger.info("PDF message received")
@ -55,21 +217,85 @@ class Processor(FlowProcessor):
logger.info(f"Decoding {v.metadata.id}...")
blob = base64.b64decode(v.data)
# Get PDF content - fetch from librarian or use inline data
if v.document_id:
logger.info(f"Fetching document {v.document_id} from librarian...")
content = await self.fetch_document_content(
document_id=v.document_id,
user=v.metadata.user,
)
if isinstance(content, str):
content = content.encode('utf-8')
blob = base64.b64decode(content)
logger.info(f"Fetched {len(blob)} bytes from librarian")
else:
blob = base64.b64decode(v.data)
# Get the source document ID
source_doc_id = v.document_id or v.metadata.id
pages = convert_from_bytes(blob)
for ix, page in enumerate(pages):
page_num = ix + 1 # 1-indexed
try:
text = pytesseract.image_to_string(page, lang='eng')
except Exception as e:
logger.warning(f"Page did not OCR: {e}")
logger.warning(f"Page {page_num} did not OCR: {e}")
continue
logger.debug(f"Processing page {page_num}")
# Generate unique page ID
pg_uri = make_page_uri()
page_doc_id = pg_uri
page_content = text.encode("utf-8")
# Save page as child document in librarian
await self.save_child_document(
doc_id=page_doc_id,
parent_id=source_doc_id,
user=v.metadata.user,
content=page_content,
document_type="page",
title=f"Page {page_num}",
)
# Emit provenance triples
doc_uri = document_uri(source_doc_id)
prov_triples = derived_entity_triples(
entity_uri=pg_uri,
parent_uri=doc_uri,
component_name=COMPONENT_NAME,
component_version=COMPONENT_VERSION,
label=f"Page {page_num}",
page_number=page_num,
)
await flow("triples").send(Triples(
metadata=Metadata(
id=pg_uri,
root=v.metadata.root,
user=v.metadata.user,
collection=v.metadata.collection,
),
triples=set_graph(prov_triples, GRAPH_SOURCE),
))
# Forward page document ID to chunker
# Chunker will fetch content from librarian
r = TextDocument(
metadata=v.metadata,
text=text.encode("utf-8"),
metadata=Metadata(
id=pg_uri,
root=v.metadata.root,
user=v.metadata.user,
collection=v.metadata.collection,
),
document_id=page_doc_id,
text=b"", # Empty, chunker will fetch from librarian
)
await flow("output").send(r)
@ -78,9 +304,21 @@ class Processor(FlowProcessor):
@staticmethod
def add_args(parser):
FlowProcessor.add_args(parser)
parser.add_argument(
'--librarian-request-queue',
default=default_librarian_request_queue,
help=f'Librarian request queue (default: {default_librarian_request_queue})',
)
parser.add_argument(
'--librarian-response-queue',
default=default_librarian_response_queue,
help=f'Librarian response queue (default: {default_librarian_response_queue})',
)
def run():
Processor.launch(default_ident, __doc__)