ragbits-core 1.4.0.dev202601170236__tar.gz → 1.4.0.dev202601300258__tar.gz

{ragbits_core-1.4.0.dev202601170236 → ragbits_core-1.4.0.dev202601300258}/CHANGELOG.md

@@ -5,6 +5,7 @@
 - Add Support for Thinking in agents (#837)
 - Add support for confirmation requests in chat (#853)
 - Add name parameter and slightly refactor HuggingFace dataloder (#829)
+- Fix trace context managers not closing properly when streaming (#922)
 
 ## 1.3.0 (2025-09-11)
 

{ragbits_core-1.4.0.dev202601170236 → ragbits_core-1.4.0.dev202601300258}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ragbits-core
-Version: 1.4.0.dev202601170236
+Version: 1.4.0.dev202601300258
 Summary: Building blocks for rapid development of GenAI applications
 Project-URL: Homepage, https://github.com/deepsense-ai/ragbits
 Project-URL: Bug Reports, https://github.com/deepsense-ai/ragbits/issues

{ragbits_core-1.4.0.dev202601170236 → ragbits_core-1.4.0.dev202601300258}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ragbits-core"
-version = "1.4.0.dev202601170236"
+version = "1.4.0.dev202601300258"
 description = "Building blocks for rapid development of GenAI applications"
 readme = "README.md"
 requires-python = ">=3.10"

{ragbits_core-1.4.0.dev202601170236 → ragbits_core-1.4.0.dev202601300258}/src/ragbits/core/llms/base.py

@@ -244,6 +244,9 @@ class LLMResultStreaming(AsyncIterator[T]):
                     self.metadata: LLMResponseWithMetadata = item
                     if item.usage:
                         self.usage += item.usage
+                    # Exhaust generator to let trace context managers close properly
+                    async for _ in self._generator:
+                        pass
                     raise StopAsyncIteration
                 case Usage():
                     self.usage += item

{ragbits_core-1.4.0.dev202601170236 → ragbits_core-1.4.0.dev202601300258}/tests/unit/llms/test_base.py

@@ -411,3 +411,20 @@ async def test_generate_with_metadata_tracing_and_metrics(
     assert metric_calls[2][1]["value"] == 90.0  # total_tokens / throughput
     assert metric_calls[2][1]["attributes"]["model"] == llm.model_name
     assert metric_calls[2][1]["metric_type"] == MetricType.HISTOGRAM
+
+
+async def test_generate_streaming_trace_handler_closes_properly(llm: MockLLM, mock_trace_handler: MagicMock):
+    set_trace_handlers(mock_trace_handler)
+
+    prompt = SimplePrompt("Hello")
+    stream = llm.generate_streaming(prompt)
+
+    responses = [response async for response in stream]
+    assert responses == ["first response", "second response"]
+
+    # Verify trace context manager was called and properly exited
+    mock_trace_handler.trace.assert_called_once()
+    # The context manager's __enter__ and __exit__ should both be called
+    trace_cm = mock_trace_handler.trace.return_value
+    trace_cm.__enter__.assert_called_once()
+    trace_cm.__exit__.assert_called_once()