Merge pull request #361 from ggozad/fix/source-uri-override

Allow uri override on create_document_from_source
This commit is contained in:
Yiorgis Gozadinos 2026-04-30 17:35:38 +03:00 committed by GitHub
commit 0e9efe448d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 195 additions and 37 deletions

View file

@ -123,6 +123,20 @@ Numbers measured under the current pinned judge (`ollama:qwen3.6`) on a recent `
*Measured on haiku.rag v0.43.1, judged by `ollama:qwen3.6` (current default), on 199 of 200 completed cases.* 28 % of cases produce a perfect citation (`cited_map` = 1.0).
### OpenRAG Bench (ORB)
[OpenRAG Bench](https://huggingface.co/datasets/vectara/open_ragbench) contains ArXiv research papers with multimodal question-answering pairs. Queries include both text-based and image-based questions, testing retrieval over visual content like figures, charts, and diagrams. Each query maps to one relevant document.
**Multimodal processing**: Picture descriptions are generated using a Vision Language Model (VLM) during document conversion, making embedded images searchable via text queries. See [Picture Description configuration](configuration/processing.md#picture-description-vlm).
#### Skill QA + citation retrieval
| Skill model | QA accuracy | Mean `cited_map` | VLM |
|------------------|-------------|------------------|----------------------|
| `ollama:gpt-oss` | 0.94 | 0.86 | Ollama / ministral-3 |
*Measured on haiku.rag v0.44.0, judged by `ollama:qwen3.6` (current default), on 2992 of 3044 completed cases.*
## Past results
These were measured under the prior pinned judge (`ollama:gpt-oss`). The pinned default has since switched to `ollama:qwen3.6` (see [Methodology — QA Accuracy](#qa-accuracy)) — under the new judge the QA accuracy numbers below typically shift up by ~510 pp.
@ -198,25 +212,3 @@ Note the significant degradation when very small models are used such as `qwen3:
| `qwen3-embedding:4b` | `gpt-oss:20b` - thinking | 0.86 |
*Measured on haiku.rag v0.20.2, judged by `ollama:gpt-oss`.*
### OpenRAG Bench (ORB)
[OpenRAG Bench](https://huggingface.co/datasets/vectara/open_ragbench) contains ArXiv research papers with multimodal question-answering pairs. Queries include both text-based and image-based questions, testing retrieval over visual content like figures, charts, and diagrams. We use MAP for retrieval evaluation since each query maps to one relevant document.
**Multimodal processing**: Picture descriptions are generated using a Vision Language Model (VLM) during document conversion, making embedded images searchable via text queries. See [Picture Description configuration](configuration/processing.md#picture-description-vlm).
#### Retrieval (MAP)
| Embedding Model | MAP | VLM |
|----------------------|--------|----------------------|
| `qwen3-embedding:4b` | 0.9626 | Ollama / ministral-3 |
*Measured on haiku.rag v0.26.8.*
#### QA Accuracy
| Embedding Model | QA Model | Accuracy | VLM |
|----------------------|-----------------------------|----------|----------------------|
| `qwen3-embedding:4b` | `gpt-oss:20b` - no thinking | 0.912 | Ollama / ministral-3 |
*Measured on haiku.rag v0.26.8, judged by `ollama:gpt-oss`.*

View file

@ -133,13 +133,12 @@ async def populate_db(
progress.advance(task)
continue
# Use the actual URI that will be stored in the database
if payload.source_path is not None:
lookup_uri = payload.source_path.absolute().as_uri()
else:
lookup_uri = payload.uri
existing = await rag.get_document_by_uri(lookup_uri)
# `payload.uri` is the canonical document identifier and is now
# honored by both `create_document` and (via the `uri=` override)
# `create_document_from_source`, so it's also the right key to
# look up an existing document, regardless of whether the source
# is a file path or inline content.
existing = await rag.get_document_by_uri(payload.uri)
if existing is not None:
assert existing.id
chunks = await rag.chunk_repository.get_by_document_id(existing.id)
@ -153,6 +152,7 @@ async def populate_db(
source=payload.source_path,
title=payload.title,
metadata=payload.metadata,
uri=payload.uri,
)
else:
assert payload.content is not None

View file

@ -195,10 +195,11 @@ class HaikuRAG:
source: str | Path,
title: str | None = None,
metadata: dict | None = None,
uri: str | None = None,
) -> Document | list[Document]:
from haiku.rag.client.documents import create_document_from_source
return await create_document_from_source(self, source, title, metadata)
return await create_document_from_source(self, source, title, metadata, uri=uri)
async def update_document(
self,

View file

@ -184,6 +184,7 @@ async def create_document_from_source(
source: str | Path,
title: str | None = None,
metadata: dict | None = None,
uri: str | None = None,
) -> Document | list[Document]:
"""Create or update document(s) from a file path, directory, or URL.
@ -192,6 +193,12 @@ async def create_document_from_source(
- If MD5 changed, updates the document
- If no document exists, creates a new one
If ``uri`` is provided, it overrides the URI auto-derived from the source
(which is normally ``file://`` for local files or the URL for remote
sources). This is useful when callers want to persist documents under a
logical identifier (e.g. an ArXiv ID) rather than the on-disk path. Not
supported for directory sources, which produce one document per file.
Returns a single Document for files/URLs, a list for directories.
"""
metadata = metadata or {}
@ -200,7 +207,7 @@ async def create_document_from_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
client, source_str, title=title, metadata=metadata, uri=uri
)
elif parsed_url.scheme == "file":
source_path = Path(parsed_url.path)
@ -208,6 +215,11 @@ async def create_document_from_source(
source_path = Path(source) if isinstance(source, str) else source
if source_path.is_dir():
if uri is not None:
raise ValueError(
"uri override is not supported for directory sources; each file "
"produces its own document with its own auto-derived URI."
)
from haiku.rag.monitor import FileFilter
documents = []
@ -224,7 +236,7 @@ async def create_document_from_source(
return documents
return await _create_document_from_file(
client, source_path, title=title, metadata=metadata
client, source_path, title=title, metadata=metadata, uri=uri
)
@ -233,8 +245,13 @@ async def _create_document_from_file(
source_path: Path,
title: str | None = None,
metadata: dict | None = None,
uri: str | None = None,
) -> Document:
"""Create or update a document from a single file path."""
"""Create or update a document from a single file path.
``uri`` overrides the auto-derived ``file://`` URI; it's used as the
canonical document identifier for lookup and storage.
"""
from haiku.rag.embeddings import embed_chunks
metadata = metadata or {}
@ -246,7 +263,8 @@ async def _create_document_from_file(
if not source_path.exists():
raise ValueError(f"File does not exist: {source_path}")
uri = source_path.absolute().as_uri()
if uri is None:
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))
@ -315,12 +333,17 @@ async def _create_or_update_document_from_url(
url: str,
title: str | None = None,
metadata: dict | None = None,
uri: str | None = None,
) -> Document:
"""Create or update a document from a URL by downloading and parsing the content."""
"""Create or update a document from a URL by downloading and parsing the content.
``uri`` overrides the URL as the stored document identifier.
"""
from haiku.rag.client.processing import get_extension_from_content_type_or_url
from haiku.rag.embeddings import embed_chunks
metadata = metadata or {}
stored_uri = uri if uri is not None else url
converter = get_converter(client._config)
supported_extensions = converter.supported_extensions
@ -334,7 +357,7 @@ async def _create_or_update_document_from_url(
content_type = response.headers.get("content-type", "").lower()
# Check if document already exists
existing_doc = await client.get_document_by_uri(url)
existing_doc = await client.get_document_by_uri(stored_uri)
if existing_doc and existing_doc.metadata.get("md5") == md5_hash:
updated = False
if title is not None and title != existing_doc.title:
@ -396,7 +419,7 @@ async def _create_or_update_document_from_url(
)
document = Document(
content=stored_content,
uri=url,
uri=stored_uri,
title=title,
metadata=metadata,
)

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

View file

@ -263,6 +263,64 @@ async def test_client_update_title_noop_behavior(temp_db_path):
assert got.title == "Title B"
@pytest.mark.vcr()
async def test_client_create_document_from_source_with_uri_override(temp_db_path):
"""A `uri` override is honored as the canonical document identifier."""
async with HaikuRAG(temp_db_path, create=True) as client:
with tempfile.TemporaryDirectory() as temp_dir:
temp_path = Path(temp_dir) / "2412.06611v2.pdf-ish.txt"
temp_path.write_text("Synthetic content for URI-override test.")
doc = await client.create_document_from_source(
source=temp_path, uri="2412.06611v2"
)
assert isinstance(doc, Document)
assert doc.uri == "2412.06611v2"
assert doc.uri != temp_path.as_uri()
# The override URI is the lookup key for subsequent reads.
looked_up = await client.get_document_by_uri("2412.06611v2")
assert looked_up is not None
assert looked_up.id == doc.id
# The original file URI is NOT a key.
assert await client.get_document_by_uri(temp_path.as_uri()) is None
@pytest.mark.vcr()
async def test_client_create_document_from_source_uri_override_dedupes(temp_db_path):
"""Re-creating from the same source with the same override is a no-op."""
async with HaikuRAG(temp_db_path, create=True) as client:
with tempfile.TemporaryDirectory() as temp_dir:
temp_path = Path(temp_dir) / "doc.txt"
temp_path.write_text("Stable content for dedup test.")
doc1 = await client.create_document_from_source(
source=temp_path, uri="paper-id-1"
)
doc2 = await client.create_document_from_source(
source=temp_path, uri="paper-id-1"
)
assert isinstance(doc1, Document) and isinstance(doc2, Document)
assert doc1.id == doc2.id
@pytest.mark.vcr()
async def test_client_create_document_from_source_uri_override_rejected_for_dir(
temp_db_path,
):
"""Directory sources reject the `uri` override (would collide on multiple files)."""
async with HaikuRAG(temp_db_path, create=True) as client:
with tempfile.TemporaryDirectory() as temp_dir:
(Path(temp_dir) / "a.txt").write_text("a")
(Path(temp_dir) / "b.txt").write_text("b")
with pytest.raises(ValueError, match="directory sources"):
await client.create_document_from_source(
source=Path(temp_dir), uri="some-uri"
)
@pytest.mark.vcr()
async def test_client_create_document_from_source_unsupported(temp_db_path):
"""Test creating a document from an unsupported file type."""