diff --git a/packages/client/src/launchdarkly_ai_server/agui.py b/packages/client/src/launchdarkly_ai_server/agui.py new file mode 100644 index 00000000..6b302e65 --- /dev/null +++ b/packages/client/src/launchdarkly_ai_server/agui.py @@ -0,0 +1,656 @@ +"""Opt-in AG-UI reader for one AI Config turn. + +Handlers may yield ``chunk`` or the text and tool subset below, then one +``done``. :func:`stream` is the opt-in entry: it brackets that subset with +``RUN_STARTED`` and ``RUN_FINISHED``. ``config().stream()`` still yields +``chunk`` then ``done``. The SDK runs the tool loop. This stream only reports it. + +Contract for one run: + +* ``RUN_STARTED`` is first and ``RUN_FINISHED`` is last. A failure or + cancellation raises and does not yield ``RUN_FINISHED`` or ``RUN_ERROR``. +* Message and tool-result ids are unique across runs. A chat UI that keeps + messages by id will not append a later answer onto an earlier one. +* ``TOOL_CALL_RESULT.content`` is always a string. Argument deltas and result + text are omitted (sent as ``""``) when ``capture_content`` is off. A failed + call still sets ``metadata.launchdarkly.toolFailed``. +* ``RUN_FINISHED`` is not sent while a tool call is open. An unfinished call + is closed first. +* Tool-result text is the string the model was given, not a second encoding. +* Optional ``thread_id`` and ``run_id`` are echoed on ``RUN_STARTED`` and + ``RUN_FINISHED``. When omitted, this reader mints them. +""" + +from __future__ import annotations + +import json +import math +import uuid +from collections.abc import AsyncGenerator +from contextlib import aclosing +from typing import Any + +from .conversation import bind_conversation_id +from .sdk_usage import report_usage +from .types import JudgeResult, LDContext + +PROTOCOL_VERSION = "1.0" + +TEXT_MESSAGE_START = "TEXT_MESSAGE_START" +TEXT_MESSAGE_CONTENT = "TEXT_MESSAGE_CONTENT" +TEXT_MESSAGE_END = "TEXT_MESSAGE_END" +TOOL_CALL_START = "TOOL_CALL_START" +TOOL_CALL_ARGS = "TOOL_CALL_ARGS" +TOOL_CALL_END = "TOOL_CALL_END" +TOOL_CALL_RESULT = "TOOL_CALL_RESULT" +RUN_STARTED = "RUN_STARTED" +RUN_FINISHED = "RUN_FINISHED" + +TEXT_TYPES = frozenset({TEXT_MESSAGE_START, TEXT_MESSAGE_CONTENT, TEXT_MESSAGE_END}) +TOOL_TYPES = frozenset( + {TOOL_CALL_START, TOOL_CALL_ARGS, TOOL_CALL_END, TOOL_CALL_RESULT} +) +HANDLER_RICH_TYPES = TEXT_TYPES | TOOL_TYPES + +_FAILURE_CONTENT = "An error occurred." + + +def argument_text(value: Any) -> str | None: + """The argument text for one ``TOOL_CALL_ARGS`` delta. + + A string is that string. Any other JSON value is compact JSON. ``None`` and + a value that is not JSON under :func:`json.dumps` with ``allow_nan=False`` + produce no delta. + """ + if value is None: + return None + if isinstance(value, str): + return value or None + if isinstance(value, (bytes, bytearray, memoryview)): + return None + if isinstance(value, float) and (math.isnan(value) or math.isinf(value)): + return None + try: + return json.dumps( + value, ensure_ascii=False, separators=(",", ":"), allow_nan=False + ) + except (TypeError, ValueError): + return None + + +def safe_tool_content(value: Any) -> tuple[str, bool]: + """``(content, failed)`` for a tool result. + + A string is not encoded again. Any other JSON value is compact JSON. + A value that is not a string and not serializable under that rule is a + failed result: the caller masks it. The model-facing value is unchanged. + """ + if isinstance(value, str): + return value, False + if isinstance(value, (bytes, bytearray, memoryview)): + return _FAILURE_CONTENT, True + if isinstance(value, float) and (math.isnan(value) or math.isinf(value)): + return _FAILURE_CONTENT, True + try: + return ( + json.dumps( + value, ensure_ascii=False, separators=(",", ":"), allow_nan=False + ), + False, + ) + except (TypeError, ValueError): + return _FAILURE_CONTENT, True + + +class TurnEmitter: + """Yields the text and tool subset for one handler run. + + ``capture_content`` gates argument deltas and result content. The tool name + and call id are always present on ``TOOL_CALL_START``. A closed call is not + reopened. + """ + + def __init__(self, *, capture_content: bool) -> None: + self.capture_content = capture_content + self._open: str | None = None + self._started: set[str] = set() + self._args: set[str] = set() + self._arg_text: dict[str, str] = {} + self._ended: set[str] = set() + self._resulted: set[str] = set() + + def _mint(self, prefix: str) -> str: + return f"{prefix}-{uuid.uuid4()}" + + def has_argument_deltas(self, tool_call_id: str) -> bool: + """True once a streamed argument slice was recorded for this call.""" + return tool_call_id in self._args + + def text(self, delta: str) -> list[dict[str, Any]]: + """Append assistant text. The first slice opens the message.""" + if not isinstance(delta, str) or not delta: + return [] + events: list[dict[str, Any]] = [] + if self._open is None: + self._open = self._mint("msg") + events.append({"type": TEXT_MESSAGE_START, "messageId": self._open}) + events.append( + { + "type": TEXT_MESSAGE_CONTENT, + "messageId": self._open, + "delta": delta, + } + ) + return events + + def end_text(self) -> list[dict[str, Any]]: + if self._open is None: + return [] + event = {"type": TEXT_MESSAGE_END, "messageId": self._open} + self._open = None + return [event] + + def tool_start( + self, tool_call_id: str, tool_call_name: str + ) -> list[dict[str, Any]]: + if not tool_call_id or not tool_call_name or tool_call_id in self._started: + return [] + self._started.add(tool_call_id) + event: dict[str, Any] = { + "type": TOOL_CALL_START, + "toolCallId": tool_call_id, + "toolCallName": tool_call_name, + } + if self._open is not None: + event["parentMessageId"] = self._open + return [event] + + def tool_args(self, tool_call_id: str, delta: str) -> list[dict[str, Any]]: + if ( + not self.capture_content + or not delta + or tool_call_id not in self._started + or tool_call_id in self._ended + ): + return [] + self._args.add(tool_call_id) + self._arg_text[tool_call_id] = self._arg_text.get(tool_call_id, "") + delta + return [{"type": TOOL_CALL_ARGS, "toolCallId": tool_call_id, "delta": delta}] + + def tool_end( + self, tool_call_id: str, *, arguments: Any = None + ) -> list[dict[str, Any]]: + """Close the argument text. + + When nothing was streamed, one delta carries the complete arguments. + When the complete text starts with what was already streamed, the + missing suffix is appended. Otherwise the streamed text wins: a + differently encoded complete value is not concatenated onto it. + """ + if tool_call_id not in self._started or tool_call_id in self._ended: + return [] + events: list[dict[str, Any]] = [] + if self.capture_content and arguments is not None: + text = argument_text(arguments) + have = self._arg_text.get(tool_call_id, "") + extra = text[len(have) :] if text and text.startswith(have) else "" + if extra: + self._args.add(tool_call_id) + self._arg_text[tool_call_id] = have + extra + events.append( + {"type": TOOL_CALL_ARGS, "toolCallId": tool_call_id, "delta": extra} + ) + self._ended.add(tool_call_id) + events.append({"type": TOOL_CALL_END, "toolCallId": tool_call_id}) + return events + + def tool_result( + self, tool_call_id: str, value: Any = None, *, failed: bool = False + ) -> list[dict[str, Any]]: + """Report a finished call, including failure. Does not run the tool.""" + if tool_call_id not in self._started or tool_call_id in self._resulted: + return [] + events = self.tool_end(tool_call_id) if tool_call_id not in self._ended else [] + self._resulted.add(tool_call_id) + event: dict[str, Any] = { + "type": TOOL_CALL_RESULT, + "messageId": self._mint("toolmsg"), + "toolCallId": tool_call_id, + } + if failed: + event["metadata"] = {"launchdarkly": {"toolFailed": True}} + event["content"] = _FAILURE_CONTENT if self.capture_content else "" + elif self.capture_content: + content, bad = safe_tool_content(value) + if bad: + event["content"] = _FAILURE_CONTENT + event["metadata"] = {"launchdarkly": {"toolFailed": True}} + else: + event["content"] = content + else: + event["content"] = "" + events.append(event) + return events + + def fail_open(self) -> list[dict[str, Any]]: + """Close every call that started and did not finish, as a masked failure.""" + events: list[dict[str, Any]] = [] + for call_id in list(self._started): + if call_id not in self._resulted: + events.extend(self.tool_result(call_id, failed=True)) + return events + + +class StableProjector: + """Projects a handler stream onto ``chunk`` then one internal ``done``. + + ``TEXT_MESSAGE_CONTENT`` becomes a chunk. Tool events and any other type are + dropped and do not close the turn. Only the handler's ``done`` does. + """ + + def __init__(self) -> None: + self.saw_text = False + self.done: dict[str, Any] | None = None + + def feed(self, event: dict[str, Any]) -> list[dict[str, Any]]: + typ = event.get("type") + if typ in TEXT_TYPES: + self.saw_text = True + if typ == "done": + self.done = event + return [] + if typ == "chunk": + if self.saw_text: + return [] + text = event.get("text", "") + if not isinstance(text, str): + return [] + return [{"type": "chunk", "text": text}] + if typ == TEXT_MESSAGE_CONTENT: + delta = event.get("delta") or "" + if not isinstance(delta, str) or not delta: + return [] + return [{"type": "chunk", "text": delta}] + return [] + + +def _token_usage( + parsed: dict[str, Any], *, provider: str, model: str +) -> dict[str, Any]: + usage: dict[str, Any] = { + "inputTokens": int(parsed.get("input") or 0), + "outputTokens": int(parsed.get("output") or 0), + "totalTokens": int(parsed.get("total") or 0), + } + details = parsed.get("input_details") + if isinstance(details, dict): + usage["cachedInputTokens"] = int(details.get("cache_read") or 0) + usage["cacheWriteInputTokens"] = int(details.get("cache_creation") or 0) + if provider: + usage["provider"] = provider + if model: + usage["model"] = model + return usage + + +def _judge_payload(results: dict[str, JudgeResult]) -> dict[str, Any] | None: + if not results: + return None + payload: dict[str, Any] = {} + for key, result in results.items(): + usage: dict[str, Any] = { + "input": result.usage.input, + "output": result.usage.output, + "total": result.usage.total, + } + details = result.usage.input_details + if details is not None: + usage["input_details"] = { + "uncached": details.uncached, + "cache_read": details.cache_read, + "cache_creation": details.cache_creation, + } + payload[key] = { + "usage": usage, + "response": result.response, + "score": result.score, + } + return payload + + +class RunReader: + """Maps one handler stream onto one AG-UI run. + + Legacy ``chunk`` events become one text message when the handler yielded no + ``TEXT_MESSAGE_*`` event. ``RUN_STARTED`` and ``RUN_FINISHED`` are minted + here. The handler's ``done`` is consumed, not forwarded. + """ + + def __init__( + self, *, thread_id: str | None = None, run_id: str | None = None + ) -> None: + self.thread_id = thread_id or str(uuid.uuid4()) + self.run_id = run_id or str(uuid.uuid4()) + self._mode: str | None = None + self._open: str | None = None + self._deltas: list[str] = [] + self._started: set[str] = set() + self._ended: set[str] = set() + self._resulted: set[str] = set() + + def start(self) -> dict[str, Any]: + return { + "type": RUN_STARTED, + "threadId": self.thread_id, + "runId": self.run_id, + "protocolVersion": PROTOCOL_VERSION, + } + + def _mint(self, prefix: str) -> str: + return f"{prefix}-{uuid.uuid4()}" + + def feed(self, event: dict[str, Any]) -> list[dict[str, Any]]: + typ = event.get("type") + if typ == "chunk": + return self._chunk(event.get("text") or "") + if typ == TEXT_MESSAGE_START: + return self._text_start(event.get("messageId")) + if typ == TEXT_MESSAGE_CONTENT: + return self._text_content(event.get("messageId"), event.get("delta")) + if typ == TEXT_MESSAGE_END: + return self._text_end(event.get("messageId")) + if typ == TOOL_CALL_START: + return self._tool_start(event) + if typ == TOOL_CALL_ARGS: + return self._tool_args(event) + if typ == TOOL_CALL_END: + return self._tool_end(event) + if typ == TOOL_CALL_RESULT: + return self._tool_result(event) + return [] + + def finish( + self, + done: dict[str, Any], + judge_results: dict[str, JudgeResult] | None, + *, + provider: str, + model: str, + ) -> list[dict[str, Any]]: + events: list[dict[str, Any]] = [] + if self._open is not None: + events.append({"type": TEXT_MESSAGE_END, "messageId": self._open}) + self._open = None + for call_id in list(self._started): + if call_id not in self._ended: + self._ended.add(call_id) + events.append({"type": TOOL_CALL_END, "toolCallId": call_id}) + response = done.get("response") + if not isinstance(response, str): + response = "".join(self._deltas) + finished: dict[str, Any] = { + "type": RUN_FINISHED, + "threadId": self.thread_id, + "runId": self.run_id, + } + if response: + finished["result"] = response + parsed = done.get("usage") + if done.get("_had_usage") and isinstance(parsed, dict): + finished["usage"] = [_token_usage(parsed, provider=provider, model=model)] + judges = _judge_payload(judge_results or {}) + if judges is not None: + finished["metadata"] = {"launchdarkly": {"judgeResults": judges}} + events.append(finished) + return events + + def _chunk(self, text: str) -> list[dict[str, Any]]: + if self._mode == "message" or not text: + return [] + self._mode = "chunk" + events: list[dict[str, Any]] = [] + if self._open is None: + self._open = self._mint("msg") + events.append({"type": TEXT_MESSAGE_START, "messageId": self._open}) + events.append( + {"type": TEXT_MESSAGE_CONTENT, "messageId": self._open, "delta": text} + ) + self._deltas.append(text) + return events + + def _close_open(self, message_id: str) -> list[dict[str, Any]]: + if self._open is None or self._open == message_id: + return [] + event = {"type": TEXT_MESSAGE_END, "messageId": self._open} + self._open = None + return [event] + + def _text_start(self, message_id: Any) -> list[dict[str, Any]]: + if not isinstance(message_id, str) or not message_id: + return [] + self._mode = "message" + if self._open == message_id: + return [] + events = self._close_open(message_id) + self._open = message_id + events.append({"type": TEXT_MESSAGE_START, "messageId": message_id}) + return events + + def _text_content(self, message_id: Any, delta: Any) -> list[dict[str, Any]]: + if not isinstance(delta, str): + return [] + self._mode = "message" + mid = message_id if isinstance(message_id, str) and message_id else None + events: list[dict[str, Any]] = [] + if self._open is None: + if mid is None: + mid = self._mint("msg") + self._open = mid + events.append({"type": TEXT_MESSAGE_START, "messageId": mid}) + elif mid is not None and mid != self._open: + events.extend(self._close_open(mid)) + self._open = mid + events.append({"type": TEXT_MESSAGE_START, "messageId": mid}) + events.append( + { + "type": TEXT_MESSAGE_CONTENT, + "messageId": self._open, + "delta": delta, + } + ) + self._deltas.append(delta) + return events + + def _text_end(self, message_id: Any) -> list[dict[str, Any]]: + if self._open is None: + return [] + if isinstance(message_id, str) and message_id and message_id != self._open: + return [] + event = {"type": TEXT_MESSAGE_END, "messageId": self._open} + self._open = None + return [event] + + def _tool_start(self, event: dict[str, Any]) -> list[dict[str, Any]]: + call_id = event.get("toolCallId") + name = event.get("toolCallName") + if ( + not isinstance(call_id, str) + or not call_id + or not isinstance(name, str) + or not name + or call_id in self._started + ): + return [] + self._started.add(call_id) + started: dict[str, Any] = { + "type": TOOL_CALL_START, + "toolCallId": call_id, + "toolCallName": name, + } + parent = event.get("parentMessageId") + if isinstance(parent, str) and parent: + started["parentMessageId"] = parent + return [started] + + def _tool_args(self, event: dict[str, Any]) -> list[dict[str, Any]]: + call_id = event.get("toolCallId") + delta = event.get("delta") + if ( + not isinstance(call_id, str) + or call_id not in self._started + or call_id in self._ended + or not isinstance(delta, str) + ): + return [] + return [{"type": TOOL_CALL_ARGS, "toolCallId": call_id, "delta": delta}] + + def _tool_end(self, event: dict[str, Any]) -> list[dict[str, Any]]: + call_id = event.get("toolCallId") + if ( + not isinstance(call_id, str) + or call_id not in self._started + or call_id in self._ended + ): + return [] + self._ended.add(call_id) + return [{"type": TOOL_CALL_END, "toolCallId": call_id}] + + def _tool_result(self, event: dict[str, Any]) -> list[dict[str, Any]]: + call_id = event.get("toolCallId") + if ( + not isinstance(call_id, str) + or call_id not in self._started + or call_id in self._resulted + ): + return [] + events: list[dict[str, Any]] = [] + if call_id not in self._ended: + self._ended.add(call_id) + events.append({"type": TOOL_CALL_END, "toolCallId": call_id}) + self._resulted.add(call_id) + message_id = event.get("messageId") + if not isinstance(message_id, str) or not message_id: + message_id = self._mint("toolmsg") + result: dict[str, Any] = { + "type": TOOL_CALL_RESULT, + "messageId": message_id, + "toolCallId": call_id, + "content": event.get("content") + if isinstance(event.get("content"), str) + else "", + } + metadata = event.get("metadata") + if isinstance(metadata, dict): + launchdarkly = metadata.get("launchdarkly") + if ( + isinstance(launchdarkly, dict) + and launchdarkly.get("toolFailed") is True + ): + result["metadata"] = {"launchdarkly": {"toolFailed": True}} + events.append(result) + return events + + +def stream( + instance: Any, + user_input: str | None, + context: LDContext, + *, + variables: dict[str, Any] | None = None, + history: list[dict[str, Any]] | None = None, + thread_id: str | None = None, + run_id: str | None = None, +) -> AsyncGenerator[dict[str, Any], None]: + """Stream one config turn as AG-UI events. + + Same arguments as ``ConfigInstance.stream``, plus optional ``thread_id`` + and ``run_id`` echoed from the caller's AG-UI run. ``RUN_STARTED`` is first + and ``RUN_FINISHED`` is last. A failure, including cancellation, raises and + does not yield ``RUN_FINISHED``. + """ + report_usage("client.agui.stream") + return bind_conversation_id( + _stream_events( + instance, + user_input, + context, + variables, + history, + thread_id=thread_id, + run_id=run_id, + ) + ) + + +async def judge_completed_stream( + instance: Any, + sink: dict[str, Any], + done: dict[str, Any], + user_input: str | None, + context: LDContext, +) -> dict[str, JudgeResult]: + """Run the turn's judges once the handler has finished. + + A non-string response is not sent to the judge. Both ``config().stream()`` + and :func:`stream` use this, so the two copies cannot drift. + """ + from .judges import run_judges + + if getattr(instance, "_skip_judges", False) or not sink: + return {} + response = done.get("response", "") + if not isinstance(response, str): + response = "" + return await run_judges( + config=sink["config"], + user_context=context, + handler=sink["handler"], + handlers=sink["handlers"], + user_input=user_input, + trajectory=done.get("trajectory", ""), + llm_response=response, + base_track_data=done.get("track_data", {}), + tool_handlers=sink["tools"], + ) + + +async def _stream_events( + instance: Any, + user_input: str | None, + context: LDContext, + variables: dict[str, Any] | None, + history: list[dict[str, Any]] | None, + *, + thread_id: str | None = None, + run_id: str | None = None, +) -> AsyncGenerator[dict[str, Any], None]: + sink: dict[str, Any] = {} + reader = RunReader(thread_id=thread_id, run_id=run_id) + started = False + done: dict[str, Any] | None = None + async with aclosing( + instance._execute_stream(user_input, context, variables, history, sink) + ) as source: + async for event in source: + if not started: + yield reader.start() + started = True + if event.get("type") == "done": + done = event + continue + for item in reader.feed(event): + yield item + if not started: + yield reader.start() + if done is None: + done = {"response": "", "usage": {}, "_had_usage": False} + judge_results = await judge_completed_stream( + instance, sink, done, user_input, context + ) + for item in reader.finish( + done, + judge_results, + provider=str(sink.get("provider") or ""), + model=str(sink.get("model") or ""), + ): + yield item diff --git a/packages/client/src/launchdarkly_ai_server/client.py b/packages/client/src/launchdarkly_ai_server/client.py index 0f77383f..da73be45 100644 --- a/packages/client/src/launchdarkly_ai_server/client.py +++ b/packages/client/src/launchdarkly_ai_server/client.py @@ -2,8 +2,10 @@ import json from collections.abc import AsyncGenerator, Callable +from contextlib import aclosing from typing import Any +from .agui import StableProjector, judge_completed_stream from .conversation import bind_conversation_id from .judges import _build_judge_tasks, run_judges from .lifecycle import extract_variation @@ -192,13 +194,20 @@ def _stream( self._stream_events(user_input, context, variables, history) ) - async def _stream_events( + async def _execute_stream( self, user_input: str | None, context: LDContext, - variables: dict[str, Any] | None = None, - history: list[dict[str, Any]] | None = None, - ) -> AsyncGenerator[StreamEvent, None]: + variables: dict[str, Any] | None, + history: list[dict[str, Any]] | None, + sink: dict[str, Any], + ) -> AsyncGenerator[dict[str, Any], None]: + """Yield handler events and the internal ``done``. + + Fills *sink* with the resolved config and handler before the first event, + so both ``config().stream()`` and the AG-UI reader run judges from one + pass. + """ resolved_handler_list = resolve_handlers( self._registry, self._normalize_handlers() ) @@ -211,41 +220,54 @@ async def _stream_events( if not resolved_handler_list: raise ValueError("No handlers provided to config()") handler = select_handler(config, meta, resolved_handler_list) - - done_event: dict[str, Any] | None = None - - async for event in execute_and_stream( - config_key=self._key, - config=config, - meta=meta, - user_context=context, - handler=handler, - user_input=user_input, - tool_handlers=resolved_tools, - variables=variables, - history=history, - ): - if event.get("type") == "chunk": + provider = "" + model = "" + if isinstance(config, dict): + provider = str((config.get("provider") or {}).get("name") or "") + model = str((config.get("model") or {}).get("name") or "") + sink["config"] = config + sink["handler"] = handler + sink["handlers"] = resolved_handler_list + sink["tools"] = resolved_tools + sink["provider"] = provider + sink["model"] = model + + async with aclosing( + execute_and_stream( + config_key=self._key, + config=config, + meta=meta, + user_context=context, + handler=handler, + user_input=user_input, + tool_handlers=resolved_tools, + variables=variables, + history=history, + ) + ) as source: + async for event in source: yield event - else: - done_event = event + + async def _stream_events( + self, + user_input: str | None, + context: LDContext, + variables: dict[str, Any] | None = None, + history: list[dict[str, Any]] | None = None, + ) -> AsyncGenerator[StreamEvent, None]: + sink: dict[str, Any] = {} + projector = StableProjector() + async with aclosing( + self._execute_stream(user_input, context, variables, history, sink) + ) as source: + async for event in source: + for stable in projector.feed(event): + yield stable + done_event = projector.done if done_event: - track_data = done_event.get("track_data", {}) - judge_results = ( - {} - if self._skip_judges - else await run_judges( - config=config, - user_context=context, - handler=handler, - handlers=resolved_handler_list, - user_input=user_input, - trajectory=done_event.get("trajectory", ""), - llm_response=done_event.get("response", ""), - base_track_data=track_data, - tool_handlers=resolved_tools, - ) + judge_results = await judge_completed_stream( + self, sink, done_event, user_input, context ) yield { "type": "done", diff --git a/packages/client/src/launchdarkly_ai_server/graph.py b/packages/client/src/launchdarkly_ai_server/graph.py index beb39d97..63e68736 100644 --- a/packages/client/src/launchdarkly_ai_server/graph.py +++ b/packages/client/src/launchdarkly_ai_server/graph.py @@ -10,6 +10,7 @@ from collections.abc import AsyncGenerator, Callable from typing import Any +from .agui import StableProjector from .conversation import bind_conversation_id, bind_span_context from .registry import resolve_handlers, resolve_tools from .types import ( @@ -509,6 +510,7 @@ async def stream_node( "graphKey": key, } + projector = StableProjector() async for event in execute_and_stream( config_key=node.key, config=node.config, @@ -521,17 +523,17 @@ async def stream_node( graph_key=key, history=opts.get("history"), ): - if event.get("type") == "chunk": + for stable in projector.feed(event): yield { "type": "chunk", - "text": event["text"], + "text": stable["text"], "nodeKey": node.key, } - else: - response = event.get("response", "") - usage = event.get("usage") or usage - track_data = event.get("track_data") or track_data - trajectory = event.get("trajectory") or trajectory + if projector.done: + response = projector.done.get("response", "") + usage = projector.done.get("usage") or usage + track_data = projector.done.get("track_data") or track_data + trajectory = projector.done.get("trajectory") or trajectory judge_results = await run_judges( config=node.config, @@ -648,6 +650,7 @@ async def stream_route( merged_tool_handlers = {**(tool_handlers or {}), **handoff_handlers} + projector = StableProjector() async for event in execute_and_stream( config_key=node.key, config=routed_config, @@ -660,17 +663,17 @@ async def stream_route( graph_key=key, history=opts.get("history"), ): - if event.get("type") == "chunk": + for stable in projector.feed(event): yield { "type": "chunk", - "text": event["text"], + "text": stable["text"], "nodeKey": node.key, } - else: - response = event.get("response", "") - usage = event.get("usage") or usage - track_data = event.get("track_data") or track_data - trajectory = event.get("trajectory") or trajectory + if projector.done: + response = projector.done.get("response", "") + usage = projector.done.get("usage") or usage + track_data = projector.done.get("track_data") or track_data + trajectory = projector.done.get("trajectory") or trajectory # Judge against the node's original config, not the routing-augmented one. judge_results = await run_judges( diff --git a/packages/client/src/launchdarkly_ai_server/tracking.py b/packages/client/src/launchdarkly_ai_server/tracking.py index 6e7a5638..acc4f8b5 100644 --- a/packages/client/src/launchdarkly_ai_server/tracking.py +++ b/packages/client/src/launchdarkly_ai_server/tracking.py @@ -4,9 +4,11 @@ import os import time import uuid -from collections.abc import AsyncGenerator, Callable +from collections.abc import AsyncGenerator, AsyncIterator, Callable +from contextlib import asynccontextmanager from typing import Any +from .agui import HANDLER_RICH_TYPES, TEXT_MESSAGE_CONTENT, TEXT_TYPES from .trajectory import HANDOFF_TOOL_PREFIX, TrajectoryRecorder, render_trajectory from .types import ( NATIVE_TOOL_KEY, @@ -21,6 +23,22 @@ from .utils import model_stamps_from_meta, parse_usage, to_ld_context +@asynccontextmanager +async def _closing_stream(stream: Any) -> AsyncIterator[Any]: + """Close a handler stream when it implements ``aclose``. + + ``contextlib.aclosing`` always awaits ``aclose``. A handler may return a + plain async iterator, which has no such method, and that call then raises + ``AttributeError`` after the loop has already succeeded. + """ + try: + yield stream + finally: + close = getattr(stream, "aclose", None) + if close is not None: + await close() + + def _try_get_environment_id() -> str | None: """ Returns the LaunchDarkly environment MongoDB ObjectId used to set @@ -250,9 +268,10 @@ async def execute_and_stream( history: list[dict[str, Any]] | None = None, ) -> AsyncGenerator[ExecuteStreamEvent, None]: """ - Streaming counterpart to ``execute_and_track``. Yields ``chunk`` events - from the handler, then a single ``done`` event once the stream completes - and LD telemetry has been flushed. + Streaming counterpart to ``execute_and_track``. Yields ``chunk`` events, + the AG-UI text and tool subset when the handler emits it, then a single + ``done`` event once the stream completes and LD telemetry has been flushed. + ``config().stream()`` projects that subset back to ``chunk`` then ``done``. Falls back to calling the blocking handler when the handler has no ``stream`` implementation. @@ -297,27 +316,51 @@ async def execute_and_stream( start_time = time.monotonic() full_text = "" raw_usage: dict[str, Any] = {} + had_usage = False + saw_text = False try: if handler.has_stream: stream_gen = await handler.stream( config, user_input, tracked_tool_handlers, merged_variables, history ) - async for event in stream_gen: - if event.get("type") == "chunk": - full_text += event.get("text", "") - yield {"type": "chunk", "text": event["text"]} - elif event.get("type") == "done": - if event.get("output") is not None: - full_text = event["output"] - raw_usage = event.get("usage", {}) + async with _closing_stream(stream_gen) as agen: + async for event in agen: + typ = event.get("type") + if typ in TEXT_TYPES: + saw_text = True + if typ == "chunk": + if saw_text: + continue + text = event.get("text", "") + if not isinstance(text, str): + continue + full_text += text + yield {"type": "chunk", "text": text} + elif typ == TEXT_MESSAGE_CONTENT: + delta = event.get("delta") or "" + if isinstance(delta, str): + full_text += delta + yield event + elif typ in HANDLER_RICH_TYPES: + yield event + elif typ == "done": + if event.get("output") is not None: + full_text = event["output"] + usage_value = event.get("usage") + if isinstance(usage_value, dict) and usage_value: + raw_usage = usage_value + had_usage = True else: result = await handler( config, user_input, tracked_tool_handlers, merged_variables, history ) out = result.get("output") full_text = str(out) if out is not None else "" - raw_usage = result.get("usage") or {} + usage_value = result.get("usage") + if isinstance(usage_value, dict) and usage_value: + raw_usage = usage_value + had_usage = True if full_text: yield {"type": "chunk", "text": full_text} except Exception: @@ -348,4 +391,5 @@ async def execute_and_stream( observable_tools=exposed_tool_keys, omitted=recorder.omitted, ), + "_had_usage": had_usage, } diff --git a/packages/client/src/launchdarkly_ai_server/types.py b/packages/client/src/launchdarkly_ai_server/types.py index 1a7f6819..75c8008a 100644 --- a/packages/client/src/launchdarkly_ai_server/types.py +++ b/packages/client/src/launchdarkly_ai_server/types.py @@ -114,7 +114,12 @@ class Message: class HandlerStreamEvent: - """Union of chunk / done events yielded by a handler's stream method.""" + """Events a handler stream yields. + + ``chunk`` and ``done`` remain valid. A handler may instead yield the AG-UI + text and tool subset (``TEXT_MESSAGE_*``, ``TOOL_CALL_*``) and one ``done``. + ``config().stream()`` still yields ``chunk`` then ``done``. + """ pass diff --git a/packages/client/tests/test_agui.py b/packages/client/tests/test_agui.py new file mode 100644 index 00000000..915e6944 --- /dev/null +++ b/packages/client/tests/test_agui.py @@ -0,0 +1,552 @@ +"""The opt-in AG-UI reader, and the stable chunk/done projection.""" + +from __future__ import annotations + +from collections.abc import AsyncGenerator +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import launchdarkly_ai_server.lifecycle as lifecycle_module +from launchdarkly_ai_server import ProviderHandler, config +from launchdarkly_ai_server.agui import ( + RUN_FINISHED, + RUN_STARTED, + TEXT_MESSAGE_CONTENT, + TEXT_MESSAGE_END, + TEXT_MESSAGE_START, + TOOL_CALL_ARGS, + TOOL_CALL_END, + TOOL_CALL_RESULT, + TOOL_CALL_START, + RunReader, + StableProjector, + TurnEmitter, + stream, +) +from launchdarkly_ai_server.types import InputTokenDetails, JudgeResult, UsageDict + +CONTEXT = {"kind": "user", "key": "u1"} + + +def _client() -> MagicMock: + client = MagicMock() + client.track = MagicMock() + client.flush = AsyncMock() + client.close = AsyncMock() + client.variation = AsyncMock( + return_value={ + "model": {"name": "gpt-4"}, + "provider": {"name": "TestProvider"}, + "instructions": "Be helpful.", + "_ldMeta": { + "enabled": True, + "variationKey": "v1", + "version": 1, + "mode": "messages", + }, + } + ) + return client + + +@pytest.fixture +def mock_ld_client() -> MagicMock: + client = _client() + lifecycle_module._set_client_for_testing(client) + yield client + lifecycle_module._reset_for_testing() + + +def _handler( + events: list[dict[str, Any]], + *, + capture_content: bool = False, +) -> ProviderHandler: + async def fn(cfg, user_input, tool_handlers, variables, history=None) -> dict: # type: ignore[override] + return {"output": "", "usage": {}} + + async def stream_fn( + cfg, user_input, tool_handlers, variables, history=None + ) -> AsyncGenerator[dict[str, Any], None]: + for event in events: + yield event + + return ProviderHandler( + fn=fn, + provides_for=("TestProvider", "messages"), + stream_fn=stream_fn, + capture_content=capture_content, + ) + + +def _types(events: list[dict[str, Any]]) -> list[str]: + return [event["type"] for event in events] + + +class TestTurnEmitter: + def test_capture_on_includes_args_and_result(self) -> None: + emitter = TurnEmitter(capture_content=True) + events: list[dict[str, Any]] = [] + events.extend(emitter.text("Looking up")) + events.extend(emitter.tool_start("a", "search")) + events.extend(emitter.tool_args("a", '{"q":')) + events.extend(emitter.tool_end("a")) + events.extend(emitter.end_text()) + events.extend(emitter.tool_result("a", {"ok": True})) + assert _types(events) == [ + TEXT_MESSAGE_START, + TEXT_MESSAGE_CONTENT, + TOOL_CALL_START, + TOOL_CALL_ARGS, + TOOL_CALL_END, + TEXT_MESSAGE_END, + TOOL_CALL_RESULT, + ] + assert events[3]["delta"] == '{"q":' + assert events[2]["parentMessageId"] == events[0]["messageId"] + assert events[6]["content"] == '{"ok":true}' + assert events[6]["messageId"] != "a" + + def test_capture_off_keeps_the_lifecycle_without_payloads(self) -> None: + emitter = TurnEmitter(capture_content=False) + events: list[dict[str, Any]] = [] + events.extend(emitter.tool_start("a", "search")) + events.extend(emitter.tool_args("a", '{"q":')) + events.extend(emitter.tool_end("a", arguments={"q": "x"})) + events.extend(emitter.tool_result("a", "secret")) + assert _types(events) == [TOOL_CALL_START, TOOL_CALL_END, TOOL_CALL_RESULT] + assert events[0]["toolCallName"] == "search" + assert "parentMessageId" not in events[0] + assert events[2]["content"] == "" + assert "delta" not in str(events) + + def test_failure_is_masked_and_a_string_is_not_reencoded(self) -> None: + emitter = TurnEmitter(capture_content=True) + events: list[dict[str, Any]] = [] + events.extend(emitter.tool_start("a", "search")) + events.extend(emitter.tool_result("a", failed=True)) + events.extend(emitter.tool_start("b", "search")) + events.extend(emitter.tool_result("b", "plain")) + failed = events[2] + assert failed["content"] == "An error occurred." + assert failed["metadata"] == {"launchdarkly": {"toolFailed": True}} + assert events[5]["content"] == "plain" + + def test_unserializable_result_is_a_failed_result(self) -> None: + emitter = TurnEmitter(capture_content=True) + events = emitter.tool_start("a", "search") + events.extend(emitter.tool_result("a", float("nan"))) + assert events[-1]["content"] == "An error occurred." + assert events[-1]["metadata"]["launchdarkly"]["toolFailed"] is True + + def test_fail_open_masks_a_call_that_never_finished(self) -> None: + emitter = TurnEmitter(capture_content=False) + emitter.tool_start("a", "search") + events = emitter.fail_open() + assert [event["type"] for event in events] == [TOOL_CALL_END, TOOL_CALL_RESULT] + assert events[-1]["content"] == "" + assert events[-1]["metadata"]["launchdarkly"]["toolFailed"] is True + + def test_a_second_start_for_the_same_id_is_dropped(self) -> None: + emitter = TurnEmitter(capture_content=True) + assert emitter.tool_start("a", "search") + assert emitter.tool_start("a", "search") == [] + + def test_message_ids_differ_across_runs(self) -> None: + first = TurnEmitter(capture_content=False).text("a")[0]["messageId"] + second = TurnEmitter(capture_content=False).text("b")[0]["messageId"] + assert first != second + + def test_tool_end_extends_a_matching_prefix_only(self) -> None: + emitter = TurnEmitter(capture_content=True) + events = emitter.tool_start("a", "search") + events.extend(emitter.tool_args("a", '{"q":')) + events.extend(emitter.tool_end("a", arguments={"q": "x"})) + deltas = [event["delta"] for event in events if event["type"] == TOOL_CALL_ARGS] + assert "".join(deltas) == '{"q":"x"}' + + spaced = TurnEmitter(capture_content=True) + spaced.tool_start("a", "search") + spaced.tool_args("a", '{"q": "') + closed = spaced.tool_end("a", arguments={"q": "xy"}) + assert not any(event["type"] == TOOL_CALL_ARGS for event in closed) + + +class TestRunReader: + def test_legacy_chunks_become_one_text_message(self) -> None: + reader = RunReader() + events = [reader.start()] + events.extend(reader.feed({"type": "chunk", "text": "Hi"})) + events.extend(reader.feed({"type": "chunk", "text": "!"})) + events.extend( + reader.finish( + { + "response": "Hi!", + "usage": {"input": 2, "output": 1, "total": 3}, + "_had_usage": True, + }, + None, + provider="OpenAI", + model="gpt-4", + ) + ) + assert _types(events) == [ + RUN_STARTED, + TEXT_MESSAGE_START, + TEXT_MESSAGE_CONTENT, + TEXT_MESSAGE_CONTENT, + TEXT_MESSAGE_END, + RUN_FINISHED, + ] + assert events[0]["protocolVersion"] == "1.0" + assert events[0]["threadId"] == events[-1]["threadId"] + assert events[-1]["result"] == "Hi!" + assert events[-1]["usage"] == [ + { + "inputTokens": 2, + "outputTokens": 1, + "totalTokens": 3, + "provider": "OpenAI", + "model": "gpt-4", + } + ] + assert "metadata" not in events[-1] + + def test_forwards_the_subset_and_drops_a_third_type(self) -> None: + reader = RunReader() + events = [reader.start()] + events.extend(reader.feed({"type": TEXT_MESSAGE_START, "messageId": "m1"})) + events.extend( + reader.feed( + {"type": TEXT_MESSAGE_CONTENT, "messageId": "m1", "delta": "Hi"} + ) + ) + events.extend(reader.feed({"type": "widget", "text": "nope"})) + events.extend(reader.feed({"type": "chunk", "text": "dropped"})) + events.extend( + reader.feed( + { + "type": TOOL_CALL_START, + "toolCallId": "a", + "toolCallName": "search", + "parentMessageId": "m1", + } + ) + ) + events.extend( + reader.feed({"type": TOOL_CALL_ARGS, "toolCallId": "a", "delta": "{}"}) + ) + events.extend(reader.feed({"type": TOOL_CALL_END, "toolCallId": "a"})) + events.extend(reader.feed({"type": TEXT_MESSAGE_END, "messageId": "m1"})) + events.extend( + reader.feed( + { + "type": TOOL_CALL_RESULT, + "messageId": "tm1", + "toolCallId": "a", + "content": "ok", + } + ) + ) + events.extend( + reader.feed( + { + "type": TOOL_CALL_START, + "toolCallId": "a", + "toolCallName": "search", + } + ) + ) + events.extend( + reader.finish( + {"response": "Hi", "_had_usage": False}, + None, + provider="", + model="", + ) + ) + assert "widget" not in _types(events) + assert "chunk" not in _types(events) + assert "dropped" not in str(events) + assert _types(events).count(TOOL_CALL_START) == 1 + assert events[-1]["type"] == RUN_FINISHED + assert "usage" not in events[-1] + assert events[-1]["result"] == "Hi" + + def test_cache_split_and_judge_metadata(self) -> None: + reader = RunReader() + finished = reader.finish( + { + "response": "", + "_had_usage": True, + "usage": { + "input": 5, + "output": 1, + "total": 6, + "input_details": { + "uncached": 3, + "cache_read": 1, + "cache_creation": 1, + }, + }, + }, + { + "quality": JudgeResult( + usage=UsageDict( + input=1, + output=1, + total=2, + input_details=InputTokenDetails( + uncached=1, cache_read=0, cache_creation=0 + ), + ), + response="fine", + score=1.0, + ) + }, + provider="", + model="", + ) + usage = finished[-1]["usage"][0] + assert usage["cachedInputTokens"] == 1 + assert usage["cacheWriteInputTokens"] == 1 + assert "provider" not in usage + assert "result" not in finished[-1] + assert ( + finished[-1]["metadata"]["launchdarkly"]["judgeResults"]["quality"]["score"] + == 1.0 + ) + + def test_finish_closes_an_open_tool_before_the_run(self) -> None: + reader = RunReader(thread_id="thread-1", run_id="run-1") + events = [reader.start()] + events.extend( + reader.feed( + { + "type": TOOL_CALL_START, + "toolCallId": "a", + "toolCallName": "search", + } + ) + ) + events.extend(reader.finish({"response": ""}, None, provider="", model="")) + assert [event["type"] for event in events] == [ + RUN_STARTED, + TOOL_CALL_START, + TOOL_CALL_END, + RUN_FINISHED, + ] + assert events[0]["threadId"] == "thread-1" + assert events[-1]["runId"] == "run-1" + + def test_a_result_without_an_end_closes_the_call_first(self) -> None: + reader = RunReader() + reader.feed( + {"type": TOOL_CALL_START, "toolCallId": "a", "toolCallName": "search"} + ) + events = reader.feed({"type": TOOL_CALL_RESULT, "toolCallId": "a"}) + assert [event["type"] for event in events] == [ + TOOL_CALL_END, + TOOL_CALL_RESULT, + ] + assert events[-1]["content"] == "" + + +class TestStableProjector: + def test_text_content_is_a_chunk_and_a_third_type_is_not_done(self) -> None: + projector = StableProjector() + assert projector.feed({"type": "widget"}) == [] + assert projector.done is None + assert projector.feed( + {"type": TEXT_MESSAGE_CONTENT, "messageId": "m", "delta": "Hi"} + ) == [{"type": "chunk", "text": "Hi"}] + assert projector.feed({"type": "chunk", "text": "nope"}) == [] + assert ( + projector.feed( + {"type": TOOL_CALL_START, "toolCallId": "a", "toolCallName": "search"} + ) + == [] + ) + done = {"type": "done", "response": "Hi", "usage": {"input": 1}} + assert projector.feed(done) == [] + assert projector.done is done + + +class TestConfigAndReader: + async def test_config_stream_stays_chunk_then_done( + self, mock_ld_client: MagicMock + ) -> None: + handler = _handler( + [ + {"type": TEXT_MESSAGE_START, "messageId": "m1"}, + {"type": TEXT_MESSAGE_CONTENT, "messageId": "m1", "delta": "Hi"}, + { + "type": TOOL_CALL_START, + "toolCallId": "a", + "toolCallName": "search", + "parentMessageId": "m1", + }, + {"type": TOOL_CALL_ARGS, "toolCallId": "a", "delta": '{"q":"x"}'}, + {"type": TOOL_CALL_END, "toolCallId": "a"}, + {"type": TEXT_MESSAGE_END, "messageId": "m1"}, + { + "type": TOOL_CALL_RESULT, + "messageId": "tm", + "toolCallId": "a", + "content": "secret", + }, + {"type": "widget", "name": "nope"}, + { + "type": "done", + "output": "Hi", + "usage": {"input_tokens": 3, "output_tokens": 1}, + }, + ] + ) + events = [ + event + async for event in config(key="flag", handler=handler).stream("q", CONTEXT) + ] + assert events == [ + {"type": "chunk", "text": "Hi"}, + { + "type": "done", + "response": "Hi", + "usage": {"input": 3, "output": 1, "total": 4}, + "judge_results": None, + }, + ] + + async def test_reader_brackets_the_subset(self, mock_ld_client: MagicMock) -> None: + handler = _handler( + [ + {"type": "chunk", "text": "Hi"}, + { + "type": "done", + "output": "Hi", + "usage": {"input_tokens": 3, "output_tokens": 1}, + }, + ] + ) + events = [ + event + async for event in stream(config(key="flag", handler=handler), "q", CONTEXT) + ] + assert _types(events) == [ + RUN_STARTED, + TEXT_MESSAGE_START, + TEXT_MESSAGE_CONTENT, + TEXT_MESSAGE_END, + RUN_FINISHED, + ] + assert events[-1]["result"] == "Hi" + assert events[-1]["usage"][0]["provider"] == "TestProvider" + assert events[-1]["usage"][0]["model"] == "gpt-4" + assert events[-1]["usage"][0]["inputTokens"] == 3 + + async def test_a_raised_turn_has_no_run_finished( + self, mock_ld_client: MagicMock + ) -> None: + async def fn(cfg, user_input, tool_handlers, variables, history=None) -> dict: # type: ignore[override] + return {"output": "", "usage": {}} + + async def stream_fn( + cfg, user_input, tool_handlers, variables, history=None + ) -> AsyncGenerator[dict[str, Any], None]: + yield {"type": "chunk", "text": "Hi"} + raise RuntimeError("stream-fail") + yield # pragma: no cover + + handler = ProviderHandler( + fn=fn, + provides_for=("TestProvider", "messages"), + stream_fn=stream_fn, + ) + seen: list[str] = [] + with pytest.raises(RuntimeError, match="stream-fail"): + async for event in stream( + config(key="flag", handler=handler), "q", CONTEXT + ): + seen.append(event["type"]) + assert seen == [RUN_STARTED, TEXT_MESSAGE_START, TEXT_MESSAGE_CONTENT] + assert RUN_FINISHED not in seen + + async def test_judge_results_land_on_run_finished( + self, mock_ld_client: MagicMock + ) -> None: + handler = _handler( + [ + {"type": "chunk", "text": "Hi"}, + { + "type": "done", + "output": "Hi", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ] + ) + judged = { + "quality": JudgeResult( + usage=UsageDict(input=1, output=1, total=2), + response="fine", + score=0.5, + ) + } + with patch( + "launchdarkly_ai_server.judges.run_judges", + new=AsyncMock(return_value=judged), + ): + events = [ + event + async for event in stream( + config(key="flag", handler=handler), "q", CONTEXT + ) + ] + assert ( + events[-1]["metadata"]["launchdarkly"]["judgeResults"]["quality"]["score"] + == 0.5 + ) + + async def test_closing_early_does_not_finish_the_run( + self, mock_ld_client: MagicMock + ) -> None: + closed = False + + async def fn_unused( + cfg, user_input, tool_handlers, variables, history=None + ) -> dict: + return {"output": "", "usage": {}} + + async def stream_fn( + cfg, user_input, tool_handlers, variables, history=None + ) -> AsyncGenerator[dict[str, Any], None]: + nonlocal closed + try: + for event in [ + {"type": "chunk", "text": "Hi"}, + {"type": "chunk", "text": " there"}, + { + "type": "done", + "output": "Hi there", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ]: + yield event + finally: + closed = True + + handler = ProviderHandler( + fn=fn_unused, + provides_for=("TestProvider", "messages"), + stream_fn=stream_fn, + capture_content=False, + ) + + gen = stream(config(key="flag", handler=handler), "q", CONTEXT) + first = await gen.__anext__() + second = await gen.__anext__() + await gen.aclose() + assert first["type"] == RUN_STARTED + assert second["type"] == TEXT_MESSAGE_START + assert closed is True diff --git a/packages/client/tests/test_client.py b/packages/client/tests/test_client.py index 1ea9f66a..07ccb3a4 100644 --- a/packages/client/tests/test_client.py +++ b/packages/client/tests/test_client.py @@ -786,6 +786,65 @@ async def _noop_fn( events_tracked = [c[0][0] for c in mock_ld_client.track.call_args_list] assert "$ld:ai:generation:error" in events_tracked + async def test_a_plain_async_iterator_exits_cleanly( + self, mock_ld_client: MagicMock + ) -> None: + class _Iterator: + def __init__(self) -> None: + self._events = [ + {"type": "chunk", "text": "hi"}, + { + "type": "done", + "output": "hi", + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ] + + def __aiter__(self) -> Any: + return self + + async def __anext__(self) -> dict[str, Any]: + if not self._events: + raise StopAsyncIteration + event = self._events.pop(0) + return event + + def _stream_fn( + cfg: Any, + user_input: Any, + tool_handlers: Any, + variables: Any, + history: Any = None, + ) -> _Iterator: + iterator = _Iterator() + assert not hasattr(iterator, "aclose") + return iterator + + async def _noop_fn( + cfg: Any, + user_input: Any, + tool_handlers: Any, + variables: Any, + history: Any = None, + ) -> dict[str, Any]: + return {"output": "", "usage": {}} + + handler = ProviderHandler( + fn=_noop_fn, + stream_fn=_stream_fn, + provides_for=("TestProvider", "messages"), + ) + events = [ + event + async for event in config(key="flag", handler=handler).stream("q", CONTEXT) + ] + chunks = [event for event in events if event["type"] == "chunk"] + assert chunks == [{"type": "chunk", "text": "hi"}] + assert events[-1]["type"] == "done" + assert events[-1]["response"] == "hi" + tracked = [call[0][0] for call in mock_ld_client.track.call_args_list] + assert "$ld:ai:generation:success" in tracked + async def test_ld_context_injected_into_stream_variables( self, mock_ld_client: MagicMock ) -> None: