diff --git a/config.yaml b/config.yaml index c4b98cb2..8f96fe1e 100644 --- a/config.yaml +++ b/config.yaml @@ -15,6 +15,8 @@ api: base_url: 'https://yunwu.ai' model: '' llm_support_json: false + # LLM provider: 'openai' (default, OpenAI-compatible API) or 'litellm' (100+ providers via LiteLLM) + provider: 'openai' # *Number of LLM multi-threaded accesses, set to 1 if using local LLM max_workers: 4 diff --git a/core/st_utils/sidebar_setting.py b/core/st_utils/sidebar_setting.py index a300de5b..4502a073 100644 --- a/core/st_utils/sidebar_setting.py +++ b/core/st_utils/sidebar_setting.py @@ -65,11 +65,25 @@ def page_setting(): # config_input(t("Cookies Path"), "youtube.cookies_path") with st.expander(t("LLM Configuration"), expanded=True): + providers = ["openai", "litellm"] + current_provider = load_key("api.provider") + if current_provider not in providers: + current_provider = "openai" + selected_provider = st.selectbox( + t("LLM Provider"), + options=providers, + index=providers.index(current_provider), + help=t("'openai' for OpenAI-compatible APIs, 'litellm' for 100+ providers (Anthropic, Google, Azure, etc.) via LiteLLM"), + ) + if selected_provider != load_key("api.provider"): + update_key("api.provider", selected_provider) + st.rerun() + config_input(t("API_KEY"), "api.key", placeholder=t("Enter your API key")) config_input( t("BASE_URL"), "api.base_url", - help=t("Openai format, will add /v1/chat/completions automatically"), + help=t("Openai format, will add /v1/chat/completions automatically") if selected_provider == "openai" else t("Optional: custom endpoint URL (leave empty to use provider defaults)"), ) # Try to use searchbox for model selection, fall back to text_input diff --git a/core/utils/ask_gpt.py b/core/utils/ask_gpt.py index 972bd20c..42411313 100644 --- a/core/utils/ask_gpt.py +++ b/core/utils/ask_gpt.py @@ -3,6 +3,7 @@ from threading import Lock import json_repair from openai import OpenAI +import litellm from core.utils.config_utils import load_key from rich import print as rprint from core.utils.decorator import except_handler @@ -40,6 +41,37 @@ def _load_cache(prompt, resp_type, log_title): # ask gpt once # ------------ +def _call_openai(model, messages, response_format, base_url, api_key): + if 'ark' in base_url: + base_url = "https://ark.cn-beijing.volces.com/api/v3" + elif 'v1' not in base_url: + base_url = base_url.strip('/') + '/v1' + client = OpenAI(api_key=api_key, base_url=base_url) + params = dict( + model=model, + messages=messages, + response_format=response_format, + timeout=300 + ) + return client.chat.completions.create(**params) + + +def _call_litellm(model, messages, response_format, base_url, api_key): + params = dict( + model=model, + messages=messages, + drop_params=True, + timeout=300, + ) + if response_format: + params["response_format"] = response_format + if api_key: + params["api_key"] = api_key + if base_url: + params["api_base"] = base_url + return litellm.completion(**params) + + @except_handler("GPT request failed", retry=5) def ask_gpt(prompt, resp_type=None, valid_def=None, log_title="default"): if not load_key("api.key"): @@ -52,22 +84,16 @@ def ask_gpt(prompt, resp_type=None, valid_def=None, log_title="default"): model = load_key("api.model") base_url = load_key("api.base_url") - if 'ark' in base_url: - base_url = "https://ark.cn-beijing.volces.com/api/v3" # huoshan base url - elif 'v1' not in base_url: - base_url = base_url.strip('/') + '/v1' - client = OpenAI(api_key=load_key("api.key"), base_url=base_url) + api_key = load_key("api.key") response_format = {"type": "json_object"} if resp_type == "json" and load_key("api.llm_support_json") else None messages = [{"role": "user", "content": prompt}] - params = dict( - model=model, - messages=messages, - response_format=response_format, - timeout=300 - ) - resp_raw = client.chat.completions.create(**params) + provider = load_key("api.provider") + if provider == "litellm": + resp_raw = _call_litellm(model, messages, response_format, base_url, api_key) + else: + resp_raw = _call_openai(model, messages, response_format, base_url, api_key) # process and return full result resp_content = resp_raw.choices[0].message.content @@ -75,13 +101,13 @@ def ask_gpt(prompt, resp_type=None, valid_def=None, log_title="default"): resp = json_repair.loads(resp_content) else: resp = resp_content - + # check if the response format is valid if valid_def: valid_resp = valid_def(resp) if valid_resp['status'] != 'success': _save_cache(model, prompt, resp_content, resp_type, resp, log_title="error", message=valid_resp['message']) - raise ValueError(f"❎ API response error: {valid_resp['message']}") + raise ValueError(f"API response error: {valid_resp['message']}") _save_cache(model, prompt, resp_content, resp_type, resp, log_title=log_title) return resp @@ -89,6 +115,6 @@ def ask_gpt(prompt, resp_type=None, valid_def=None, log_title="default"): if __name__ == '__main__': from rich import print as rprint - + result = ask_gpt("""test respond ```json\n{\"code\": 200, \"message\": \"success\"}\n```""", resp_type="json") rprint(f"Test json output result: {result}") diff --git a/requirements.txt b/requirements.txt index 2050d02c..8b5f539c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,6 +5,7 @@ transformers>=4.48.0 moviepy==1.0.3 numpy>=2.0.2 openai>=1.55.3,<2 +litellm>=1.80.0,<1.87.0 opencv-python==4.11.0.86 openpyxl==3.1.5 pandas>=2.2.3 diff --git a/tests/test_litellm_provider.py b/tests/test_litellm_provider.py new file mode 100644 index 00000000..72e24d10 --- /dev/null +++ b/tests/test_litellm_provider.py @@ -0,0 +1,167 @@ +import sys +import os +import json +import types +from unittest import mock +from types import SimpleNamespace + +import pytest + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + + +def _make_mock_response(content="test response"): + return SimpleNamespace( + choices=[SimpleNamespace(message=SimpleNamespace(content=content))], + usage=SimpleNamespace(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ) + + +@pytest.fixture(autouse=True) +def tmp_config(tmp_path, monkeypatch): + config = { + "api": { + "key": "test-key-123", + "base_url": "https://api.example.com", + "model": "anthropic/claude-haiku-4-5", + "llm_support_json": False, + "provider": "litellm", + }, + "display_language": "en", + "language_split_with_space": ["en"], + "language_split_without_space": ["zh"], + } + config_file = tmp_path / "config.yaml" + from ruamel.yaml import YAML + + yaml = YAML() + with open(config_file, "w") as f: + yaml.dump(config, f) + + monkeypatch.setattr("core.utils.config_utils.CONFIG_PATH", str(config_file)) + return config_file + + +class TestCallLiteLLM: + def test_basic_completion(self): + mock_resp = _make_mock_response("hello world") + with mock.patch("litellm.completion", return_value=mock_resp) as mock_comp: + from core.utils.ask_gpt import _call_litellm + + result = _call_litellm( + "anthropic/claude-haiku-4-5", + [{"role": "user", "content": "say hi"}], + None, + "", + "test-key", + ) + assert result.choices[0].message.content == "hello world" + mock_comp.assert_called_once() + call_kwargs = mock_comp.call_args[1] + assert call_kwargs["model"] == "anthropic/claude-haiku-4-5" + assert call_kwargs["drop_params"] is True + assert call_kwargs["api_key"] == "test-key" + + def test_drop_params_always_true(self): + mock_resp = _make_mock_response() + with mock.patch("litellm.completion", return_value=mock_resp) as mock_comp: + from core.utils.ask_gpt import _call_litellm + + _call_litellm("gemini/gemini-pro", [{"role": "user", "content": "test"}], None, "", "key") + assert mock_comp.call_args[1]["drop_params"] is True + + def test_api_key_omitted_when_empty(self): + mock_resp = _make_mock_response() + with mock.patch("litellm.completion", return_value=mock_resp) as mock_comp: + from core.utils.ask_gpt import _call_litellm + + _call_litellm("openai/gpt-4o", [{"role": "user", "content": "test"}], None, "", "") + assert "api_key" not in mock_comp.call_args[1] + + def test_base_url_forwarded_when_set(self): + mock_resp = _make_mock_response() + with mock.patch("litellm.completion", return_value=mock_resp) as mock_comp: + from core.utils.ask_gpt import _call_litellm + + _call_litellm( + "openai/gpt-4o", + [{"role": "user", "content": "test"}], + None, + "https://custom.endpoint.com", + "key", + ) + assert mock_comp.call_args[1]["api_base"] == "https://custom.endpoint.com" + + def test_base_url_omitted_when_empty(self): + mock_resp = _make_mock_response() + with mock.patch("litellm.completion", return_value=mock_resp) as mock_comp: + from core.utils.ask_gpt import _call_litellm + + _call_litellm("anthropic/claude-haiku-4-5", [{"role": "user", "content": "test"}], None, "", "key") + assert "api_base" not in mock_comp.call_args[1] + + def test_json_response_format_forwarded(self): + mock_resp = _make_mock_response() + with mock.patch("litellm.completion", return_value=mock_resp) as mock_comp: + from core.utils.ask_gpt import _call_litellm + + rf = {"type": "json_object"} + _call_litellm("openai/gpt-4o", [{"role": "user", "content": "test"}], rf, "", "key") + assert mock_comp.call_args[1]["response_format"] == rf + + def test_response_format_omitted_when_none(self): + mock_resp = _make_mock_response() + with mock.patch("litellm.completion", return_value=mock_resp) as mock_comp: + from core.utils.ask_gpt import _call_litellm + + _call_litellm("openai/gpt-4o", [{"role": "user", "content": "test"}], None, "", "key") + assert "response_format" not in mock_comp.call_args[1] + + +class TestAskGptProviderRouting: + def test_litellm_provider_routes_to_litellm(self, tmp_config): + mock_resp = _make_mock_response('{"code": 200}') + with mock.patch("core.utils.ask_gpt._call_litellm", return_value=mock_resp) as mock_lit, \ + mock.patch("core.utils.ask_gpt._call_openai") as mock_oai: + from core.utils.ask_gpt import ask_gpt + + ask_gpt("test prompt", log_title="test_routing") + mock_lit.assert_called_once() + mock_oai.assert_not_called() + + def test_openai_provider_routes_to_openai(self, tmp_config): + from core.utils.config_utils import update_key + + update_key("api.provider", "openai") + mock_resp = _make_mock_response("ok") + with mock.patch("core.utils.ask_gpt._call_openai", return_value=mock_resp) as mock_oai, \ + mock.patch("core.utils.ask_gpt._call_litellm") as mock_lit: + from core.utils.ask_gpt import ask_gpt + + ask_gpt("test prompt", log_title="test_routing_oai") + mock_oai.assert_called_once() + mock_lit.assert_not_called() + + +class TestCallOpenAI: + def test_ark_base_url_override(self): + with mock.patch("core.utils.ask_gpt.OpenAI") as MockClient: + mock_instance = mock.MagicMock() + mock_instance.chat.completions.create.return_value = _make_mock_response() + MockClient.return_value = mock_instance + + from core.utils.ask_gpt import _call_openai + + _call_openai("model", [{"role": "user", "content": "hi"}], None, "https://ark.example.com/api", "key") + assert MockClient.call_args[1]["base_url"] == "https://ark.cn-beijing.volces.com/api/v3" + + def test_v1_appended_when_missing(self): + with mock.patch("core.utils.ask_gpt.OpenAI") as MockClient: + mock_instance = mock.MagicMock() + mock_instance.chat.completions.create.return_value = _make_mock_response() + MockClient.return_value = mock_instance + + from core.utils.ask_gpt import _call_openai + + _call_openai("model", [{"role": "user", "content": "hi"}], None, "https://api.example.com", "key") + assert MockClient.call_args[1]["base_url"] == "https://api.example.com/v1"