xiaoshiai-hub 1.1.1__py3-none-any.whl → 1.1.3__py3-none-any.whl

xiaoshiai_hub/__init__.py

@@ -28,6 +28,10 @@ from .auth import (
     load_token,
     delete_token,
 )
+from .envelope_crypto import (
+    Algorithm,
+    envelope_enc_file,
+)
 
 # Upload functionality (requires GitPython)
 try:
@@ -58,6 +62,9 @@ __all__ = [
     "save_token",
     "load_token",
     "delete_token",
+    # Encryption
+    "Algorithm",
+    "envelope_enc_file",
     # Exceptions
     "HubException",
     "RepositoryNotFoundError",

xiaoshiai_hub/cli.py

@@ -1,5 +1,5 @@
 """
-小时 AI Hub SDK 命令行工具
+XiaoShi AI Hub SDK 命令行工具
 """
 
 import argparse
@@ -62,6 +62,9 @@ def cmd_upload_folder(args):
                   "请设置 MOHA_ENCRYPTION_PASSWORD 环境变量或使用 --encryption-password 参数", file=sys.stderr)
             return 1
 
+    # 获取加密算法
+    algorithm = args.algorithm if encryption_password else None
+
     try:
         upload_folder(
             folder_path=args.folder,
@@ -76,6 +79,7 @@ def cmd_upload_folder(args):
             encryption_password=encryption_password,
             ignore_patterns=ignore_patterns,
             temp_dir=args.temp_dir,
+            algorithm=algorithm,
         )
         print("上传成功！")
         return 0
@@ -102,6 +106,9 @@ def cmd_upload_file(args):
     # 确定仓库中的路径
     path_in_repo = args.path_in_repo if args.path_in_repo else os.path.basename(args.file)
 
+    # 获取加密算法
+    algorithm = args.algorithm if encryption_password else None
+
     try:
         upload_file(
             path_file=args.file,
@@ -115,6 +122,7 @@ def cmd_upload_file(args):
             password=password,
             token=token,
             encryption_password=encryption_password,
+            algorithm=algorithm,
         )
         print(f"上传成功: {path_in_repo}")
         return 0
@@ -255,6 +263,248 @@ def cmd_whoami(args):
         return 1
 
 
