Model portability: swapping Bedrock for the Mistral API

2
# The model registry: Terraform owns it, renders it to JSON, and parks it in
3
# an SSM String parameter the handler reads at startup. Each entry names a
4
# provider and a model id; Bedrock entries also carry the inference-profile
5
# ARN. DEFAULT_MODEL (set on the Lambda) selects the active one, so switching
6
# the agent's model is a parameter change, not a code change.
7
locals {
8
  metrics_namespace = "TerraformPrAgent/Models"
9

10
  models = {
11
    haiku = {
12
      provider              = "bedrock"
13
      model_id              = local.bedrock_model_id
14
      inference_profile_arn = aws_bedrock_inference_profile.agent.arn
15
    }
16
    "mistral-large" = {
17
      provider = "mistral"
18
      model_id = "mistral-large-latest"
19
    }
20
    "devstral-small" = {
21
      provider = "mistral"
22
      model_id = "devstral-small-2507"
23
    }
24
  }
25

26
  mistral_key_wired = var.mistral_api_key != ""
27
}
28

29
resource "aws_ssm_parameter" "models" {
30
  name        = "/terraform-pr-agent/models"
31
  description = "Model registry for the terraform-pr-agent Lambda (provider + model id per entry)."
32
  type        = "String"
33
  value       = jsonencode(local.models)
34
}

38
# The Mistral API key, SecureString, fetched by the handler through the same
39
# Parameters and Secrets extension path as the Logfire token. Only created
40
# when TF_VAR_mistral_api_key is set, mirroring the Logfire token wiring; with
41
# it unset the Mistral providers are simply unreachable and a Bedrock default
42
# still works.
43
resource "aws_ssm_parameter" "mistral_api_key" {
44
  count = local.mistral_key_wired ? 1 : 0
45

46
  name        = "/terraform-pr-agent/mistral-api-key"
47
  description = "Mistral API key. Consumed by the terraform-pr-agent Lambda."
48
  type        = "SecureString"
49
  value       = var.mistral_api_key
50
}

149
@cache
150
def _build_model(name: str) -> Model:
151
    """Build the pydantic-ai model registered under ``name``.
152

153
    The registry lives in an SSM String parameter, so this runs on the first
154
    INVOKE (the extension is not ready during INIT) and is memoised per model
155
    name for warm invocations. Bedrock models authenticate via the Lambda
156
    role; Mistral models read an API key from a SecureString parameter,
157
    fetched the same way as the Logfire token.
158
    """
159
    registry = json.loads(_fetch_ssm_parameter(os.environ["MODELS_PARAMETER"]))
160
    config = registry[name]
161
    provider = config["provider"]
162
    if provider == "bedrock":
163
        return BedrockConverseModel(
164
            config["model_id"],
165
            settings={"bedrock_inference_profile": config["inference_profile_arn"]},
166
        )
167
    if provider == "mistral":
168
        key_param = os.environ.get("MISTRAL_API_KEY_PARAMETER")
169
        if not key_param:
170
            raise RuntimeError(
171
                f"model {name!r} uses the Mistral API, but MISTRAL_API_KEY_PARAMETER "
172
                "is not set. Set MISTRAL_API_KEY and re-apply so the key is wired, or "
173
                "select a Bedrock model via DEFAULT_MODEL or the event's model field."
174
            )
175
        return MistralModel(
176
            config["model_id"],
177
            provider=MistralProvider(
178
                api_key=_fetch_ssm_parameter(key_param),
179
                http_client=_retrying_http_client(),
180
            ),
181
        )
182
    raise ValueError(f"unknown provider {provider!r} for model {name!r}")

258
def _emit_emf(spans: Sequence[ReadableSpan]) -> None:
259
    """Emit one EMF metric line for the trace, read off the root span.
260

261
    pydantic-ai records gen_ai.usage.* on the root agent span as the run total
262
    (the sum of its child chat spans), so a single read is the correct total,
263
    not a sum across every span. The model dimension is the registry key the
264
    handler passed as run metadata; pydantic-ai serialises that to the root
265
    span's `metadata` attribute (even on a failed run), so it is read back here
266
    rather than carried in module state. That key is exactly what the dashboard
267
    iterates, so a Bedrock run and a Mistral run share one set of widgets.
268
    Logging the _aws envelope to stdout is enough; CloudWatch Logs extracts the
269
    metrics from the structured line.
270
    """
271
    root = next((span for span in spans if span.parent is None), None)
272
    if root is None:
273
        return
274
    attributes = root.attributes or {}
275
    model = json.loads(attributes.get("metadata", "{}")).get("model", "unknown")
276
    errored = root.status.status_code is StatusCode.ERROR
277
    record = {
278
        "_aws": {
279
            "Timestamp": root.end_time // 1_000_000,
280
            "CloudWatchMetrics": [
281
                {
282
                    "Namespace": os.environ["METRICS_NAMESPACE"],
283
                    "Dimensions": [["Model"]],
284
                    "Metrics": [
285
                        {"Name": "InputTokens", "Unit": "Count"},
286
                        {"Name": "OutputTokens", "Unit": "Count"},
287
                        {"Name": "CacheReadTokens", "Unit": "Count"},
288
                        {"Name": "CacheWriteTokens", "Unit": "Count"},
289
                        {"Name": "Latency", "Unit": "Milliseconds"},
290
                        {"Name": "Invocations", "Unit": "Count"},
291
                        {"Name": "Errors", "Unit": "Count"},
292
                    ],
293
                }
294
            ],
295
        },
296
        "Model": model,
297
        "InputTokens": attributes.get("gen_ai.usage.input_tokens", 0),
298
        "OutputTokens": attributes.get("gen_ai.usage.output_tokens", 0),
299
        # pydantic-ai sets these only when non-zero, so default to 0. Providers
300
        # without prompt caching (e.g. the Mistral API) simply never report them.
301
        "CacheReadTokens": attributes.get("gen_ai.usage.cache_read.input_tokens", 0),
302
        "CacheWriteTokens": attributes.get("gen_ai.usage.cache_creation.input_tokens", 0),
303
        "Latency": (root.end_time - root.start_time) / 1_000_000,
304
        "Invocations": 1,
305
        "Errors": 1 if errored else 0,
306
    }
307
    log.info("trace_metrics", **record)
308

309

310
def _on_trace_complete(spans: Sequence[ReadableSpan]) -> None:
311
    """Ship the audit copy, then emit metrics: one hook, two sinks."""
312
    _ship_trace(spans)
313
    _emit_emf(spans)

61
# JSON logs to stdout, which CloudWatch Logs ingests as-is. The same stream also
62
# carries the EMF metric envelope (see _emit_emf), so one structured sink covers
63
# both application logs and metrics. Logging has no extension dependency, so it
64
# is configured at import rather than on the first INVOKE.
65
structlog.configure(
66
    processors=[
67
        structlog.processors.add_log_level,
68
        structlog.processors.TimeStamper(fmt="iso"),
69
        structlog.processors.EventRenamer("message"),
70
        structlog.processors.JSONRenderer(),
71
    ],
72
    logger_factory=structlog.PrintLoggerFactory(),
73
    cache_logger_on_first_use=True,
74
)
75
log = structlog.get_logger()