diff --git a/app/engine/providers/glm.py b/app/engine/providers/glm.py index 78b328b..9b9f9ed 100644 --- a/app/engine/providers/glm.py +++ b/app/engine/providers/glm.py @@ -13,9 +13,16 @@ DEFAULT_MODEL = "glm-4-flash" class GLMProvider(AIProvider): - def __init__(self, api_key: str, model: str = "glm-4-flash", base_url: str = ""): + def __init__( + self, + api_key: str, + model: str = "glm-4-flash", + base_url: str = "", + max_tokens: int | None = None, + ): self.client = ZhipuAI(api_key=api_key, base_url=base_url or None) self.model = model + self.max_tokens = max_tokens async def chat(self, messages: list[dict], system_prompt: str) -> str: full_messages = [{"role": "system", "content": system_prompt}] + messages @@ -74,6 +81,7 @@ class GLMProvider(AIProvider): model=self.model, messages=messages, temperature=temperature, + **({"max_tokens": self.max_tokens} if self.max_tokens is not None else {}), ) return resp.choices[0].message.content diff --git a/tests/test_glm_provider.py b/tests/test_glm_provider.py new file mode 100644 index 0000000..f5afe82 --- /dev/null +++ b/tests/test_glm_provider.py @@ -0,0 +1,34 @@ +"""GLMProvider max_tokens 透传(site-import spec §3.2)。""" +from unittest.mock import MagicMock, patch + +from app.engine.providers.glm import GLMProvider + + +def _provider(**kw): + with patch("app.engine.providers.glm.ZhipuAI"): + return GLMProvider(api_key="k", **kw) + + +def _ok(cli): + cli.chat.completions.create.return_value.choices = [ + MagicMock(message=MagicMock(content="ok"))] + + +class TestMaxTokensPassthrough: + def test_默认不传_调用无max_tokens参数(self): + import asyncio + p = _provider() + with patch.object(p, "client") as cli: + _ok(cli) + asyncio.run(p.chat([{"role": "user", "content": "hi"}], "sys")) + _, kwargs = cli.chat.completions.create.call_args + assert "max_tokens" not in kwargs + + def test_传入max_tokens_透传到create(self): + import asyncio + p = _provider(max_tokens=4096) + with patch.object(p, "client") as cli: + _ok(cli) + asyncio.run(p.chat([{"role": "user", "content": "hi"}], "sys")) + _, kwargs = cli.chat.completions.create.call_args + assert kwargs.get("max_tokens") == 4096