Remove private imports from client, move util function only used by CLI to cli/_utils.py

Author: XkunW

tests/vec_inf/cli/test_utils.py

@@ -0,0 +1,17 @@
+"""Tests for the utils functions in the vec-inf cli."""
+
+from vec_inf.cli._utils import create_table
+
+
+def test_create_table_with_header():
+    """Test that create_table creates a table with the correct header."""
+    table = create_table("Key", "Value")
+    assert table.columns[0].header == "Key"
+    assert table.columns[1].header == "Value"
+    assert table.show_header is True
+
+
+def test_create_table_without_header():
+    """Test create_table without header."""
+    table = create_table(show_header=False)
+    assert table.show_header is False

tests/vec_inf/client/test_utils.py

@@ -1,4 +1,4 @@
-"""Tests for the utility functions in the CLI module."""
+"""Tests for the utility functions in the vec-inf client."""
 
 import os
 from unittest.mock import MagicMock, patch
@@ -8,7 +8,6 @@
 
 from vec_inf.client._utils import (
     MODEL_READY_SIGNATURE,
-    create_table,
     get_base_url,
     is_server_running,
     load_config,
@@ -134,20 +133,6 @@ def test_model_health_check_request_exception():
         assert result == ("FAILED", "Connection error")
 
 
-def test_create_table_with_header():
-    """Test that create_table creates a table with the correct header."""
-    table = create_table("Key", "Value")
-    assert table.columns[0].header == "Key"
-    assert table.columns[1].header == "Value"
-    assert table.show_header is True
-
-
-def test_create_table_without_header():
-    """Test create_table without header."""
-    table = create_table(show_header=False)
-    assert table.show_header is False
-
-
 def test_load_config_default_only():
     """Test loading the actual default configuration file from the filesystem."""
     configs = load_config()

vec_inf/cli/_cli.py

@@ -13,8 +13,7 @@
     MetricsResponseFormatter,
     StatusResponseFormatter,
 )
-from vec_inf.client._models import LaunchOptions, LaunchOptionsDict
-from vec_inf.client.api import VecInfClient
+from vec_inf.client import LaunchOptions, LaunchOptionsDict, VecInfClient
 
 
 CONSOLE = Console()

vec_inf/cli/_helper.py

@@ -9,10 +9,9 @@
 from rich.panel import Panel
 from rich.table import Table
 
-import vec_inf.client._utils as utils
 from vec_inf.cli._models import MODEL_TYPE_COLORS, MODEL_TYPE_PRIORITY
-from vec_inf.client._config import ModelConfig
-from vec_inf.client._models import ModelInfo, StatusResponse
+from vec_inf.cli._utils import create_table
+from vec_inf.client import ModelConfig, ModelInfo, StatusResponse
 
 
 class LaunchResponseFormatter:
@@ -24,7 +23,7 @@ def __init__(self, model_name: str, params: dict[str, Any]):
 
     def format_table_output(self) -> Table:
         """Format output as rich Table."""
-        table = utils.create_table(key_title="Job Config", value_title="Value")
+        table = create_table(key_title="Job Config", value_title="Value")
 
         # Add key information with consistent styling
         table.add_row("Slurm Job ID", self.params["slurm_job_id"], style="blue")
@@ -89,7 +88,7 @@ def output_json(self) -> None:
 
     def output_table(self) -> Table:
         """Create and display rich table."""
-        table = utils.create_table(key_title="Job Status", value_title="Value")
+        table = create_table(key_title="Job Status", value_title="Value")
         table.add_row("Model Name", self.status_info.model_name)
         table.add_row("Model Status", self.status_info.server_status, style="blue")
 
@@ -107,7 +106,7 @@ class MetricsResponseFormatter:
 
     def __init__(self, metrics: Union[dict[str, float], str]):
         self.metrics = self._set_metrics(metrics)
-        self.table = utils.create_table("Metric", "Value")
+        self.table = create_table("Metric", "Value")
         self.enabled_prefix_caching = self._check_prefix_caching()
 
     def _set_metrics(self, metrics: Union[dict[str, float], str]) -> dict[str, float]:
@@ -211,7 +210,7 @@ def _format_single_model_output(
             )
             return config_dict
 
-        table = utils.create_table(key_title="Model Config", value_title="Value")
+        table = create_table(key_title="Model Config", value_title="Value")
         for field, value in config.model_dump().items():
             if field not in {"venv", "log_dir"}:
                 table.add_row(field, str(value))

vec_inf/cli/_utils.py

@@ -0,0 +1,13 @@
+"""Helper functions for the CLI."""
+
+from rich.table import Table
+
+
+def create_table(
+    key_title: str = "", value_title: str = "", show_header: bool = True
+) -> Table:
+    """Create a table for displaying model status."""
+    table = Table(show_header=show_header, header_style="bold magenta")
+    table.add_column(key_title, style="dim")
+    table.add_column(value_title)
+    return table

vec_inf/client/__init__.py

@@ -5,6 +5,7 @@
 users direct control over the lifecycle of inference servers via python scripts.
 """
 
+from vec_inf.client._config import ModelConfig
 from vec_inf.client._models import (
     LaunchOptions,
     LaunchOptionsDict,
@@ -28,4 +29,5 @@
     "ModelType",
     "LaunchOptions",
     "LaunchOptionsDict",
+    "ModelConfig",
 ]

vec_inf/client/_utils.py

@@ -9,7 +9,6 @@
 
 import requests
 import yaml
-from rich.table import Table
 
 from vec_inf.client._config import ModelConfig
 from vec_inf.client._models import ModelStatus
@@ -118,16 +117,6 @@ def model_health_check(
         return (ModelStatus.FAILED, str(e))
 
 
-def create_table(
-    key_title: str = "", value_title: str = "", show_header: bool = True
-) -> Table:
-    """Create a table for displaying model status."""
-    table = Table(show_header=show_header, header_style="bold magenta")
-    table.add_column(key_title, style="dim")
-    table.add_column(value_title)
-    return table
-
-
 def load_config() -> list[ModelConfig]:
     """Load the model configuration."""
     default_path = (