Model portability: swapping Bedrock for the Mistral API

 """AWS Lambda handler for the terraform-pr-agent.

Model portability: the agent runs on whichever model DEFAULT_MODEL names in
the registry, which Terraform parks in an SSM String parameter (MODELS_PARAMETER).
Each registry entry declares a provider (``bedrock`` or ``mistral``) and a model
id; the handler builds the matching pydantic-ai model. Nothing else about the
agent changes when the model does, which is the whole point of the abstraction.

 First invocation:
-- Reads the Logfire token from SSM via the Parameters and Secrets Lambda
  Extension at http://localhost:2773 when LOGFIRE_TOKEN_PARAMETER is set
  in the env; sets LOGFIRE_TOKEN so the logfire SDK picks it up.
-- Configures logfire with send_to_logfire="if-token-present" and registers
  PerTraceAuditProcessor, a custom OTel SpanProcessor that buffers spans
  by trace_id and ships one OTLP-JSON Firehose record when the trace's
  root span ends.
-- Instruments pydantic-ai with version=5 for the spec-compliant span
  names (invoke_agent, chat, execute_tool) and current GenAI semantic
  conventions. Pinned explicitly so the schema readers see in the
  audit copy stays stable across pydantic-ai releases.
+- Reads the Logfire token and (for Mistral models) the Mistral API key from SSM
  via the Parameters and Secrets Lambda Extension at http://localhost:2773.
+- Configures logfire and registers PerTraceAuditProcessor, the custom OTel
  SpanProcessor that buffers spans by trace_id and, when the root span ends,
  ships one OTLP-JSON Firehose record AND emits one EMF metric line. The EMF
  line is what makes the CloudWatch dashboard provider-agnostic: off Bedrock
  there are no AWS/Bedrock metrics, so the handler emits its own from the same
  gen_ai.* span attributes pydantic-ai records for every provider.

The extension's HTTP server rejects requests during INIT with a
"not ready to serve traffic" 400, so this work runs on the first
INVOKE and is memoised with @cache for subsequent warm invocations.

Handler:
-- Reads `prompt` from the invocation event.
-- Runs the agent synchronously and returns the structured output.

The audit copy lands in Firehose from inside the processor's on_end
when the agent root span closes, so the handler has no flush logic.
A Firehose-side failure raises on the same thread as agent.run_sync
and propagates as a Lambda 5xx; the system-of-record copy is never
silently dropped.
The extension's HTTP server rejects requests during INIT, so the SSM reads (and
the model build that depends on them) run on the first INVOKE and are memoised
with @cache for subsequent warm invocations.
 """

 import json
7 collapsed lines
 import os
 import threading
 import urllib.parse
 import urllib.request
 from collections.abc import Callable, Sequence
 from functools import cache
 from typing import NotRequired, TypedDict

 import boto3
 import logfire
import structlog
 from google.protobuf import json_format
from httpx import AsyncClient, HTTPStatusError, Response
 from logfire.sampling import SamplingOptions
 from opentelemetry.context import (
     _SUPPRESS_INSTRUMENTATION_KEY,
5 collapsed lines
     attach,
     detach,
     set_value,
 )
 from opentelemetry.exporter.otlp.proto.common._internal.trace_encoder import (
     encode_spans,
 )
 from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor
from opentelemetry.trace.status import StatusCode
 from pydantic_ai import Agent
from pydantic_ai.models import Model
 from pydantic_ai.models.bedrock import BedrockConverseModel
from pydantic_ai.models.mistral import MistralModel
from pydantic_ai.providers.mistral import MistralProvider
from pydantic_ai.retries import AsyncTenacityTransport, RetryConfig, wait_retry_after
from tenacity import retry_if_exception_type, stop_after_attempt, wait_exponential

# 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()

def _fetch_logfire_token() -> str | None:
    """Read the Logfire token from the Parameters and Secrets extension.

    Returns None when no token parameter is configured, so the function
    runs fine without the Logfire integration. The extension caches the
    value across invocations, so this is cheap on warm starts.
    """
    parameter_name = os.environ.get("LOGFIRE_TOKEN_PARAMETER")
    if not parameter_name:
        return None
def _fetch_ssm_parameter(name: str) -> str:
    """Read one parameter through the Parameters and Secrets extension.

    Works for plain String and SecureString parameters alike; withDecryption
    is a no-op on a String. safe="" forces urllib to percent-encode the leading
    and embedded slashes in a hierarchical name (/a/b/c -> %2Fa%2Fb%2Fc), which
    the extension requires.
    """
     session_token = os.environ["AWS_SESSION_TOKEN"]

    # safe="" forces urllib to percent-encode the leading and embedded
    # slashes in a hierarchical parameter name (e.g. /a/b/c -> %2Fa%2Fb%2Fc).
    # The Parameters and Secrets extension rejects unencoded slashes with
    # HTTP 400; AWS' own Python sample shows the same %2F-encoded form.
     url = (
         "http://localhost:2773/systemsmanager/parameters/get"
        f"?name={urllib.parse.quote(parameter_name, safe='')}&withDecryption=true"
        f"?name={urllib.parse.quote(name, safe='')}&withDecryption=true"
     )
     req = urllib.request.Request(
         url,
4 collapsed lines
         headers={"X-Aws-Parameters-Secrets-Token": session_token},
     )
     with urllib.request.urlopen(req, timeout=2) as resp:
         payload = json.load(resp)
     return payload["Parameter"]["Value"]


