From 2bdff9294e430b070163a9e0d5bc15ee0c893204 Mon Sep 17 00:00:00 2001 From: Yiiii0 Date: Thu, 26 Mar 2026 01:33:07 -0400 Subject: [PATCH] feat: Add Forge LLM provider support ## Changes - Add Forge as LLM provider for ACE-Step-1.5 Files modified: .env.example acestep/text_tasks/external_ai_request_helpers.py acestep/text_tasks/external_ai_request_helpers_test.py acestep/text_tasks/external_lm_providers.py acestep/text_tasks/external_lm_providers_test.py --- .env.example | 6 +++- .../text_tasks/external_ai_request_helpers.py | 2 +- .../external_ai_request_helpers_test.py | 18 ++++++++++++ acestep/text_tasks/external_lm_providers.py | 16 +++++++++- .../text_tasks/external_lm_providers_test.py | 29 +++++++++++++++++++ 5 files changed, 68 insertions(+), 3 deletions(-) diff --git a/.env.example b/.env.example index afe18af0a..7b6187e8e 100644 --- a/.env.example +++ b/.env.example @@ -58,6 +58,10 @@ ACESTEP_INIT_LLM=auto # API key for authentication (optional) # ACESTEP_API_KEY=sk-your-secret-key +# Forge external LM settings (optional) +# FORGE_API_KEY=your-forge-api-key +# FORGE_API_BASE=https://api.forge.tensorblock.co/v1 + # ==================== Gradio UI Settings ==================== # Server port (default: 7860) # PORT=7860 @@ -75,4 +79,4 @@ ACESTEP_INIT_LLM=auto # ==================== Startup Settings ==================== # By default models are lazy-loaded on first request (fast server startup). # Set to false to force eager model loading at startup. -# ACESTEP_NO_INIT=true \ No newline at end of file +# ACESTEP_NO_INIT=true diff --git a/acestep/text_tasks/external_ai_request_helpers.py b/acestep/text_tasks/external_ai_request_helpers.py index c0512945e..afe9a2530 100644 --- a/acestep/text_tasks/external_ai_request_helpers.py +++ b/acestep/text_tasks/external_ai_request_helpers.py @@ -134,7 +134,7 @@ def build_request_for_protocol( "max_tokens": max_tokens or int(os.getenv("ACESTEP_OPENAI_MAX_TOKENS", "3072")), "temperature": 0.4, } - if require_json_output and provider in {"openai", "zai"}: + if require_json_output and provider in {"openai", "forge", "zai"}: payload["response_format"] = {"type": "json_object"} payload["stop"] = ["```"] if disable_thinking and provider == "zai": diff --git a/acestep/text_tasks/external_ai_request_helpers_test.py b/acestep/text_tasks/external_ai_request_helpers_test.py index 92c7e6188..f480ee861 100644 --- a/acestep/text_tasks/external_ai_request_helpers_test.py +++ b/acestep/text_tasks/external_ai_request_helpers_test.py @@ -127,6 +127,24 @@ def test_build_request_for_protocol_requests_json_output_for_openai_format(self) self.assertEqual(payload["stop"], ["```"]) self.assertNotIn("thinking", payload) + def test_build_request_for_protocol_requests_json_output_for_forge_format(self) -> None: + """Forge format-mode requests should use OpenAI-compatible JSON output flags.""" + + payload, _headers = build_request_for_protocol( + protocol="openai_chat", + provider="forge", + api_key="test-key", + model="OpenAI/gpt-4o-mini", + messages=[{"role": "system", "content": "s"}, {"role": "user", "content": "u"}], + base_url="https://api.forge.tensorblock.co/v1/chat/completions", + max_tokens=768, + require_json_output=True, + ) + + self.assertEqual(payload["response_format"], {"type": "json_object"}) + self.assertEqual(payload["stop"], ["```"]) + self.assertNotIn("thinking", payload) + def test_build_request_for_protocol_disables_zai_thinking_and_requests_json(self) -> None: """Z.ai format calls should disable thinking and request JSON output.""" diff --git a/acestep/text_tasks/external_lm_providers.py b/acestep/text_tasks/external_lm_providers.py index e1ad9bd65..1196925f4 100644 --- a/acestep/text_tasks/external_lm_providers.py +++ b/acestep/text_tasks/external_lm_providers.py @@ -53,6 +53,20 @@ class ExternalProviderProfile: ("OpenAI chat completions", "https://api.openai.com/v1/chat/completions"), ), ), + "forge": ExternalProviderProfile( + provider_id="forge", + label="Forge", + protocol="openai_chat", + default_model="OpenAI/gpt-4o-mini", + default_base_url="https://api.forge.tensorblock.co/v1/chat/completions", + api_key_env="FORGE_API_KEY", + api_key_required=True, + secret_path_env="ACESTEP_FORGE_SECRET_PATH", + secret_file_name="forge_api_key.enc", + base_url_presets=( + ("Forge chat completions", "https://api.forge.tensorblock.co/v1/chat/completions"), + ), + ), "claude": ExternalProviderProfile( provider_id="claude", label="Anthropic Claude", @@ -101,7 +115,7 @@ def get_external_provider_profile(provider: str | None) -> ExternalProviderProfi def get_external_provider_choices() -> list[tuple[str, str]]: """Return provider dropdown choices as ``(label, value)`` pairs.""" - order = ("zai", "openai", "claude", "ollama") + order = ("zai", "openai", "forge", "claude", "ollama") return [ (_EXTERNAL_PROVIDER_PROFILES[provider_id].label, provider_id) for provider_id in order diff --git a/acestep/text_tasks/external_lm_providers_test.py b/acestep/text_tasks/external_lm_providers_test.py index 40a8f5970..e3ecf90c8 100644 --- a/acestep/text_tasks/external_lm_providers_test.py +++ b/acestep/text_tasks/external_lm_providers_test.py @@ -6,8 +6,10 @@ from acestep.text_tasks.external_lm_providers import ( CUSTOM_BASE_URL_PRESET, + build_external_model_choice, get_external_base_url_preset_choices, get_external_base_url_preset_value, + get_external_provider_choices, get_external_provider_profile, ) @@ -15,12 +17,39 @@ class ExternalLmProvidersTests(unittest.TestCase): """Verify provider lookup and base-URL preset helpers stay explicit.""" + def test_get_external_provider_profile_returns_forge_defaults(self) -> None: + """Forge should use its dedicated OpenAI-compatible defaults.""" + + profile = get_external_provider_profile("forge") + + self.assertEqual(profile.protocol, "openai_chat") + self.assertEqual(profile.default_model, "OpenAI/gpt-4o-mini") + self.assertEqual( + profile.default_base_url, + "https://api.forge.tensorblock.co/v1/chat/completions", + ) + self.assertEqual(profile.api_key_env, "FORGE_API_KEY") + def test_get_external_provider_profile_rejects_unknown_provider(self) -> None: """Unknown providers should fail fast instead of silently defaulting.""" with self.assertRaises(ValueError): get_external_provider_profile("mystery") + def test_get_external_provider_choices_includes_forge(self) -> None: + """Provider choices should expose Forge to the external LM picker.""" + + choices = get_external_provider_choices() + + self.assertIn(("Forge", "forge"), choices) + + def test_build_external_model_choice_defaults_forge_model(self) -> None: + """Missing Forge model names should fall back to the provider default.""" + + choice = build_external_model_choice("forge", "") + + self.assertEqual(choice, "external:forge:OpenAI/gpt-4o-mini") + def test_base_url_preset_helpers_use_shared_custom_token(self) -> None: """Custom base-URL selection should use the centralized custom token."""