+def _parse_repo_id(repo_id: str):
+    """解析仓库 ID，返回 (organization, repo_name)"""
+    parts = repo_id.split("/")
+    if len(parts) != 2:
+        raise ValueError(f"无效的仓库 ID: {repo_id}，格式应为: 组织/仓库名")
+    return parts[0], parts[1]
+
+
+def _get_client(args):
+    """创建 HubClient 实例"""
+    from xiaoshiai_hub.client import HubClient
+    token, username, password = get_auth(args)
+    return HubClient(
+        base_url=args.base_url,
+        username=username,
+        password=password,
+        token=token,
+    )
+
+
+def cmd_repo_create(args):
+    """处理 repo-create 命令"""
+    try:
+        org, repo_name = _parse_repo_id(args.repo_id)
+        client = _get_client(args)
+
+        # 解析 metadata
+        metadata = {}
+        if args.license:
+            metadata["license"] = args.license
+        if args.tasks:
+            metadata["tasks"] = args.tasks
+        if args.languages:
+            metadata["languages"] = args.languages
+        if args.tags:
+            metadata["tags"] = args.tags
+        if args.frameworks:
+            metadata["frameworks"] = args.frameworks
+
+        client.create_repository(
+            organization=org,
+            repo_type=args.repo_type,
+            repo_name=repo_name,
+            description=args.description,
+            visibility=args.visibility,
+            metadata=metadata if metadata else None,
+            base_model=args.base_model,
+            relationship=args.relationship,
+        )
+        print(f"仓库创建成功: {org}/{repo_name}")
+        return 0
+    except Exception as e:
+        print(f"错误: {e}", file=sys.stderr)
+        return 1
+
+
+def cmd_repo_update(args):
+    """处理 repo-update 命令"""
+    try:
+        org, repo_name = _parse_repo_id(args.repo_id)
+        client = _get_client(args)
+
+        # 先获取仓库当前信息
+        repo = client.get_repository_info(
+            organization=org,
+            repo_type=args.repo_type,
+            repo_name=repo_name,
+        )
+
+        # 使用现有值作为默认值，只更新用户指定的字段
+        description = args.description if args.description is not None else repo.description
+        visibility = args.visibility if args.visibility is not None else repo.visibility
+
+        # 合并 metadata：用户指定的覆盖现有的
+        metadata = dict(repo.metadata) if repo.metadata else {}
+        if args.license:
+            metadata["license"] = args.license
+        if args.tasks:
+            metadata["tasks"] = args.tasks
+        if args.languages:
+            metadata["languages"] = args.languages
+        if args.tags:
+            metadata["tags"] = args.tags
+        if args.frameworks:
+            metadata["frameworks"] = args.frameworks
+
+        # 处理 genealogy：用户指定的覆盖现有的
+        base_model = args.base_model
+        relationship = args.relationship
+        if repo.genealogy:
+            if base_model is None and 'baseModel' in repo.genealogy:
+                base_model = repo.genealogy['baseModel']
+            if relationship is None and 'relationship' in repo.genealogy:
+                relationship = repo.genealogy['relationship']
+
+        client.update_repository(
+            organization=org,
+            repo_type=args.repo_type,
+            repo_name=repo_name,
+            description=description,
+            visibility=visibility,
+            metadata=metadata if metadata else None,
+            annotations=repo.annotations,
+            base_model=base_model,
+            relationship=relationship,
+        )
+        print(f"仓库更新成功: {org}/{repo_name}")
+        return 0
+    except Exception as e:
+        print(f"错误: {e}", file=sys.stderr)
+        return 1
+
+
+def cmd_repo_delete(args):
+    """处理 repo-delete 命令"""
+    try:
+        org, repo_name = _parse_repo_id(args.repo_id)
+        client = _get_client(args)
+
+        # 确认删除
+        if not args.yes:
+            confirm = input(f"确定要删除仓库 {org}/{repo_name} 吗？此操作不可恢复！[y/N]: ")
+            if confirm.lower() != 'y':
+                print("已取消删除")
+                return 0
+
+        client.delete_repository(
+            organization=org,
+            repo_type=args.repo_type,
+            repo_name=repo_name,
+        )
+        print(f"仓库删除成功: {org}/{repo_name}")
+        return 0
+    except Exception as e:
+        print(f"错误: {e}", file=sys.stderr)
+        return 1
+
+
+def cmd_repo_info(args):
+    """处理 repo-info 命令"""
+    try:
+        org, repo_name = _parse_repo_id(args.repo_id)
+        client = _get_client(args)
+
+        repo = client.get_repository_info(
+            organization=org,
+            repo_type=args.repo_type,
+            repo_name=repo_name,
+        )
+        print(f"仓库名称: {repo.name}")
+        print(f"组织: {repo.organization}")
+        print(f"所有者: {repo.owner}")
+        print(f"创建者: {repo.creator}")
+        print(f"类型: {repo.type}")
+        print(f"可见性: {repo.visibility}")
+        if repo.description:
+            print(f"描述: {repo.description}")
+        if repo.genealogy:
+            print(f"模型血缘: {repo.genealogy}")
+        if repo.metadata:
+            print(f"元数据: {repo.metadata}")
+        if repo.annotations:
+            print(f"注解: {repo.annotations}")
+        return 0
+    except Exception as e:
+        print(f"错误: {e}", file=sys.stderr)
+        return 1
+
+
+def cmd_branch_create(args):
+    """处理 branch-create 命令"""
+    try:
+        org, repo_name = _parse_repo_id(args.repo_id)
+        client = _get_client(args)
+
+        client.create_branch(
+            organization=org,
+            repo_type=args.repo_type,
+            repo_name=repo_name,
+            branch_name=args.branch,
+            from_branch=args.from_branch,
+        )
+        print(f"分支创建成功: {args.branch}")
+        return 0
+    except Exception as e:
+        print(f"错误: {e}", file=sys.stderr)
+        return 1
+
+
+def cmd_branch_delete(args):
+    """处理 branch-delete 命令"""
+    try:
+        org, repo_name = _parse_repo_id(args.repo_id)
+        client = _get_client(args)
+
+        # 确认删除
+        if not args.yes:
+            confirm = input(f"确定要删除分支 {args.branch} 吗？[y/N]: ")
+            if confirm.lower() != 'y':
+                print("已取消删除")
+                return 0
+
+        client.delete_branch(
+            organization=org,
+            repo_type=args.repo_type,
+            repo_name=repo_name,
+            branch_name=args.branch,
+        )
+        print(f"分支删除成功: {args.branch}")
+        return 0
+    except Exception as e:
+        print(f"错误: {e}", file=sys.stderr)
+        return 1
+
+
+def cmd_branch_list(args):
+    """处理 branch-list 命令"""
+    try:
+        org, repo_name = _parse_repo_id(args.repo_id)
+        client = _get_client(args)
+
+        refs = client.get_repository_refs(
+            organization=org,
+            repo_type=args.repo_type,
+            repo_name=repo_name,
+        )
+
+        if not refs:
+            print("没有找到分支")
+            return 0
+
+        print(f"仓库 {org}/{repo_name} 的分支列表：")
+        for ref in refs:
+            commit_hash = ref.hash[:8] if ref.hash else 'N/A'
+            default_mark = " (默认)" if ref.is_default else ""
+            print(f"  - {ref.name} [{ref.type}] (commit: {commit_hash}){default_mark}")
+        return 0
+    except Exception as e:
+        print(f"错误: {e}", file=sys.stderr)
+        return 1
+
+
 def _add_common_args(parser):
     """添加通用参数"""
     parser.add_argument(
@@ -286,21 +536,26 @@ def _add_encryption_args(parser):
         "--encryption-password",
         help="加密密码（或设置 MOHA_ENCRYPTION_PASSWORD 环境变量）",
     )