def _fetch_logfire_token() -> str | None:
    """Logfire token, or None when the integration is not wired."""
    name = os.environ.get("LOGFIRE_TOKEN_PARAMETER")
    if not name:
        return None
    return _fetch_ssm_parameter(name)




# Status codes worth retrying: rate limit and transient gateway errors. Others
# (auth, bad request) fail fast so a real problem is not retried five times.
_RETRYABLE_STATUS = frozenset({429, 502, 503, 504})


def _raise_for_retryable(response: Response) -> None:
    if response.status_code in _RETRYABLE_STATUS:
        response.raise_for_status()


def _retrying_http_client() -> AsyncClient:
    """An httpx client that retries rate-limit and transient errors.

    pydantic-ai raises ModelHTTPError on transport errors and does not retry
    them itself, so a rate-limited Mistral call (the agent loop can burst past
    the free-tier per-second cap) would otherwise fail the whole run.
    wait_retry_after honours the Retry-After header Mistral sends on a 429,
    falling back to exponential backoff. Bedrock uses boto3 rather than httpx
    and has its own retry config, so this is the Mistral client only.
    """
    transport = AsyncTenacityTransport(
        config=RetryConfig(
            retry=retry_if_exception_type(HTTPStatusError),
            wait=wait_retry_after(
                fallback_strategy=wait_exponential(multiplier=1, max=60),
                max_wait=300,
            ),
            stop=stop_after_attempt(5),
            reraise=True,
        ),
        validate_response=_raise_for_retryable,
    )
    return AsyncClient(transport=transport)


@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}")




 class PerTraceAuditProcessor(SpanProcessor):
    """Buffer spans by trace_id, ship as one batch when the root ends.
    """Buffer spans by trace_id, hand them off as one batch when the root ends.

    The OTel SDK has no `OnTraceComplete` hook, so this implements it
    against the only signal available: `on_end` fires synchronously and
    `span.parent is None` on a root. Late children (spans ended on a
    transport thread after the root has already shipped) are dropped,
    mirroring logfire's tail sampler. See pydantic/logfire#1034.
    The OTel SDK has no `OnTraceComplete` hook, so this implements it against
    the only signal available: `on_end` fires synchronously and `span.parent
    is None` on a root. Late children (spans ended on a transport thread after
    the root has already shipped) are dropped, mirroring logfire's tail sampler.
    See pydantic/logfire#1034.
     """

     def __init__(
42 collapsed lines
         self,
         on_trace_complete: Callable[[Sequence[ReadableSpan]], None],
     ) -> None:
         self._on_trace_complete = on_trace_complete
         self._buffers: dict[int, list[ReadableSpan]] = {}
         self._shipped: set[int] = set()
         self._lock = threading.Lock()

     def on_end(self, span: ReadableSpan) -> None:
         if not (span.context and span.context.trace_flags.sampled):
             return
         trace_id = span.context.trace_id
         with self._lock:
             if trace_id in self._shipped:
                 return
             self._buffers.setdefault(trace_id, []).append(span)
             if span.parent is not None:
                 return
             spans = self._buffers.pop(trace_id)
             self._shipped.add(trace_id)
         self._ship(spans)

     def force_flush(self, timeout_millis: int = 30000) -> bool:
         with self._lock:
             pending = list(self._buffers.values())
             self._shipped.update(self._buffers)
             self._buffers.clear()
         for spans in pending:
             self._ship(spans)
         return True

     def shutdown(self) -> None:
         self.force_flush()

     def _ship(self, spans: Sequence[ReadableSpan]) -> None:
         # Suppress instrumentation around the callback so an instrumented
         # boto3/requests client inside it does not emit a span that
         # re-enters on_end for a sibling trace.
         token = attach(set_value(_SUPPRESS_INSTRUMENTATION_KEY, True))
         try:
             self._on_trace_complete(spans)
         finally:
             detach(token)




 _firehose = boto3.client("firehose")
 _DELIVERY_STREAM = os.environ["FIREHOSE_DELIVERY_STREAM"]

7 collapsed lines

 def _ship_trace(spans: Sequence[ReadableSpan]) -> None:
     """Serialise one trace as OTLP-JSON and ship it as a single Firehose record."""
     payload = json_format.MessageToJson(encode_spans(spans), indent=None) + "\n"
     _firehose.put_record(
         DeliveryStreamName=_DELIVERY_STREAM,
         Record={"Data": payload.encode("utf-8")},
     )


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)




 @cache
 def _init_logfire() -> None:
    """Wire Logfire and the audit span processor once per warm
    """Wire Logfire and the audit/metrics span processor once per warm
     container, on the first INVOKE.

    The Parameters and Secrets extension is not ready to serve traffic
    during the Lambda INIT phase, so the token fetch (and the matching
    logfire setup) cannot run at module import time. @cache memoises
    on the empty argument tuple, so this runs exactly once per
    The Parameters and Secrets extension is not ready to serve traffic during
    the Lambda INIT phase, so the token fetch cannot run at module import time.
    @cache memoises on the empty argument tuple, so this runs exactly once per
     container and is a no-op on every subsequent invocation.
     """
     token = _fetch_logfire_token()
     if token:
         os.environ["LOGFIRE_TOKEN"] = token
    # head=1.0 and tail=None are today's Logfire defaults; pinned here
    # because this is an audit pipeline, so every trace must reach S3.
    # Volume is low (one trace per Lambda invocation) and the audit
    # requirement outweighs Logfire ingest cost. Splitting the rates
    # (e.g. 1% to Logfire, 100% to S3) is possible with a small extra
    # sampler; see the post.
     logfire.configure(
         send_to_logfire="if-token-present",
         sampling=SamplingOptions(head=1.0, tail=None),
        additional_span_processors=[PerTraceAuditProcessor(_ship_trace)],
        additional_span_processors=[PerTraceAuditProcessor(_on_trace_complete)],
     )
    # include_content=True is the pydantic-ai default; pinned because the
    # audit copy needs the actual prompts, tool args, and responses to be
    # useful for after-the-fact forensics. If that ever becomes a
    # compliance problem (PII, secrets in prompts), flip to False and
    # accept a metadata-only audit trail.
     logfire.instrument_pydantic_ai(version=5, include_content=True)




