Coverage for src / qdrant_loader_core / llm / providers / bedrock_utils.py: 59%

1from __future__ import annotations 

2 

3import logging 

4import math 

5from typing import Any 

6 

7from ..errors import ( 

8 AuthError, 

9 InvalidRequestError, 

10 LLMError, 

11 RateLimitedError, 

12 ServerError, 

13) 

14from ..types import TokenCounter 

15 

16logger = logging.getLogger(__name__) 

17 

18try: 

19 from botocore.exceptions import ( 

20 BotoCoreError, 

21 ClientError, 

22 EndpointConnectionError, 

23 NoCredentialsError, 

24 ) 

25except ImportError: 

26 

27 class _BedrockBaseError(Exception): 

28 pass 

29 

30 class _BedrockClientError(_BedrockBaseError): 

31 pass 

32 

33 class _BedrockNoCredentialsError(_BedrockBaseError): 

34 pass 

35 

36 class _BedrockEndpointConnectionError(_BedrockBaseError): 

37 pass 

38 

39 class _BedrockBotoCoreError(_BedrockBaseError): 

40 pass 

41 

42 BotoCoreError = _BedrockBotoCoreError 

43 ClientError = _BedrockClientError 

44 EndpointConnectionError = _BedrockEndpointConnectionError 

45 NoCredentialsError = _BedrockNoCredentialsError 

46 

47 

48def _map_bedrock_exception(exc: Exception) -> LLMError: 

49 """Map a botocore/boto3 exception into a qdrant_loader_core LLMError.""" 

50 if isinstance(exc, NoCredentialsError): 

51 return AuthError(str(exc)) 

52 

53 error_code = "" 

54 status_code = None 

55 if hasattr(exc, "response"): 

56 try: 

57 error_code = exc.response.get("Error", {}).get("Code", "") 

58 status_code = exc.response.get("ResponseMetadata", {}).get("HTTPStatusCode") 

59 except Exception: 

60 pass 

61 

62 if error_code in ("ThrottlingException", "TooManyRequestsException", "Throttling"): 

63 return RateLimitedError(str(exc)) 

64 if error_code in ( 

65 "AccessDeniedException", 

66 "UnrecognizedClientException", 

67 "InvalidSignatureException", 

68 ): 

69 return AuthError(str(exc)) 

70 if error_code in ( 

71 "ValidationException", 

72 "InvalidParameterException", 

73 "ResourceNotFoundException", 

74 "UnsupportedOperationException", 

75 ): 

76 return InvalidRequestError(str(exc)) 

77 if error_code in ("ModelErrorException", "ModelTimeoutException"): 

78 return ServerError(str(exc)) 

79 if isinstance(exc, ClientError): 

80 if status_code in (408, 424): 

81 return ServerError(str(exc)) 

82 if status_code == 429: 

83 return RateLimitedError(str(exc)) 

84 if status_code in (401, 403): 

85 return AuthError(str(exc)) 

86 if isinstance(status_code, int) and status_code >= 500: 

87 return ServerError(str(exc)) 

88 if isinstance(status_code, int) and status_code >= 400: 

89 return InvalidRequestError(str(exc)) 

90 return ServerError(str(exc)) 

91 

92 if isinstance(exc, EndpointConnectionError): 

93 return ServerError(str(exc)) 

94 

95 if isinstance(exc, BotoCoreError): 

96 return ServerError(str(exc)) 

97 

98 return ServerError(str(exc)) 

99 

100 

101def _extract_embeddings(response_payload: Any) -> list[list[float]]: 

102 """Normalize Bedrock embedding response payloads into a list of float vectors.""" 

103 if not isinstance(response_payload, (dict, list)): 

104 raise InvalidRequestError("Bedrock response has unexpected format") 

105 

106 raw_embeddings: list[Any] 

107 

108 if isinstance(response_payload, list): 

109 raw_embeddings = response_payload 

110 

111 elif "embedding" in response_payload: 

112 raw_embeddings = [response_payload["embedding"]] 

113 

114 elif "embeddings" in response_payload: 

115 raw_embeddings = response_payload["embeddings"] 

116 

117 elif "data" in response_payload and isinstance(response_payload["data"], list): 

118 raw_embeddings = [ 

119 item.get("embedding", item) if isinstance(item, dict) else item 

120 for item in response_payload["data"] 

121 ] 

122 

123 else: 

124 raise InvalidRequestError("Bedrock response did not contain embeddings") 

125 

126 normalized: list[list[float]] = [] 

127 

128 for vector in raw_embeddings: 

129 if isinstance(vector, dict): 

130 vector = vector.get("embedding", vector) 

131 

132 if not isinstance(vector, list): 

133 raise InvalidRequestError( 

134 "Bedrock embedding payload must be a list of floats" 

135 ) 

136 

137 try: 

138 parsed_vector = [float(value) for value in vector] 

139 except (TypeError, ValueError) as exc: 

140 raise InvalidRequestError( 

141 f"Invalid embedding element from Bedrock: {exc}" 

142 ) from exc 

143 if not all(math.isfinite(value) for value in parsed_vector): 

144 raise InvalidRequestError( 

145 "Invalid embedding element from Bedrock: non-finite value" 

146 ) 

147 normalized.append(parsed_vector) 

148 

149 return normalized 

150 

151 

152class BedrockTokenizer(TokenCounter): 

153 """ 

154 Temporary fallback tokenizer for Bedrock providers. 

155 

156 This implementation approximates token counts using character length 

157 and is NOT model-accurate. Counts may differ significantly from 

158 actual Bedrock model token usage. 

159 """ 

160 

161 def __init__(self) -> None: 

162 logger.warning( 

163 "BedrockTokenizer is using a character-count fallback. " 

164 "Token counts are approximate and may not match actual " 

165 "Bedrock model tokenization." 

166 ) 

167 

168 def count(self, text: str) -> int: 

169 return len(text)