diff --git a/packages/ai/src/ai/common/chat.py b/packages/ai/src/ai/common/chat.py index 64c403914..35c769ae1 100644 --- a/packages/ai/src/ai/common/chat.py +++ b/packages/ai/src/ai/common/chat.py @@ -38,6 +38,59 @@ def _stop_kwargs() -> dict: return {'stop': stop} if stop else {} +def _is_stop_rejection(e: Exception, stop_kwargs: dict) -> bool: + """Determine if an exception is due to an unsupported stop parameter.""" + if not stop_kwargs: + return False + + status = getattr(e, 'status_code', None) + if status == 400: + body = getattr(e, 'body', None) or {} + code = getattr(e, 'code', None) + param = getattr(e, 'param', None) + if isinstance(body, dict): + code = code or body.get('code') + param = param or body.get('param') + if code == 'unsupported_parameter' or param == 'stop': + return True + + err_str = str(e).lower() + return 'stop' in err_str and ('unsupported' in err_str or 'invalid' in err_str) + + +def _invoke_with_stop_fallback(llm, prompt, stop_kwargs): + """Invoke the LLM, but fallback to no stop kwargs if the API rejects them.""" + try: + return llm.invoke(prompt, **stop_kwargs) + except Exception as e: + if _is_stop_rejection(e, stop_kwargs): + warning(f"LLM rejected 'stop' parameter: {e}. Retrying without stop.") + # Note: Retrying without the stop parameter means the model might hallucinate the + # '\\nObservation:' sequence and return it before agent-level truncation trims it. + return llm.invoke(prompt) + raise + + +def _stream_with_stop_fallback(llm, prompt, stop_kwargs): + """Stream from the LLM, but fallback to no stop kwargs if the API rejects them.""" + try: + it = llm.stream(prompt, **stop_kwargs) + first = next(it) + except StopIteration: + return + except Exception as e: + if _is_stop_rejection(e, stop_kwargs): + warning(f"LLM rejected 'stop' parameter: {e}. Retrying stream without stop.") + # Note: Retrying without the stop parameter means the model might hallucinate the + # '\\nObservation:' sequence and stream it to the UI before agent truncation trims it. + yield from llm.stream(prompt) + return + raise + + yield first + yield from it + + def _make_think_tag_splitter(): """Split ``...`` CoT out of the content stream (Ollama, Perplexity). @@ -240,7 +293,7 @@ def _chat(self, prompt: str) -> str: """ # Ask the LLM. The stop kwarg is only added when the agent set stop sequences, # so non-agent callers (and backends/mocks without a stop param) are unaffected. - results = self._llm.invoke(prompt, **_stop_kwargs()) + results = _invoke_with_stop_fallback(self._llm, prompt, _stop_kwargs()) # Return the results return results.content @@ -496,7 +549,7 @@ def _chat_string_responses( # Only retry non-streaming if nothing has reached the UI; otherwise # the full fallback would arrive on top of the partial we already streamed. if emitted is None or not emitted['any']: - results = self._llm.invoke(prompt, **_stop_kwargs()) + results = _invoke_with_stop_fallback(self._llm, prompt, _stop_kwargs()) content = getattr(results, 'content', '') or '' content_text = content if isinstance(content, str) else str(content) text_parts = [content_text] @@ -597,7 +650,7 @@ def on_chunk_w(t): finish_reason: Optional[str] = None _signature_only_note_sent = False _think_split = _make_think_tag_splitter() - for piece in _llm.stream(prompt, **_stop_kwargs()): + for piece in _stream_with_stop_fallback(_llm, prompt, _stop_kwargs()): # content: str for OpenAI-style, list of typed blocks for Anthropic. content = piece.content text = '' diff --git a/packages/ai/tests/ai/common/test_chat_fallback.py b/packages/ai/tests/ai/common/test_chat_fallback.py new file mode 100644 index 000000000..0f5c11266 --- /dev/null +++ b/packages/ai/tests/ai/common/test_chat_fallback.py @@ -0,0 +1,161 @@ +import pytest + +from ai.common.chat import _invoke_with_stop_fallback, _stream_with_stop_fallback + +class MockInvokeLLM: + def __init__(self, should_fail_on_stop=True): + self.should_fail_on_stop = should_fail_on_stop + self.invoked_kwargs = [] + + def invoke(self, prompt, **kwargs): + self.invoked_kwargs.append(kwargs) + if self.should_fail_on_stop and 'stop' in kwargs: + # Simulate the OpenAI 400 Bad Request error for unsupported stop parameter + raise ValueError("unsupported_parameter: 'stop'") + return 'success' + +class MockStreamLLM: + def __init__(self, should_fail_on_stop=True): + self.should_fail_on_stop = should_fail_on_stop + self.streamed_kwargs = [] + + def stream(self, prompt, **kwargs): + self.streamed_kwargs.append(kwargs) + if self.should_fail_on_stop and 'stop' in kwargs: + def failing_generator(): + # The error usually occurs on the first yield of the lazy stream + raise ValueError("unsupported_parameter: 'stop' is invalid") + yield 'never' + return failing_generator() + + def success_generator(): + yield 'chunk1' + yield 'chunk2' + return success_generator() + +def test_invoke_with_stop_fallback_recovers(): + """Verify that a 400 unsupported parameter error on invoke is caught and retried.""" + llm = MockInvokeLLM(should_fail_on_stop=True) + stop_kwargs = {'stop': ['\nObservation:']} + + result = _invoke_with_stop_fallback(llm, 'Hello', stop_kwargs) + + assert result == 'success' + assert len(llm.invoked_kwargs) == 2 + assert 'stop' in llm.invoked_kwargs[0] + assert 'stop' not in llm.invoked_kwargs[1] + +def test_invoke_with_stop_fallback_success(): + """Verify that if the model supports stop, it works on the first try.""" + llm = MockInvokeLLM(should_fail_on_stop=False) + stop_kwargs = {'stop': ['\nObservation:']} + + result = _invoke_with_stop_fallback(llm, 'Hello', stop_kwargs) + + assert result == 'success' + assert len(llm.invoked_kwargs) == 1 + assert 'stop' in llm.invoked_kwargs[0] + +def test_stream_with_stop_fallback_recovers(): + """Verify that a 400 unsupported parameter error on stream is caught and retried.""" + llm = MockStreamLLM(should_fail_on_stop=True) + stop_kwargs = {'stop': ['\nObservation:']} + + chunks = list(_stream_with_stop_fallback(llm, 'Hello', stop_kwargs)) + + assert chunks == ['chunk1', 'chunk2'] + assert len(llm.streamed_kwargs) == 2 + assert 'stop' in llm.streamed_kwargs[0] + assert 'stop' not in llm.streamed_kwargs[1] + +def test_stream_with_stop_fallback_success(): + """Verify that if the model supports stop, streaming works on the first try.""" + llm = MockStreamLLM(should_fail_on_stop=False) + stop_kwargs = {'stop': ['\nObservation:']} + + chunks = list(_stream_with_stop_fallback(llm, 'Hello', stop_kwargs)) + + assert chunks == ['chunk1', 'chunk2'] + assert len(llm.streamed_kwargs) == 1 + assert 'stop' in llm.streamed_kwargs[0] + +def test_invoke_with_stop_fallback_empty_stop_kwargs(): + """Verify that an error does not trigger a retry if stop_kwargs is empty.""" + llm = MockInvokeLLM() + def failing_invoke(prompt, **kwargs): + llm.invoked_kwargs.append(kwargs) + raise ValueError("unsupported_parameter: 'stop'") + llm.invoke = failing_invoke + + with pytest.raises(ValueError, match="unsupported_parameter: 'stop'"): + _invoke_with_stop_fallback(llm, 'Hello', {}) + + assert len(llm.invoked_kwargs) == 1 # No retry occurred + +def test_invoke_with_stop_fallback_unrelated_error(): + """Verify that an unrelated error does not trigger a retry even with stop_kwargs.""" + llm = MockInvokeLLM() + def failing_invoke(prompt, **kwargs): + llm.invoked_kwargs.append(kwargs) + raise ValueError('Rate limit exceeded') + llm.invoke = failing_invoke + + stop_kwargs = {'stop': ['\nObservation:']} + with pytest.raises(ValueError, match='Rate limit exceeded'): + _invoke_with_stop_fallback(llm, 'Hello', stop_kwargs) + + assert len(llm.invoked_kwargs) == 1 # No retry occurred + +def test_stream_with_stop_fallback_empty_stop_kwargs(): + """Verify that a stream error does not trigger a retry if stop_kwargs is empty.""" + llm = MockStreamLLM() + def failing_stream(prompt, **kwargs): + llm.streamed_kwargs.append(kwargs) + def gen(): + raise ValueError("unsupported_parameter: 'stop'") + yield 'never' + return gen() + llm.stream = failing_stream + + with pytest.raises(ValueError, match="unsupported_parameter: 'stop'"): + list(_stream_with_stop_fallback(llm, 'Hello', {})) + + assert len(llm.streamed_kwargs) == 1 # No retry occurred + + +def test_stream_with_stop_fallback_unrelated_error(): + """Verify that an unrelated stream error does not trigger a retry even with stop_kwargs.""" + llm = MockStreamLLM() + def failing_stream(prompt, **kwargs): + llm.streamed_kwargs.append(kwargs) + def gen(): + raise ValueError('Rate limit exceeded') + yield 'never' + return gen() + llm.stream = failing_stream + + stop_kwargs = {'stop': ['\nObservation:']} + with pytest.raises(ValueError, match='Rate limit exceeded'): + list(_stream_with_stop_fallback(llm, 'Hello', stop_kwargs)) + + assert len(llm.streamed_kwargs) == 1 # No retry occurred + + +def test_stream_with_stop_fallback_mid_stream_error(): + """Verify that an error raised after the first chunk does not trigger a retry, locking in no-duplicate behavior.""" + llm = MockStreamLLM() + def failing_stream(prompt, **kwargs): + llm.streamed_kwargs.append(kwargs) + def gen(): + yield 'chunk1' + raise ValueError("unsupported_parameter: 'stop'") + return gen() + llm.stream = failing_stream + + stop_kwargs = {'stop': ['\nObservation:']} + with pytest.raises(ValueError, match="unsupported_parameter: 'stop'"): + # The generator will yield 'chunk1', then raise the error. + # It should NOT retry because the try/except block does not wrap the main yield loop. + list(_stream_with_stop_fallback(llm, 'Hello', stop_kwargs)) + + assert len(llm.streamed_kwargs) == 1 # No retry occurred