_INFERENCE_PROFILE_ARN = os.environ["BEDROCK_INFERENCE_PROFILE_ARN"]
_MODEL_ID = os.environ["BEDROCK_MODEL_ID"]


def _build_model() -> BedrockConverseModel:
    return BedrockConverseModel(
        _MODEL_ID,
        settings={"bedrock_inference_profile": _INFERENCE_PROFILE_ARN},
    )




 SYSTEM_PROMPT = (
     "You are the terraform-pr-agent. For now you are a placeholder; "
     "respond briefly to whatever prompt you are given."
 )


agent = Agent(_build_model(), system_prompt=SYSTEM_PROMPT)
# The agent carries no default model: the model is built at INVOKE from the
# registry and passed per run. Tests override the model on this instance, which
# takes precedence over the per-run model.
agent = Agent(system_prompt=SYSTEM_PROMPT)


 class HandlerEvent(TypedDict):
     prompt: NotRequired[str]
    model: NotRequired[str]


 class HandlerResponse(TypedDict):
     status: str
    model: str
     output: str


 def handler(event: HandlerEvent, context: object) -> HandlerResponse:
     """Lambda entry point.

    Falls back to a default prompt so the function can be smoke-tested
    with an empty payload.

    The audit copy ships from inside PerTraceAuditProcessor.on_end when
    the agent's root span closes, so the handler does not need a finally
    block: a Firehose failure raises on the same thread as agent.run_sync
    and propagates as a Lambda 5xx. A failed agent run also closes its
    root span (with status=ERROR) before the exception unwinds, so the
    partial trace still ships.
    Falls back to a default prompt so the function can be smoke-tested with an
    empty payload. ``model`` in the event overrides DEFAULT_MODEL for the run,
    so a single deployment can be exercised against any registry entry without
    redeploying. The audit copy and the EMF metric line both ship from
    PerTraceAuditProcessor.on_end when the agent's root span closes, so the
    handler has no flush logic.
     """
     _init_logfire()
     prompt = event.get("prompt", "Say hello.")
    result = agent.run_sync(prompt)
    model_name = event.get("model", os.environ["DEFAULT_MODEL"])
    result = agent.run_sync(prompt, model=_build_model(model_name), metadata={"model": model_name})
     return {
         "status": "ok",
        "model": model_name,
         "output": str(result.output),
     }

1
"""AWS Lambda handler for the terraform-pr-agent.
2

3
Model portability: the agent runs on whichever model DEFAULT_MODEL names in
4
the registry, which Terraform parks in an SSM String parameter (MODELS_PARAMETER).
5
Each registry entry declares a provider (``bedrock`` or ``mistral``) and a model
6
id; the handler builds the matching pydantic-ai model. Nothing else about the
7
agent changes when the model does, which is the whole point of the abstraction.
8

9
First invocation:
10
- Reads the Logfire token and (for Mistral models) the Mistral API key from SSM
11
  via the Parameters and Secrets Lambda Extension at http://localhost:2773.
12
- Configures logfire and registers PerTraceAuditProcessor, the custom OTel
13
  SpanProcessor that buffers spans by trace_id and, when the root span ends,
14
  ships one OTLP-JSON Firehose record AND emits one EMF metric line. The EMF
15
  line is what makes the CloudWatch dashboard provider-agnostic: off Bedrock
16
  there are no AWS/Bedrock metrics, so the handler emits its own from the same
17
  gen_ai.* span attributes pydantic-ai records for every provider.
18

19
The extension's HTTP server rejects requests during INIT, so the SSM reads (and
20
the model build that depends on them) run on the first INVOKE and are memoised
21
with @cache for subsequent warm invocations.
22
"""
23

