feat(llm): add streaming tool call accumulation and LLMResponse parity (#1789)

addresses #1760 (comment)
part of #1760

Author: Pouyanpi

nemoguardrails/actions/llm/utils.py

@@ -29,7 +29,7 @@
 from nemoguardrails.exceptions import LLMCallException
 from nemoguardrails.logging.explain import LLMCallInfo
 from nemoguardrails.logging.llm_tracker import track_llm_call
-from nemoguardrails.types import ChatMessage, LLMModel, LLMResponse, LLMResponseChunk
+from nemoguardrails.types import ChatMessage, LLMModel, LLMResponse, LLMResponseChunk, UsageInfo
 
 if TYPE_CHECKING:
     from nemoguardrails.streaming import StreamingHandler
@@ -98,43 +98,84 @@ async def _stream_llm_call(
     llm_params: Optional[dict] = None,
 ) -> LLMResponse:
     handler.stop = stop or []
-    accumulated_metadata: Dict[str, Any] = {}
-    last_chunk: Optional[LLMResponseChunk] = None
+    streaming_handler_metadata: Dict[str, Any] = {}
+    accumulated_provider_metadata: Dict[str, Any] = {}
+    accumulated_reasoning: List[str] = []
+    tool_calls = None
+    model_name: Optional[str] = None
+    finish_reason: Optional[str] = None
+    request_id: Optional[str] = None
+    usage: Optional[UsageInfo] = None
 
     try:
         async for chunk in model.stream_async(prompt, stop=stop, **(llm_params or {})):
-            last_chunk = chunk
             content = chunk.delta_content or ""
 
+            if chunk.delta_reasoning:
+                accumulated_reasoning.append(chunk.delta_reasoning)
+            if chunk.delta_tool_calls:
+                tool_calls = chunk.delta_tool_calls
+            if chunk.model:
+                model_name = chunk.model
+            if chunk.finish_reason:
+                finish_reason = chunk.finish_reason
+            if chunk.request_id:
+                request_id = chunk.request_id
+            if chunk.usage:
+                usage = chunk.usage
+            if chunk.provider_metadata:
+                accumulated_provider_metadata.update(chunk.provider_metadata)
+
             chunk_metadata = _extract_chunk_metadata(chunk)
             if chunk_metadata:
-                accumulated_metadata.update(chunk_metadata)
+                streaming_handler_metadata.update(chunk_metadata)
 
             await handler.push_chunk(content, chunk_metadata)
 
-        if accumulated_metadata:
-            llm_response_metadata_var.set(accumulated_metadata)
+        llm_response_metadata_var.set(accumulated_provider_metadata or None)
 
         await handler.finish()
 
         llm_call_info = llm_call_info_var.get()
         if llm_call_info:
             llm_call_info.completion = handler.completion
 
-        if last_chunk is not None:
-            _update_token_stats_from_chunk(last_chunk)
+        if usage:
+            fake_chunk = LLMResponseChunk(usage=usage)
+            _update_token_stats_from_chunk(fake_chunk)
+
+        if tool_calls:
+            tool_calls_var.set([tc.to_dict() for tc in tool_calls])
+        else:
+            tool_calls_var.set(None)
+
+        reasoning_content = "".join(accumulated_reasoning) if accumulated_reasoning else None
+        # TODO: call _extract_and_remove_think_tags on the completed response
+        # to handle models that stream reasoning via <think> tags in content
+        # rather than via delta_reasoning. Pre-existing gap, not introduced here.
+        reasoning_trace_var.set(reasoning_content)
 
         return LLMResponse(
             content=handler.completion,
-            usage=last_chunk.usage if last_chunk else None,
-            provider_metadata=accumulated_metadata if accumulated_metadata else None,
+            reasoning=reasoning_content,
+            tool_calls=tool_calls,
+            model=model_name,
+            finish_reason=finish_reason,
+            request_id=request_id,
+            usage=usage,
+            provider_metadata=accumulated_provider_metadata or None,
         )
 
     except Exception as e:
         _raise_llm_call_exception(e, model)
 
 
 def _extract_chunk_metadata(chunk: LLMResponseChunk) -> Optional[Dict[str, Any]]:
+    # This feeds handler.push_chunk() for the StreamingHandler consumer path
+    # (API responses, output rails). Separate from the field accumulation in
+    # _stream_llm_call which builds the returned LLMResponse for the pipeline.
+    # TODO(Pouyanpi): consider pushing tool_calls and reasoning through the handler too,
+    # so output rails and streaming consumers can see them in real-time.
     metadata: Dict[str, Any] = {}
     if chunk.provider_metadata:
         metadata["provider_metadata"] = chunk.provider_metadata

nemoguardrails/integrations/langchain/llm_adapter.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
 import logging
 import uuid
 from typing import Any, AsyncIterator, Dict, List, NamedTuple, Optional, Union
@@ -205,8 +206,33 @@ async def stream_async(
     ) -> AsyncIterator[LLMResponseChunk]:
         llm = self._prepare_llm(kwargs)
         messages = self._to_langchain_input(prompt)
+
+        tool_call_acc: Dict[int, Dict[str, Any]] = {}
+
         async for chunk in llm.astream(messages, stop=stop):
-            yield _langchain_chunk_to_llm_response_chunk(chunk)
+            for tc_chunk in getattr(chunk, "tool_call_chunks", None) or []:
+                idx = tc_chunk.get("index", 0)
+                if idx not in tool_call_acc:
+                    tool_call_acc[idx] = {
+                        "id": tc_chunk.get("id") or "",
+                        "name": tc_chunk.get("name") or "",
+                        "arguments_buffer": "",
+                    }
+                else:
+                    if tc_chunk.get("id"):
+                        tool_call_acc[idx]["id"] = tc_chunk["id"]
+                    if tc_chunk.get("name"):
+                        tool_call_acc[idx]["name"] = tc_chunk["name"]
+                arg_fragment = tc_chunk.get("args") or ""
+                if arg_fragment:
+                    tool_call_acc[idx]["arguments_buffer"] += arg_fragment
+
+            response_chunk = _langchain_chunk_to_llm_response_chunk(chunk)
+
+            if response_chunk.finish_reason == "tool_calls" and tool_call_acc:
+                response_chunk.delta_tool_calls = _finalize_tool_call_acc(tool_call_acc)
+
+            yield response_chunk
 
 
 class LangChainFramework:
@@ -354,6 +380,29 @@ def _extract_tool_calls(response: Any) -> Optional[List[ToolCall]]:
     ]
 
 
