Coverage for gco_mcp / mission / embeddings.py: 100.00%

1"""Directive embedding for mission memory. 

2 

3One free function, :func:`embed_text`, turns a Mission directive into the 

4vector that the ``{project}-mission-memory`` DynamoDB vector index stores 

5and queries. The Bedrock client is constructed exactly the way 

6:class:`mcp.mission.sampling.BedrockSamplingBackend` builds its client 

7(``boto3.Session().client("bedrock-runtime", config=Config(read_timeout= 

8BEDROCK_READ_TIMEOUT_SECONDS))``), and the model id resolves through 

9:func:`gco.bedrock.get_default_embedding_model_id` so the runtime and the 

10deployed index share one configuration source. 

11 

12Failure contract: every failure raises the typed :class:`EmbeddingError` 

13so callers decide whether to swallow. The engine's memory write/read 

14paths are best-effort and do swallow; the CLI surfaces the message. The 

15one deliberate exception is the Bedrock first-time-use gate — 

16:func:`gco.bedrock.raise_if_bedrock_ftu_form_error` escalates that case 

17to :class:`gco.bedrock.BedrockFTUFormNotAcceptedError` because it is a 

18permanent, account-scoped misconfiguration with a one-line fix, and the 

19established posture (see ``sampling.py``) is that it must never be 

20absorbed by a graceful-degradation handler. 

21 

22Request shape note: the body follows the Amazon Titan Text Embeddings 

23contract (``inputText`` plus, when supplied, the V2-only ``dimensions`` 

24key). ``dimensions`` is omitted from the request when the caller passes 

25``None`` so V1-family models — which reject the key — keep working. 

26 

27``boto3`` / ``botocore`` are imported lazily inside the functions so 

28pure-Python consumers of the mission package do not pay for the SDK 

29import, and so tests can inject a fake ``boto3`` module through 

30``sys.modules`` (the same pattern ``sampling.py`` uses). 

31""" 

32 

33from __future__ import annotations 

34 

35import json 

36import logging 

37from typing import Any 

38 

39from gco.bedrock import ( 

40 BEDROCK_READ_TIMEOUT_SECONDS, 

41 get_default_embedding_model_id, 

42 raise_if_bedrock_ftu_form_error, 

43) 

44 

45logger = logging.getLogger(__name__) 

46 

47__all__ = [ 

48 "EmbeddingError", 

49 "embed_text", 

50] 

51 

52 

53class EmbeddingError(RuntimeError): 

54 """A directive embedding could not be produced. 

55 

56 The message is a short machine-matchable code, mirroring the 

57 ``SamplingTransportError`` convention in ``sampling.py``: 

58 

59 * ``embedding_empty_text`` — the input was empty or whitespace-only. 

60 * ``embedding_no_credentials`` — ``boto3`` could not resolve 

61 credentials at client-construction time. 

62 * ``embedding_bedrock_<ErrorCode>`` — ``InvokeModel`` raised a 

63 ``ClientError``; ``<ErrorCode>`` is the AWS error code. 

64 * ``embedding_transport_failure`` — a non-``ClientError`` botocore 

65 transport fault (endpoint unreachable, read timeout, ...). 

66 * ``embedding_malformed_response`` — the response body was not JSON 

67 or did not carry a non-empty numeric ``embedding`` list. 

68 """ 

69 

70 

71def _build_client() -> Any: 

72 """Return a fresh ``bedrock-runtime`` client. 

73 

74 Split out of :func:`embed_text` so tests can monkeypatch this one 

75 seam instead of faking the whole ``boto3`` module (both work). The 

76 imports are local — see the module docstring. 

77 """ 

78 import boto3 

79 from botocore.config import Config 

80 from botocore.exceptions import ( 

81 NoCredentialsError, 

82 PartialCredentialsError, 

83 ) 

84 

85 try: 

86 return boto3.Session().client( 

87 "bedrock-runtime", 

88 config=Config(read_timeout=BEDROCK_READ_TIMEOUT_SECONDS), 

89 ) 

90 except (NoCredentialsError, PartialCredentialsError) as err: 

91 raise EmbeddingError("embedding_no_credentials") from err 

92 

93 

94def embed_text( 

95 text: str, 

96 *, 

97 model_id: str | None = None, 

98 dimensions: int | None = None, 

99) -> list[float]: 

100 """Embed ``text`` and return the vector as a plain list of floats. 

101 

102 Args: 

103 text: The text to embed — for mission memory, the verbatim 

104 operator directive. Must be non-empty. 

105 model_id: Bedrock model id override. ``None`` resolves the 

106 checked-in default via 

107 :func:`gco.bedrock.get_default_embedding_model_id`. 

108 dimensions: Requested output width, passed through to the model 

109 (Titan Text Embeddings V2 accepts 256/512/1024). ``None`` 

110 omits the key so the model's own default width applies. 

111 The deployed vector index width is a one-way door — query 

112 vectors must match it, so callers that know the configured 

113 width should pass it. 

114 

115 Raises: 

116 EmbeddingError: On any embedding failure; see the class 

117 docstring for the code taxonomy. 

118 gco.bedrock.BedrockFTUFormNotAcceptedError: The account has not 

119 submitted Anthropic's first-time-use form (only reachable 

120 when the embedding model is Anthropic-gated). 

121 gco.bedrock.BedrockModelConfigurationError: The canonical 

122 ``cdk.json`` could not supply a default model id. 

123 """ 

124 from botocore.exceptions import BotoCoreError, ClientError 

125 

126 if not isinstance(text, str) or not text.strip(): 

127 raise EmbeddingError("embedding_empty_text") 

128 

129 resolved_model_id = model_id or get_default_embedding_model_id() 

130 

131 body: dict[str, Any] = {"inputText": text} 

132 if dimensions is not None: 

133 body["dimensions"] = int(dimensions) 

134 

135 client = _build_client() 

136 try: 

137 response = client.invoke_model( 

138 modelId=resolved_model_id, 

139 body=json.dumps(body), 

140 contentType="application/json", 

141 accept="application/json", 

142 ) 

143 except ClientError as err: 

144 raise_if_bedrock_ftu_form_error(err) 

145 code = err.response.get("Error", {}).get("Code") or "ClientError" 

146 raise EmbeddingError(f"embedding_bedrock_{code}") from err 

147 except BotoCoreError as err: 

148 raise EmbeddingError("embedding_transport_failure") from err 

149 

150 try: 

151 payload = json.loads(response["body"].read()) 

152 except (AttributeError, KeyError, TypeError, ValueError) as err: 

153 raise EmbeddingError("embedding_malformed_response") from err 

154 

155 vector = payload.get("embedding") if isinstance(payload, dict) else None 

156 if ( 

157 not isinstance(vector, list) 

158 or not vector 

159 or not all(isinstance(v, (int, float)) and not isinstance(v, bool) for v in vector) 

160 ): 

161 raise EmbeddingError("embedding_malformed_response") 

162 return [float(v) for v in vector]