24
import json
25
import os
26
import threading
27
import urllib.parse
28
import urllib.request
29
from collections.abc import Callable, Sequence
30
from functools import cache
31
from typing import NotRequired, TypedDict
32

33
import boto3
34
import logfire
35
import structlog
36
from google.protobuf import json_format
37
from httpx import AsyncClient, HTTPStatusError, Response
38
from logfire.sampling import SamplingOptions
39
from opentelemetry.context import (
40
    _SUPPRESS_INSTRUMENTATION_KEY,
41
    attach,
42
    detach,
43
    set_value,
44
)
45
from opentelemetry.exporter.otlp.proto.common._internal.trace_encoder import (
46
    encode_spans,
47
)
48
from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor
49
from opentelemetry.trace.status import StatusCode
50
from pydantic_ai import Agent
51
from pydantic_ai.models import Model
52
from pydantic_ai.models.bedrock import BedrockConverseModel
53
from pydantic_ai.models.mistral import MistralModel
54
from pydantic_ai.providers.mistral import MistralProvider
55
from pydantic_ai.retries import AsyncTenacityTransport, RetryConfig, wait_retry_after
56
from tenacity import retry_if_exception_type, stop_after_attempt, wait_exponential
57

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

74

75
def _fetch_ssm_parameter(name: str) -> str:
76
    """Read one parameter through the Parameters and Secrets extension.
77

78
    Works for plain String and SecureString parameters alike; withDecryption
79
    is a no-op on a String. safe="" forces urllib to percent-encode the leading
80
    and embedded slashes in a hierarchical name (/a/b/c -> %2Fa%2Fb%2Fc), which
81
    the extension requires.
82
    """
83
    session_token = os.environ["AWS_SESSION_TOKEN"]
84
    url = (
85
        "http://localhost:2773/systemsmanager/parameters/get"
86
        f"?name={urllib.parse.quote(name, safe='')}&withDecryption=true"
87
    )
88
    req = urllib.request.Request(
89
        url,
90
        headers={"X-Aws-Parameters-Secrets-Token": session_token},
91
    )
92
    with urllib.request.urlopen(req, timeout=2) as resp:
93
        payload = json.load(resp)
94
    return payload["Parameter"]["Value"]
95

96

97
def _fetch_logfire_token() -> str | None:
98
    """Logfire token, or None when the integration is not wired."""
99
    name = os.environ.get("LOGFIRE_TOKEN_PARAMETER")
100
    if not name:
101
        return None
102
    return _fetch_ssm_parameter(name)
103

104

105

106

107
# Status codes worth retrying: rate limit and transient gateway errors. Others
108
# (auth, bad request) fail fast so a real problem is not retried five times.
109
_RETRYABLE_STATUS = frozenset({429, 502, 503, 504})
110

111

112
def _raise_for_retryable(response: Response) -> None:
113
    if response.status_code in _RETRYABLE_STATUS:
114
        response.raise_for_status()
115

116

117
def _retrying_http_client() -> AsyncClient:
118
    """An httpx client that retries rate-limit and transient errors.
119

120
    pydantic-ai raises ModelHTTPError on transport errors and does not retry
121
    them itself, so a rate-limited Mistral call (the agent loop can burst past
122
    the free-tier per-second cap) would otherwise fail the whole run.
123
    wait_retry_after honours the Retry-After header Mistral sends on a 429,
124
    falling back to exponential backoff. Bedrock uses boto3 rather than httpx
125
    and has its own retry config, so this is the Mistral client only.
126
    """
127
    transport = AsyncTenacityTransport(
128
        config=RetryConfig(
129
            retry=retry_if_exception_type(HTTPStatusError),
130
            wait=wait_retry_after(
131
                fallback_strategy=wait_exponential(multiplier=1, max=60),
132
                max_wait=300,
133
            ),
134
            stop=stop_after_attempt(5),
135
            reraise=True,
136
        ),
137
        validate_response=_raise_for_retryable,
138
    )
139
    return AsyncClient(transport=transport)
140

141

142
@cache
143
def _build_model(name: str) -> Model:
144
    """Build the pydantic-ai model registered under ``name``.
145

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

177

178

179

180
class PerTraceAuditProcessor(SpanProcessor):
181
    """Buffer spans by trace_id, hand them off as one batch when the root ends.
182

183
    The OTel SDK has no `OnTraceComplete` hook, so this implements it against
