Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-llm/src/lexigram/ai/llm/clients/_bedrock_mappers.py: 12%

1"""Mapping helpers for Bedrock request and response payloads.""" 

2 

3from __future__ import annotations 

4 

5import asyncio 

6from typing import TYPE_CHECKING, Any 

7 

8from lexigram.ai.llm.types import ( 

9 Completion, 

10 FunctionCall, 

11 StreamChunk, 

12 ThinkingResult, 

13 TokenUsage, 

14 ToolCall, 

15) 

16from lexigram.serialization import dumps_str 

17 

18if TYPE_CHECKING: 

19 from collections.abc import AsyncIterator 

20 from concurrent.futures import ThreadPoolExecutor 

21 

22 

23def extract_system(messages: list[Any]) -> str: 

24 """Extract the system prompt from a message list. 

25 

26 Args: 

27 messages: OpenAI-compatible message list (dicts or ChatMessage). 

28 

29 Returns: 

30 Combined system prompt text, or empty string when absent. 

31 """ 

32 parts: list[str] = [] 

33 for msg in messages: 

34 role = ( 

35 msg.get("role", "") if isinstance(msg, dict) else getattr(msg, "role", "") 

36 ) 

37 role_str = role.value if hasattr(role, "value") else str(role) 

38 if role_str == "system": 

39 content = ( 

40 msg.get("content", "") 

41 if isinstance(msg, dict) 

42 else getattr(msg, "content", "") 

43 ) 

44 if isinstance(content, list): 

45 parts.extend(p.get("text", "") for p in content if isinstance(p, dict)) 

46 else: 

47 parts.append(str(content)) 

48 return " ".join(parts) 

49 

50 

51def parse_bedrock_response(raw: dict[str, Any], model: str) -> Completion: 

52 """Parse a Bedrock ``Converse`` response to a :class:`Completion`. 

53 

54 Args: 

55 raw: Raw response dict from ``boto3`` ``converse()``. 

56 model: Model ID for the ``Completion`` metadata. 

57 

58 Returns: 

59 Normalised :class:`~lexigram.ai.llm.types.Completion`. 

60 """ 

61 output = raw.get("output", {}).get("message", {}) 

62 content_blocks = output.get("content", []) 

63 

64 text_parts: list[str] = [] 

65 thinking_parts: list[str] = [] 

66 thinking_signature: str | None = None 

67 tool_calls: list[ToolCall] = [] 

68 

69 for block in content_blocks: 

70 if "text" in block: 

71 text_parts.append(block["text"]) 

72 elif "reasoningContent" in block: 

73 rc = block["reasoningContent"] 

74 # Bedrock Claude: {reasoningText: {text: "...", signature: "..."}} 

75 rt = rc.get("reasoningText", {}) 

76 if rt.get("text"): 

77 thinking_parts.append(rt["text"]) 

78 if rt.get("signature"): 

79 thinking_signature = rt["signature"] 

80 elif "toolUse" in block: 

81 tu = block["toolUse"] 

82 tool_calls.append( 

83 ToolCall( 

84 id=tu.get("toolUseId", tu.get("name", "")), 

85 type="function", 

86 function=FunctionCall( 

87 name=tu.get("name", ""), 

88 arguments=dumps_str(tu.get("input", {})), 

89 ), 

90 ) 

91 ) 

92 

93 usage_raw = raw.get("usage", {}) 

94 usage = TokenUsage( 

95 prompt_tokens=usage_raw.get("inputTokens", 0), 

96 completion_tokens=usage_raw.get("outputTokens", 0), 

97 total_tokens=usage_raw.get("totalTokens", 0), 

98 ) 

99 thinking_text = "".join(thinking_parts) or None 

100 thinking: ThinkingResult | None = ( 

101 ThinkingResult(content=thinking_text, signature=thinking_signature) 

102 if thinking_text 

103 else None 

104 ) 

105 

106 return Completion( 

107 content="".join(text_parts), 

108 model=model, 

109 finish_reason=raw.get("stopReason"), 

110 thinking=thinking, 

111 usage=usage, 

112 tool_calls=tool_calls or None, 

113 ) 

114 

115 

116async def bedrock_stream_chunks( 

117 raw_stream: Any, 

118 model: str, 

119 thread_pool: ThreadPoolExecutor, 

120) -> AsyncIterator[StreamChunk]: 

121 """Yield :class:`StreamChunk` objects from a Bedrock ``ConverseStream`` response. 

122 

123 Bedrock returns an event-stream dict with a ``stream`` key containing an 

124 iterator of typed event dicts. Events of type ``contentBlockDelta`` carry 

125 incremental text in ``delta.text``. 

126 

127 Args: 

128 raw_stream: Raw ``converse_stream()`` response from boto3. 

129 model: Model label for each chunk. 

130 

131 Yields: 

132 :class:`StreamChunk` with incremental text deltas. 

133 """ 

134 stream = raw_stream.get("stream", []) 

135 loop = asyncio.get_event_loop() 

136 index = 0 

137 

138 def _next_event(it: Any) -> Any: 

139 try: 

140 return next(it) 

141 except StopIteration: 

142 return None 

143 

144 it = iter(stream) 

145 while True: 

146 event = await loop.run_in_executor(thread_pool, _next_event, it) 

147 if event is None: 

148 break 

149 if "contentBlockDelta" in event: 

150 delta_obj = event["contentBlockDelta"].get("delta", {}) 

151 if "text" in delta_obj: 

152 yield StreamChunk( 

153 delta=delta_obj["text"], 

154 model=model, 

155 finish_reason=None, 

156 index=index, 

157 ) 

158 index += 1 

159 elif "reasoningContent" in delta_obj: 

160 # Bedrock Claude thinking delta 

161 thinking_text = delta_obj["reasoningContent"].get("text", "") 

162 if thinking_text: 

163 yield StreamChunk( 

164 thinking_delta=thinking_text, 

165 is_thinking=True, 

166 model=model, 

167 finish_reason=None, 

168 index=index, 

169 ) 

170 index += 1 

171 elif "messageStop" in event: 

172 stop_reason = event["messageStop"].get("stopReason") 

173 yield StreamChunk( 

174 delta="", model=model, finish_reason=stop_reason, index=index 

175 ) 

176 break 

177 

178 

179def tool_to_bedrock(tool: Any) -> dict[str, Any]: 

180 """Convert a tool descriptor to Bedrock ``ToolSpec`` format. 

181 

182 Args: 

183 tool: A class with ``__tool_schema__``, or a dict in OpenAI tool format. 

184 

185 Returns: 

186 Bedrock ``ToolSpec`` dict wrapped in a ``{"toolSpec": {...}}`` envelope. 

187 """ 

188 if hasattr(tool, "__tool_schema__"): 

189 schema: dict[str, Any] = tool.__tool_schema__ 

190 return { 

191 "toolSpec": { 

192 "name": schema["name"], 

193 "description": schema.get("description", ""), 

194 "inputSchema": {"json": schema.get("parameters", {})}, 

195 } 

196 } 

197 if isinstance(tool, dict): 

198 func = tool.get("function", tool) 

199 return { 

200 "toolSpec": { 

201 "name": func.get("name", ""), 

202 "description": func.get("description", ""), 

203 "inputSchema": {"json": func.get("parameters", {})}, 

204 } 

205 } 

206 return { 

207 "toolSpec": { 

208 "name": getattr(tool, "name", str(tool)), 

209 "description": getattr(tool, "description", ""), 

210 "inputSchema": { 

211 "json": getattr(tool, "parameters", None) 

212 or { 

213 "type": "object", 

214 "properties": {}, 

215 } 

216 }, 

217 } 

218 }