+    parser.add_argument(
+        "--algorithm", "-a",
+        choices=["AES", "SM4"],
+        default="AES",
+        help="加密算法（默认: AES）",
+    )
 
 
 def create_parser():
     """创建命令行参数解析器"""
     parser = argparse.ArgumentParser(
         prog="moha",
-        description="小时 AI Hub 命令行工具 - 上传和下载模型及数据集",
+        description="晓石 AI Hub 命令行工具 - 模型及数据集管理",
     )
 
     subparsers = parser.add_subparsers(dest="command", help="可用命令")
 
     # ========== 上传文件夹命令 ==========
     upload_folder_parser = subparsers.add_parser(
-        "upload-folder",
-        aliases=["upload"],
+        "upload",
         help="上传文件夹到仓库",
     )
     upload_folder_parser.add_argument("folder", help="要上传的文件夹路径")
@@ -366,7 +621,6 @@ def create_parser():
     # ========== 下载仓库命令 ==========
     download_repo_parser = subparsers.add_parser(
         "download",
-        aliases=["download-repo"],
         help="下载整个仓库",
     )
     download_repo_parser.add_argument("repo_id", help="仓库 ID（格式: 组织/仓库名）")
@@ -476,6 +730,218 @@ def create_parser():
     )
     whoami_parser.set_defaults(func=cmd_whoami)
 
