Coverage for src / lexigram / ai / relay / stream / gemini.py: 93%

1"""Google Gemini ``generateContent`` target stream emitter. 

2 

3Maps one canonical :class:`StreamDelta` into zero or more 

4:class:`GeminiResponse` stream chunks. Thinking parts carry the 

5``thought`` flag and optional ``thoughtSignature``; tool calls are 

6emitted as ``functionCall`` parts on the terminal chunk, where the full 

7accumulated argument JSON can be parsed into the target's object args. 

8""" 

9 

10from __future__ import annotations 

11 

12from typing import Any 

13 

14from lexigram.ai.relay.errors import stream_state_invalid 

15from lexigram.ai.relay.finish_reasons import ( 

16 finish_reason_to_wire, 

17 normalize_finish_reason, 

18) 

19from lexigram.ai.relay.stream.state import StreamSnapshot 

20from lexigram.contracts.ai.exceptions import RelayError 

21from lexigram.contracts.ai.relay.dto import ( 

22 GeminiCandidate, 

23 GeminiContent, 

24 GeminiPart, 

25 GeminiResponse, 

26 GeminiUsageMetadata, 

27) 

28from lexigram.contracts.ai.relay.ir import StreamDelta 

29from lexigram.contracts.ai.relay.types import RelayFormat, RelayUsage 

30from lexigram.contracts.core.result import Err, Ok, Result 

31from lexigram.serialization import loads_str 

32 

33__all__ = ["gemini_emitter"] 

34 

35 

36def _gemini_finish_reason(reason: str) -> str: 

37 """Map a canonical finish reason onto Gemini's wire values.""" 

38 return finish_reason_to_wire(normalize_finish_reason(reason), RelayFormat.GEMINI) 

39 

40 

41def _usage_to_wire(usage: RelayUsage) -> GeminiUsageMetadata: 

42 """Serialize canonical usage into the Gemini usage shape.""" 

43 return GeminiUsageMetadata( 

44 prompt_token_count=usage.prompt_tokens, 

45 candidates_token_count=usage.completion_tokens, 

46 total_token_count=usage.total_tokens, 

47 cached_content_token_count=usage.cache_read_tokens or 0, 

48 thoughts_token_count=usage.reasoning_tokens or 0, 

49 ) 

50 

51 

52def _zero_usage() -> RelayUsage: 

53 """A canonical usage of zero, padded across all counters.""" 

54 return RelayUsage(prompt_tokens=0, completion_tokens=0) 

55 

56 

57def _chunk( 

58 state: StreamSnapshot, 

59 *, 

60 parts: list[GeminiPart] | None = None, 

61 finish_reason: str | None = None, 

62 usage: GeminiUsageMetadata | None = None, 

63) -> GeminiResponse: 

64 """Build one wire chunk with a single candidate at index zero. 

65 

66 The goldens record relaykit serializing ``finishReason`` explicitly 

67 (``null`` on in-progress chunks), so the candidate passes that 

68 artifact through verbatim. 

69 """ 

70 candidate = GeminiCandidate( 

71 content=GeminiContent(role="model", parts=parts or []), 

72 finish_reason=finish_reason, 

73 index=0, 

74 safety_ratings=[], 

75 passthrough={"finishReason": None} if finish_reason is None else {}, 

76 ) 

77 return GeminiResponse(candidates=[candidate], usage_metadata=usage) 

78 

79 

80def _terminal_usage(state: StreamSnapshot, delta: StreamDelta) -> RelayUsage: 

81 """Usage stamped on the terminal chunk for the source's relay hop. 

82 

83 The recorded goldens carry no usage across the *responses* hop, so 

84 that source relays a zeroed usage; every other source relays the 

85 latest usage seen (which the transcription attaches to the correct 

86 chunk). 

87 """ 

88 if state.source == RelayFormat.OPENAI_RESPONSES: 

89 return _zero_usage() 

90 return state.usage or (delta.usage or _zero_usage()) 

91 

92 

93def _function_call_parts(state: StreamSnapshot) -> list[GeminiPart]: 

94 """Build ``functionCall`` parts from the accumulated tool calls.""" 

95 parts: list[GeminiPart] = [] 

96 for record in state.tool_calls: 

97 args: dict[str, Any] 

98 try: 

99 parsed = loads_str(record.arguments) 

100 args = parsed if isinstance(parsed, dict) else {} 

101 except (ValueError, TypeError): 

102 args = {} 

103 parts.append( 

104 GeminiPart( 

105 function_call={"name": record.name, "args": args}, 

106 ) 

107 ) 

108 return parts 

109 

110 

111def gemini_emitter( 

112 delta: StreamDelta, *, state: StreamSnapshot 

113) -> Result[tuple[GeminiResponse, ...], RelayError]: 

114 """Map one canonical delta into Gemini stream chunks. 

115 

116 Args: 

117 delta: One canonical stream delta. 

118 state: Accumulated session snapshot. 

119 

120 Returns: 

121 Ok(tuple of chunks) on success; ``stream_state_invalid`` for an 

122 unknown delta kind. 

123 """ 

124 if delta.kind == "role": 

125 return Ok(()) 

126 if delta.kind == "content": 

127 if not delta.content: 

128 return Ok(()) 

129 return Ok( 

130 ( 

131 _chunk( 

132 state, 

133 parts=[GeminiPart(text=delta.content)], 

134 usage=_usage_to_wire(_zero_usage()), 

135 ), 

136 ) 

137 ) 

138 if delta.kind == "thinking": 

139 if not delta.thinking_delta: 

140 return Ok(()) 

141 signature = delta.passthrough.get("signature") 

142 thought_signature = ( 

143 signature if isinstance(signature, str) and signature else None 

144 ) 

145 return Ok( 

146 ( 

147 _chunk( 

148 state, 

149 parts=[ 

150 GeminiPart( 

151 text=delta.thinking_delta, 

152 thought=True, 

153 thought_signature=thought_signature, 

154 ) 

155 ], 

156 usage=_usage_to_wire(_zero_usage()), 

157 ), 

158 ) 

159 ) 

160 if delta.kind == "tool_call": 

161 return Ok(()) 

162 if delta.kind == "finish": 

163 parts = _function_call_parts(state) 

164 finish_reason = _gemini_finish_reason(delta.finish_reason or "stop") 

165 usage = _usage_to_wire(_terminal_usage(state, delta)) 

166 return Ok( 

167 (_chunk(state, parts=parts, finish_reason=finish_reason, usage=usage),) 

168 ) 

169 if delta.kind == "usage": 

170 return Ok(()) 

171 if delta.kind == "status": 

172 return Ok(()) 

173 return Err(stream_state_invalid(f"unknown delta kind {delta.kind!r} for gemini"))