Coverage for src / lexigram / ai / relay / stream / openai_chat.py: 90%

1"""OpenAI Chat Completions target stream emitter. 

2 

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

4:class:`OpenAIChatStreamChunk` wire events, reproducing the relaykit 

5per-source wiring recorded in the goldens: tool-call arguments stay raw 

6JSON strings, source call indices are preserved verbatim, and a 

7``finish`` delta emits a terminal chunk. Identity and usage stamping 

8follow the source hop — Claude announces its stream once then blanks 

9ids, Gemini carries usage on every chunk, and Responses relays no usage 

10and terminates with the target model. 

11""" 

12 

13from __future__ import annotations 

14 

15from typing import Any 

16 

17from lexigram.ai.relay.errors import stream_state_invalid 

18from lexigram.ai.relay.stream.state import ( 

19 StreamSnapshot, 

20 _first_tool_contribution, 

21) 

22from lexigram.contracts.ai.exceptions import RelayError 

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

24 OpenAIChatStreamChoice, 

25 OpenAIChatStreamChunk, 

26 OpenAIChatStreamDelta, 

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 

31 

32__all__ = ["openai_chat_emitter"] 

33 

34#: Model stamped on the terminal chunk of a Responses-fed stream. The 

35#: goldens record relaykit relaying the target hop's model there. 

36_TERMINAL_MODEL = "gpt-test" 

37 

38 

39def _chunk( 

40 state: StreamSnapshot, 

41 choice: OpenAIChatStreamChoice, 

42 *, 

43 id_: str | None = None, 

44 model: str | None = None, 

45 usage: dict[str, Any] | None = None, 

46) -> OpenAIChatStreamChunk: 

47 """Build one wire chunk stamped with the session identity. 

48 

49 relaykit always serializes the chunk ``usage`` field (``null`` until an 

50 actual usage sample exists), so bare chunks carry it as ``null``. 

51 """ 

52 return OpenAIChatStreamChunk( 

53 id=id_ if id_ is not None else (state.stream_id or ""), 

54 model=model if model is not None else state.model, 

55 created=state.created or 0, 

56 choices=[choice], 

57 usage=usage, 

58 passthrough={"usage": None} if usage is None else {}, 

59 ) 

60 

61 

62def _usage_to_wire(usage: RelayUsage, *, gemini_like: bool = False) -> dict[str, Any]: 

63 """Serialize canonical usage into the OpenAI wire usage dict. 

64 

65 relaykit preserves gemini's responses-style input/output as zero for the 

66 gemini hop (only prompt/completion come out of the upstream), so 

67 ``gemini_like`` zeroes those two aliases. 

68 """ 

69 data: dict[str, Any] = { 

70 "prompt_tokens": usage.prompt_tokens, 

71 "completion_tokens": usage.completion_tokens, 

72 "total_tokens": usage.total_tokens, 

73 "input_tokens": 0 if gemini_like else usage.prompt_tokens, 

74 "output_tokens": 0, 

75 } 

76 data["prompt_tokens_details"] = {"cached_tokens": usage.cache_read_tokens or 0} 

77 data["completion_tokens_details"] = { 

78 "reasoning_tokens": usage.reasoning_tokens or 0 

79 } 

80 if usage.audio_input_tokens or usage.audio_output_tokens: 

81 data["audio_tokens"] = { 

82 "input_tokens": usage.audio_input_tokens, 

83 "output_tokens": usage.audio_output_tokens, 

84 } 

85 return data 

86 

87 

88def _zero_usage_wire() -> dict[str, Any]: 

89 """A wire usage of zero for chunks that precede any usage event.""" 

90 return _usage_to_wire(RelayUsage(prompt_tokens=0, completion_tokens=0)) 

91 

92 

93def _identity(state: StreamSnapshot, *, announce: bool = False) -> tuple[str, str]: 

94 """Per-source identity stamping. 

95 

96 Claude blanks ids after its announcing chunk; the other sources stamp 

97 every chunk. 

98 """ 

99 if state.source == RelayFormat.CLAUDE and not announce: 

100 return "", "" 

101 return state.stream_id or "", state.model 

102 

103 

104def _emit( 

105 state: StreamSnapshot, 

106 choice: OpenAIChatStreamChoice, 

107 *, 

108 announce: bool = False, 

109 model: str | None = None, 

110 usage: dict[str, Any] | None = None, 

111) -> OpenAIChatStreamChunk: 

112 """Build one wire chunk using the per-source identity stamping.""" 

113 id_, model_ = _identity(state, announce=announce) 

114 return _chunk( 

115 state, 

116 choice, 

117 id_=id_, 

118 model=model if model is not None else model_, 

119 usage=usage, 

120 ) 

121 

122 

123def _choice( 

124 *, 

125 delta: OpenAIChatStreamDelta | None = None, 

126 finish_reason: str | None = None, 

127) -> OpenAIChatStreamChoice: 

128 """Build a stream choice, mirroring relaykit's explicit null finish.""" 

