Coverage for src / quber / playground / answers / streaming.py: 38%

1"""Streaming a typed answer. 

2 

3The playground streams today because a prose answer is worth watching arrive. 

4A figure is not — by the time a scalar validates there is nothing to animate — 

5but the question is not whether streaming is useful for a scalar, it is 

6whether adding a shape breaks the streaming path that prose still needs. 

7 

8So the stream yields the same two events as before. `delta` carries whatever 

9text is readable so far, which is the prose payload's text while it grows and 

10nothing at all for a scalar; `final` carries the validated `Answer`. A caller 

11that only wants the value ignores the deltas and waits for the final event. 

12""" 

13 

14from __future__ import annotations 

15 

16from typing import AsyncIterator, List, Optional, Tuple 

17 

18from quber.agents.langsmith_tracer import usage_metadata_from 

19from quber.playground.agent import DEFAULT_MODEL 

20from quber.playground.answers.document import DocumentIdentity 

21from quber.playground.answers.fixed import build_agent, prompt_for, trace_inputs 

22from quber.playground.answers.shapes import Answer 

23from quber.playground.retrieval import RetrievedChunk 

24from quber.playground.tracing import run_metadata, tracer 

25 

26 

27def readable(partial: object) -> str: 

28 """The text worth showing from a partially built answer, if any. 

29 

30 Only the prose payload has text that grows. A partial whose payload is not 

31 yet determined, or is a figure, has nothing to show. 

32 """ 

33 payload = getattr(partial, "payload", None) 

34 if payload is None: 

35 return "" 

36 if getattr(payload, "kind", None) != "prose": 

37 return "" 

38 return getattr(payload, "text", "") or "" 

39 

40 

41async def answer_stream( 

42 question: str, chunks: List[RetrievedChunk], document: Optional[DocumentIdentity] = None 

43) -> AsyncIterator[Tuple[str, object]]: 

44 """Yield ("delta", text) as prose accumulates, then ("final", Answer).""" 

45 prompt = prompt_for(question, chunks, document) 

46 async with tracer().llm_run( 

47 "answer_fixed", trace_inputs(prompt, chunks), model=DEFAULT_MODEL, extra_metadata=run_metadata() 

48 ) as run: 

49 async with build_agent().run_stream(prompt) as result: 

50 last = "" 

51 async for partial in result.stream_output(debounce_by=0.05): 

52 text = readable(partial) 

53 if text and text != last: 

54 last = text 

55 yield "delta", text 

56 final: Answer = await result.get_output() 

57 run.outputs = { 

58 "messages": [{"role": "assistant", "content": final.model_dump_json()}], 

59 "usage_metadata": usage_metadata_from(result.usage), 

60 } 

61 yield "final", final