Add support for code and MCP repository type (#118)

- Introduced MCP repository type
- Updated API client to handle MCP
- Enhanced logging for API operations
- Modified download/upload commands for MCP
- Updated constants and utility functions for MCP

Co-authored-by: Haihui.Wang <wanghh2000@163.com>

Author: HaiHui886

launch.py

@@ -1,6 +1,15 @@
 import re
 import sys
+import logging
 from pycsghub.cli import app
+
 if __name__ == '__main__':
+    log_level = "INFO"
+    logging.basicConfig(
+        level=getattr(logging, log_level),
+        format='%(asctime)s - %(levelname)s - %(message)s',
+        datefmt='%Y-%m-%d %H:%M:%S',
+        handlers=[logging.StreamHandler()]
+    )
     sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0])
-    sys.exit(app())
+    sys.exit(app())

pycsghub/api_client/__init__.py

@@ -1,15 +1,28 @@
 from typing import Optional
-
+import logging
+from pycsghub.constants import REPO_TYPE_CODE, REPO_TYPE_MCPSERVER
 from pycsghub.utils import disable_xnet, get_endpoint, get_token_to_send
 from .api_client import CsghubApi
 from .api_client_interface import HubApi
 from .api_client_xet import CsgXnetApi
+from pycsghub.cmd.repo_types import RepoType
+
+logger = logging.getLogger(__name__)
 
-def get_csghub_api(token: Optional[str] = None,
+def get_csghub_api(repo_type: Optional[RepoType] = None,
+                   token: Optional[str] = None,
                    endpoint: Optional[str] = None,
                    user_name: Optional[str] = None) -> HubApi:
     token = get_token_to_send(token)
     endpoint = get_endpoint(endpoint=endpoint)
-    return CsghubApi(token=token, endpoint=endpoint, user_name=user_name) if disable_xnet() else CsgXnetApi(token=token,
-                                                                                                            endpoint=endpoint,
-                                                                                                            user_name=user_name)
+
+    if repo_type in [REPO_TYPE_CODE, REPO_TYPE_MCPSERVER]:
+        logger.debug(f"Use CsghubApi for repo_type {repo_type}")
+        return CsghubApi(token=token, endpoint=endpoint, user_name=user_name)
+
+    if disable_xnet():
+        logger.debug("xnet disabled")
+        return CsghubApi(token=token, endpoint=endpoint, user_name=user_name)  
+    else:
+        logger.debug("xnet enabled")
+        return CsgXnetApi(token=token, endpoint=endpoint, user_name=user_name)

pycsghub/api_client/api_client_xet.py

@@ -1,10 +1,13 @@
 import os
+import logging
 from typing import Optional
 
 from huggingface_hub import constants, HfApi
 
 from pycsghub.utils import get_default_cache_dir, get_xnet_endpoint, get_cache_dir
 
+logger = logging.getLogger(__name__)
+
 class CsgXnetApi(HfApi):
     def __init__(self, token: Optional[str] = None, endpoint: Optional[str] = None, user_name: Optional[str] = None):
         endpoint = get_xnet_endpoint(endpoint=endpoint)
@@ -102,5 +105,6 @@ def snapshot_download(self, *args, **kwargs):
         # Ensure token is passed
         if 'cache_dir' not in kwargs:
             kwargs['cache_dir'] = self._cache_dir
-            
+
+        logging.debug(f"invoke origin_snapshot_download")
         return origin_snapshot_download(*args, **kwargs)

pycsghub/cli.py

@@ -21,7 +21,7 @@
 from pycsghub.cmd.repo_types import RepoType
 from pycsghub.constants import DEFAULT_CSGHUB_DOMAIN, DEFAULT_REVISION, REPO_SOURCE_CSG
 from pycsghub.api_client import get_csghub_api
-from .utils import print_download_result
+from .utils import print_download_result, disable_xnet
 from .lfs import LfsEnableCommand, LfsUploadCommand
 from .upload_large_folder.main import upload_large_folder_internal
 
@@ -89,7 +89,7 @@ def version_callback(value: bool):
     "path"              : typer.Argument(help="Local path to repository you want to configure."),
 }
 
