Skip to content

tdfs4ds.utils.credential_renewal

tdfs4ds.utils.credential_renewal

Renew an expired LLM credential in the middle of a long run (v0.3.1.215).

rebuild_documentation() and vector indexing make hundreds of LLM / embedding calls. With a short-lived credential -- typically a temporary Bedrock bearer token (INSTRUCT_MODEL_API_KEY, copied to AWS_BEARER_TOKEN_BEDROCK by _inject_provider_env) -- the token can expire half-way through, and the run used to die on the first call after expiry.

Renewal is cheap because every LLM call site rebuilds its client from tdfs4ds.INSTRUCT_MODEL_API_KEY (build_llm is called per call). So once the session holds a new token, retrying the failed call is enough -- nothing has to be resumed. :func:call_with_credential_renewal does exactly that: run the call, and on an expiry error ask for a new token (hidden input) and retry, up to tdfs4ds.CREDENTIAL_RENEWAL_ATTEMPTS times.

Two rules, each with a test:

  • The callable must build its client inside itself. A client constructed before the renewal still holds the old credential.
  • The classifier matches AWS's own wording, never a bare word or number. The reconnect classifier once matched "session", which every teradatasql message carries; a credential prompt on an unrelated 400 would be the same mistake with a password-shaped prompt attached.

When there is nobody to ask (a server, papermill, a script) nothing is prompted and the original error propagates unchanged.

_EXPIRED_MARKERS module-attribute

_EXPIRED_MARKERS = ('expiredtokenexception', 'expiredtoken', 'unrecognizedclientexception', 'security token included in the request is expired', 'security token included in the request is invalid', 'the provided token has expired', 'bearer token has expired', 'bearer token is expired', 'bearer token is invalid', 'invalid bearer token', 'api key has expired', 'api key is expired', 'please make sure your api key is valid')

_DECLINED_FOR module-attribute

_DECLINED_FOR = None

logger_safe

logger_safe(level, message, *args, **kwargs)

Wrapper around the global logger that only emits logs when tdfs4ds.DISPLAY_LOGS is True. level is a string like "info", "error", etc.

_texts

_texts(exc)

Yield "<Class>: <message>" for exc and its explicit causes.

__cause__ only, never __context__: an exception raised inside an except block carries an unrelated failure as its context.

is_credential_expired_error

is_credential_expired_error(exc)

True when exc (or an explicit cause) says the LLM credential expired or was rejected as invalid.

reset_renewal_state

reset_renewal_state()

Forget a previous "stop" answer, so the next expiry prompts again.

_declined

_declined()

_max_attempts

_max_attempts()

apply_new_token

apply_new_token(token)

Install token in the session so the next client built picks it up.

renew_credentials

renew_credentials(exc, attempt, max_attempts, what='')

Ask for a new token. Returns True when one was installed.

Never raises and never logs the token (only its length).

call_with_credential_renewal

call_with_credential_renewal(fn, *, what='', on_renewal=None)

Call fn(); on an expired-credential error, ask for a new token and retry. fn must build its LLM / embedding client inside itself.

on_renewal (optional) runs after each successful renewal, before the retry -- e.g. to rebuild a client the caller holds on to.

Any other error, a declined prompt, the absence of anybody to ask, or exhausting tdfs4ds.CREDENTIAL_RENEWAL_ATTEMPTS re-raises the error unchanged.