Allow uri override on create_document_from_source

evaluations/evaluations/benchmark.py

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

haiku_rag_slim/haiku/rag/client/__init__.py

@@ -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,

haiku_rag_slim/haiku/rag/client/documents.py

@@ -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,
             )

tests/test_client.py

@@ -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."""