diff --git a/src/ccproxy/cli.py b/src/ccproxy/cli.py index 5586d96..d3c8c24 100644 --- a/src/ccproxy/cli.py +++ b/src/ccproxy/cli.py @@ -326,6 +326,16 @@ def start_litellm(config_dir: Path, args: list[str] | None = None, detach: bool cmd = [str(litellm_path), "--config", str(config_path)] + # Pass num_workers from ccproxy.yaml litellm section if configured + ccproxy_config_path = config_dir / "ccproxy.yaml" + if ccproxy_config_path.exists(): + with ccproxy_config_path.open() as f: + ccproxy_config = yaml.safe_load(f) + if ccproxy_config: + num_workers = ccproxy_config.get("litellm", {}).get("num_workers") + if num_workers is not None: + cmd.extend(["--num_workers", str(num_workers)]) + # Add any additional arguments if args: cmd.extend(args) diff --git a/tests/test_num_workers.py b/tests/test_num_workers.py new file mode 100644 index 0000000..4dbff82 --- /dev/null +++ b/tests/test_num_workers.py @@ -0,0 +1,28 @@ +"""Tests for num_workers configuration passthrough.""" + +from pathlib import Path +from unittest.mock import Mock, patch + +import pytest + +from ccproxy.cli import start_litellm + + +class TestNumWorkers: + """Test suite for num_workers in ccproxy.yaml.""" + + @patch("subprocess.run") + def test_num_workers_passed_to_litellm(self, mock_run: Mock, tmp_path: Path) -> None: + """Test num_workers from ccproxy.yaml is passed as --num_workers to litellm.""" + (tmp_path / "config.yaml").write_text("model_list: []") + (tmp_path / "ccproxy.yaml").write_text( + "ccproxy:\n handler: 'ccproxy.handler:CCProxyHandler'\nlitellm:\n num_workers: 8\n" + ) + mock_run.return_value = Mock(returncode=0) + + with pytest.raises(SystemExit): + start_litellm(tmp_path) + + cmd = mock_run.call_args[0][0] + assert "--num_workers" in cmd, f"--num_workers missing from command: {cmd}" + assert cmd[cmd.index("--num_workers") + 1] == "8"