fix: stop max_retries leaking to bare SDK client in _run_tool_loop (#1850)

Author: ghostiee-11

lumen/ai/llm.py

@@ -431,6 +431,10 @@ async def _run_tool_loop(
         max_tool_rounds: int = 16,
         **kwargs
     ) -> BaseModel | str:
+        # ``max_retries`` is consumed by the instructor wrapper; on bare-client
+        # paths it leaks to the SDK and raises ``TypeError``. Pop once and
+        # re-attach only to the instructor (response_model) round-trip.
+        max_retries = kwargs.pop("max_retries", None)
         requested_stream = bool(kwargs.get("stream", False))
         if requested_stream and structured_model is None:
             # In stream mode we can't know upfront whether a tool will be called.
@@ -442,6 +446,8 @@ async def _run_tool_loop(
         stream = False
         if structured_model is not None and not tool_instances:
             kwargs["response_model"] = structured_model
+            if max_retries is not None:
+                kwargs["max_retries"] = max_retries
         else:
             if tool_instances:
                 stream = kwargs.pop("stream", False)
@@ -470,6 +476,8 @@ async def _run_tool_loop(
 
         if structured_model:
             kwargs["response_model"] = structured_model
+            if max_retries is not None:
+                kwargs["max_retries"] = max_retries
             output = await self.run_client(model_spec, messages_curr, stream=stream, **kwargs)
         return output
 
@@ -1290,6 +1298,7 @@ async def _run_tool_loop(
                 messages, structured_model, tool_instances, tool_contexts, model_spec, max_tool_rounds, **kwargs
             )
 
+        max_retries = kwargs.pop("max_retries", None)
         requested_stream = bool(kwargs.get("stream", False))
         if requested_stream and structured_model is None:
             kwargs.pop("response_model", None)
@@ -1304,6 +1313,8 @@ async def _run_tool_loop(
         # can ask for the structured response in a single round-trip.
         if structured_model is not None and not tool_instances and not has_inbuilt:
             kwargs["response_model"] = structured_model
+            if max_retries is not None:
+                kwargs["max_retries"] = max_retries
         else:
             kwargs.pop("response_model", None)
 
@@ -1332,6 +1343,8 @@ async def _run_tool_loop(
         if structured_model:
             final_kwargs = dict(kwargs)
             final_kwargs["response_model"] = structured_model
+            if max_retries is not None:
+                final_kwargs["max_retries"] = max_retries
             response_id = getattr(output, "id", None)
             if response_id:
                 final_kwargs["previous_response_id"] = response_id

lumen/tests/ai/test_llm.py

@@ -20,6 +20,7 @@
     pytest.skip("lumen.ai could not be imported, skipping tests.", allow_module_level=True)
 
 from instructor.processing.multimodal import Image
+from pydantic import BaseModel
 
 # ---------------------------------------------------------------------------
 # Helpers
@@ -331,6 +332,53 @@ def test_openai_responses_stream_tool_call_id_preserved_from_added_event():
     OpenAI._accumulate_tool_calls(accum, order, OpenAI._extract_stream_tool_calls(done_event))
     tool_calls = OpenAI._tool_calls_from_accum(accum, order)
     assert tool_calls[0]["id"] == "call_abc"
+
+
+async def test_run_tool_loop_drops_max_retries_on_bare_client(monkeypatch):
+    """Regression: ``max_retries`` is an instructor-only kwarg.
+
+    When both ``response_model`` and ``tools`` are supplied (the planner's
+    path), ``_run_tool_loop`` drops ``response_model`` and routes to the bare
+    SDK client.  Before the fix, ``max_retries`` was left in ``kwargs`` and
+    leaked through to ``AsyncMessages.create`` / ``AsyncCompletions.create``,
+    raising ``TypeError: got an unexpected keyword argument 'max_retries'``.
+    """
+
+    class _Plan(BaseModel):
+        next_step: str
+
+    llm = OpenAI(model_kwargs={"default": {"model": "gpt-4.1-mini"}})
+    calls: list[dict] = []
+
+    async def fake_run_client(_model_spec, _messages, **kwargs):
+        calls.append(dict(kwargs))
+        return SimpleNamespace(content="done", tool_calls=None)
+
+    monkeypatch.setattr(llm, "run_client", fake_run_client)
+    monkeypatch.setattr(llm, "_extract_tool_calls", lambda _output: [])
+
+    await llm._run_tool_loop(
+        messages=[{"role": "user", "content": "hi"}],
+        structured_model=_Plan,
+        tool_instances={"lookup": object()},
+        tool_contexts={},
+        max_retries=3,
+    )
+    # The first call uses the bare client (no response_model). It must not
+    # carry max_retries through, since the bare SDK rejects it.
+    bare_call = calls[0]
+    assert "response_model" not in bare_call
+    assert "max_retries" not in bare_call, (
+        "max_retries must be popped from kwargs on the bare-client path; "
+        "it is consumed by the instructor wrapper, not the underlying SDK."
+    )
+    # Final structured-output call (instructor) should still carry the
+    # user-passed max_retries so retry semantics are preserved end to end.
+    final_call = calls[-1]
+    assert final_call.get("response_model") is _Plan
+    assert final_call.get("max_retries") == 3
+
+
 # ---------------------------------------------------------------------------
 # _normalize_multimodal_messages tests
 # ---------------------------------------------------------------------------