+    # ========== 创建仓库命令 ==========
+    repo_create_parser = subparsers.add_parser(
+        "repo-create",
+        help="创建仓库",
+    )
+    repo_create_parser.add_argument("repo_id", help="仓库 ID（格式: 组织/仓库名）")
+    repo_create_parser.add_argument(
+        "--repo-type", "-t",
+        default="models",
+        choices=["models", "datasets"],
+        help="仓库类型（默认: models）",
+    )
+    repo_create_parser.add_argument(
+        "--description", "-d",
+        help="仓库描述",
+    )
+    repo_create_parser.add_argument(
+        "--visibility", "-v",
+        default="internal",
+        choices=["public", "internal", "private"],
+        help="可见性（默认: internal）",
+    )
+    repo_create_parser.add_argument(
+        "--license",
+        action="append",
+        help="许可证（可多次指定）",
+    )
+    repo_create_parser.add_argument(
+        "--tasks",
+        action="append",
+        help="任务类型（可多次指定）",
+    )
+    repo_create_parser.add_argument(
+        "--languages",
+        action="append",
+        help="语言（可多次指定）",
+    )
+    repo_create_parser.add_argument(
+        "--tags",
+        action="append",
+        help="标签（可多次指定）",
+    )
+    repo_create_parser.add_argument(
+        "--frameworks",
+        action="append",
+        help="框架（可多次指定）",
+    )
+    repo_create_parser.add_argument(
+        "--base-model",
+        action="append",
+        help="基础模型（可多次指定，格式: 组织/模型名）",
+    )
+    repo_create_parser.add_argument(
+        "--relationship",
+        choices=["adapter", "finetune", "quantized", "merge", "repackage"],
+        help="与基础模型的关系",
+    )
+    _add_common_args(repo_create_parser)
+    repo_create_parser.set_defaults(func=cmd_repo_create)
+
+    # ========== 更新仓库命令 ==========
+    repo_update_parser = subparsers.add_parser(
+        "repo-update",
+        help="更新仓库",
+    )
+    repo_update_parser.add_argument("repo_id", help="仓库 ID（格式: 组织/仓库名）")
+    repo_update_parser.add_argument(
+        "--repo-type", "-t",
+        default="models",
+        choices=["models", "datasets"],
+        help="仓库类型（默认: models）",
+    )
+    repo_update_parser.add_argument(
+        "--description", "-d",
+        help="仓库描述",
+    )
+    repo_update_parser.add_argument(
+        "--visibility", "-v",
+        choices=["public", "internal", "private"],
+        help="可见性",
+    )
+    repo_update_parser.add_argument(
+        "--license",
+        action="append",
+        help="许可证（可多次指定）",
+    )
+    repo_update_parser.add_argument(
+        "--tasks",
+        action="append",
+        help="任务类型（可多次指定）",
+    )
+    repo_update_parser.add_argument(
+        "--languages",
+        action="append",
+        help="语言（可多次指定）",
+    )
+    repo_update_parser.add_argument(
+        "--tags",
+        action="append",
+        help="标签（可多次指定）",
+    )
+    repo_update_parser.add_argument(
+        "--frameworks",
+        action="append",
+        help="框架（可多次指定）",
+    )
+    repo_update_parser.add_argument(
+        "--base-model",
+        action="append",
+        help="基础模型（可多次指定，格式: 组织/模型名）",
+    )
+    repo_update_parser.add_argument(
+        "--relationship",
+        choices=["adapter", "finetune", "quantized", "merge", "repackage"],
+        help="与基础模型的关系",
+    )
+    _add_common_args(repo_update_parser)
+    repo_update_parser.set_defaults(func=cmd_repo_update)
+
+    # ========== 删除仓库命令 ==========
+    repo_delete_parser = subparsers.add_parser(
+        "repo-delete",
+        help="删除仓库",
+    )
+    repo_delete_parser.add_argument("repo_id", help="仓库 ID（格式: 组织/仓库名）")
+    repo_delete_parser.add_argument(
+        "--repo-type", "-t",
+        default="models",
+        choices=["models", "datasets"],
+        help="仓库类型（默认: models）",
+    )
+    repo_delete_parser.add_argument(
+        "--yes", "-y",
+        action="store_true",
+        help="跳过确认提示",
+    )
+    _add_common_args(repo_delete_parser)
+    repo_delete_parser.set_defaults(func=cmd_repo_delete)
+
+    # ========== 查看仓库信息命令 ==========
+    repo_info_parser = subparsers.add_parser(
+        "repo-info",
+        help="查看仓库信息",
+    )
+    repo_info_parser.add_argument("repo_id", help="仓库 ID（格式: 组织/仓库名）")
+    repo_info_parser.add_argument(
+        "--repo-type", "-t",
+        default="models",
+        choices=["models", "datasets"],
+        help="仓库类型（默认: models）",
+    )
+    _add_common_args(repo_info_parser)
+    repo_info_parser.set_defaults(func=cmd_repo_info)
+
+    # ========== 创建分支命令 ==========
+    branch_create_parser = subparsers.add_parser(
+        "branch-create",
+        help="创建分支",
+    )
+    branch_create_parser.add_argument("repo_id", help="仓库 ID（格式: 组织/仓库名）")
+    branch_create_parser.add_argument("branch", help="要创建的分支名称")
+    branch_create_parser.add_argument(
+        "--from", "-f",
+        dest="from_branch",
+        default="main",
+        help="基于哪个分支创建（默认: main）",
+    )
+    branch_create_parser.add_argument(
+        "--repo-type", "-t",
+        default="models",
+        choices=["models", "datasets"],
+        help="仓库类型（默认: models）",
+    )
+    _add_common_args(branch_create_parser)
+    branch_create_parser.set_defaults(func=cmd_branch_create)
+
+    # ========== 删除分支命令 ==========
+    branch_delete_parser = subparsers.add_parser(
+        "branch-delete",
+        help="删除分支",
+    )
+    branch_delete_parser.add_argument("repo_id", help="仓库 ID（格式: 组织/仓库名）")
+    branch_delete_parser.add_argument("branch", help="要删除的分支名称")
+    branch_delete_parser.add_argument(
+        "--repo-type", "-t",
+        default="models",
+        choices=["models", "datasets"],
+        help="仓库类型（默认: models）",
+    )
+    branch_delete_parser.add_argument(
+        "--yes", "-y",
+        action="store_true",
+        help="跳过确认提示",
+    )
+    _add_common_args(branch_delete_parser)
+    branch_delete_parser.set_defaults(func=cmd_branch_delete)
+
+    # ========== 列出分支命令 ==========
+    branch_list_parser = subparsers.add_parser(
+        "branch-list",
+        help="列出仓库的所有分支",
+    )
+    branch_list_parser.add_argument("repo_id", help="仓库 ID（格式: 组织/仓库名）")
+    branch_list_parser.add_argument(
+        "--repo-type", "-t",
+        default="models",
+        choices=["models", "datasets"],
+        help="仓库类型（默认: models）",
+    )
+    _add_common_args(branch_list_parser)
+    branch_list_parser.set_defaults(func=cmd_branch_list)
+
     return parser