diff --git a/apps/agent/src/model.py b/apps/agent/src/model.py index 33725ab..c1e7062 100644 --- a/apps/agent/src/model.py +++ b/apps/agent/src/model.py @@ -6,6 +6,7 @@ from langchain.agents.middleware import AgentMiddleware from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult from langchain_core.messages import AIMessageChunk +from langchain_core.runnables import RunnableBinding from pydantic import Field from langchain_anthropic import ChatAnthropic from langchain_core.language_models.chat_models import BaseChatModel @@ -146,31 +147,48 @@ def _target(self, kwargs): tool_options = options.pop("_credential_tool_options", {}) if tools is not None: model = model.bind_tools(tools, **tool_options) + # Call the provider's private methods so only this wrapper's run reports + # events. A public nested call inherits the graph's callbacks and streams + # every token and tool-call delta to AG-UI a second time. + if isinstance(model, RunnableBinding): + options = {**model.kwargs, **options} + model = model.bound return model, options def _generate(self, messages, stop=None, run_manager=None, **kwargs): with provider_error_boundary(): model, options = self._target(kwargs) - message = model.invoke(messages, stop=stop, **options) - return ChatResult(generations=[ChatGeneration(message=message)]) + return model._generate(messages, stop=stop, **options) async def _agenerate(self, messages, stop=None, run_manager=None, **kwargs): with provider_error_boundary(): model, options = self._target(kwargs) - message = await model.ainvoke(messages, stop=stop, **options) - return ChatResult(generations=[ChatGeneration(message=message)]) + return await model._agenerate(messages, stop=stop, **options) def _stream(self, messages, stop=None, run_manager=None, **kwargs): with provider_error_boundary(): model, options = self._target(kwargs) - for chunk in model.stream(messages, stop=stop, **options): - yield ChatGenerationChunk(message=as_chunk(chunk)) + if not streams(model): + yield as_generation_chunk(model._generate(messages, stop=stop, **options)) + return + yield from model._stream(messages, stop=stop, **options) async def _astream(self, messages, stop=None, run_manager=None, **kwargs): with provider_error_boundary(): model, options = self._target(kwargs) - async for chunk in model.astream(messages, stop=stop, **options): - yield ChatGenerationChunk(message=as_chunk(chunk)) + if not streams(model): + yield as_generation_chunk(await model._agenerate(messages, stop=stop, **options)) + return + async for chunk in model._astream(messages, stop=stop, **options): + yield chunk + + +def streams(model): + return type(model)._stream is not BaseChatModel._stream or type(model)._astream is not BaseChatModel._astream + + +def as_generation_chunk(result): + return ChatGenerationChunk(message=as_chunk(result.generations[0].message)) def as_chunk(message): diff --git a/apps/agent/tests/test_model.py b/apps/agent/tests/test_model.py index b9fc29f..7b860a6 100644 --- a/apps/agent/tests/test_model.py +++ b/apps/agent/tests/test_model.py @@ -84,3 +84,32 @@ def test_missing_provider_key_has_actionable_error(monkeypatch, model, key, valu monkeypatch.setenv(key, value) with pytest.raises(ValueError, match=key): build_model() + + +def test_credential_scoped_model_streams_each_chunk_once(): + import asyncio + from langchain_core.language_models.fake_chat_models import GenericFakeChatModel + from langchain_core.messages import AIMessage + from langchain_core.runnables import RunnableLambda + from src.model import CredentialScopedModel + + class ToolFake(GenericFakeChatModel): + def bind_tools(self, tools, **kwargs): + return self.bind(tools=tools) + + inner = ToolFake(messages=iter([AIMessage(content="a b c")])) + # Bound tools wrap the provider in a RunnableBinding, as in the real agent. + wrapped = CredentialScopedModel(fallback=inner).bind_tools([{"name": "t"}]) + + # Like a LangGraph node: the parent run's callbacks are inherited by calls inside it. + async def node(_): + return await wrapped.ainvoke("hi") + + async def stream_events(): + return [ + e async for e in RunnableLambda(node).astream_events("x", version="v2") + if e["event"] == "on_chat_model_stream" + ] + + tokens = [e["data"]["chunk"].content for e in asyncio.run(stream_events())] + assert "".join(tokens) == "a b c"