-@app.command(name="download", help="Download model/dataset/space from OpenCSG Hub", no_args_is_help=True)
+@app.command(name="download", help="Download model/dataset/space/code/mcp from OpenCSG Hub", no_args_is_help=True)
 def download(
     repo_id: Annotated[str, OPTIONS["repoID"]],
     filenames: Annotated[
@@ -112,7 +112,7 @@ def download(
     force_download: Annotated[Optional[bool], OPTIONS["force_download"]] = False,
     max_workers: Annotated[Optional[int], OPTIONS["max_workers"]] = 8,
 ):
-    api = get_csghub_api(token=token, endpoint=endpoint)
+    api = get_csghub_api(repo_type=repo_type, token=token, endpoint=endpoint)
 
     # Handle single/multiple file args similar to HF
     filenames_list = filenames if filenames is not None else []
@@ -186,7 +186,8 @@ def download(
     except Exception as e:
         if quiet:
             raise e
-        print(f"Download failed: {e}")
+        
+        logger.error(f"Download with xnet-disabled {disable_xnet()} failed: {e} ")
         raise typer.Exit(code=1)
 
     if quiet:
@@ -196,7 +197,7 @@ def download(
     else:
         print_download_result(result)
 
-@app.command(name="upload", help="Upload repository files to OpenCSG Hub", no_args_is_help=True)
+@app.command(name="upload", help="Upload model/dataset/space/code/mcp files to OpenCSG Hub", no_args_is_help=True)
 def upload(
     repo_id: Annotated[str, OPTIONS["repoID"]],
     local_path: Annotated[str, OPTIONS["localPath"]],
@@ -216,8 +217,9 @@ def upload(
     every: Annotated[Optional[float], OPTIONS["every"]] = None,
     quiet: Annotated[bool, OPTIONS["quiet"]] = False,
 ):
+    
     repo_type_str = repo_type.value
-    api = get_csghub_api(token=token, endpoint=endpoint, user_name=user_name)
+    api = get_csghub_api(repo_type=repo_type, token=token, endpoint=endpoint, user_name=user_name)
 
     resolved_local_path = local_path
     resolved_path_in_repo = path_in_repo
@@ -236,7 +238,7 @@ def run_upload() -> dict:
             # Placeholder for Scheduler
             # Scheduler not implemented yet in csghub-sdk
             # Similar to HF CommitScheduler logic
-            print(f"Scheduling commits every {every} minutes... (Not fully implemented)")
+            logger.info(f"Scheduling commits every {every} minutes... (Not fully implemented)")
             try:
                 while True:
                     time.sleep(100)
@@ -246,6 +248,7 @@ def run_upload() -> dict:
         if not os.path.isfile(resolved_local_path) and not os.path.isdir(resolved_local_path):
             raise FileNotFoundError(f"No such file or directory: '{resolved_local_path}'.")
 
+        logger.debug(f"Uploading {repo_type_str} {repo_id} from {resolved_local_path} to {resolved_path_in_repo}")
         # Create repo if needed
         api.create_repo(
             repo_id=repo_id,
@@ -259,7 +262,7 @@ def run_upload() -> dict:
             try:
                 api.repo_info(repo_id=repo_id, repo_type=repo_type_str, revision=revision)
             except Exception:  # RevisionNotFoundError
-                logger.info(f"Branch '{revision}' not found. Creating it...")
+                logger.warning(f"Branch '{revision}' not found. Creating it...")
                 api.create_branch(repo_id=repo_id, repo_type=repo_type_str, branch=revision, exist_ok=True)
 
         if os.path.isfile(resolved_local_path):
@@ -292,10 +295,12 @@ def run_upload() -> dict:
         disable_progress_bars()
         with warnings.catch_warnings():
             warnings.simplefilter("ignore")
-            print(run_upload())
+            result = run_upload()
+            logger.debug(f"Upload result {result}")
         enable_progress_bars()
     else:
-        print(run_upload())
+        result = run_upload()
+        logger.debug(f"Upload result {result}")
 
 @app.command(name="upload-large-folder", help="Upload large folder to OpenCSG Hub using multiple workers",
              no_args_is_help=True)

pycsghub/cmd/repo_types.py

@@ -1,8 +1,9 @@
 from enum import Enum
-from pycsghub.constants import REPO_TYPE_MODEL, REPO_TYPE_DATASET, REPO_TYPE_SPACE, REPO_TYPE_CODE
+from pycsghub.constants import REPO_TYPE_MODEL, REPO_TYPE_DATASET, REPO_TYPE_SPACE, REPO_TYPE_CODE, REPO_TYPE_MCPSERVER
 
 class RepoType(str, Enum):
     MODEL = REPO_TYPE_MODEL
     DATASET = REPO_TYPE_DATASET
     SPACE = REPO_TYPE_SPACE
     CODE = REPO_TYPE_CODE
+    MCP = REPO_TYPE_MCPSERVER

pycsghub/constants.py

@@ -7,7 +7,8 @@
 REPO_TYPE_MODEL = "model"
 REPO_TYPE_SPACE = "space"
 REPO_TYPE_CODE = "code"
-REPO_TYPES = [None, REPO_TYPE_MODEL, REPO_TYPE_DATASET, REPO_TYPE_SPACE]
+REPO_TYPE_MCPSERVER = "mcp"
+REPO_TYPES = [None, REPO_TYPE_MODEL, REPO_TYPE_DATASET, REPO_TYPE_SPACE, REPO_TYPE_CODE, REPO_TYPE_MCPSERVER]
 
 REPO_SOURCE_CSG = "csg"
 REPO_SOURCE_HF = "hf"

pycsghub/file_download.py

@@ -122,7 +122,7 @@ def file_download(
             if force_download or not cache.exists(repo_file_info):
                 # get download url
                 url = get_file_download_url(
-                    model_id=repo_id,
+                    repo_id=repo_id,
                     file_path=file_name,
                     revision=revision,
                     endpoint=download_endpoint,

pycsghub/file_upload.py

@@ -2,7 +2,7 @@
 import requests
 from typing import Optional
 from pycsghub.constants import (DEFAULT_REVISION)
-from pycsghub.utils import (build_csg_headers, get_endpoint)
+from pycsghub.utils import (build_csg_headers, get_endpoint, get_repo_url_prefix)
 import logging
 
 logger = logging.getLogger(__name__)
@@ -23,7 +23,7 @@ def http_upload_file(
     http_endpoint = endpoint if endpoint is not None else get_endpoint()
     if not http_endpoint.endswith("/"):
         http_endpoint += "/"
-    http_url = http_endpoint + "api/v1/" + repo_type + "s/" + repo_id + "/upload_file"
+    http_url = http_endpoint + "api/v1/" + get_repo_url_prefix(repo_type=repo_type) + "/" + repo_id + "/upload_file"
     post_headers = build_csg_headers(token=token)
     file_data = {'file': open(file_path, 'rb')}
     form_data = {'file_path': destination_path, 'branch': revision, 'message': (commit_message or ('upload ' + os.path.basename(file_path)))}
@@ -34,7 +34,7 @@ def http_upload_file(
     elif response.status_code != 200 and exist_msg in response.text:
         logger.info(f"file '{file_path}' already exists.")
     else:
-        msg = f"fail to upload {file_path} with response code: {response.status_code}, error: {response.content.decode()}"
+        msg = f"failed to upload {file_path} to branch {revision} on {http_url} with response code: {response.status_code}, error: {response.content.decode()}"
         logger.error(msg)
         raise RuntimeError(msg)
 

pycsghub/repository.py

@@ -17,7 +17,9 @@
 from pycsghub.constants import (GIT_HIDDEN_DIR, GIT_ATTRIBUTES_FILE)
 from pycsghub.utils import (build_csg_headers,
                             model_id_to_group_owner_name,
-                            get_endpoint)
+                            get_endpoint,
+                            get_repo_url_prefix,
+                            get_repo_git_prefix)
 import logging
 
 logger = logging.getLogger(__name__)
@@ -88,14 +90,7 @@ def __init__(
         self.delete_patterns = delete_patterns
 
     def get_url_prefix(self):
-        if self.repo_type == REPO_TYPE_DATASET:
-            return "datasets"
-        if self.repo_type == REPO_TYPE_SPACE:
-            return "spaces"
-        if self.repo_type == REPO_TYPE_CODE:
-            return "codes"
-        else:
-            return "models"
+        return get_repo_url_prefix(repo_type=self.repo_type)
 
     def upload(self) -> None:
         if not os.path.exists(self.upload_path):
@@ -267,17 +262,18 @@ def create_new_repo(self):
             "Content-Type": "application/json"
         })
         response = requests.post(url, json=data, headers=headers)
-        exist_msg = "SYS-ERR-4: Duplicate entry for key"
+        exist_msg = "duplicate key value violates unique constraint"
         if response.status_code != 200 and exist_msg not in response.text :
-            logger.info(f"create repo on {url} response: {response.text}")
+            logger.error(f"create repo on {url} response: {response.text}")
         response.raise_for_status()
         return response
 
     def generate_repo_clone_url(self) -> str:
         clone_endpoint = get_endpoint(endpoint=self.endpoint, operation=OPERATION_ACTION_GIT)
-        clone_url = f"{clone_endpoint}/{self.repo_url_prefix}/{self.repo_id}.git"
+        clone_url = f"{clone_endpoint}/{get_repo_git_prefix(repo_type=self.repo_type)}/{self.repo_id}.git"
         scheme = urlparse(clone_url).scheme
         clone_url = clone_url.replace(f"{scheme}://", f"{scheme}://{self.user_name}:{self.token}@")
+        logger.debug(f"generate clone url: {clone_url} on branch {self.branch_name}")
         return clone_url
 
     def git_clone(

pycsghub/snapshot_download.py

@@ -124,7 +124,7 @@ def _download_one(repo_file: str):
                     logger.info(f"File {file_name} already in '{cache.get_root_location()}', skip downloading!")
                     return
                 url = get_file_download_url(
-                        model_id=repo_id,
+                        repo_id=repo_id,
                         file_path=repo_file,
                         repo_type=repo_type,
                         revision=revision,