Extract document orchestration into client/documents.py
This commit is contained in:
Yiorgis Gozadinos
parent
cb02a2c3d8
commit
3224d1f20c
7 changed files with 561 additions and 607 deletions
|
|
@ -183,120 +183,6 @@ class HaikuRAG:
|
|||
|
||||
return await generate_title(self._config, document)
|
||||
|
||||
async def _store_document_with_chunks(
|
||||
self,
|
||||
document: Document,
|
||||
chunks: list[Chunk],
|
||||
docling_document: "DoclingDocument",
|
||||
) -> Document:
|
||||
"""Store a document with chunks, embedding any that lack embeddings.
|
||||
|
||||
Handles versioning/rollback on failure.
|
||||
|
||||
Args:
|
||||
document: The document to store (will be created).
|
||||
chunks: Chunks to store (will be embedded if lacking embeddings).
|
||||
docling_document: The DoclingDocument to extract items from.
|
||||
|
||||
Returns:
|
||||
The created Document instance with ID set.
|
||||
"""
|
||||
# Ensure all chunks have embeddings before storing
|
||||
chunks = await self._ensure_chunks_embedded(chunks)
|
||||
|
||||
# Snapshot table versions for versioned rollback (if supported)
|
||||
versions = await self.store.current_table_versions()
|
||||
|
||||
# Create the document
|
||||
created_doc = await self.document_repository.create(document)
|
||||
|
||||
try:
|
||||
assert created_doc.id is not None, (
|
||||
"Document ID should not be None after creation"
|
||||
)
|
||||
# Set document_id and order for all chunks
|
||||
for order, chunk in enumerate(chunks):
|
||||
chunk.document_id = created_doc.id
|
||||
chunk.order = order
|
||||
|
||||
# Batch create all chunks in a single operation
|
||||
await self.chunk_repository.create(chunks)
|
||||
|
||||
# Extract and store document items for context expansion
|
||||
items = extract_items(created_doc.id, docling_document)
|
||||
await self.document_item_repository.create_items(created_doc.id, items)
|
||||
|
||||
# Vacuum old versions in background (non-blocking) if auto_vacuum enabled
|
||||
if self._config.storage.auto_vacuum:
|
||||
self._schedule_vacuum()
|
||||
|
||||
return created_doc
|
||||
except Exception:
|
||||
# Roll back to the captured versions and re-raise
|
||||
await self.store.restore_table_versions(versions)
|
||||
raise
|
||||
|
||||
async def _update_document_with_chunks(
|
||||
self,
|
||||
document: Document,
|
||||
chunks: list[Chunk],
|
||||
docling_document: "DoclingDocument | None" = None,
|
||||
) -> Document:
|
||||
"""Update a document and replace its chunks, embedding any that lack embeddings.
|
||||
|
||||
Handles versioning/rollback on failure.
|
||||
|
||||
Args:
|
||||
document: The document to update (must have ID set).
|
||||
chunks: Chunks to replace existing (will be embedded if lacking embeddings).
|
||||
docling_document: The DoclingDocument to extract items from.
|
||||
When None, existing items are preserved.
|
||||
|
||||
Returns:
|
||||
The updated Document instance.
|
||||
"""
|
||||
assert document.id is not None, "Document ID is required for update"
|
||||
|
||||
# Ensure all chunks have embeddings before storing
|
||||
chunks = await self._ensure_chunks_embedded(chunks)
|
||||
|
||||
# Snapshot table versions for versioned rollback
|
||||
versions = await self.store.current_table_versions()
|
||||
|
||||
# Delete existing chunks before writing new ones
|
||||
await self.chunk_repository.delete_by_document_id(document.id)
|
||||
|
||||
try:
|
||||
# Update the document
|
||||
updated_doc = await self.document_repository.update(document)
|
||||
|
||||
# Set document_id and order for all chunks
|
||||
assert updated_doc.id is not None
|
||||
for order, chunk in enumerate(chunks):
|
||||
chunk.document_id = updated_doc.id
|
||||
chunk.order = order
|
||||
|
||||
# Batch create all chunks in a single operation
|
||||
await self.chunk_repository.create(chunks)
|
||||
|
||||
# Replace document items when a new DoclingDocument is provided
|
||||
if docling_document is not None:
|
||||
await self.document_item_repository.delete_by_document_id(
|
||||
updated_doc.id
|
||||
)
|
||||
items = extract_items(updated_doc.id, docling_document)
|
||||
await self.document_item_repository.create_items(updated_doc.id, items)
|
||||
|
||||
# Vacuum old versions in background (non-blocking) if auto_vacuum enabled
|
||||
if self._config.storage.auto_vacuum:
|
||||
self._schedule_vacuum()
|
||||
|
||||
return updated_doc
|
||||
except Exception:
|
||||
# Roll back to the captured versions and re-raise
|
||||
await self.store.restore_table_versions(versions)
|
||||
raise
|
||||
|
||||
async def create_document(
|
||||
self,
|
||||
content: str,
|
||||
|
|
@ -305,49 +191,9 @@ class HaikuRAG:
|
|||
metadata: dict | None = None,
|
||||
format: str = "md",
|
||||
) -> Document:
|
||||
"""Create a new document from text content.
|
||||
from haiku.rag.client.documents import create_document
|
||||
|
||||
Converts the content, chunks it, and generates embeddings.
|
||||
|
||||
Args:
|
||||
content: The text content of the document.
|
||||
uri: Optional URI identifier for the document.
|
||||
title: Optional title for the document.
|
||||
metadata: Optional metadata dictionary.
|
||||
format: The format of the content ("md", "html", or "plain").
|
||||
Defaults to "md". Use "plain" for plain text without parsing.
|
||||
|
||||
Returns:
|
||||
The created Document instance.
|
||||
"""
|
||||
from haiku.rag.embeddings import embed_chunks
|
||||
|
||||
# Convert → Chunk → Embed using primitives
|
||||
converter = get_converter(self._config)
|
||||
docling_document = await converter.convert_text(content, format=format)
|
||||
chunks = await self.chunk(docling_document)
|
||||
embedded_chunks = await embed_chunks(chunks, self._config)
|
||||
|
||||
# Store markdown export as content for better display/readability
|
||||
# The original content is preserved in docling_document
|
||||
stored_content = docling_document.export_to_markdown()
|
||||
|
||||
if title is None:
|
||||
title = await self._resolve_title(docling_document, stored_content)
|
||||
|
||||
# Create document model
|
||||
document = Document(
|
||||
content=stored_content,
|
||||
uri=uri,
|
||||
title=title,
|
||||
metadata=metadata or {},
|
||||
)
|
||||
document.set_docling(docling_document)
|
||||
|
||||
# Store document and chunks
|
||||
return await self._store_document_with_chunks(
|
||||
document, embedded_chunks, docling_document
|
||||
)
|
||||
return await create_document(self, content, uri, title, metadata, format)
|
||||
|
||||
async def import_document(
|
||||
self,
|
||||
|
|
@ -357,316 +203,43 @@ class HaikuRAG:
|
|||
title: str | None = None,
|
||||
metadata: dict | None = None,
|
||||
) -> Document:
|
||||
"""Import a pre-processed document with chunks.
|
||||
from haiku.rag.client.documents import import_document
|
||||
|
||||
Use this when document conversion, chunking, and embedding were done
|
||||
externally and you want to store the results in haiku.rag.
|
||||
|
||||
Args:
|
||||
docling_document: The DoclingDocument to import.
|
||||
chunks: Pre-created chunks. Chunks without embeddings will be
|
||||
automatically embedded.
|
||||
uri: Optional URI identifier for the document.
|
||||
title: Optional title for the document.
|
||||
metadata: Optional metadata dictionary.
|
||||
|
||||
Returns:
|
||||
The created Document instance.
|
||||
"""
|
||||
content = docling_document.export_to_markdown()
|
||||
if title is None:
|
||||
title = await self._resolve_title(docling_document, content)
|
||||
|
||||
document = Document(
|
||||
content=content,
|
||||
uri=uri,
|
||||
title=title,
|
||||
metadata=metadata or {},
|
||||
)
|
||||
document.set_docling(docling_document)
|
||||
|
||||
return await self._store_document_with_chunks(
|
||||
document, chunks, docling_document
|
||||
return await import_document(
|
||||
self, docling_document, chunks, uri, title, metadata
|
||||
)
|
||||
|
||||
async def create_document_from_source(
|
||||
self, source: str | Path, title: str | None = None, metadata: dict | None = None
|
||||
self,
|
||||
source: str | Path,
|
||||
title: str | None = None,
|
||||
metadata: dict | None = None,
|
||||
) -> Document | list[Document]:
|
||||
"""Create or update document(s) from a file path, directory, or URL.
|
||||
from haiku.rag.client.documents import create_document_from_source
|
||||
|
||||
Checks if a document with the same URI already exists:
|
||||
- If MD5 is unchanged, returns existing document
|
||||
- If MD5 changed, updates the document
|
||||
- If no document exists, creates a new one
|
||||
return await create_document_from_source(self, source, title, metadata)
|
||||
|
||||
Args:
|
||||
source: File path, directory (as string or Path), or URL to parse
|
||||
title: Optional title (only used for single files, not directories)
|
||||
metadata: Optional metadata dictionary
|
||||
async def update_document(
|
||||
self,
|
||||
document_id: str,
|
||||
content: str | None = None,
|
||||
metadata: dict | None = None,
|
||||
chunks: list[Chunk] | None = None,
|
||||
title: str | None = None,
|
||||
docling_document: "DoclingDocument | None" = None,
|
||||
) -> Document:
|
||||
from haiku.rag.client.documents import update_document
|
||||
|
||||
Returns:
|
||||
Document instance (created, updated, or existing) for single files/URLs
|
||||
List of Document instances for directories
|
||||
|
||||
Raises:
|
||||
ValueError: If the file/URL cannot be parsed or doesn't exist
|
||||
httpx.RequestError: If URL request fails
|
||||
"""
|
||||
# Normalize metadata
|
||||
metadata = metadata or {}
|
||||
|
||||
# Check if it's a URL
|
||||
source_str = str(source)
|
||||
parsed_url = urlparse(source_str)
|
||||
if parsed_url.scheme in ("http", "https"):
|
||||
return await self._create_or_update_document_from_url(
|
||||
source_str, title=title, metadata=metadata
|
||||
)
|
||||
elif parsed_url.scheme == "file":
|
||||
# Handle file:// URI by converting to path
|
||||
source_path = Path(parsed_url.path)
|
||||
else:
|
||||
# Handle as regular file path
|
||||
source_path = Path(source) if isinstance(source, str) else source
|
||||
|
||||
# Handle directories
|
||||
if source_path.is_dir():
|
||||
from haiku.rag.monitor import FileFilter
|
||||
|
||||
documents = []
|
||||
filter = FileFilter(
|
||||
ignore_patterns=self._config.monitor.ignore_patterns or None,
|
||||
include_patterns=self._config.monitor.include_patterns or None,
|
||||
)
|
||||
for path in source_path.rglob("*"):
|
||||
if path.is_file() and filter.include_file(str(path)):
|
||||
doc = await self._create_document_from_file(
|
||||
path, title=None, metadata=metadata
|
||||
)
|
||||
documents.append(doc)
|
||||
return documents
|
||||
|
||||
# Handle single file
|
||||
return await self._create_document_from_file(
|
||||
source_path, title=title, metadata=metadata
|
||||
return await update_document(
|
||||
self,
|
||||
document_id,
|
||||
content,
|
||||
metadata,
|
||||
chunks,
|
||||
title,
|
||||
docling_document,
|
||||
)
|
||||
|
||||
async def _create_document_from_file(
|
||||
self, source_path: Path, title: str | None = None, metadata: dict | None = None
|
||||
) -> Document:
|
||||
"""Create or update a document from a single file path.
|
||||
|
||||
Args:
|
||||
source_path: Path to the file
|
||||
title: Optional title
|
||||
metadata: Optional metadata dictionary
|
||||
|
||||
Returns:
|
||||
Document instance (created, updated, or existing)
|
||||
|
||||
Raises:
|
||||
ValueError: If the file cannot be parsed or doesn't exist
|
||||
"""
|
||||
from haiku.rag.embeddings import embed_chunks
|
||||
|
||||
metadata = metadata or {}
|
||||
|
||||
converter = get_converter(self._config)
|
||||
if source_path.suffix.lower() not in converter.supported_extensions:
|
||||
raise ValueError(f"Unsupported file extension: {source_path.suffix}")
|
||||
|
||||
if not source_path.exists():
|
||||
raise ValueError(f"File does not exist: {source_path}")
|
||||
|
||||
uri = source_path.absolute().as_uri()
|
||||
md5_hash = hashlib.md5(
|
||||
source_path.read_bytes(), usedforsecurity=False
|
||||
).hexdigest()
|
||||
|
||||
# Get content type from file extension (do before early return)
|
||||
content_type, _ = mimetypes.guess_type(str(source_path))
|
||||
if not content_type:
|
||||
content_type = "application/octet-stream"
|
||||
# Merge metadata with contentType and md5
|
||||
metadata.update({"contentType": content_type, "md5": md5_hash})
|
||||
|
||||
# Check if document already exists
|
||||
existing_doc = await self.get_document_by_uri(uri)
|
||||
if existing_doc and existing_doc.metadata.get("md5") == md5_hash:
|
||||
# MD5 unchanged; update title/metadata if provided
|
||||
updated = False
|
||||
if title is not None and title != existing_doc.title:
|
||||
existing_doc.title = title
|
||||
updated = True
|
||||
|
||||
# Check if metadata actually changed (beyond contentType and md5)
|
||||
merged_metadata = {**(existing_doc.metadata or {}), **metadata}
|
||||
if merged_metadata != existing_doc.metadata:
|
||||
existing_doc.metadata = merged_metadata
|
||||
updated = True
|
||||
|
||||
if updated:
|
||||
return await self.document_repository.update(existing_doc)
|
||||
return existing_doc
|
||||
|
||||
# Convert → Chunk → Embed using primitives
|
||||
docling_document = await self.convert(source_path)
|
||||
chunks = await self.chunk(docling_document)
|
||||
embedded_chunks = await embed_chunks(chunks, self._config)
|
||||
|
||||
stored_content = docling_document.export_to_markdown()
|
||||
|
||||
if existing_doc:
|
||||
# Update existing document and rechunk
|
||||
existing_doc.content = stored_content
|
||||
existing_doc.metadata = metadata
|
||||
existing_doc.set_docling(docling_document)
|
||||
if title is not None:
|
||||
existing_doc.title = title
|
||||
elif existing_doc.title is None:
|
||||
existing_doc.title = await self._resolve_title(
|
||||
docling_document, stored_content
|
||||
)
|
||||
return await self._update_document_with_chunks(
|
||||
existing_doc, embedded_chunks, docling_document
|
||||
)
|
||||
else:
|
||||
# Create new document
|
||||
if title is None:
|
||||
title = await self._resolve_title(docling_document, stored_content)
|
||||
document = Document(
|
||||
content=stored_content,
|
||||
uri=uri,
|
||||
title=title,
|
||||
metadata=metadata,
|
||||
)
|
||||
document.set_docling(docling_document)
|
||||
return await self._store_document_with_chunks(
|
||||
document, embedded_chunks, docling_document
|
||||
)
|
||||
|
||||
async def _create_or_update_document_from_url(
|
||||
self, url: str, title: str | None = None, metadata: dict | None = None
|
||||
) -> Document:
|
||||
"""Create or update a document from a URL by downloading and parsing the content.
|
||||
|
||||
Checks if a document with the same URI already exists:
|
||||
- If MD5 is unchanged, returns existing document
|
||||
- If MD5 changed, updates the document
|
||||
- If no document exists, creates a new one
|
||||
|
||||
Args:
|
||||
url: URL to download and parse
|
||||
metadata: Optional metadata dictionary
|
||||
|
||||
Returns:
|
||||
Document instance (created, updated, or existing)
|
||||
|
||||
Raises:
|
||||
ValueError: If the content cannot be parsed
|
||||
httpx.RequestError: If URL request fails
|
||||
"""
|
||||
from haiku.rag.embeddings import embed_chunks
|
||||
|
||||
metadata = metadata or {}
|
||||
|
||||
converter = get_converter(self._config)
|
||||
supported_extensions = converter.supported_extensions
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
response = await client.get(url)
|
||||
response.raise_for_status()
|
||||
|
||||
md5_hash = hashlib.md5(response.content).hexdigest()
|
||||
|
||||
# Get content type early (used for potential no-op update)
|
||||
content_type = response.headers.get("content-type", "").lower()
|
||||
|
||||
# Check if document already exists
|
||||
existing_doc = await self.get_document_by_uri(url)
|
||||
if existing_doc and existing_doc.metadata.get("md5") == md5_hash:
|
||||
# MD5 unchanged; update title/metadata if provided
|
||||
updated = False
|
||||
if title is not None and title != existing_doc.title:
|
||||
existing_doc.title = title
|
||||
updated = True
|
||||
|
||||
metadata.update({"contentType": content_type, "md5": md5_hash})
|
||||
# Check if metadata actually changed (beyond contentType and md5)
|
||||
merged_metadata = {**(existing_doc.metadata or {}), **metadata}
|
||||
if merged_metadata != existing_doc.metadata:
|
||||
existing_doc.metadata = merged_metadata
|
||||
updated = True
|
||||
|
||||
if updated:
|
||||
return await self.document_repository.update(existing_doc)
|
||||
return existing_doc
|
||||
file_extension = self._get_extension_from_content_type_or_url(
|
||||
url, content_type
|
||||
)
|
||||
|
||||
if file_extension not in supported_extensions:
|
||||
raise ValueError(
|
||||
f"Unsupported content type/extension: {content_type}/{file_extension}"
|
||||
)
|
||||
|
||||
# Create a temporary file with the appropriate 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:
|
||||
# Convert → Chunk → Embed using primitives
|
||||
docling_document = await self.convert(temp_path)
|
||||
chunks = await self.chunk(docling_document)
|
||||
embedded_chunks = await embed_chunks(chunks, self._config)
|
||||
finally:
|
||||
temp_path.unlink(missing_ok=True)
|
||||
|
||||
# Merge metadata with contentType and md5
|
||||
metadata.update({"contentType": content_type, "md5": md5_hash})
|
||||
|
||||
stored_content = docling_document.export_to_markdown()
|
||||
|
||||
if existing_doc:
|
||||
# Update existing document and rechunk
|
||||
existing_doc.content = stored_content
|
||||
existing_doc.metadata = metadata
|
||||
existing_doc.set_docling(docling_document)
|
||||
if title is not None:
|
||||
existing_doc.title = title
|
||||
elif existing_doc.title is None:
|
||||
existing_doc.title = await self._resolve_title(
|
||||
docling_document, stored_content
|
||||
)
|
||||
return await self._update_document_with_chunks(
|
||||
existing_doc, embedded_chunks, docling_document
|
||||
)
|
||||
else:
|
||||
# Create new document
|
||||
if title is None:
|
||||
title = await self._resolve_title(docling_document, stored_content)
|
||||
document = Document(
|
||||
content=stored_content,
|
||||
uri=url,
|
||||
title=title,
|
||||
metadata=metadata,
|
||||
)
|
||||
document.set_docling(docling_document)
|
||||
return await self._store_document_with_chunks(
|
||||
document, embedded_chunks, docling_document
|
||||
)
|
||||
|
||||
def _get_extension_from_content_type_or_url(
|
||||
self, url: str, content_type: str
|
||||
) -> str:
|
||||
from haiku.rag.client.processing import get_extension_from_content_type_or_url
|
||||
|
||||
return get_extension_from_content_type_or_url(url, content_type)
|
||||
|
||||
async def get_document_by_id(self, document_id: str) -> Document | None:
|
||||
"""Get a document by its ID.
|
||||
|
||||
|
|
@ -724,99 +297,6 @@ class HaikuRAG:
|
|||
|
||||
return None
|
||||
|
||||
async def update_document(
|
||||
self,
|
||||
document_id: str,
|
||||
content: str | None = None,
|
||||
metadata: dict | None = None,
|
||||
chunks: list[Chunk] | None = None,
|
||||
title: str | None = None,
|
||||
docling_document: "DoclingDocument | None" = None,
|
||||
) -> Document:
|
||||
"""Update a document by ID.
|
||||
|
||||
Updates specified fields. When content or docling_document is provided,
|
||||
the document is rechunked and re-embedded. Updates to only metadata or title
|
||||
skip rechunking for efficiency.
|
||||
|
||||
Args:
|
||||
document_id: The ID of the document to update.
|
||||
content: New content (mutually exclusive with docling_document).
|
||||
metadata: New metadata dict.
|
||||
chunks: Custom chunks (will be embedded if missing embeddings).
|
||||
title: New title.
|
||||
docling_document: DoclingDocument to replace content (mutually exclusive with content).
|
||||
|
||||
Returns:
|
||||
The updated Document instance.
|
||||
|
||||
Raises:
|
||||
ValueError: If document not found, or if both content and docling_document
|
||||
are provided.
|
||||
"""
|
||||
from haiku.rag.embeddings import embed_chunks
|
||||
|
||||
# Validate: content and docling_document are mutually exclusive
|
||||
if content is not None and docling_document is not None:
|
||||
raise ValueError(
|
||||
"content and docling_document are mutually exclusive. "
|
||||
"Provide one or the other, not both."
|
||||
)
|
||||
|
||||
# Fetch the existing document
|
||||
existing_doc = await self.get_document_by_id(document_id)
|
||||
if existing_doc is None:
|
||||
raise ValueError(f"Document with ID {document_id} not found")
|
||||
|
||||
# Update metadata/title fields
|
||||
if title is not None:
|
||||
existing_doc.title = title
|
||||
if metadata is not None:
|
||||
existing_doc.metadata = metadata
|
||||
|
||||
# Only metadata/title update - no rechunking needed
|
||||
if content is None and chunks is None and docling_document is None:
|
||||
return await self.document_repository.update(existing_doc)
|
||||
|
||||
# Custom chunks provided - use them as-is
|
||||
if chunks is not None:
|
||||
# Store docling data if provided
|
||||
if docling_document is not None:
|
||||
existing_doc.content = docling_document.export_to_markdown()
|
||||
existing_doc.set_docling(docling_document)
|
||||
elif content is not None:
|
||||
existing_doc.content = content
|
||||
|
||||
return await self._update_document_with_chunks(
|
||||
existing_doc, chunks, docling_document
|
||||
)
|
||||
|
||||
# DoclingDocument provided without chunks - chunk and embed using primitives
|
||||
if docling_document is not None:
|
||||
existing_doc.content = docling_document.export_to_markdown()
|
||||
existing_doc.set_docling(docling_document)
|
||||
|
||||
new_chunks = await self.chunk(docling_document)
|
||||
embedded_chunks = await embed_chunks(new_chunks, self._config)
|
||||
return await self._update_document_with_chunks(
|
||||
existing_doc, embedded_chunks, docling_document
|
||||
)
|
||||
|
||||
# Content provided without chunks - convert, chunk, and embed using primitives
|
||||
assert content is not None
|
||||
existing_doc.content = content
|
||||
converter = get_converter(self._config)
|
||||
converted_docling = await converter.convert_text(
|
||||
existing_doc.content, format="md"
|
||||
)
|
||||
existing_doc.set_docling(converted_docling)
|
||||
|
||||
new_chunks = await self.chunk(converted_docling)
|
||||
embedded_chunks = await embed_chunks(new_chunks, self._config)
|
||||
return await self._update_document_with_chunks(
|
||||
existing_doc, embedded_chunks, converted_docling
|
||||
)
|
||||
|
||||
async def delete_document(self, document_id: str) -> bool:
|
||||
"""Delete a document by its ID."""
|
||||
return await self.document_repository.delete(document_id)
|
||||
|
|
@ -920,18 +400,6 @@ class HaikuRAG:
|
|||
async for doc_id in rebuild_database(self, mode):
|
||||
yield doc_id
|
||||
|
||||
def _check_source_accessible(self, uri: str) -> bool:
|
||||
"""Check if a document's source URI is accessible."""
|
||||
parsed_url = urlparse(uri)
|
||||
try:
|
||||
if parsed_url.scheme == "file":
|
||||
return Path(parsed_url.path).exists()
|
||||
elif parsed_url.scheme in ("http", "https"):
|
||||
return True
|
||||
return False
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
async def vacuum(self) -> None:
|
||||
"""Optimize and clean up old versions across all tables."""
|
||||
await self.store.vacuum()
|
||||
|
|
|
|||
491
haiku_rag_slim/haiku/rag/client/documents.py
Normal file
491
haiku_rag_slim/haiku/rag/client/documents.py
Normal file
|
|
@ -0,0 +1,491 @@
|
|||
import hashlib
|
||||
import mimetypes
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import httpx
|
||||
|
||||
from haiku.rag.converters import get_converter
|
||||
from haiku.rag.store.models.chunk import Chunk
|
||||
from haiku.rag.store.models.document import Document
|
||||
from haiku.rag.store.models.document_item import extract_items
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from docling_core.types.doc.document import DoclingDocument
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
|
||||
async def _store_document_with_chunks(
|
||||
client: "HaikuRAG",
|
||||
document: Document,
|
||||
chunks: list[Chunk],
|
||||
docling_document: "DoclingDocument",
|
||||
) -> Document:
|
||||
"""Store a document with chunks, embedding any that lack embeddings.
|
||||
|
||||
Handles versioning/rollback on failure.
|
||||
"""
|
||||
# Ensure all chunks have embeddings before storing
|
||||
chunks = await client._ensure_chunks_embedded(chunks)
|
||||
|
||||
# Snapshot table versions for versioned rollback (if supported)
|
||||
versions = await client.store.current_table_versions()
|
||||
|
||||
# Create the document
|
||||
created_doc = await client.document_repository.create(document)
|
||||
|
||||
try:
|
||||
assert created_doc.id is not None, (
|
||||
"Document ID should not be None after creation"
|
||||
)
|
||||
# Set document_id and order for all chunks
|
||||
for order, chunk in enumerate(chunks):
|
||||
chunk.document_id = created_doc.id
|
||||
chunk.order = order
|
||||
|
||||
# Batch create all chunks in a single operation
|
||||
await client.chunk_repository.create(chunks)
|
||||
|
||||
# Extract and store document items for context expansion
|
||||
items = extract_items(created_doc.id, docling_document)
|
||||
await client.document_item_repository.create_items(created_doc.id, items)
|
||||
|
||||
# Vacuum old versions in background (non-blocking) if auto_vacuum enabled
|
||||
if client._config.storage.auto_vacuum:
|
||||
client._schedule_vacuum()
|
||||
|
||||
return created_doc
|
||||
except Exception:
|
||||
# Roll back to the captured versions and re-raise
|
||||
await client.store.restore_table_versions(versions)
|
||||
raise
|
||||
|
||||
|
||||
async def _update_document_with_chunks(
|
||||
client: "HaikuRAG",
|
||||
document: Document,
|
||||
chunks: list[Chunk],
|
||||
docling_document: "DoclingDocument | None" = None,
|
||||
) -> Document:
|
||||
"""Update a document and replace its chunks, embedding any that lack embeddings.
|
||||
|
||||
Handles versioning/rollback on failure. When `docling_document` is None,
|
||||
existing items are preserved.
|
||||
"""
|
||||
assert document.id is not None, "Document ID is required for update"
|
||||
|
||||
chunks = await client._ensure_chunks_embedded(chunks)
|
||||
|
||||
versions = await client.store.current_table_versions()
|
||||
|
||||
# Delete existing chunks before writing new ones
|
||||
await client.chunk_repository.delete_by_document_id(document.id)
|
||||
|
||||
try:
|
||||
updated_doc = await client.document_repository.update(document)
|
||||
|
||||
assert updated_doc.id is not None
|
||||
for order, chunk in enumerate(chunks):
|
||||
chunk.document_id = updated_doc.id
|
||||
chunk.order = order
|
||||
|
||||
await client.chunk_repository.create(chunks)
|
||||
|
||||
# Replace document items when a new DoclingDocument is provided
|
||||
if docling_document is not None:
|
||||
await client.document_item_repository.delete_by_document_id(updated_doc.id)
|
||||
items = extract_items(updated_doc.id, docling_document)
|
||||
await client.document_item_repository.create_items(updated_doc.id, items)
|
||||
|
||||
if client._config.storage.auto_vacuum:
|
||||
client._schedule_vacuum()
|
||||
|
||||
return updated_doc
|
||||
except Exception:
|
||||
await client.store.restore_table_versions(versions)
|
||||
raise
|
||||
|
||||
|
||||
async def create_document(
|
||||
client: "HaikuRAG",
|
||||
content: str,
|
||||
uri: str | None = None,
|
||||
title: str | None = None,
|
||||
metadata: dict | None = None,
|
||||
format: str = "md",
|
||||
) -> Document:
|
||||
"""Create a new document from text content.
|
||||
|
||||
Converts the content, chunks it, and generates embeddings.
|
||||
"""
|
||||
from haiku.rag.embeddings import embed_chunks
|
||||
|
||||
# Convert → Chunk → Embed using primitives
|
||||
converter = get_converter(client._config)
|
||||
docling_document = await converter.convert_text(content, format=format)
|
||||
chunks = await client.chunk(docling_document)
|
||||
embedded_chunks = await embed_chunks(chunks, client._config)
|
||||
|
||||
# Store markdown export as content for better display/readability.
|
||||
# The original is preserved in docling_document.
|
||||
stored_content = docling_document.export_to_markdown()
|
||||
|
||||
if title is None:
|
||||
title = await client._resolve_title(docling_document, stored_content)
|
||||
|
||||
document = Document(
|
||||
content=stored_content,
|
||||
uri=uri,
|
||||
title=title,
|
||||
metadata=metadata or {},
|
||||
)
|
||||
document.set_docling(docling_document)
|
||||
|
||||
return await _store_document_with_chunks(
|
||||
client, document, embedded_chunks, docling_document
|
||||
)
|
||||
|
||||
|
||||
async def import_document(
|
||||
client: "HaikuRAG",
|
||||
docling_document: "DoclingDocument",
|
||||
chunks: list[Chunk],
|
||||
uri: str | None = None,
|
||||
title: str | None = None,
|
||||
metadata: dict | None = None,
|
||||
) -> Document:
|
||||
"""Import a pre-processed document with chunks.
|
||||
|
||||
Use this when conversion, chunking, and embedding were done externally.
|
||||
Chunks without embeddings will be automatically embedded.
|
||||
"""
|
||||
content = docling_document.export_to_markdown()
|
||||
if title is None:
|
||||
title = await client._resolve_title(docling_document, content)
|
||||
|
||||
document = Document(
|
||||
content=content,
|
||||
uri=uri,
|
||||
title=title,
|
||||
metadata=metadata or {},
|
||||
)
|
||||
document.set_docling(docling_document)
|
||||
|
||||
return await _store_document_with_chunks(client, document, chunks, docling_document)
|
||||
|
||||
|
||||
async def create_document_from_source(
|
||||
client: "HaikuRAG",
|
||||
source: str | Path,
|
||||
title: str | None = None,
|
||||
metadata: dict | None = None,
|
||||
) -> Document | list[Document]:
|
||||
"""Create or update document(s) from a file path, directory, or URL.
|
||||
|
||||
Checks if a document with the same URI already exists:
|
||||
- If MD5 is unchanged, returns existing document
|
||||
- If MD5 changed, updates the document
|
||||
- If no document exists, creates a new one
|
||||
|
||||
Returns a single Document for files/URLs, a list for directories.
|
||||
"""
|
||||
metadata = metadata or {}
|
||||
|
||||
source_str = str(source)
|
||||
parsed_url = urlparse(source_str)
|
||||
if parsed_url.scheme in ("http", "https"):
|
||||
return await _create_or_update_document_from_url(
|
||||
client, source_str, title=title, metadata=metadata
|
||||
)
|
||||
elif parsed_url.scheme == "file":
|
||||
source_path = Path(parsed_url.path)
|
||||
else:
|
||||
source_path = Path(source) if isinstance(source, str) else source
|
||||
|
||||
if source_path.is_dir():
|
||||
from haiku.rag.monitor import FileFilter
|
||||
|
||||
documents = []
|
||||
filter = FileFilter(
|
||||
ignore_patterns=client._config.monitor.ignore_patterns or None,
|
||||
include_patterns=client._config.monitor.include_patterns or None,
|
||||
)
|
||||
for path in source_path.rglob("*"):
|
||||
if path.is_file() and filter.include_file(str(path)):
|
||||
doc = await _create_document_from_file(
|
||||
client, path, title=None, metadata=metadata
|
||||
)
|
||||
documents.append(doc)
|
||||
return documents
|
||||
|
||||
return await _create_document_from_file(
|
||||
client, source_path, title=title, metadata=metadata
|
||||
)
|
||||
|
||||
|
||||
async def _create_document_from_file(
|
||||
client: "HaikuRAG",
|
||||
source_path: Path,
|
||||
title: str | None = None,
|
||||
metadata: dict | None = None,
|
||||
) -> Document:
|
||||
"""Create or update a document from a single file path."""
|
||||
from haiku.rag.embeddings import embed_chunks
|
||||
|
||||
metadata = metadata or {}
|
||||
|
||||
converter = get_converter(client._config)
|
||||
if source_path.suffix.lower() not in converter.supported_extensions:
|
||||
raise ValueError(f"Unsupported file extension: {source_path.suffix}")
|
||||
|
||||
if not source_path.exists():
|
||||
raise ValueError(f"File does not exist: {source_path}")
|
||||
|
||||
uri = source_path.absolute().as_uri()
|
||||
md5_hash = hashlib.md5(source_path.read_bytes(), usedforsecurity=False).hexdigest()
|
||||
|
||||
content_type, _ = mimetypes.guess_type(str(source_path))
|
||||
if not content_type:
|
||||
content_type = "application/octet-stream"
|
||||
metadata.update({"contentType": content_type, "md5": md5_hash})
|
||||
|
||||
# Check if document already exists
|
||||
existing_doc = await client.get_document_by_uri(uri)
|
||||
if existing_doc and existing_doc.metadata.get("md5") == md5_hash:
|
||||
# MD5 unchanged; update title/metadata if provided
|
||||
updated = False
|
||||
if title is not None and title != existing_doc.title:
|
||||
existing_doc.title = title
|
||||
updated = True
|
||||
|
||||
merged_metadata = {**(existing_doc.metadata or {}), **metadata}
|
||||
if merged_metadata != existing_doc.metadata:
|
||||
existing_doc.metadata = merged_metadata
|
||||
updated = True
|
||||
|
||||
if updated:
|
||||
return await client.document_repository.update(existing_doc)
|
||||
return existing_doc
|
||||
|
||||
# Convert → Chunk → Embed
|
||||
docling_document = await client.convert(source_path)
|
||||
chunks = await client.chunk(docling_document)
|
||||
embedded_chunks = await embed_chunks(chunks, client._config)
|
||||
|
||||
stored_content = docling_document.export_to_markdown()
|
||||
|
||||
if existing_doc:
|
||||
# Update existing document and rechunk
|
||||
existing_doc.content = stored_content
|
||||
existing_doc.metadata = metadata
|
||||
existing_doc.set_docling(docling_document)
|
||||
if title is not None:
|
||||
existing_doc.title = title
|
||||
elif existing_doc.title is None:
|
||||
existing_doc.title = await client._resolve_title(
|
||||
docling_document, stored_content
|
||||
)
|
||||
return await _update_document_with_chunks(
|
||||
client, existing_doc, embedded_chunks, docling_document
|
||||
)
|
||||
else:
|
||||
if title is None:
|
||||
title = await client._resolve_title(docling_document, stored_content)
|
||||
document = Document(
|
||||
content=stored_content,
|
||||
uri=uri,
|
||||
title=title,
|
||||
metadata=metadata,
|
||||
)
|
||||
document.set_docling(docling_document)
|
||||
return await _store_document_with_chunks(
|
||||
client, document, embedded_chunks, docling_document
|
||||
)
|
||||
|
||||
|
||||
async def _create_or_update_document_from_url(
|
||||
client: "HaikuRAG",
|
||||
url: str,
|
||||
title: str | None = None,
|
||||
metadata: dict | None = None,
|
||||
) -> Document:
|
||||
"""Create or update a document from a URL by downloading and parsing the content."""
|
||||
from haiku.rag.client.processing import get_extension_from_content_type_or_url
|
||||
from haiku.rag.embeddings import embed_chunks
|
||||
|
||||
metadata = metadata or {}
|
||||
|
||||
converter = get_converter(client._config)
|
||||
supported_extensions = converter.supported_extensions
|
||||
|
||||
async with httpx.AsyncClient() as http:
|
||||
response = await http.get(url)
|
||||
response.raise_for_status()
|
||||
|
||||
md5_hash = hashlib.md5(response.content).hexdigest()
|
||||
|
||||
content_type = response.headers.get("content-type", "").lower()
|
||||
|
||||
# Check if document already exists
|
||||
existing_doc = await client.get_document_by_uri(url)
|
||||
if existing_doc and existing_doc.metadata.get("md5") == md5_hash:
|
||||
updated = False
|
||||
if title is not None and title != existing_doc.title:
|
||||
existing_doc.title = title
|
||||
updated = True
|
||||
|
||||
metadata.update({"contentType": content_type, "md5": md5_hash})
|
||||
merged_metadata = {**(existing_doc.metadata or {}), **metadata}
|
||||
if merged_metadata != existing_doc.metadata:
|
||||
existing_doc.metadata = merged_metadata
|
||||
updated = True
|
||||
|
||||
if updated:
|
||||
return await client.document_repository.update(existing_doc)
|
||||
return existing_doc
|
||||
|
||||
file_extension = get_extension_from_content_type_or_url(url, content_type)
|
||||
|
||||
if file_extension not in 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:
|
||||
docling_document = await client.convert(temp_path)
|
||||
chunks = await client.chunk(docling_document)
|
||||
embedded_chunks = await embed_chunks(chunks, client._config)
|
||||
finally:
|
||||
temp_path.unlink(missing_ok=True)
|
||||
|
||||
metadata.update({"contentType": content_type, "md5": md5_hash})
|
||||
|
||||
stored_content = docling_document.export_to_markdown()
|
||||
|
||||
if existing_doc:
|
||||
existing_doc.content = stored_content
|
||||
existing_doc.metadata = metadata
|
||||
existing_doc.set_docling(docling_document)
|
||||
if title is not None:
|
||||
existing_doc.title = title
|
||||
elif existing_doc.title is None:
|
||||
existing_doc.title = await client._resolve_title(
|
||||
docling_document, stored_content
|
||||
)
|
||||
return await _update_document_with_chunks(
|
||||
client, existing_doc, embedded_chunks, docling_document
|
||||
)
|
||||
else:
|
||||
if title is None:
|
||||
title = await client._resolve_title(docling_document, stored_content)
|
||||
document = Document(
|
||||
content=stored_content,
|
||||
uri=url,
|
||||
title=title,
|
||||
metadata=metadata,
|
||||
)
|
||||
document.set_docling(docling_document)
|
||||
return await _store_document_with_chunks(
|
||||
client, document, embedded_chunks, docling_document
|
||||
)
|
||||
|
||||
|
||||
async def update_document(
|
||||
client: "HaikuRAG",
|
||||
document_id: str,
|
||||
content: str | None = None,
|
||||
metadata: dict | None = None,
|
||||
chunks: list[Chunk] | None = None,
|
||||
title: str | None = None,
|
||||
docling_document: "DoclingDocument | None" = None,
|
||||
) -> Document:
|
||||
"""Update a document by ID.
|
||||
|
||||
Updates specified fields. When content or docling_document is provided, the
|
||||
document is rechunked and re-embedded. Updates to only metadata or title
|
||||
skip rechunking for efficiency.
|
||||
|
||||
Raises:
|
||||
ValueError: If document not found, or if both content and
|
||||
docling_document are provided.
|
||||
"""
|
||||
from haiku.rag.embeddings import embed_chunks
|
||||
|
||||
# Validate: content and docling_document are mutually exclusive
|
||||
if content is not None and docling_document is not None:
|
||||
raise ValueError(
|
||||
"content and docling_document are mutually exclusive. "
|
||||
"Provide one or the other, not both."
|
||||
)
|
||||
|
||||
existing_doc = await client.get_document_by_id(document_id)
|
||||
if existing_doc is None:
|
||||
raise ValueError(f"Document with ID {document_id} not found")
|
||||
|
||||
if title is not None:
|
||||
existing_doc.title = title
|
||||
if metadata is not None:
|
||||
existing_doc.metadata = metadata
|
||||
|
||||
# Only metadata/title update - no rechunking needed
|
||||
if content is None and chunks is None and docling_document is None:
|
||||
return await client.document_repository.update(existing_doc)
|
||||
|
||||
# Custom chunks provided - use them as-is
|
||||
if chunks is not None:
|
||||
if docling_document is not None:
|
||||
existing_doc.content = docling_document.export_to_markdown()
|
||||
existing_doc.set_docling(docling_document)
|
||||
elif content is not None:
|
||||
existing_doc.content = content
|
||||
|
||||
return await _update_document_with_chunks(
|
||||
client, existing_doc, chunks, docling_document
|
||||
)
|
||||
|
||||
# DoclingDocument provided without chunks - chunk and embed
|
||||
if docling_document is not None:
|
||||
existing_doc.content = docling_document.export_to_markdown()
|
||||
existing_doc.set_docling(docling_document)
|
||||
|
||||
new_chunks = await client.chunk(docling_document)
|
||||
embedded_chunks = await embed_chunks(new_chunks, client._config)
|
||||
return await _update_document_with_chunks(
|
||||
client, existing_doc, embedded_chunks, docling_document
|
||||
)
|
||||
|
||||
# Content provided without chunks - convert, chunk, and embed
|
||||
assert content is not None
|
||||
existing_doc.content = content
|
||||
converter = get_converter(client._config)
|
||||
converted_docling = await converter.convert_text(existing_doc.content, format="md")
|
||||
existing_doc.set_docling(converted_docling)
|
||||
|
||||
new_chunks = await client.chunk(converted_docling)
|
||||
embedded_chunks = await embed_chunks(new_chunks, client._config)
|
||||
return await _update_document_with_chunks(
|
||||
client, existing_doc, embedded_chunks, converted_docling
|
||||
)
|
||||
|
||||
|
||||
def check_source_accessible(uri: str) -> bool:
|
||||
"""Check if a document's source URI is accessible."""
|
||||
parsed_url = urlparse(uri)
|
||||
try:
|
||||
if parsed_url.scheme == "file":
|
||||
return Path(parsed_url.path).exists()
|
||||
elif parsed_url.scheme in ("http", "https"):
|
||||
return True
|
||||
return False
|
||||
except Exception:
|
||||
return False
|
||||
|
|
@ -4,6 +4,7 @@ from collections.abc import AsyncGenerator
|
|||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from haiku.rag.client.documents import check_source_accessible
|
||||
from haiku.rag.converters import get_converter
|
||||
from haiku.rag.store.models.chunk import Chunk
|
||||
from haiku.rag.store.models.document import Document
|
||||
|
|
@ -261,7 +262,7 @@ async def _rebuild_full(
|
|||
assert doc.id is not None
|
||||
|
||||
# Try to rebuild from source if available
|
||||
if doc.uri and client._check_source_accessible(doc.uri):
|
||||
if doc.uri and check_source_accessible(doc.uri):
|
||||
try:
|
||||
# Flush pending batch before source rebuild (creates new doc)
|
||||
if pending_docs:
|
||||
|
|
|
|||
|
|
@ -1,6 +1,10 @@
|
|||
import pytest
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.client.documents import (
|
||||
_store_document_with_chunks,
|
||||
_update_document_with_chunks,
|
||||
)
|
||||
from haiku.rag.store.engine import Store
|
||||
from haiku.rag.store.models.document_item import (
|
||||
DocumentItem,
|
||||
|
|
@ -217,7 +221,7 @@ class TestDocumentItemPopulation:
|
|||
|
||||
# Use _store_document_with_chunks directly with empty chunks
|
||||
# to avoid needing embeddings
|
||||
created = await rag._store_document_with_chunks(document, [], docling_doc)
|
||||
created = await _store_document_with_chunks(rag, document, [], docling_doc)
|
||||
assert created.id is not None
|
||||
|
||||
count = await rag.document_item_repository.get_item_count(created.id)
|
||||
|
|
@ -245,7 +249,7 @@ class TestDocumentItemPopulation:
|
|||
uri="test://doc",
|
||||
)
|
||||
document.set_docling(docling_doc)
|
||||
created = await rag._store_document_with_chunks(document, [], docling_doc)
|
||||
created = await _store_document_with_chunks(rag, document, [], docling_doc)
|
||||
assert created.id is not None
|
||||
assert await rag.document_item_repository.get_item_count(created.id) == 6
|
||||
|
||||
|
|
@ -254,7 +258,7 @@ class TestDocumentItemPopulation:
|
|||
new_doc.add_text(label=DocItemLabel.PARAGRAPH, text="Only one item now.")
|
||||
created.set_docling(new_doc)
|
||||
|
||||
await rag._update_document_with_chunks(created, [], new_doc)
|
||||
await _update_document_with_chunks(rag, created, [], new_doc)
|
||||
assert await rag.document_item_repository.get_item_count(created.id) == 1
|
||||
|
||||
async def test_delete_document_cascades_items(self, temp_db_path):
|
||||
|
|
@ -269,7 +273,7 @@ class TestDocumentItemPopulation:
|
|||
uri="test://doc",
|
||||
)
|
||||
document.set_docling(docling_doc)
|
||||
created = await rag._store_document_with_chunks(document, [], docling_doc)
|
||||
created = await _store_document_with_chunks(rag, document, [], docling_doc)
|
||||
assert created.id is not None
|
||||
assert await rag.document_item_repository.get_item_count(created.id) == 6
|
||||
|
||||
|
|
|
|||
|
|
@ -481,49 +481,35 @@ async def test_client_create_document_from_url_http_error(temp_db_path):
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
async def test_get_extension_from_content_type_or_url(temp_db_path):
|
||||
"""Test the helper method for determining file extensions."""
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
# Test content type mappings
|
||||
assert (
|
||||
client._get_extension_from_content_type_or_url("", "text/html") == ".html"
|
||||
)
|
||||
assert (
|
||||
client._get_extension_from_content_type_or_url("", "application/pdf")
|
||||
== ".pdf"
|
||||
)
|
||||
assert (
|
||||
client._get_extension_from_content_type_or_url("", "text/plain") == ".txt"
|
||||
)
|
||||
def test_get_extension_from_content_type_or_url():
|
||||
"""Test the helper function for determining file extensions."""
|
||||
from haiku.rag.client.processing import get_extension_from_content_type_or_url
|
||||
|
||||
# Test URL extension detection
|
||||
assert (
|
||||
client._get_extension_from_content_type_or_url(
|
||||
"https://example.com/doc.pdf", ""
|
||||
)
|
||||
== ".pdf"
|
||||
)
|
||||
assert (
|
||||
client._get_extension_from_content_type_or_url(
|
||||
"https://example.com/data.json", ""
|
||||
)
|
||||
== ".json"
|
||||
)
|
||||
# Content type mappings
|
||||
assert get_extension_from_content_type_or_url("", "text/html") == ".html"
|
||||
assert get_extension_from_content_type_or_url("", "application/pdf") == ".pdf"
|
||||
assert get_extension_from_content_type_or_url("", "text/plain") == ".txt"
|
||||
|
||||
# Test default fallback
|
||||
assert (
|
||||
client._get_extension_from_content_type_or_url("https://example.com/", "")
|
||||
== ".html"
|
||||
)
|
||||
# URL extension detection
|
||||
assert (
|
||||
get_extension_from_content_type_or_url("https://example.com/doc.pdf", "")
|
||||
== ".pdf"
|
||||
)
|
||||
assert (
|
||||
get_extension_from_content_type_or_url("https://example.com/data.json", "")
|
||||
== ".json"
|
||||
)
|
||||
|
||||
# Test content type priority over URL extension
|
||||
assert (
|
||||
client._get_extension_from_content_type_or_url(
|
||||
"https://example.com/file.txt", "application/pdf"
|
||||
)
|
||||
== ".pdf"
|
||||
# Default fallback
|
||||
assert get_extension_from_content_type_or_url("https://example.com/", "") == ".html"
|
||||
|
||||
# Content type priority over URL extension
|
||||
assert (
|
||||
get_extension_from_content_type_or_url(
|
||||
"https://example.com/file.txt", "application/pdf"
|
||||
)
|
||||
== ".pdf"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import pytest
|
||||
|
||||
from haiku.rag.client.documents import _store_document_with_chunks
|
||||
from haiku.rag.context import (
|
||||
_expand_outward,
|
||||
_find_expansion_range,
|
||||
|
|
@ -249,7 +250,8 @@ class TestExpandWithItems:
|
|||
from haiku.rag.store.models.document import Document
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as rag:
|
||||
doc = await rag._store_document_with_chunks(
|
||||
doc = await _store_document_with_chunks(
|
||||
rag,
|
||||
Document(content="test"),
|
||||
[],
|
||||
__import__(
|
||||
|
|
@ -262,6 +264,7 @@ class TestExpandWithItems:
|
|||
document_id=doc.id,
|
||||
doc_item_refs=["#/texts/999999"],
|
||||
)
|
||||
assert doc.id is not None
|
||||
expanded = await expand_with_items(
|
||||
rag.document_item_repository, doc.id, [result], 5000
|
||||
)
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ from docling_core.types.doc.document import DoclingDocument, TableData
|
|||
from docling_core.types.doc.labels import DocItemLabel
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.client.documents import _store_document_with_chunks
|
||||
from haiku.rag.config.models import AppConfig
|
||||
from haiku.rag.store.models import SearchResult
|
||||
|
||||
|
|
@ -314,7 +315,7 @@ async def test_expand_context_single_item_document(temp_db_path):
|
|||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
document = Document(content="Simple test content")
|
||||
document.set_docling(docling_doc)
|
||||
doc = await client._store_document_with_chunks(document, [], docling_doc)
|
||||
doc = await _store_document_with_chunks(client, document, [], docling_doc)
|
||||
assert doc.id is not None
|
||||
|
||||
# Create a search result with a doc_item_ref pointing to the item
|
||||
|
|
|
|||
Loading…
Reference in a new issue