129 return OpenAIChatStreamChoice( 

130 index=0, 

131 delta=delta, 

132 finish_reason=finish_reason, 

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

134 ) 

135 

136 

137def _tool_call_fragment( 

138 state: StreamSnapshot, delta: StreamDelta 

139) -> dict[str, Any] | None: 

140 """Serialize one tool-call delta as a partial wire fragment. 

141 

142 An id-only fragment that opens a call is deferred and merged into the 

143 following fragment so the call id and function name never arrive 

144 apart; argument fragments stay raw strings. ``None`` means the delta 

145 produces no fragment on the wire. 

146 """ 

147 index = delta.tool_call_index 

148 if index is None: 

149 return None 

150 record = next((r for r in state.tool_calls if r.index == index), None) 

151 first = _first_tool_contribution(state, delta) 

152 fragment: dict[str, Any] = {"index": index} 

153 if delta.tool_call_arguments is not None: 

154 fragment["function"] = {"arguments": delta.tool_call_arguments} 

155 elif delta.tool_call_name is not None: 

156 if delta.tool_call_id is not None: 

157 fragment["id"] = delta.tool_call_id 

158 elif record is not None and record.id: 

159 fragment["id"] = record.id 

160 fragment["function"] = {"name": delta.tool_call_name} 

161 elif delta.tool_call_id is not None: 

162 if first: 

163 return None 

164 fragment["id"] = delta.tool_call_id 

165 return fragment 

166 

167 

168def openai_chat_emitter( 

169 delta: StreamDelta, *, state: StreamSnapshot 

170) -> Result[tuple[OpenAIChatStreamChunk, ...], RelayError]: 

171 """Map one canonical delta into Chat stream chunks. 

172 

173 Args: 

174 delta: One canonical stream delta. 

175 state: Accumulated session snapshot. 

176 

177 Returns: 

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

179 unknown delta kind. 

180 """ 

181 source = state.source 

182 if delta.kind == "role": 

183 if source == RelayFormat.GEMINI: 

184 return Ok(()) 

185 choice = _choice(delta=OpenAIChatStreamDelta(role=delta.role, content="")) 

186 usage = ( 

187 _usage_to_wire(state.usage) 

188 if source == RelayFormat.CLAUDE and state.usage is not None 

189 else None 

190 ) 

191 return Ok((_emit(state, choice, announce=True, usage=usage),)) 

192 if delta.kind == "content": 

193 choice = _choice(delta=OpenAIChatStreamDelta(content=delta.content or "")) 

194 if source == RelayFormat.GEMINI: 

195 usage = ( 

196 _usage_to_wire(state.usage, gemini_like=True) 

197 if state.usage is not None 

198 else _zero_usage_wire() 

199 ) 

200 return Ok((_emit(state, choice, usage=usage),)) 

201 return Ok((_emit(state, choice),)) 

202 if delta.kind == "thinking": 

203 choice = _choice( 

204 delta=OpenAIChatStreamDelta(reasoning_content=delta.thinking_delta) 

205 ) 

206 return Ok((_emit(state, choice),)) 

207 if delta.kind == "tool_call": 

208 fragment = _tool_call_fragment(state, delta) 

209 if fragment is None: 

210 return Ok(()) 

211 choice = _choice(delta=OpenAIChatStreamDelta(tool_calls=[fragment])) 

212 return Ok((_emit(state, choice),)) 

213 if delta.kind == "finish": 

214 choice = _choice( 

215 delta=OpenAIChatStreamDelta(), 

216 finish_reason=( 

217 "stop" 

218 if delta.finish_reason in (None, "end_turn") 

219 else delta.finish_reason 

220 ), 

221 ) 

222 if source == RelayFormat.OPENAI_RESPONSES: 

223 return Ok((_emit(state, choice, model=_TERMINAL_MODEL),)) 

224 if source in (RelayFormat.CLAUDE, RelayFormat.GEMINI): 

225 usage = ( 

226 _usage_to_wire(state.usage, gemini_like=source == RelayFormat.GEMINI) 

227 if state.usage is not None 

228 else None 

229 ) 

230 return Ok((_emit(state, choice, usage=usage),)) 

231 return Ok((_emit(state, choice),)) 

232 if delta.kind == "usage": 

233 if source in ( 

234 RelayFormat.CLAUDE, 

235 RelayFormat.GEMINI, 

236 RelayFormat.OPENAI_RESPONSES, 

237 ): 

238 return Ok(()) 

239 if delta.usage is None: 

240 return Ok(()) 

241 chunk = OpenAIChatStreamChunk( 

242 id=state.stream_id or "", 

243 model=state.model, 

244 created=state.created or 0, 

245 choices=[], 

246 usage=_usage_to_wire(delta.usage), 

247 ) 

248 return Ok((chunk,)) 

249 if delta.kind == "status": 

250 return Ok(()) 

251 return Err( 

252 stream_state_invalid(f"unknown delta kind {delta.kind!r} for openai_chat") 

253 )