184
    the only signal available: `on_end` fires synchronously and `span.parent
185
    is None` on a root. Late children (spans ended on a transport thread after
186
    the root has already shipped) are dropped, mirroring logfire's tail sampler.
187
    See pydantic/logfire#1034.
188
    """
189

190
    def __init__(
191
        self,
192
        on_trace_complete: Callable[[Sequence[ReadableSpan]], None],
193
    ) -> None:
194
        self._on_trace_complete = on_trace_complete
195
        self._buffers: dict[int, list[ReadableSpan]] = {}
196
        self._shipped: set[int] = set()
197
        self._lock = threading.Lock()
198

199
    def on_end(self, span: ReadableSpan) -> None:
200
        if not (span.context and span.context.trace_flags.sampled):
201
            return
202
        trace_id = span.context.trace_id
203
        with self._lock:
204
            if trace_id in self._shipped:
205
                return
206
            self._buffers.setdefault(trace_id, []).append(span)
207
            if span.parent is not None:
208
                return
209
            spans = self._buffers.pop(trace_id)
210
            self._shipped.add(trace_id)
211
        self._ship(spans)
212

213
    def force_flush(self, timeout_millis: int = 30000) -> bool:
214
        with self._lock:
215
            pending = list(self._buffers.values())
216
            self._shipped.update(self._buffers)
217
            self._buffers.clear()
218
        for spans in pending:
219
            self._ship(spans)
220
        return True
221

222
    def shutdown(self) -> None:
223
        self.force_flush()
224

225
    def _ship(self, spans: Sequence[ReadableSpan]) -> None:
226
        # Suppress instrumentation around the callback so an instrumented
227
        # boto3/requests client inside it does not emit a span that
228
        # re-enters on_end for a sibling trace.
229
        token = attach(set_value(_SUPPRESS_INSTRUMENTATION_KEY, True))
230
        try:
231
            self._on_trace_complete(spans)
232
        finally:
233
            detach(token)
234

235

236
_firehose = boto3.client("firehose")
237
_DELIVERY_STREAM = os.environ["FIREHOSE_DELIVERY_STREAM"]
238

239

240
def _ship_trace(spans: Sequence[ReadableSpan]) -> None:
241
    """Serialise one trace as OTLP-JSON and ship it as a single Firehose record."""
242
    payload = json_format.MessageToJson(encode_spans(spans), indent=None) + "\n"
243
    _firehose.put_record(
244
        DeliveryStreamName=_DELIVERY_STREAM,
245
        Record={"Data": payload.encode("utf-8")},
246
    )
247

248

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

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

300

301
def _on_trace_complete(spans: Sequence[ReadableSpan]) -> None:
302
    """Ship the audit copy, then emit metrics: one hook, two sinks."""
303
    _ship_trace(spans)
304
    _emit_emf(spans)
305

306

307

308

309
@cache
310
def _init_logfire() -> None:
311
    """Wire Logfire and the audit/metrics span processor once per warm
312
    container, on the first INVOKE.
313

314
    The Parameters and Secrets extension is not ready to serve traffic during
315
    the Lambda INIT phase, so the token fetch cannot run at module import time.
316
    @cache memoises on the empty argument tuple, so this runs exactly once per
317
    container and is a no-op on every subsequent invocation.
318
    """
319
    token = _fetch_logfire_token()
320
    if token:
321
        os.environ["LOGFIRE_TOKEN"] = token
322
    logfire.configure(
323
        send_to_logfire="if-token-present",
324
        sampling=SamplingOptions(head=1.0, tail=None),
325
        additional_span_processors=[PerTraceAuditProcessor(_on_trace_complete)],
326
    )
327
    logfire.instrument_pydantic_ai(version=5, include_content=True)
328

329

330
SYSTEM_PROMPT = (
331
    "You are the terraform-pr-agent. For now you are a placeholder; "
332
    "respond briefly to whatever prompt you are given."
333
)
334

335

336
# The agent carries no default model: the model is built at INVOKE from the
337
# registry and passed per run. Tests override the model on this instance, which
338
# takes precedence over the per-run model.
339
agent = Agent(system_prompt=SYSTEM_PROMPT)
340

341

342
class HandlerEvent(TypedDict):
343
    prompt: NotRequired[str]
344
    model: NotRequired[str]
345

346

347
class HandlerResponse(TypedDict):
348
    status: str
349
    model: str
350
    output: str
351

352

353
def handler(event: HandlerEvent, context: object) -> HandlerResponse:
354
    """Lambda entry point.
355

356
    Falls back to a default prompt so the function can be smoke-tested with an
357
    empty payload. ``model`` in the event overrides DEFAULT_MODEL for the run,
358
    so a single deployment can be exercised against any registry entry without
359
    redeploying. The audit copy and the EMF metric line both ship from
360
    PerTraceAuditProcessor.on_end when the agent's root span closes, so the
361
    handler has no flush logic.
