From 25335fce88d4c8b9261b1b9cb99b88e4f31a6ecd Mon Sep 17 00:00:00 2001
From: Eric Curtin <ecurtin@redhat.com>
Date: Thu, 17 Oct 2024 13:34:02 +0100
Subject: [PATCH] Add hf:// as an alias to huggingface://

Also check start of string contains :// , by checking if start of
the string simply starts with things like "hf", "ollama", etc. we
are kinda polluting the namespace and not allowing models to start
with these terms.

Signed-off-by: Eric Curtin <ecurtin@redhat.com>
---
 ramalama/cli.py           | 10 +++++-----
 ramalama/huggingface.py   |  4 +++-
 test/system/050-pull.bats |  5 +++++
 3 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/ramalama/cli.py b/ramalama/cli.py
index 479eb15f..c64f4a41 100644
--- a/ramalama/cli.py
+++ b/ramalama/cli.py
@@ -193,7 +193,7 @@ def login_parser(subparsers):
 def login_cli(args):
     transport = args.TRANSPORT
     if transport != "":
-        transport = os.getenv("RAMALAMA_TRANSPORT")
+        transport = os.getenv("RAMALAMA_TRANSPORT") + "://"
     model = New(str(transport))
     return model.login(args)
 
@@ -211,7 +211,7 @@ def logout_parser(subparsers):
 def logout_cli(args):
     transport = args.TRANSPORT
     if transport != "":
-        transport = os.getenv("RAMALAMA_TRANSPORT")
+        transport = os.getenv("RAMALAMA_TRANSPORT") + "://"
     model = New(str(transport))
     return model.logout(args)
 
@@ -679,11 +679,11 @@ def dry_run(args):
 
 
 def New(model):
-    if model.startswith("huggingface"):
+    if model.startswith("huggingface://") or model.startswith("hf://"):
         return Huggingface(model)
-    if model.startswith("ollama"):
+    if model.startswith("ollama://"):
         return Ollama(model)
-    if model.startswith("oci") | model.startswith("docker"):
+    if model.startswith("oci://") or model.startswith("docker://"):
         return OCI(model)
 
     transport = os.getenv("RAMALAMA_TRANSPORT")
diff --git a/ramalama/huggingface.py b/ramalama/huggingface.py
index 1887ec90..e3eb5b2c 100644
--- a/ramalama/huggingface.py
+++ b/ramalama/huggingface.py
@@ -12,7 +12,9 @@
 
 class Huggingface(Model):
     def __init__(self, model):
-        super().__init__(model.removeprefix("huggingface://"))
+        model = model.removeprefix("huggingface://")
+        model = model.removeprefix("hf://")
+        super().__init__(model)
         self.type = "HuggingFace"
         split = self.model.rsplit("/", 1)
         self.directory = ""
diff --git a/test/system/050-pull.bats b/test/system/050-pull.bats
index 28cdb925..bcabd067 100644
--- a/test/system/050-pull.bats
+++ b/test/system/050-pull.bats
@@ -25,6 +25,11 @@ load helpers
 
 # bats test_tags=distro-integration
 @test "ramalama pull huggingface" {
+    run_ramalama pull hf://afrideva/Tiny-Vicuna-1B-GGUF/tiny-vicuna-1b.q2_k.gguf
+    run_ramalama list
+    is "$output" ".*afrideva/Tiny-Vicuna-1B-GGUF/tiny-vicuna-1b.q2_k" "image was actually pulled locally"
+    run_ramalama rm hf://afrideva/Tiny-Vicuna-1B-GGUF/tiny-vicuna-1b.q2_k.gguf
+
     run_ramalama pull huggingface://afrideva/Tiny-Vicuna-1B-GGUF/tiny-vicuna-1b.q2_k.gguf
     run_ramalama list
     is "$output" ".*afrideva/Tiny-Vicuna-1B-GGUF/tiny-vicuna-1b.q2_k" "image was actually pulled locally"