+def _finalize_tool_call_acc(acc: Dict[int, Dict[str, Any]]) -> List[ToolCall]:
+    result = []
+    for idx in sorted(acc.keys()):
+        entry = acc[idx]
+        raw_args = entry["arguments_buffer"]
+        try:
+            args_dict = json.loads(raw_args) if raw_args else {}
+        except json.JSONDecodeError:
+            log.warning("Failed to parse tool call arguments for '%s' (index %d): %r", entry["name"], idx, raw_args)
+            args_dict = {}
+        result.append(
+            ToolCall(
+                id=entry["id"] or str(uuid.uuid4()),
+                type="function",
+                function=ToolCallFunction(
+                    name=entry["name"],
+                    arguments=args_dict,
+                ),
+            )
+        )
+    return result
+
+
 def _extract_usage(response: Any) -> Optional[UsageInfo]:
     usage = _build_usage_info(getattr(response, "usage_metadata", None))
     if usage is not None:

tests/test_actions_llm_utils.py

@@ -21,17 +21,25 @@
     _log_completion,
     _store_reasoning_traces,
     _store_tool_calls,
+    _stream_llm_call,
     _update_token_stats_from_chunk,
     llm_call,
 )
-from nemoguardrails.context import llm_call_info_var, llm_stats_var, reasoning_trace_var, tool_calls_var
+from nemoguardrails.context import (
+    llm_call_info_var,
+    llm_response_metadata_var,
+    llm_stats_var,
+    reasoning_trace_var,
+    tool_calls_var,
+)
 from nemoguardrails.exceptions import LLMCallException
 from nemoguardrails.integrations.langchain.llm_adapter import (
     LangChainLLMAdapter,
     _infer_provider_from_module,
 )
 from nemoguardrails.logging.explain import LLMCallInfo
 from nemoguardrails.logging.stats import LLMStats
+from nemoguardrails.streaming import StreamingHandler
 from nemoguardrails.types import ChatMessage, LLMResponse, LLMResponseChunk, Role, ToolCall, ToolCallFunction, UsageInfo
 
 
@@ -482,3 +490,150 @@ def provider_url(self):
         await llm_call(model, [])
 
         assert received_prompt == []
