|
| 1 | +import pytest |
| 2 | +import requests |
| 3 | +import os |
| 4 | +import yaml |
| 5 | + |
| 6 | +from pathlib import Path |
| 7 | +from test_runner import ( |
| 8 | + run, |
| 9 | + start_server, |
| 10 | + stop_server, |
| 11 | + wait_for_websocket_download_success_event, |
| 12 | +) |
| 13 | + |
| 14 | +collection_id = "cortexso/local-models-6683a6e29e8f3018845b16db" |
| 15 | +token = os.getenv("HF_TOKEN") |
| 16 | +if not token: |
| 17 | + raise ValueError("HF_TOKEN environment variable not set") |
| 18 | + |
| 19 | +def get_repos_in_collection(collection_id, token): |
| 20 | + # API endpoint to get list of repos in the collection |
| 21 | + url = f"https://huggingface.co/api/collections/{collection_id}" |
| 22 | + headers = {"Authorization": f"Bearer {token}"} |
| 23 | + response = requests.get(url, headers=headers) |
| 24 | + |
| 25 | + # Check response and retrieve repo IDs if successful |
| 26 | + if response.status_code == 200: |
| 27 | + return [repo['id'] for repo in response.json()["items"]] |
| 28 | + else: |
| 29 | + print("Error fetching repos:", response.status_code, response.json()) |
| 30 | + return [] |
| 31 | + |
| 32 | +def get_repo_default_branch(repo_id, token): |
| 33 | + # Direct link to metadata.yaml on the main branch |
| 34 | + url = f"https://huggingface.co/{repo_id}/resolve/main/metadata.yml" |
| 35 | + headers = {"Authorization": f"Bearer {token}"} |
| 36 | + response = requests.get(url, headers=headers) |
| 37 | + |
| 38 | + # Check response and retrieve the 'default' field value |
| 39 | + if response.status_code == 200: |
| 40 | + # Read YAML content from response text |
| 41 | + metadata = yaml.safe_load(response.text) |
| 42 | + return metadata.get("default") |
| 43 | + else: |
| 44 | + print(f"Error fetching metadata for {repo_id}:", response.status_code, response.json()) |
| 45 | + return None |
| 46 | + |
| 47 | +def get_all_repos_and_default_branches_from_metadata(collection_id, token): |
| 48 | + # Get list of repos from the collection |
| 49 | + repos = get_repos_in_collection(collection_id, token) |
| 50 | + combined_list = [] |
| 51 | + |
| 52 | + # Iterate over each repo and fetch the default branch from metadata |
| 53 | + for repo_id in repos: |
| 54 | + default_branch = get_repo_default_branch(repo_id, token) |
| 55 | + if default_branch and "gguf" in default_branch: |
| 56 | + combined_list.append(f"{repo_id.split('/')[1]}:{default_branch}") |
| 57 | + |
| 58 | + return combined_list |
| 59 | + |
| 60 | +#Call the function and print the results |
| 61 | +repo_branches = get_all_repos_and_default_branches_from_metadata(collection_id, token) |
| 62 | + |
| 63 | +class TestCortexsoModels: |
| 64 | + |
| 65 | + @pytest.fixture(autouse=True) |
| 66 | + def setup_and_teardown(self, request): |
| 67 | + # Setup |
| 68 | + success = start_server() |
| 69 | + if not success: |
| 70 | + raise Exception("Failed to start server") |
| 71 | + # Delete model if exists |
| 72 | + for model_url in repo_branches: |
| 73 | + run( |
| 74 | + "Delete model", |
| 75 | + [ |
| 76 | + "models", |
| 77 | + "delete", |
| 78 | + model_url, |
| 79 | + ], |
| 80 | + ) |
| 81 | + yield |
| 82 | + |
| 83 | + # Teardown |
| 84 | + for model_url in repo_branches: |
| 85 | + run( |
| 86 | + "Delete model", |
| 87 | + [ |
| 88 | + "models", |
| 89 | + "delete", |
| 90 | + model_url, |
| 91 | + ], |
| 92 | + ) |
| 93 | + stop_server() |
| 94 | + |
| 95 | + @pytest.mark.parametrize("model_url", repo_branches) |
| 96 | + @pytest.mark.asyncio |
| 97 | + async def test_models_on_cortexso_hub(self, model_url): |
| 98 | + |
| 99 | + # Pull model from cortexso hub |
| 100 | + json_body = { |
| 101 | + "model": model_url |
| 102 | + } |
| 103 | + response = requests.post("http://localhost:3928/models/pull", json=json_body) |
| 104 | + assert response.status_code == 200, f"Failed to pull model: {model_url}" |
| 105 | + |
| 106 | + await wait_for_websocket_download_success_event(timeout=None) |
| 107 | + |
| 108 | + # Check if the model was pulled successfully |
| 109 | + get_model_response = requests.get( |
| 110 | + f"http://127.0.0.1:3928/models/{model_url}" |
| 111 | + ) |
| 112 | + assert get_model_response.status_code == 200, f"Failed to fetch model: {model_url}" |
| 113 | + assert ( |
| 114 | + get_model_response.json()["model"] == model_url |
| 115 | + ), f"Unexpected model name for: {model_url}" |
| 116 | + |
| 117 | + # Check if the model is available in the list of models |
| 118 | + response = requests.get("http://localhost:3928/models") |
| 119 | + assert response.status_code == 200 |
| 120 | + models = [i["id"] for i in response.json()["data"]] |
| 121 | + assert model_url in models, f"Model not found in list: {model_url}" |
| 122 | + |
| 123 | + # Install Engine |
| 124 | + exit_code, output, error = run( |
| 125 | + "Install Engine", ["engines", "install", "llama-cpp"], timeout=None, capture = False |
| 126 | + ) |
| 127 | + root = Path.home() |
| 128 | + assert os.path.exists(root / "cortexcpp" / "engines" / "cortex.llamacpp" / "version.txt") |
| 129 | + assert exit_code == 0, f"Install engine failed with error: {error}" |
| 130 | + |
| 131 | + # Start the model |
| 132 | + response = requests.post("http://localhost:3928/models/start", json=json_body) |
| 133 | + assert response.status_code == 200, f"status_code: {response.status_code}" |
| 134 | + |
| 135 | + # Send an inference request |
| 136 | + inference_json_body = { |
| 137 | + "frequency_penalty": 0.2, |
| 138 | + "max_tokens": 4096, |
| 139 | + "messages": [ |
| 140 | + { |
| 141 | + "content": "", |
| 142 | + "role": "user" |
| 143 | + } |
| 144 | + ], |
| 145 | + "model": model_url, |
| 146 | + "presence_penalty": 0.6, |
| 147 | + "stop": [ |
| 148 | + "End" |
| 149 | + ], |
| 150 | + "stream": False, |
| 151 | + "temperature": 0.8, |
| 152 | + "top_p": 0.95 |
| 153 | + } |
| 154 | + response = requests.post("http://localhost:3928/v1/chat/completions", json=inference_json_body, headers={"Content-Type": "application/json"}) |
| 155 | + assert response.status_code == 200, f"status_code: {response.status_code} response: {response.json()}" |
| 156 | + |
| 157 | + # Stop the model |
| 158 | + response = requests.post("http://localhost:3928/models/stop", json=json_body) |
| 159 | + assert response.status_code == 200, f"status_code: {response.status_code}" |
| 160 | + |
| 161 | + # Uninstall Engine |
| 162 | + exit_code, output, error = run( |
| 163 | + "Uninstall engine", ["engines", "uninstall", "llama-cpp"] |
| 164 | + ) |
| 165 | + assert "Engine llama-cpp uninstalled successfully!" in output |
| 166 | + assert exit_code == 0, f"Install engine failed with error: {error}" |
0 commit comments