260 lines
8.8 KiB
Python
260 lines
8.8 KiB
Python
import tempfile
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING
|
|
from urllib.parse import urlparse
|
|
|
|
import httpx
|
|
|
|
from haiku.rag.config import AppConfig
|
|
from haiku.rag.converters import get_converter
|
|
from haiku.rag.store.models.chunk import Chunk
|
|
|
|
if TYPE_CHECKING:
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
|
|
|
|
async def convert(
|
|
config: AppConfig, source: Path | str, *, format: str = "md"
|
|
) -> "DoclingDocument":
|
|
"""Convert a file, URL, or text to DoclingDocument.
|
|
|
|
Args:
|
|
config: Application configuration.
|
|
source: One of:
|
|
- Path: Local file path to convert
|
|
- str (URL): HTTP/HTTPS URL to download and convert
|
|
- str (text): Raw text content to convert
|
|
format: The format of text content ("md", "html", or "plain").
|
|
Defaults to "md". Use "plain" for plain text without parsing.
|
|
Only used when source is raw text (not a file path or URL).
|
|
Files and URLs determine format from extension/content-type.
|
|
|
|
Returns:
|
|
DoclingDocument from the converted source.
|
|
|
|
Raises:
|
|
ValueError: If the file doesn't exist or has unsupported extension.
|
|
httpx.RequestError: If URL download fails.
|
|
"""
|
|
converter = get_converter(config)
|
|
|
|
# Path object - convert file directly
|
|
if isinstance(source, Path):
|
|
if not source.exists():
|
|
raise ValueError(f"File does not exist: {source}")
|
|
if source.suffix.lower() not in converter.supported_extensions:
|
|
raise ValueError(f"Unsupported file extension: {source.suffix}")
|
|
return await converter.convert_file(source)
|
|
|
|
# String - check if URL or text
|
|
parsed = urlparse(source)
|
|
|
|
if parsed.scheme in ("http", "https"):
|
|
# URL - download and convert
|
|
async with httpx.AsyncClient() as http:
|
|
response = await http.get(source)
|
|
response.raise_for_status()
|
|
|
|
content_type = response.headers.get("content-type", "").lower()
|
|
file_extension = get_extension_from_content_type_or_url(
|
|
source, content_type
|
|
)
|
|
|
|
if file_extension not in converter.supported_extensions:
|
|
raise ValueError(
|
|
f"Unsupported content type/extension: {content_type}/{file_extension}"
|
|
)
|
|
|
|
with tempfile.NamedTemporaryFile(
|
|
mode="wb", suffix=file_extension, delete=False
|
|
) as temp_file:
|
|
temp_file.write(response.content)
|
|
temp_file.flush()
|
|
temp_path = Path(temp_file.name)
|
|
|
|
try:
|
|
return await converter.convert_file(temp_path)
|
|
finally:
|
|
temp_path.unlink(missing_ok=True)
|
|
|
|
elif parsed.scheme == "file":
|
|
# file:// URI
|
|
file_path = Path(parsed.path)
|
|
if not file_path.exists():
|
|
raise ValueError(f"File does not exist: {file_path}")
|
|
if file_path.suffix.lower() not in converter.supported_extensions:
|
|
raise ValueError(f"Unsupported file extension: {file_path.suffix}")
|
|
return await converter.convert_file(file_path)
|
|
|
|
else:
|
|
# Treat as text content
|
|
return await converter.convert_text(source, format=format)
|
|
|
|
|
|
async def chunk(
|
|
config: AppConfig,
|
|
docling_document: "DoclingDocument",
|
|
*,
|
|
existing_picture_data: dict[str, bytes] | None = None,
|
|
document_id: str | None = None,
|
|
) -> list[Chunk]:
|
|
"""Chunk a DoclingDocument into Chunks.
|
|
|
|
When the configured embedder supports images, also emit one synthetic
|
|
Chunk per ``PictureItem`` with available bytes (see ``build_picture_chunks``)
|
|
and merge them with text chunks in structural (``iterate_items()``) order.
|
|
``chunk.order`` is the index in the merged list.
|
|
|
|
``existing_picture_data`` (snapshot keyed by ``self_ref``) supplies bytes
|
|
for pictures whose ``image.uri`` has been stripped — used by the rebuild
|
|
path where the docling is loaded from the stored blob.
|
|
"""
|
|
from haiku.rag.chunkers import get_chunker
|
|
from haiku.rag.embeddings import get_embedder
|
|
|
|
chunker = get_chunker(config)
|
|
text_chunks = await chunker.chunk(docling_document)
|
|
|
|
if not get_embedder(config).supports_images:
|
|
for i, c in enumerate(text_chunks):
|
|
c.order = i
|
|
return text_chunks
|
|
|
|
picture_chunks = build_picture_chunks(
|
|
docling_document,
|
|
document_id=document_id,
|
|
existing_picture_data=existing_picture_data,
|
|
)
|
|
|
|
if not picture_chunks:
|
|
for i, c in enumerate(text_chunks):
|
|
c.order = i
|
|
return text_chunks
|
|
|
|
positions = {
|
|
item.self_ref: pos
|
|
for pos, (item, _level) in enumerate(docling_document.iterate_items())
|
|
}
|
|
|
|
def first_pos(c: Chunk) -> int:
|
|
refs = (c.metadata or {}).get("doc_item_refs") or []
|
|
return positions.get(refs[0], len(positions)) if refs else len(positions)
|
|
|
|
merged = sorted(text_chunks + picture_chunks, key=first_pos)
|
|
for i, c in enumerate(merged):
|
|
c.order = i
|
|
return merged
|
|
|
|
|
|
def build_picture_chunks(
|
|
docling_document: "DoclingDocument",
|
|
*,
|
|
document_id: str | None = None,
|
|
existing_picture_data: dict[str, bytes] | None = None,
|
|
) -> list[Chunk]:
|
|
"""Emit one synthetic ``Chunk`` per ``PictureItem`` with available bytes.
|
|
|
|
Bytes come from ``picture.image.uri`` (live data URI on a freshly-converted
|
|
docling) or from ``existing_picture_data`` keyed by ``self_ref`` (snapshot
|
|
taken before a delete-and-re-extract cycle, when the live docling has had
|
|
its picture URIs stripped). Pictures with no available bytes are skipped.
|
|
|
|
The bytes ride on ``Chunk._picture_data`` (a PrivateAttr — not serialized)
|
|
so ``embed_chunks`` can route them through ``embed_image_query``. The
|
|
``order`` field is left at its default (0); the caller (``chunk()``)
|
|
reassigns it after merging with text chunks in structural order.
|
|
"""
|
|
from haiku.rag.store.models.document_item import (
|
|
_decode_picture_bytes,
|
|
extract_item_text,
|
|
)
|
|
|
|
existing = existing_picture_data or {}
|
|
chunks: list[Chunk] = []
|
|
|
|
for picture in docling_document.pictures:
|
|
picture_data = _decode_picture_bytes(picture)
|
|
if picture_data is None:
|
|
picture_data = existing.get(picture.self_ref)
|
|
if picture_data is None:
|
|
continue
|
|
|
|
text = extract_item_text(picture, docling_document) or ""
|
|
|
|
page_numbers: list[int] = []
|
|
for p in picture.prov:
|
|
if p.page_no not in page_numbers:
|
|
page_numbers.append(p.page_no)
|
|
|
|
metadata = {
|
|
"doc_item_refs": [picture.self_ref],
|
|
"labels": ["picture"],
|
|
"page_numbers": sorted(page_numbers),
|
|
"headings": None,
|
|
}
|
|
chunk = Chunk(
|
|
document_id=document_id,
|
|
content=text,
|
|
metadata=metadata,
|
|
)
|
|
chunk._picture_data = picture_data
|
|
chunks.append(chunk)
|
|
|
|
return chunks
|
|
|
|
|
|
async def ensure_chunks_embedded(config: AppConfig, chunks: list[Chunk]) -> list[Chunk]:
|
|
"""Ensure all chunks have embeddings, embedding any that don't.
|
|
|
|
Chunks that already have embeddings are passed through unchanged; missing
|
|
embeddings are filled in in-place in the returned list (preserving order).
|
|
"""
|
|
from haiku.rag.embeddings import embed_chunks
|
|
|
|
chunks_to_embed = [c for c in chunks if c.embedding is None]
|
|
|
|
if not chunks_to_embed:
|
|
return chunks
|
|
|
|
embedded = await embed_chunks(chunks_to_embed, config)
|
|
|
|
# Build result maintaining original order
|
|
embedded_map = {(c.content, c.order): c for c in embedded}
|
|
result = []
|
|
for ch in chunks:
|
|
if ch.embedding is not None:
|
|
result.append(ch)
|
|
else:
|
|
result.append(embedded_map[(ch.content, ch.order)])
|
|
|
|
return result
|
|
|
|
|
|
def get_extension_from_content_type_or_url(url: str, content_type: str) -> str:
|
|
"""Determine file extension from HTTP Content-Type header or URL suffix.
|
|
|
|
Returns the mapped extension for known content types, falling back to the
|
|
URL path suffix, and finally `.html` for generic web content.
|
|
"""
|
|
content_type_map = {
|
|
"text/html": ".html",
|
|
"text/plain": ".txt",
|
|
"text/markdown": ".md",
|
|
"application/pdf": ".pdf",
|
|
"application/json": ".json",
|
|
"text/csv": ".csv",
|
|
"application/vnd.openxmlformats-officedocument.wordprocessingml.document": ".docx",
|
|
"application/vnd.openxmlformats-officedocument.presentationml.presentation": ".pptx",
|
|
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": ".xlsx",
|
|
}
|
|
|
|
for ct, ext in content_type_map.items():
|
|
if ct in content_type:
|
|
return ext
|
|
|
|
parsed_url = urlparse(url)
|
|
path = Path(parsed_url.path)
|
|
if path.suffix:
|
|
return path.suffix.lower()
|
|
|
|
return ".html"
|