+
+
+def _make_chunk_model(chunks):
+    class _Model:
+        model_name = "test-model"
+        provider_name = "test"
+        provider_url = None
+
+        async def generate_async(self, prompt, *, stop=None, **kwargs):
+            return LLMResponse(content="")
+
+        async def stream_async(self, prompt, *, stop=None, **kwargs):
+            for c in chunks:
+                yield c
+
+    return _Model()
+
+
+class TestStreamLlmCallAccumulation:
+    @pytest.mark.asyncio
+    async def test_accumulates_tool_calls(self):
+        tc = [ToolCall(id="call_1", function=ToolCallFunction(name="get_weather", arguments={"city": "Paris"}))]
+        model = _make_chunk_model(
+            [
+                LLMResponseChunk(model="gpt-4o"),
+                LLMResponseChunk(delta_tool_calls=tc, finish_reason="tool_calls"),
+                LLMResponseChunk(usage=UsageInfo(input_tokens=10, output_tokens=5, total_tokens=15)),
+            ]
+        )
+
+        result = await _stream_llm_call(model, "test", StreamingHandler(), stop=None)
+
+        assert result.tool_calls == tc
+        assert result.model == "gpt-4o"
+        assert result.finish_reason == "tool_calls"
+        assert result.usage.total_tokens == 15
+        assert tool_calls_var.get() is not None
+
+    @pytest.mark.asyncio
+    async def test_accumulates_reasoning(self):
+        model = _make_chunk_model(
+            [
+                LLMResponseChunk(delta_reasoning="Let me ", model="gpt-4o"),
+                LLMResponseChunk(delta_reasoning="think..."),
+                LLMResponseChunk(delta_content="42", finish_reason="stop"),
+                LLMResponseChunk(usage=UsageInfo(input_tokens=5, output_tokens=3, total_tokens=8)),
+            ]
+        )
+
+        result = await _stream_llm_call(model, "test", StreamingHandler(), stop=None)
+
+        assert result.content == "42"
+        assert result.reasoning == "Let me think..."
+        assert result.model == "gpt-4o"
+        assert result.finish_reason == "stop"
+        assert reasoning_trace_var.get() == "Let me think..."
+
+    @pytest.mark.asyncio
+    async def test_text_only(self):
+        model = _make_chunk_model(
+            [
+                LLMResponseChunk(delta_content="Hello", model="gpt-4o"),
+                LLMResponseChunk(delta_content=" world", finish_reason="stop"),
+                LLMResponseChunk(usage=UsageInfo(input_tokens=5, output_tokens=2, total_tokens=7)),
+            ]
+        )
+
+        result = await _stream_llm_call(model, "test", StreamingHandler(), stop=None)
+
+        assert result.content == "Hello world"
+        assert result.tool_calls is None
+        assert result.reasoning is None
+        assert result.model == "gpt-4o"
+        assert result.finish_reason == "stop"
+        assert result.usage.total_tokens == 7
+
+    @pytest.mark.asyncio
+    async def test_request_id_accumulated(self):
+        model = _make_chunk_model(
+            [
+                LLMResponseChunk(delta_content="hi", request_id="req-123", model="gpt-4o"),
+                LLMResponseChunk(finish_reason="stop"),
+            ]
+        )
+
+        result = await _stream_llm_call(model, "test", StreamingHandler(), stop=None)
+
+        assert result.request_id == "req-123"
+
+    @pytest.mark.asyncio
+    async def test_clears_tool_calls_var_when_none(self):
+        tool_calls_var.set([{"id": "stale", "type": "function", "function": {"name": "old", "arguments": {}}}])
+
+        model = _make_chunk_model(
+            [
+                LLMResponseChunk(delta_content="no tools here", finish_reason="stop"),
+            ]
+        )
+
+        await _stream_llm_call(model, "test", StreamingHandler(), stop=None)
+
+        assert tool_calls_var.get() is None
+
+    @pytest.mark.asyncio
+    async def test_clears_reasoning_var_when_none(self):
+        reasoning_trace_var.set("stale reasoning")
+
+        model = _make_chunk_model(
+            [
+                LLMResponseChunk(delta_content="no reasoning", finish_reason="stop"),
+            ]
+        )
+
+        await _stream_llm_call(model, "test", StreamingHandler(), stop=None)
+
+        assert reasoning_trace_var.get() is None
+
+    @pytest.mark.asyncio
+    async def test_provider_metadata_stored_flat(self):
+        model = _make_chunk_model(
+            [
+                LLMResponseChunk(
+                    delta_content="hi",
+                    provider_metadata={"system_fingerprint": "fp_abc"},
+                    finish_reason="stop",
+                ),
+            ]
+        )
+
+        await _stream_llm_call(model, "test", StreamingHandler(), stop=None)
+
+        metadata = llm_response_metadata_var.get()
+        assert metadata == {"system_fingerprint": "fp_abc"}
+
+    @pytest.mark.asyncio
+    async def test_clears_metadata_var_when_none(self):
+        llm_response_metadata_var.set({"stale": True})
+
+        model = _make_chunk_model(
+            [
+                LLMResponseChunk(delta_content="no metadata", finish_reason="stop"),
+            ]
+        )
+
+        await _stream_llm_call(model, "test", StreamingHandler(), stop=None)
+
+        assert llm_response_metadata_var.get() is None