mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
359004646b
## Description When a wrapped `ChatOpenAI` model is configured with `streaming=True`, calling `ainvoke()` (the non-streaming async API) on the resulting `HeadroomChatModel` crashes with `AttributeError: 'AsyncStream' object has no attribute 'model_dump'`. This happens because `_agenerate()` passes through to the wrapped model's `_agenerate()`, which — when `streaming=True` — returns a raw OpenAI SDK `AsyncStream` object instead of a LangChain `ChatResult`. The caller then tries to call `.model_dump()` on the stream, which doesn't have that method. `_agenerate()` now detects `streaming=True` on the wrapped model and temporarily disables it for the duration of the non-streaming call, then restores it in a `finally` block. Closes #1285 ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `headroom/integrations/langchain/chat_model.py`: Modified `_agenerate()` to detect `streaming=True` on the wrapped model, temporarily set it to `False` for the duration of the non-streaming call, and restore it in a `finally` block (even on exceptions). Gracefully handles models without a `streaming` attribute or immutable fields. - `tests/test_integrations/langchain/test_chat_model.py`: Added `TestAinvokeStreamingTrue` with 5 test cases covering the core fix, streaming state restoration, exception safety, and passthrough for models without `streaming`. - `CHANGELOG.md`: Added bug fix entry under Unreleased → Bug Fixes. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [ ] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text $ python -m pytest tests/test_integrations/langchain/test_chat_model.py -k TestAinvokeStreamingTrue 5 passed, 39 deselected in 4.14s $ python -m pytest tests/test_integrations/langchain/test_chat_model.py -k "not Ollama and not RealLangChain" 35 passed, 9 deselected in 4.62s $ ruff check headroom/integrations/langchain/chat_model.py tests/test_integrations/langchain/test_chat_model.py All checks passed! $ ruff format --check headroom/integrations/langchain/chat_model.py tests/test_integrations/langchain/test_chat_model.py 2 files already formatted ``` Verification that tests catch the bug (reverted only `chat_model.py`, ran tests): ```text test_agenerate_returns_chatresult_with_streaming_true FAILED assert False = isinstance(<FakeAsyncStream object>, ChatResult) test_streaming_disabled_during_agenerate_call FAILED assert [True] == [False] # streaming was NOT disabled during the call ``` ## Real Behavior Proof - Environment: Linux 6.17.0, Python 3.11.14, langchain-core 1.4.8, pytest 9.1.1, pytest-asyncio 1.4.0 - Exact command / steps: `uv pip install -e ".[dev,langchain]"` then `python -m pytest tests/test_integrations/langchain/test_chat_model.py -k TestAinvokeStreamingTrue` then full module suite with `-k "not Ollama and not RealLangChain"` - Observed result: 5/5 new tests pass, 35/35 existing tests pass, lint clean. Tests fail without the fix (2 failures matching the bug). - Not tested: Real OpenAI API calls (no API key available). Mock-based test simulates `ChatOpenAI`'s streaming behavior faithfully — when `streaming=True`, `_agenerate` returns an `AsyncStream`-like object; when `streaming=False`, it returns a proper `ChatResult`. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [x] I have updated the CHANGELOG.md if applicable ## Screenshots (if applicable) N/A ## Additional Notes - `mypy` was not run as it is not part of the local dev dependencies in this environment. The fix is straightforward attribute access with `getattr`/`setattr` and does not introduce new type complexities. - The fix is minimal: `ainvoke()` is the non-streaming API, so it should never trigger streaming. Temporarily disabling `streaming` on the wrapped model is the safest approach — the setting is always restored in a `finally` block. - If `streaming` is an immutable (frozen pydantic) field, the code catches the exception and falls through without crashing. The caller would need to disable `streaming` on the wrapped model directly in that case. |
||
|---|---|---|
| .. | ||
| __init__.py | ||
| test_agents.py | ||
| test_chat_model.py | ||
| test_evals.py | ||
| test_extended.py | ||
| test_langchain_live.py | ||
| test_langgraph.py | ||
| test_memory.py | ||
| test_retriever.py | ||
| test_streaming.py | ||