362
    """
363
    _init_logfire()
364
    prompt = event.get("prompt", "Say hello.")
365
    model_name = event.get("model", os.environ["DEFAULT_MODEL"])
366
    result = agent.run_sync(prompt, model=_build_model(model_name), metadata={"model": model_name})
367
    return {
368
        "status": "ok",
369
        "model": model_name,
370
        "output": str(result.output),
371
    }

# alias/aws/sns is the AWS-managed SNS key; a CMK is unnecessary for an
# operational alarm topic that publishes alarm-fired events, not audit data.
#trivy:ignore:avd-aws-0136
11 collapsed lines
resource "aws_sns_topic" "agent_alerts" {
  name              = "terraform-pr-agent-alerts"
  kms_master_key_id = "alias/aws/sns"
}

# The subscription stays PENDING until the recipient clicks the AWS
# confirmation email. Alarms fire either way, but no email goes out until the
# subscription is confirmed.
resource "aws_sns_topic_subscription" "agent_alerts_email" {
  topic_arn = aws_sns_topic.agent_alerts.arn
  protocol  = "email"
  endpoint  = var.alert_email
}

# Reads the EMF tokens the handler emits (Model dimension), not AWS/Bedrock,
# so the budget guard works whichever provider DEFAULT_MODEL points at. Scoped
# to the active model.
resource "aws_cloudwatch_metric_alarm" "daily_tokens" {
  alarm_name          = "${aws_bedrock_inference_profile.agent.name}-daily-tokens"
  alarm_description   = "Daily input + output token usage for the agent inference profile exceeded the configured threshold."
  alarm_name          = "terraform-pr-agent-daily-tokens"
  alarm_description   = "Daily input + output token usage for the active model exceeded the configured threshold."
  comparison_operator = "GreaterThanThreshold"
  evaluation_periods  = 1
  threshold           = var.daily_token_alarm_threshold
9 collapsed lines
  treat_missing_data  = "notBreaching"

  metric_query {
    id          = "total"
    expression  = "input + output"
    label       = "Total tokens (input + output)"
    return_data = true
  }

  metric_query {
    id = "input"
    metric {
      namespace   = "AWS/Bedrock"
      metric_name = "InputTokenCount"
      dimensions  = { ModelId = local.profile_id }
      namespace   = local.metrics_namespace
      metric_name = "InputTokens"
      dimensions  = { Model = var.default_model }
      stat        = "Sum"
      period      = 86400
    }
2 collapsed lines
  }

  metric_query {
    id = "output"
    metric {
      namespace   = "AWS/Bedrock"
      metric_name = "OutputTokenCount"
      dimensions  = { ModelId = local.profile_id }
      namespace   = local.metrics_namespace
      metric_name = "OutputTokens"
      dimensions  = { Model = var.default_model }
      stat        = "Sum"
      period      = 86400
    }
  }

  alarm_actions = [aws_sns_topic.agent_alerts.arn]
}

1
# alias/aws/sns is the AWS-managed SNS key; a CMK is unnecessary for an
2
# operational alarm topic that publishes alarm-fired events, not audit data.
3
#trivy:ignore:avd-aws-0136
4
resource "aws_sns_topic" "agent_alerts" {
5
  name              = "terraform-pr-agent-alerts"
6
  kms_master_key_id = "alias/aws/sns"
7
}
8

9
# The subscription stays PENDING until the recipient clicks the AWS
10
# confirmation email. Alarms fire either way, but no email goes out until the
11
# subscription is confirmed.
12
resource "aws_sns_topic_subscription" "agent_alerts_email" {
13
  topic_arn = aws_sns_topic.agent_alerts.arn
14
  protocol  = "email"
15
  endpoint  = var.alert_email
16
}
17

18
# Reads the EMF tokens the handler emits (Model dimension), not AWS/Bedrock,
19
# so the budget guard works whichever provider DEFAULT_MODEL points at. Scoped
20
# to the active model.
21
resource "aws_cloudwatch_metric_alarm" "daily_tokens" {
22
  alarm_name          = "terraform-pr-agent-daily-tokens"
23
  alarm_description   = "Daily input + output token usage for the active model exceeded the configured threshold."
24
  comparison_operator = "GreaterThanThreshold"
25
  evaluation_periods  = 1
26
  threshold           = var.daily_token_alarm_threshold
27
  treat_missing_data  = "notBreaching"
28

29
  metric_query {
30
    id          = "total"
31
    expression  = "input + output"
32
    label       = "Total tokens (input + output)"
33
    return_data = true
34
  }
35

36
  metric_query {
37
    id = "input"
38
    metric {
39
      namespace   = local.metrics_namespace
40
      metric_name = "InputTokens"
41
      dimensions  = { Model = var.default_model }
42
      stat        = "Sum"
43
      period      = 86400
44
    }
45
  }
46

47
  metric_query {
48
    id = "output"
49
    metric {
50
      namespace   = local.metrics_namespace
51
      metric_name = "OutputTokens"
52
      dimensions  = { Model = var.default_model }
53
      stat        = "Sum"
54
      period      = 86400
55
    }
56
  }
57

58
  alarm_actions = [aws_sns_topic.agent_alerts.arn]
59
}

1
locals {
2
  audit_bucket_name = "terraform-pr-agent-audit-${data.aws_caller_identity.current.account_id}-${data.aws_region.current.region}"
3
}
4

5
# Access logging would require a second bucket and is out of scope: the audit
6
# copy here is the system of record for what the agent did, not for who read it.
7
#trivy:ignore:avd-aws-0089
8
resource "aws_s3_bucket" "audit" {
9
  bucket              = local.audit_bucket_name
10
  object_lock_enabled = true
11
}
12

13
resource "aws_s3_bucket_versioning" "audit" {
14
  bucket = aws_s3_bucket.audit.id
15

16
  versioning_configuration {
17
    status = "Enabled"
18
  }
19
}
20

21
resource "aws_s3_bucket_public_access_block" "audit" {
22
  bucket                  = aws_s3_bucket.audit.id
23
  block_public_acls       = true
24
  block_public_policy     = true
25
  ignore_public_acls      = true
26
  restrict_public_buckets = true
27
}
28

29
resource "aws_s3_bucket_server_side_encryption_configuration" "audit" {
30
  bucket = aws_s3_bucket.audit.id
31

32
  rule {
33
    apply_server_side_encryption_by_default {
34
      sse_algorithm     = "aws:kms"
35
      kms_master_key_id = aws_kms_key.audit.arn
36
    }
37

38
    bucket_key_enabled = true
39
  }
40
}
41

42
data "aws_iam_policy_document" "audit_bucket_resource_policy" {
43
  statement {
44
    sid     = "DenyInsecureTransport"
45
    effect  = "Deny"
46
    actions = ["s3:*"]
47
    resources = [
48
      aws_s3_bucket.audit.arn,
49
      "${aws_s3_bucket.audit.arn}/*",
50
    ]
51
    principals {
52
      type        = "*"
53
      identifiers = ["*"]
54
    }
55
    condition {
56
      test     = "Bool"
57
      variable = "aws:SecureTransport"
58
      values   = ["false"]
59
    }
60
  }
61
}
62

63
resource "aws_s3_bucket_policy" "audit" {
64
  bucket = aws_s3_bucket.audit.id
65
  policy = data.aws_iam_policy_document.audit_bucket_resource_policy.json
66
}
67

68
# GOVERNANCE leaves the bucket deletable for the tutorial; switch to "COMPLIANCE" in production.
69
resource "aws_s3_bucket_object_lock_configuration" "audit" {
70
  bucket = aws_s3_bucket.audit.id
71

72
  rule {
73
    default_retention {
74
      mode = "GOVERNANCE"
75
      days = var.audit_retention_days
76
    }
77
  }
78
}
79

80
resource "aws_s3_bucket_lifecycle_configuration" "audit" {
81
  bucket = aws_s3_bucket.audit.id
82

83
  rule {
84
    id     = "transition-to-glacier-ir"
85
    status = "Enabled"
86

87
    filter {}
88

89
    transition {
90
      days          = 90
91
      storage_class = "GLACIER_IR"
92
    }
93

94
    noncurrent_version_transition {
95
      noncurrent_days = 90
96
      storage_class   = "GLACIER_IR"
97
    }
98
  }
99

100
  rule {
101
    id     = "expire-after-retention"
102
    status = "Enabled"
103

104
    filter {}
105

106
    expiration {
107
      days = var.audit_retention_days
108
    }
109

110
    noncurrent_version_expiration {
111
      noncurrent_days = 1
112
    }
113
  }
114
}
115

116
output "audit_bucket_name" {
117
  value = aws_s3_bucket.audit.id
118
}
119

120
output "audit_bucket_arn" {
121
  value = aws_s3_bucket.audit.arn
122
}

1
locals {
2
  bedrock_model_id            = "anthropic.claude-haiku-4-5-20251001-v1:0"
3
  bedrock_cross_region_prefix = "eu"
4

5
  system_inference_profile_arn = format(
6
    "arn:aws:bedrock:%s:%s:inference-profile/%s.%s",
7
    data.aws_region.current.region,
8
    data.aws_caller_identity.current.account_id,
9
    local.bedrock_cross_region_prefix,
10
    local.bedrock_model_id,
11
  )
12
}
13

14
resource "aws_bedrock_inference_profile" "agent" {
15
  name        = "terraform-pr-agent"
16
  description = "Application inference profile for the Terraform PR Agent series."
17

18
  model_source {
19
    copy_from = local.system_inference_profile_arn
20
  }
21
}
22

23
output "inference_profile_arn" {
24
  value = aws_bedrock_inference_profile.agent.arn
25
}

locals {
  cloudwatch_region = data.aws_region.current.region
  dashboard_name    = "terraform-pr-agent"

  dashboard_label = aws_bedrock_inference_profile.agent.name
  profile_id      = aws_bedrock_inference_profile.agent.id
  # The dashboard now reads the EMF metrics the handler emits (namespace
  # local.metrics_namespace, dimensioned by Model), not AWS/Bedrock. That is
  # the whole point: a Bedrock model and a Mistral-API model land in the same
  # namespace, so one set of widgets covers both. Each registry model gets its
  # own line, built by iterating keys(local.models).
  model_keys = keys(local.models)
}

resource "aws_cloudwatch_dashboard" "agent" {
  dashboard_name = aws_bedrock_inference_profile.agent.name
  dashboard_name = local.dashboard_name
  dashboard_body = jsonencode({
    widgets = [
      {
9 collapsed lines
        type   = "metric"
        x      = 0
        y      = 0
        width  = 12
        height = 6
        properties = {
          title  = "Tokens"
          region = local.cloudwatch_region
          view   = "timeSeries"
          stat   = "Sum"
          period = 60
          metrics = [
            # here we plug in the labels so this looks nicer
            ["AWS/Bedrock", "InputTokenCount", "ModelId", local.profile_id, { label = "${local.dashboard_label} / input" }],
            # the dot syntax allows not to repeat the same as above
            [".", "OutputTokenCount", ".", ".", { label = "${local.dashboard_label} / output" }],
            for m in flatten([
              for key in local.model_keys : [
                { metric = "InputTokens", label = "${key} / input", key = key },
                { metric = "OutputTokens", label = "${key} / output", key = key },
              ]
            ]) : [local.metrics_namespace, m.metric, "Model", m.key, { label = m.label }]
          ]
        }
      },
4 collapsed lines
      {
        type   = "metric"
        x      = 12
        y      = 0
        width  = 12
        height = 6
        properties = {
          title  = "Cache tokens"
          title  = "Invocations and errors"
          region = local.cloudwatch_region
          view   = "timeSeries"
          stat   = "Sum"
          period = 60
          metrics = [
            ["AWS/Bedrock", "CacheReadInputTokenCount", "ModelId", local.profile_id, { label = "${local.dashboard_label} / cache read" }],
            [".", "CacheWriteInputTokenCount", ".", ".", { label = "${local.dashboard_label} / cache write" }],
            for m in flatten([
              for key in local.model_keys : [
                { metric = "Invocations", label = "${key} / invocations", key = key },
                { metric = "Errors", label = "${key} / errors", key = key },
              ]
            ]) : [local.metrics_namespace, m.metric, "Model", m.key, { label = m.label }]
          ]
        }
      },
4 collapsed lines
      {
        type   = "metric"
        x      = 0
        y      = 6
        width  = 12
        height = 6
        properties = {
          title  = "Invocations and errors"
          title  = "Latency (ms)"
          region = local.cloudwatch_region
          view   = "timeSeries"
          stat   = "Sum"
          period = 60
          metrics = [
            ["AWS/Bedrock", "Invocations", "ModelId", local.profile_id, { label = "${local.dashboard_label} / invocations" }],
            [".", "InvocationClientErrors", ".", ".", { label = "${local.dashboard_label} / 4xx" }],
            [".", "InvocationServerErrors", ".", ".", { label = "${local.dashboard_label} / 5xx" }],
            [".", "InvocationThrottles", ".", ".", { label = "${local.dashboard_label} / throttles" }],
            for m in flatten([
              for key in local.model_keys : [
                { label = "${key} / avg", key = key, stat = "Average" },
                { label = "${key} / p99", key = key, stat = "p99" },
              ]
            ]) : [local.metrics_namespace, "Latency", "Model", m.key, { label = m.label, stat = m.stat }]
          ]
        }
      },
4 collapsed lines
      {
        type   = "metric"
        x      = 12
        y      = 6
        width  = 12
        height = 6
        properties = {
          title  = "Latency (ms)"
          title  = "Cache tokens"
          region = local.cloudwatch_region
          view   = "timeSeries"
          stat   = "Sum"
          period = 60
          metrics = [
            ["AWS/Bedrock", "InvocationLatency", "ModelId", local.profile_id, { label = "${local.dashboard_label} / avg", stat = "Average" }],
            [".", ".", ".", ".", { label = "${local.dashboard_label} / p99", stat = "p99" }],
            for m in flatten([
              for key in local.model_keys : [
                { metric = "CacheReadTokens", label = "${key} / cache read", key = key },
                { metric = "CacheWriteTokens", label = "${key} / cache write", key = key },
              ]
            ]) : [local.metrics_namespace, m.metric, "Model", m.key, { label = m.label }]
          ]
        }
      },
9 collapsed lines
    ]
  })
}

output "cloudwatch_dashboard_url" {
  value = format(
    "https://%s.console.aws.amazon.com/cloudwatch/home?region=%s#dashboards/dashboard/%s",
    local.cloudwatch_region,
    local.cloudwatch_region,
    aws_cloudwatch_dashboard.agent.dashboard_name,
  )
}