diff --git a/client.py b/client.py index be0b585..f394ebb 100644 --- a/client.py +++ b/client.py @@ -7,6 +7,7 @@ Functions for creating and configuring the Claude Agent SDK client. import json import os +import re import shutil import sys from pathlib import Path @@ -68,7 +69,7 @@ EXTRA_READ_PATHS_BLOCKLIST = { ".netrc", } -def compute_mode(model: str) -> str: +def convert_model_for_vertex(model: str) -> str: """ Convert model name format for Vertex AI compatibility. @@ -89,8 +90,7 @@ def compute_mode(model: str) -> str: # Pattern: claude-{name}-{version}-{date} -> claude-{name}-{version}@{date} # Example: claude-opus-4-5-20251101 -> claude-opus-4-5@20251101 # The date is always 8 digits at the end - import re - match = re.match(r'^(claude-[a-z0-9-]+?)-(\d{8})$', model) + match = re.match(r'^(claude-.+)-(\d{8})$', model) if match: base_name, date = match.groups() return f"{base_name}@{date}" @@ -208,6 +208,7 @@ def get_extra_read_paths() -> list[Path]: return validated_paths + # Feature MCP tools for feature/test management FEATURE_MCP_TOOLS = [ # Core feature operations @@ -438,7 +439,7 @@ def create_client( is_vertex = sdk_env.get("CLAUDE_CODE_USE_VERTEX") == "1" is_alternative_api = bool(base_url) or is_vertex is_ollama = "localhost:11434" in base_url or "127.0.0.1:11434" in base_url - model = compute_mode(model) + model = convert_model_for_vertex(model) if sdk_env: print(f" - API overrides: {', '.join(sdk_env.keys())}") if is_vertex: diff --git a/test_client.py b/test_client.py new file mode 100644 index 0000000..48f52c4 --- /dev/null +++ b/test_client.py @@ -0,0 +1,105 @@ +#!/usr/bin/env python3 +""" +Client Utility Tests +==================== + +Tests for the client module utility functions. +Run with: python test_client.py +""" + +import os +import unittest + +from client import convert_model_for_vertex + + +class TestConvertModelForVertex(unittest.TestCase): + """Tests for convert_model_for_vertex function.""" + + def setUp(self): + """Save original env state.""" + self._orig_vertex = os.environ.get("CLAUDE_CODE_USE_VERTEX") + + def tearDown(self): + """Restore original env state.""" + if self._orig_vertex is None: + os.environ.pop("CLAUDE_CODE_USE_VERTEX", None) + else: + os.environ["CLAUDE_CODE_USE_VERTEX"] = self._orig_vertex + + # --- Vertex AI disabled (default) --- + + def test_returns_model_unchanged_when_vertex_disabled(self): + os.environ.pop("CLAUDE_CODE_USE_VERTEX", None) + self.assertEqual( + convert_model_for_vertex("claude-opus-4-5-20251101"), + "claude-opus-4-5-20251101", + ) + + def test_returns_model_unchanged_when_vertex_set_to_zero(self): + os.environ["CLAUDE_CODE_USE_VERTEX"] = "0" + self.assertEqual( + convert_model_for_vertex("claude-opus-4-5-20251101"), + "claude-opus-4-5-20251101", + ) + + def test_returns_model_unchanged_when_vertex_set_to_empty(self): + os.environ["CLAUDE_CODE_USE_VERTEX"] = "" + self.assertEqual( + convert_model_for_vertex("claude-sonnet-4-5-20250929"), + "claude-sonnet-4-5-20250929", + ) + + # --- Vertex AI enabled: standard conversions --- + + def test_converts_opus_model(self): + os.environ["CLAUDE_CODE_USE_VERTEX"] = "1" + self.assertEqual( + convert_model_for_vertex("claude-opus-4-5-20251101"), + "claude-opus-4-5@20251101", + ) + + def test_converts_sonnet_model(self): + os.environ["CLAUDE_CODE_USE_VERTEX"] = "1" + self.assertEqual( + convert_model_for_vertex("claude-sonnet-4-5-20250929"), + "claude-sonnet-4-5@20250929", + ) + + def test_converts_haiku_model(self): + os.environ["CLAUDE_CODE_USE_VERTEX"] = "1" + self.assertEqual( + convert_model_for_vertex("claude-3-5-haiku-20241022"), + "claude-3-5-haiku@20241022", + ) + + # --- Vertex AI enabled: already converted or non-matching --- + + def test_already_vertex_format_unchanged(self): + os.environ["CLAUDE_CODE_USE_VERTEX"] = "1" + self.assertEqual( + convert_model_for_vertex("claude-opus-4-5@20251101"), + "claude-opus-4-5@20251101", + ) + + def test_non_claude_model_unchanged(self): + os.environ["CLAUDE_CODE_USE_VERTEX"] = "1" + self.assertEqual( + convert_model_for_vertex("gpt-4o"), + "gpt-4o", + ) + + def test_model_without_date_suffix_unchanged(self): + os.environ["CLAUDE_CODE_USE_VERTEX"] = "1" + self.assertEqual( + convert_model_for_vertex("claude-opus-4-5"), + "claude-opus-4-5", + ) + + def test_empty_string_unchanged(self): + os.environ["CLAUDE_CODE_USE_VERTEX"] = "1" + self.assertEqual(convert_model_for_vertex(""), "") + + +if __name__ == "__main__": + unittest.main()