New to the series? Tooling, AWS access, and project setup are covered in Part 1 (linked above).
What this post covers
Recently the US government decided to put export controls in place for Anthropic Mythos and Fable models. See here for details. While this is only for the recently released Fable/Mythos models, it did get me thinking about the increasing risk of reliance on US only foundation models. While I am obviously aware that this post is still running on AWS, I wanted to at least make a move to a European foundational model.
Admittedly, there is not a grand deal of choice and it also meant moving away from AWS Bedrock. Bedrock does have a few Mistral models, but regions are extremely inflexible and the specific one I wanted to use (Mistral Large 3) was not available in the EU at all (the model card says so, but it is not). Losing Bedrock also meant losing direct integration with CloudWatch, but luckily the decision to go with OTLP for audit meant I already had the code hooked up to extract these metrics out of the trace. That in combination with EMF (Embedded Metrics Format) meant I could easily send these as custom metrics to CloudWatch without a great deal of code changes.
Originally I had planned to only add the ability to switch between the models later when we get to evaluation, but with the recent events I changed the order, so the new code does still support Haiku via Bedrock, but added also the ability to use Mistral models via Mistral's API.
The final tree. + is new in post 3, ~ extends a post 2 file, blank carries unchanged. Click any changed or new file to read it; the download below fast-forwards to this state.
terraform-pr-agent/
agent/
__init__.py
~ handler.py
infra/
~ alerts.tf
audit-bucket.tf
bedrock.tf
~ cloudwatch.tf
firehose.tf
iam.tf
kms.tf
~ lambda.tf
logfire.tf
main.tf
+ models.tf
~ variables.tf
scripts/
build-lambda.sh
chat.py
queries.sql
traces.sql
tests/
+ conftest.py
+ test_handler.py
.envrc
~ .envrc.local
.gitignore
AGENTS.md
~ pyproject.toml
Browse these files interactively on the original post.
Fast-forward to the final code of this post:
mkdir -p ~/projects
cd ~/projects
curl -fsSL https://andreaslang.dev/terraform-pr-agent/terraform-pr-agent-03.tar.gz | tar xz
To use Mistral models, you will need to create an API key and configure it in your .envrc.local file. Sign up here and create an API key here. For this post's usage the free tier is fine, but you may as well load 10 Euros on it and switch to the "Scale" plan of the API. Otherwise you will very quickly receive 429 errors.
Architecture
Post 2 ran a single Bedrock model behind the Lambda and shipped spans to Logfire and the S3 audit copy. Post 3 keeps that intact and turns the model into a runtime choice: Terraform renders a model registry into SSM Parameter Store, the handler builds the pydantic-ai model on first invoke by reading that registry, and a Mistral API entry sits alongside the Bedrock one (with the Mistral key fetched from SSM the same way as the Logfire token). Metrics move to EMF, so a Bedrock model and a Mistral-API model land in the same CloudWatch namespace and one dashboard covers both.
See this diagram full-size on the original post.
The model registry
To support both models I am passing a simple config via AWS SSM Parameter Store into the Lambda. It defines provider model id and if on bedrock inference profile to be used.
infra/models.tf
# The model registry: Terraform owns it, renders it to JSON, and parks it in
# an SSM String parameter the handler reads at startup. Each entry names a
# provider and a model id; Bedrock entries also carry the inference-profile
# ARN. DEFAULT_MODEL (set on the Lambda) selects the active one, so switching
# the agent's model is a parameter change, not a code change.
locals {
metrics_namespace = "TerraformPrAgent/Models"
models = {
haiku = {
provider = "bedrock"
model_id = local.bedrock_model_id
inference_profile_arn = aws_bedrock_inference_profile.agent.arn
}
"mistral-large" = {
provider = "mistral"
model_id = "mistral-large-latest"
}
"devstral-small" = {
provider = "mistral"
model_id = "devstral-small-2507"
}
}
mistral_key_wired = var.mistral_api_key != ""
}
resource "aws_ssm_parameter" "models" {
name = "/terraform-pr-agent/models"
description = "Model registry for the terraform-pr-agent Lambda (provider + model id per entry)."
type = "String"
value = jsonencode(local.models)
}
In addition we need a Mistral API key wired and retrieved the same way as the Logfire key via SSM Parameter Store (encrypted).
infra/models.tf
# The Mistral API key, SecureString, fetched by the handler through the same
# Parameters and Secrets extension path as the Logfire token. Only created
# when TF_VAR_mistral_api_key is set, mirroring the Logfire token wiring; with
# it unset the Mistral providers are simply unreachable and a Bedrock default
# still works.
resource "aws_ssm_parameter" "mistral_api_key" {
count = local.mistral_key_wired ? 1 : 0
name = "/terraform-pr-agent/mistral-api-key"
description = "Mistral API key. Consumed by the terraform-pr-agent Lambda."
type = "SecureString"
value = var.mistral_api_key
}
Building the model at invoke time
Now that we support Bedrock and Mistral models, we just need to create the right pydantic-ai model object with the matching configuration. The handler has also been modified so the model to be used can be provided via the event payload. The default is Mistral Large 3 if nothing is provided.
agent/handler.py
@cache
def _build_model(name: str) -> Model:
"""Build the pydantic-ai model registered under ``name``.
The registry lives in an SSM String parameter, so this runs on the first
INVOKE (the extension is not ready during INIT) and is memoised per model
name for warm invocations. Bedrock models authenticate via the Lambda
role; Mistral models read an API key from a SecureString parameter,
fetched the same way as the Logfire token.
"""
registry = json.loads(_fetch_ssm_parameter(os.environ["MODELS_PARAMETER"]))
config = registry[name]
provider = config["provider"]
if provider == "bedrock":
return BedrockConverseModel(
config["model_id"],
settings={"bedrock_inference_profile": config["inference_profile_arn"]},
)
if provider == "mistral":
key_param = os.environ.get("MISTRAL_API_KEY_PARAMETER")
if not key_param:
raise RuntimeError(
f"model {name!r} uses the Mistral API, but MISTRAL_API_KEY_PARAMETER "
"is not set. Set MISTRAL_API_KEY and re-apply so the key is wired, or "
"select a Bedrock model via DEFAULT_MODEL or the event's model field."
)
return MistralModel(
config["model_id"],
provider=MistralProvider(
api_key=_fetch_ssm_parameter(key_param),
http_client=_retrying_http_client(),
),
)
raise ValueError(f"unknown provider {provider!r} for model {name!r}")
Provider-agnostic metrics with EMF
To avoid having one model via the inference profile and the Mistral models via a different mechanism, we switch all models to use EMF logged metrics, so we can build a clean dashboard (check it in the code you can download above).
agent/handler.py
def _emit_emf(spans: Sequence[ReadableSpan]) -> None:
"""Emit one EMF metric line for the trace, read off the root span.
pydantic-ai records gen_ai.usage.* on the root agent span as the run total
(the sum of its child chat spans), so a single read is the correct total,
not a sum across every span. The model dimension is the registry key the
handler passed as run metadata; pydantic-ai serialises that to the root
span's `metadata` attribute (even on a failed run), so it is read back here
rather than carried in module state. That key is exactly what the dashboard
iterates, so a Bedrock run and a Mistral run share one set of widgets.
Logging the _aws envelope to stdout is enough; CloudWatch Logs extracts the
metrics from the structured line.
"""
root = next((span for span in spans if span.parent is None), None)
if root is None:
return
attributes = root.attributes or {}
model = json.loads(attributes.get("metadata", "{}")).get("model", "unknown")
errored = root.status.status_code is StatusCode.ERROR
record = {
"_aws": {
"Timestamp": root.end_time // 1_000_000,
"CloudWatchMetrics": [
{
"Namespace": os.environ["METRICS_NAMESPACE"],
"Dimensions": [["Model"]],
"Metrics": [
{"Name": "InputTokens", "Unit": "Count"},
{"Name": "OutputTokens", "Unit": "Count"},
{"Name": "CacheReadTokens", "Unit": "Count"},
{"Name": "CacheWriteTokens", "Unit": "Count"},
{"Name": "Latency", "Unit": "Milliseconds"},
{"Name": "Invocations", "Unit": "Count"},
{"Name": "Errors", "Unit": "Count"},
],
}
],
},
"Model": model,
"InputTokens": attributes.get("gen_ai.usage.input_tokens", 0),
"OutputTokens": attributes.get("gen_ai.usage.output_tokens", 0),
# pydantic-ai sets these only when non-zero, so default to 0. Providers
# without prompt caching (e.g. the Mistral API) simply never report them.
"CacheReadTokens": attributes.get("gen_ai.usage.cache_read.input_tokens", 0),
"CacheWriteTokens": attributes.get("gen_ai.usage.cache_creation.input_tokens", 0),
"Latency": (root.end_time - root.start_time) / 1_000_000,
"Invocations": 1,
"Errors": 1 if errored else 0,
}
log.info("trace_metrics", **record)
def _on_trace_complete(spans: Sequence[ReadableSpan]) -> None:
"""Ship the audit copy, then emit metrics: one hook, two sinks."""
_ship_trace(spans)
_emit_emf(spans)
You might also wonder about log.info("trace_metrics", **record) and how this logs in the right format for EMF. Well, the answer is I sneaked in structlog. It is an amazing Python logging library that has all the things and ease of use the standard logging library misses.
agent/handler.py
# JSON logs to stdout, which CloudWatch Logs ingests as-is. The same stream also
# carries the EMF metric envelope (see _emit_emf), so one structured sink covers
# both application logs and metrics. Logging has no extension dependency, so it
# is configured at import rather than on the first INVOKE.
structlog.configure(
processors=[
structlog.processors.add_log_level,
structlog.processors.TimeStamper(fmt="iso"),
structlog.processors.EventRenamer("message"),
structlog.processors.JSONRenderer(),
],
logger_factory=structlog.PrintLoggerFactory(),
cache_logger_on_first_use=True,
)
log = structlog.get_logger()
End State
Ease of switching between models and EMF logging/monitoring configured and the ability to run a (good) European foundation model 🇪🇺!
Coming next: workspace and small toolkit for the agent to get to work.
Top comments (0)