diff --git a/src/rosetta-cli/rosetta_cli/cli.py b/src/rosetta-cli/rosetta_cli/cli.py index 5dbb6c5e..16e8a0cf 100644 --- a/src/rosetta-cli/rosetta_cli/cli.py +++ b/src/rosetta-cli/rosetta_cli/cli.py @@ -378,7 +378,8 @@ def main() -> int: chunk_method=config.chunk_method, parser_config=config.parser_config, page_size=config.page_size, - timeout=config.timeout + dataset_default=config.dataset_default, + timeout=config.timeout, ) return execute_command(args.command, args, client, config) diff --git a/src/rosetta-cli/rosetta_cli/ragflow_client.py b/src/rosetta-cli/rosetta_cli/ragflow_client.py index 9286211a..7efc7ab9 100644 --- a/src/rosetta-cli/rosetta_cli/ragflow_client.py +++ b/src/rosetta-cli/rosetta_cli/ragflow_client.py @@ -116,7 +116,8 @@ def __init__( embedding_model: str | None = None, chunk_method: str = "naive", parser_config: JsonDict | None = None, - page_size: int = 1000 + page_size: int = 1000, + dataset_default: str = "aia", ): """ Initialize RAGFlow client. @@ -130,6 +131,7 @@ def __init__( chunk_method: Chunking method (default: naive) parser_config: Parser configuration dict for chunk_method settings page_size: Default page size for list operations (default: 1000) + dataset_default: Default dataset name when no release-specific or explicit dataset name is available Raises: ValueError: If api_key or base_url is empty @@ -147,6 +149,7 @@ def __init__( self.chunk_method = chunk_method self.parser_config = parser_config or {} self.page_size = page_size + self.dataset_default = dataset_default # Initialize RAGFlow SDK client self._client = RAGFlow(api_key=api_key, base_url=base_url, version=version) @@ -469,26 +472,38 @@ def _ensure_dataset(self, name: str, description: str = "", dry_run: bool = Fals # Dataset doesn't exist, create it (gated by dry_run) return self.create_dataset(name, description, dry_run=dry_run) - def _resolve_dataset_name(self, template: str, release: str | None) -> str: + def _resolve_dataset_name( + self, + dataset_template: str, + release: str | None, + dataset_name: str | None, + ) -> str: """ - Resolve dataset name from template. + Resolve dataset name using release + dataset_template, explicit dataset_name or default. Args: - template: Name template (e.g., "aia-{release}") + dataset_template: Name template (e.g., "aia-{release}") release: Release identifier (e.g., "r1") + dataset_name: Explicit dataset name, if provided (e.g., "aia-r1") Returns: Resolved dataset name Examples: - >>> _resolve_dataset_name("aia-{release}", "r1") + >>> _resolve_dataset_name("aia-{release}", "r1", "aia-r2") "aia-r1" - >>> _resolve_dataset_name("aia", None) - "aia" + >>> _resolve_dataset_name("aia-{release}", None, "aia-r2") + "aia-r2" """ - if release and "{release}" in template: - return template.format(release=release) - return template + if release and "{release}" in dataset_template: + resolved_name = dataset_template.format(release=release) + else: + resolved_name = dataset_name or self.dataset_default + + if "{release}" in resolved_name: + raise ValueError(f"Unresolved release placeholder in dataset name: {resolved_name}") + + return resolved_name def _build_title_with_tags(self, tags: list[str], filename: str) -> str: """ @@ -521,7 +536,7 @@ def upload_document( OPTIMIZED: Now accepts pre-read content to avoid redundant file I/O. This method: - 1. Resolves dataset name from template + release + 1. Resolves the target dataset name 2. Ensures dataset exists 3. Builds title with tag prefixes 4. Checks if document exists (by ims_doc_id) @@ -529,6 +544,12 @@ def upload_document( 6. Deletes existing document if changed 7. Uploads new document with metadata + Dataset name resolution precedence: + 1. If metadata.release is present, and dataset_template contains "{release}", + resolve dataset_template with the release. + 2. Otherwise, use dataset_name when provided. + 3. Otherwise, fall back to dataset_default. + Args: file_path: Path to file (for filename, backward compatibility) metadata: Document metadata with pre-calculated hash @@ -567,10 +588,7 @@ def upload_document( actual_hash = metadata.content_hash # Resolve dataset name - resolved_name = self._resolve_dataset_name( - dataset_template if "{release}" in dataset_template else (dataset_name or dataset_template), - metadata.release - ) + resolved_name = self._resolve_dataset_name(dataset_template, metadata.release, dataset_name) # Ensure dataset exists (dry_run gates the underlying create_dataset call) dataset = self._ensure_dataset( diff --git a/src/rosetta-cli/tests/test_ragflow_client_dataset_name_resolution.py b/src/rosetta-cli/tests/test_ragflow_client_dataset_name_resolution.py new file mode 100644 index 00000000..c1319418 --- /dev/null +++ b/src/rosetta-cli/tests/test_ragflow_client_dataset_name_resolution.py @@ -0,0 +1,93 @@ +from unittest.mock import Mock + +import pytest + +from rosetta_cli.ragflow_client import DocumentMetadata, RAGFlowClient + + +def _make_client() -> RAGFlowClient: + client = object.__new__(RAGFlowClient) + client._client = None + client.dataset_default = "aia" + client._ensure_dataset = Mock(return_value=None) + return client + + +def _make_metadata(**overrides) -> DocumentMetadata: + base = dict( + tags=["a", "b"], + domain="d", + release="r2", + content_hash="h" * 32, + ims_doc_id="ims-1", + ) + base.update(overrides) + return DocumentMetadata(**base) + + +def test_upload_document_uses_default_dataset_name_when_release_is_empty(): + client = _make_client() + + client.upload_document( + file_path=None, + metadata=_make_metadata(release=""), + dataset_name="aia", + dataset_template="aia-{release}", + content=b"hello", + dry_run=True, + ) + + resolved_name = client._ensure_dataset.call_args.args[0] + + assert resolved_name == "aia" + + +def test_upload_document_resolves_dataset_name_when_release_is_present(): + client = _make_client() + + client.upload_document( + file_path=None, + metadata=_make_metadata(), + dataset_name="aia", + dataset_template="aia-{release}", + content=b"hello", + dry_run=True, + ) + + resolved_name = client._ensure_dataset.call_args.args[0] + + assert resolved_name == "aia-r2" + + +def test_upload_document_uses_dataset_default_when_dataset_name_is_none(): + client = _make_client() + + client.upload_document( + file_path=None, + metadata=_make_metadata(release=""), + dataset_name=None, + dataset_template="aia-{release}", + content=b"hello", + dry_run=True, + ) + + resolved_name = client._ensure_dataset.call_args.args[0] + + assert resolved_name == "aia" + + +def test_upload_document_rejects_unresolved_release_placeholder(): + client = _make_client() + + with pytest.raises( + ValueError, + match="Unresolved release placeholder", + ): + client.upload_document( + file_path=None, + metadata=_make_metadata(release=""), + dataset_name="custom-{release}", + dataset_template="aia-{release}", + content=b"hello", + dry_run=True, + ) diff --git a/src/rosetta-cli/tests/test_ragflow_client_meta_fields_v25.py b/src/rosetta-cli/tests/test_ragflow_client_meta_fields_v25.py index 71b4ba26..06ff61ac 100644 --- a/src/rosetta-cli/tests/test_ragflow_client_meta_fields_v25.py +++ b/src/rosetta-cli/tests/test_ragflow_client_meta_fields_v25.py @@ -49,6 +49,7 @@ def _make_client(dataset: _FakeDataset, list_documents_returns=()): client = object.__new__(RAGFlowClient) client._client = None client.page_size = 1000 + client.dataset_default = "aia" client._doc_index_by_dataset = {} client._ensure_dataset = lambda *_a, **_k: dataset client.list_documents = lambda *_a, **_k: list(list_documents_returns) diff --git a/src/rosetta-cli/tests/test_ragflow_client_upload_exception_handling.py b/src/rosetta-cli/tests/test_ragflow_client_upload_exception_handling.py index 4f4ceede..2b8c9fc9 100644 --- a/src/rosetta-cli/tests/test_ragflow_client_upload_exception_handling.py +++ b/src/rosetta-cli/tests/test_ragflow_client_upload_exception_handling.py @@ -29,6 +29,7 @@ def _make_client(): client = object.__new__(RAGFlowClient) client._client = None client.page_size = 1000 + client.dataset_default = "aia" client._doc_index_by_dataset = {} return client