Venkatakrishna Reddy Avula
2da639e59b
Feature/DAIP2-2314 DAIP2 1687 hybrid * remove -files from s3 prefix requirements * Resolve input paths * fix: VendorProcessor.process_file returns (df, None) tuple runner.safe_process_file unpacks the result as (cc_df, dashboard_df), so returning a single DataFrame caused every vendor/generic file to fail with "too many values to unpack (expected 2)" — Python iterates DataFrame columns during unpacking. Vendor pipelines have no dashboard variant; second slot is None and the existing `dashboard_result is not None` guard in runner.py already handles it. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * DAIP2-2314 + DAIP2-1687: pad DYNAMIC_PRIMARY + DYNAMIC_PRIMARY_ENTITY_CLASSIFICATION over 1024-token cache floor - Pad DYNAMIC_PRIMARY_INSTRUCTION with three new sections: [SCOPE BOUNDARIES], [SOURCE TEXT INTERPRETATION], [REASONING DISCIPLINE], plus a [WORKED EXAMPLES] block. Estimated tokens: 447 -> 1117 (Sonnet 4.5 1024-min, +93 margin). All additions reinforce existing rules (sibling-field separation, alias mapping, pricing-vs-LOB distinction, contrastive-clause exclusion, exhibit-header binding) — no new directives that could bias extraction. - Pad DYNAMIC_PRIMARY_ENTITY_CLASSIFICATION_INSTRUCTION with a [FINAL CHECKLIST BEFORE OUTPUT] block. Estimated tokens: 956 -> 1101 (Sonnet 4.5 1024-min, +77 margin). Reinforces the existing 4-step anti-duplication protocol and JSON shape requirements. - Register both new entries in cache_registry: DYNAMIC_PRIMARY_ENTITY_CLASSIFICATION as INSTRUCTION_PLUS_CONTEXT (caches at warm-up), DYNAMIC_PRIMARY_ENTITIES as CONTEXT (instruction is intentionally short; CONTEXT c… * black format fix * Merged dev into feature/DAIP2-2314-DAIP2-1687-hybrid * fixed raw lob values in base lob field mapping and composite entities fix * black format fix * fixed LOB Program output issues * issue fixes * remove debugging code * Updated prompts * updated additional instructions * Update Program-->LOB * LLM-based AD-Program/Product mapping to LOB even when there is a crosswalk * black format fix * Merged dev into feature/DAIP2-2314-DAIP2-1687-hybrid * added logging in prompt call tracking * added updated logging in prompt call tracking * aaded min cache token per usage label * added cache registry for dynamic primary mapping prompt calls * reolved mapping prompts ambiguities * black format fix * Phase 2 modifications added * reverted phase 2 modifications Approved-by: Katon Minhas
1499 lines
57 KiB
Python
1499 lines
57 KiB
Python
import base64
|
|
import hashlib
|
|
import inspect
|
|
import json
|
|
import logging
|
|
import os
|
|
import random
|
|
import threading
|
|
import time
|
|
from collections import OrderedDict
|
|
|
|
import anthropic
|
|
import src.utils.string_utils as string_utils
|
|
import src.utils.prompt_call_tracking as prompt_call_tracking
|
|
import src.utils.usage_tracking as usage_tracking
|
|
from botocore.exceptions import ClientError, ReadTimeoutError
|
|
from src import config
|
|
|
|
# Cache warming state
|
|
_cache_warmed = {}
|
|
_cache_warming_lock = threading.Lock()
|
|
|
|
|
|
# ============================================================================
|
|
# MODEL SUPPORT CONFIGURATION
|
|
# ============================================================================
|
|
# Centralized configuration for supported and deprecated models.
|
|
# Model IDs are defined in config.py - update there first, then add to these lists.
|
|
|
|
# Supported models (from config.py)
|
|
# Add new models here when they're added to config.py
|
|
_SUPPORTED_MODELS = {
|
|
config.MODEL_ID_CLAUDE45_HAIKU,
|
|
config.MODEL_ID_CLAUDE35_SONNET,
|
|
config.MODEL_ID_CLAUDE37_SONNET,
|
|
config.MODEL_ID_CLAUDE4_SONNET,
|
|
config.MODEL_ID_CLAUDE45_SONNET,
|
|
}
|
|
|
|
# Deprecated model patterns (fuzzy matching)
|
|
# Add patterns here to deprecate model families (e.g., "claude-2" matches any claude-2 variant)
|
|
# Patterns are case-insensitive and match if found anywhere in the model ID
|
|
_DEPRECATED_PATTERNS = [
|
|
"claude-instant", # Deprecates all Claude Instant/Claude 2 models
|
|
"claude-2",
|
|
]
|
|
|
|
|
|
def _validate_model_support(
|
|
model_id: str, resolved_model_id: str, context: str = ""
|
|
) -> None:
|
|
"""Validate that a model is supported and not deprecated.
|
|
|
|
Checks both exact model IDs and pattern matching for deprecated models.
|
|
Raises ValueError with warning if model is deprecated or unsupported.
|
|
|
|
Args:
|
|
model_id: Original model ID (may be alias)
|
|
resolved_model_id: Resolved model ID from config.resolve_model_id()
|
|
context: Additional context for error message (e.g., "for image array")
|
|
"""
|
|
model_id_lower = resolved_model_id.lower()
|
|
|
|
# Check for deprecated patterns (fuzzy match)
|
|
for pattern in _DEPRECATED_PATTERNS:
|
|
if pattern in model_id_lower:
|
|
# Determine appropriate error message based on pattern
|
|
if "claude-2" in pattern or "claude-instant" in pattern:
|
|
deprecation_msg = (
|
|
"Claude 2 is no longer supported. Please use Claude 3+ models."
|
|
)
|
|
supported_models = "Haiku, Sonnet 3.5/3.7/4/4.5"
|
|
else:
|
|
deprecation_msg = (
|
|
"This model is no longer supported. Please use Claude 3+ models."
|
|
)
|
|
supported_models = "Haiku, Sonnet 3.5/3.7/4/4.5"
|
|
|
|
warning_msg = (
|
|
f"Deprecated model detected: {model_id} (resolved: {resolved_model_id}). "
|
|
f"{deprecation_msg}"
|
|
)
|
|
if context:
|
|
warning_msg = (
|
|
f"Deprecated model detected {context}: {model_id} (resolved: {resolved_model_id}). "
|
|
f"{deprecation_msg}"
|
|
)
|
|
logging.warning(warning_msg)
|
|
raise ValueError(
|
|
f"Unsupported model_id{(' ' + context) if context else ''}: {model_id}. "
|
|
f"{deprecation_msg} Supported models: {supported_models}."
|
|
)
|
|
|
|
# Check if supported (exact match)
|
|
supported_set = {m for m in _SUPPORTED_MODELS if m is not None}
|
|
if resolved_model_id not in supported_set:
|
|
raise ValueError(
|
|
f"Unsupported model_id{(' ' + context) if context else ''}: {model_id} "
|
|
f"(resolved: {resolved_model_id}). "
|
|
f"Only Claude 3+ models are supported (Haiku, Sonnet 3.5/3.7/4/4.5)."
|
|
)
|
|
|
|
|
|
def _supports_prompt_cache(model_id: str) -> bool:
|
|
"""Return True if the given model supports prompt caching via system cache_control.
|
|
|
|
Enable for models that are known to support prompt caching:
|
|
- Claude 3.5 Sonnet v2+ (20241022-v2:0)
|
|
- Claude 3.7 Sonnet
|
|
- Claude 4 Sonnet
|
|
- Claude 4.5 Sonnet
|
|
- Claude 4.5 Haiku
|
|
"""
|
|
supported_ids = {
|
|
getattr(config, "MODEL_ID_CLAUDE35_SONNET_V2", None),
|
|
getattr(config, "MODEL_ID_CLAUDE37_SONNET", None),
|
|
getattr(config, "MODEL_ID_CLAUDE4_SONNET", None),
|
|
getattr(config, "MODEL_ID_CLAUDE45_SONNET", None),
|
|
getattr(config, "MODEL_ID_CLAUDE45_HAIKU", None),
|
|
}
|
|
# Filter out None values in case any model IDs are not configured
|
|
return model_id in {m for m in supported_ids if m is not None}
|
|
|
|
|
|
# Per Anthropic docs (https://docs.anthropic.com/en/docs/prompt-caching),
|
|
# the minimum cacheable prompt length differs by model family:
|
|
# 1024 tokens — Sonnet 3.7 / 4 / 4.5, Opus 4 / 4.1
|
|
# 2048 tokens — Sonnet 4.6, Haiku 3.5 (deprecated)
|
|
# 4096 tokens — Sonnet/Opus Mythos, Opus 4.5/4.6/4.7, Haiku 4.5
|
|
# Caches against blocks below the threshold are silently dropped — the
|
|
# request still succeeds and returns content, but `cache_creation_tokens`
|
|
# and `cache_read_tokens` will both be 0 forever. That is the
|
|
# `no_cache_event` signature surfaced in runtime tracking.
|
|
_MODEL_MIN_CACHE_TOKENS_DEFAULT = 1024
|
|
_MODEL_MIN_CACHE_TOKENS: dict[str, int] = {
|
|
# Sonnet 4.5 — 1024
|
|
getattr(config, "MODEL_ID_CLAUDE45_SONNET", "_unset_45_sonnet"): 1024,
|
|
# Sonnet 4 / 3.7 / 3.5 v2 — 1024
|
|
getattr(config, "MODEL_ID_CLAUDE4_SONNET", "_unset_4_sonnet"): 1024,
|
|
getattr(config, "MODEL_ID_CLAUDE37_SONNET", "_unset_37_sonnet"): 1024,
|
|
getattr(config, "MODEL_ID_CLAUDE35_SONNET_V2", "_unset_35_sonnet_v2"): 1024,
|
|
# Haiku 4.5 — 4096
|
|
getattr(config, "MODEL_ID_CLAUDE45_HAIKU", "_unset_45_haiku"): 4096,
|
|
}
|
|
|
|
|
|
def _min_cache_tokens(model_id: str) -> int:
|
|
"""Return the minimum prompt-cache block size honored by the model.
|
|
|
|
Falls back to 1024 (the historical Sonnet minimum) when the model is
|
|
unknown — that is the most permissive choice; callers only get a
|
|
surprise `no_cache_event` if the actual minimum is higher.
|
|
"""
|
|
return _MODEL_MIN_CACHE_TOKENS.get(model_id, _MODEL_MIN_CACHE_TOKENS_DEFAULT)
|
|
|
|
|
|
def _classify_cache_outcome(
|
|
*,
|
|
has_context_cache: bool,
|
|
cache_read_tokens: int,
|
|
cache_creation_tokens: int,
|
|
context_chars: int,
|
|
cache_key: str | None,
|
|
) -> str:
|
|
"""Classify why a cache-attempting call did or didn't hit the Anthropic cache.
|
|
|
|
Values: hit, first_call, ttl_expired, under_min_tokens, silent_miss, not_attempted.
|
|
Anthropic's minimum cacheable block is 1024 tokens; we approximate as
|
|
context_chars / 3.5. Updates _cache_key_seen on every classify for future calls.
|
|
"""
|
|
if not has_context_cache:
|
|
return "not_attempted"
|
|
if cache_read_tokens > 0:
|
|
if cache_key is not None:
|
|
with _cache_key_seen_lock:
|
|
_cache_key_seen[cache_key] = time.monotonic()
|
|
return "hit"
|
|
if cache_creation_tokens > 0:
|
|
# Fresh write — either this is the first call with this key, or a prior
|
|
# entry expired (TTL ~5 min) and had to be re-written.
|
|
if cache_key is not None:
|
|
with _cache_key_seen_lock:
|
|
prior_seen = cache_key in _cache_key_seen
|
|
_cache_key_seen[cache_key] = time.monotonic()
|
|
return "ttl_expired" if prior_seen else "first_call"
|
|
return "first_call"
|
|
# has_context_cache=True but both token counts are 0 — silent failure mode.
|
|
approx_tokens = int(context_chars / 3.5) if context_chars else 0
|
|
if approx_tokens and approx_tokens < 1024:
|
|
return "under_min_tokens"
|
|
return "silent_miss"
|
|
|
|
|
|
def _ctx_minus_explicit_keys(ctx: dict, *explicit_keys: str) -> dict:
|
|
"""Filter a ctx dict so explicit kwargs aren't duplicated via **ctx splat."""
|
|
skip = {"_usage_label", "segment", *explicit_keys}
|
|
return {k: v for k, v in ctx.items() if not k.startswith("_") and k not in skip}
|
|
|
|
|
|
def _log_llm_retry(
|
|
filename: str,
|
|
model_id: str,
|
|
attempt: int,
|
|
error_class: str,
|
|
error_msg: str,
|
|
backoff_sec: float,
|
|
) -> None:
|
|
"""Emit one llm_retry instrumentation event (no-op when disabled)."""
|
|
from src.utils import instrumentation, instrumentation_context
|
|
|
|
if not instrumentation.is_enabled():
|
|
return
|
|
ctx = dict(instrumentation_context.get_ctx())
|
|
usage_label = ctx.get("_usage_label")
|
|
ctx_segment = ctx.get("segment")
|
|
instrumentation.log(
|
|
"llm_retry",
|
|
filename=filename,
|
|
model=model_id,
|
|
usage_label=usage_label,
|
|
segment=instrumentation.resolve_segment(usage_label, ctx_segment),
|
|
attempt=attempt,
|
|
error_class=error_class,
|
|
error_msg=(error_msg or "")[:200],
|
|
backoff_sec=round(backoff_sec, 4),
|
|
**_ctx_minus_explicit_keys(ctx, "filename", "model", "attempt"),
|
|
)
|
|
|
|
|
|
def _log_llm_error(
|
|
filename: str,
|
|
model_id: str,
|
|
attempt: int,
|
|
error_class: str,
|
|
error_msg: str,
|
|
) -> None:
|
|
"""Emit one llm_error instrumentation event (no-op when disabled).
|
|
|
|
Fires when retries are exhausted and the call is about to raise.
|
|
"""
|
|
from src.utils import instrumentation, instrumentation_context
|
|
|
|
if not instrumentation.is_enabled():
|
|
return
|
|
ctx = dict(instrumentation_context.get_ctx())
|
|
usage_label = ctx.get("_usage_label")
|
|
ctx_segment = ctx.get("segment")
|
|
instrumentation.log(
|
|
"llm_error",
|
|
filename=filename,
|
|
model=model_id,
|
|
usage_label=usage_label,
|
|
segment=instrumentation.resolve_segment(usage_label, ctx_segment),
|
|
attempt=attempt,
|
|
error_class=error_class,
|
|
error_msg=(error_msg or "")[:200],
|
|
**_ctx_minus_explicit_keys(ctx, "filename", "model", "attempt"),
|
|
)
|
|
|
|
|
|
def _track_usage_from_response(
|
|
response_body: dict, filename: str, model_id: str
|
|
) -> None:
|
|
"""Extract and track usage from API response.
|
|
|
|
Args:
|
|
response_body: Parsed JSON response from Bedrock API
|
|
filename: Name of file being processed
|
|
model_id: Model ID used for the API call
|
|
"""
|
|
if not filename:
|
|
return
|
|
|
|
(
|
|
input_tokens,
|
|
output_tokens,
|
|
cache_creation_tokens,
|
|
cache_read_tokens,
|
|
) = usage_tracking.extract_usage_from_response(response_body)
|
|
|
|
if (
|
|
input_tokens > 0
|
|
or output_tokens > 0
|
|
or cache_creation_tokens > 0
|
|
or cache_read_tokens > 0
|
|
):
|
|
usage_tracking.track_usage(
|
|
filename,
|
|
model_id,
|
|
input_tokens,
|
|
output_tokens,
|
|
cache_creation_tokens,
|
|
cache_read_tokens,
|
|
)
|
|
# One llm_call event per Bedrock response with non-zero tokens.
|
|
from src.utils import instrumentation, instrumentation_context
|
|
|
|
if instrumentation.is_enabled():
|
|
from src.utils.usage_tracking import get_cost_per_token
|
|
|
|
input_cpt, output_cpt, cache_read_cpt = get_cost_per_token(model_id)
|
|
if "sonnet" in model_id.lower():
|
|
cache_write_cpt = input_cpt * 1.25
|
|
elif "haiku" in model_id.lower():
|
|
cache_write_cpt = input_cpt * 1.20
|
|
else:
|
|
cache_write_cpt = input_cpt
|
|
cost = (
|
|
input_tokens * input_cpt
|
|
+ output_tokens * output_cpt
|
|
+ cache_read_tokens * cache_read_cpt
|
|
+ cache_creation_tokens * cache_write_cpt
|
|
)
|
|
ctx = dict(instrumentation_context.get_ctx())
|
|
usage_label = ctx.pop("_usage_label", None)
|
|
has_context_cache = ctx.pop("_has_context_cache", False)
|
|
context_chars = ctx.pop("_context_chars", 0)
|
|
prompt_chars = ctx.pop("_prompt_chars", 0)
|
|
cache_key = ctx.pop("_cache_key", None)
|
|
segment = instrumentation.resolve_segment(
|
|
usage_label, ctx.pop("segment", None)
|
|
)
|
|
cache_miss_reason = _classify_cache_outcome(
|
|
has_context_cache=has_context_cache,
|
|
cache_read_tokens=cache_read_tokens,
|
|
cache_creation_tokens=cache_creation_tokens,
|
|
context_chars=context_chars,
|
|
cache_key=cache_key,
|
|
)
|
|
instrumentation.log(
|
|
"llm_call",
|
|
filename=filename,
|
|
model=model_id,
|
|
usage_label=usage_label,
|
|
segment=segment,
|
|
input_tokens=input_tokens,
|
|
output_tokens=output_tokens,
|
|
cache_creation_tokens=cache_creation_tokens,
|
|
cache_read_tokens=cache_read_tokens,
|
|
total_tokens=(
|
|
input_tokens
|
|
+ output_tokens
|
|
+ cache_creation_tokens
|
|
+ cache_read_tokens
|
|
),
|
|
cost_usd=cost,
|
|
cache_hit_inmem=False,
|
|
cache_miss_reason=cache_miss_reason,
|
|
has_context_cache=has_context_cache,
|
|
context_chars=context_chars,
|
|
prompt_chars=prompt_chars,
|
|
**_ctx_minus_explicit_keys(
|
|
ctx,
|
|
"filename",
|
|
"model",
|
|
"usage_label",
|
|
"input_tokens",
|
|
"output_tokens",
|
|
"cache_creation_tokens",
|
|
"cache_read_tokens",
|
|
"total_tokens",
|
|
"cost_usd",
|
|
"cache_hit_inmem",
|
|
"cache_miss_reason",
|
|
"has_context_cache",
|
|
"context_chars",
|
|
"prompt_chars",
|
|
),
|
|
)
|
|
|
|
|
|
def _extract_text_from_claude_response(response_body: dict) -> str:
|
|
"""Extract text from a Bedrock Claude response body.
|
|
|
|
Supports the messages API content blocks and includes a completion fallback
|
|
for older payload formats. Raises ValueError when no text is available.
|
|
"""
|
|
if not isinstance(response_body, dict):
|
|
raise ValueError("Claude response body is not a dictionary")
|
|
|
|
content = response_body.get("content")
|
|
if isinstance(content, list):
|
|
text_blocks = []
|
|
for block in content:
|
|
if isinstance(block, dict) and isinstance(block.get("text"), str):
|
|
text_blocks.append(block["text"])
|
|
if text_blocks:
|
|
return "".join(text_blocks)
|
|
|
|
# Fallback for legacy response shape.
|
|
completion = response_body.get("completion")
|
|
if isinstance(completion, str) and completion:
|
|
return completion
|
|
|
|
stop_reason = response_body.get("stop_reason", "unknown")
|
|
content_len = len(content) if isinstance(content, list) else "n/a"
|
|
raise ValueError(
|
|
f"Claude response missing text content (stop_reason={stop_reason}, content_len={content_len})"
|
|
)
|
|
|
|
|
|
def _build_claude_3_request_body(
|
|
prompt: str,
|
|
max_tokens: int,
|
|
instruction: str = None,
|
|
cache: bool = False,
|
|
model_id: str = None,
|
|
context_for_caching: str = None,
|
|
) -> dict:
|
|
"""Build request body for Claude 3+ API calls.
|
|
|
|
Args:
|
|
prompt: Input text prompt
|
|
max_tokens: Maximum tokens to generate
|
|
instruction: Optional system instruction
|
|
cache: Whether to enable prompt caching
|
|
model_id: Model ID to check for cache support
|
|
context_for_caching: Optional context text to cache separately (placed before prompt)
|
|
|
|
Returns:
|
|
Dictionary ready for JSON serialization
|
|
"""
|
|
request_body = {
|
|
"anthropic_version": "bedrock-2023-05-31",
|
|
"max_tokens": max_tokens,
|
|
"temperature": 0.0,
|
|
}
|
|
|
|
if instruction:
|
|
if cache and model_id and _supports_prompt_cache(model_id):
|
|
request_body["system"] = [
|
|
{
|
|
"type": "text",
|
|
"text": instruction,
|
|
"cache_control": {"type": "ephemeral"},
|
|
}
|
|
]
|
|
else:
|
|
request_body["system"] = instruction
|
|
|
|
# Build user message content with optional cached context
|
|
user_content = []
|
|
|
|
# Add cacheable context first (if provided)
|
|
if context_for_caching:
|
|
if cache and model_id and _supports_prompt_cache(model_id):
|
|
token_estimate = len(context_for_caching.split()) * 1.3
|
|
min_tokens = _min_cache_tokens(model_id)
|
|
if token_estimate < min_tokens:
|
|
logging.warning(
|
|
f"context_for_caching is ~{token_estimate:.0f} tokens, below the "
|
|
f"Anthropic {min_tokens}-token minimum for prompt caching on {model_id}; "
|
|
"context will be sent but the cache_control block may not actually be cached."
|
|
)
|
|
user_content.append(
|
|
{
|
|
"type": "text",
|
|
"text": context_for_caching,
|
|
"cache_control": {"type": "ephemeral"},
|
|
}
|
|
)
|
|
else:
|
|
user_content.append(
|
|
{
|
|
"type": "text",
|
|
"text": context_for_caching,
|
|
}
|
|
)
|
|
|
|
# Add main prompt (not cached)
|
|
user_content.append({"type": "text", "text": prompt})
|
|
|
|
request_body["messages"] = [{"role": "user", "content": user_content}]
|
|
|
|
return request_body
|
|
|
|
|
|
def _store_in_cache(cache_key: str, response: str) -> None:
|
|
"""Thread-safe cache storage with LRU eviction.
|
|
|
|
Args:
|
|
cache_key: Unique cache key for this response
|
|
response: Response text to cache
|
|
"""
|
|
with _cache_lock:
|
|
if len(claude_cache) >= CACHE_LIMIT:
|
|
claude_cache.popitem(last=False) # Remove oldest (least recently used)
|
|
claude_cache[cache_key] = response
|
|
claude_cache.move_to_end(cache_key) # Mark as recently used
|
|
|
|
|
|
# In-memory cache
|
|
claude_cache = OrderedDict()
|
|
CACHE_LIMIT = 20000
|
|
_cache_lock = threading.Lock() # Thread safety for cache operations
|
|
|
|
# Tracks which Anthropic cache-key hashes have been seen in this process and
|
|
# when (monotonic seconds). Used to classify cache_miss_reason=first_call vs
|
|
# ttl_expired on llm_call events. O(distinct cache keys) — tiny.
|
|
_cache_key_seen: dict[str, float] = {}
|
|
_cache_key_seen_lock = threading.Lock()
|
|
|
|
# Model routing set (derived from _SUPPORTED_MODELS for consistency)
|
|
_CLAUDE_3_AND_UP_MODELS = {m for m in _SUPPORTED_MODELS if m is not None}
|
|
|
|
|
|
def get_cache_key(
|
|
prompt,
|
|
model_id,
|
|
instruction=None,
|
|
max_tokens=None,
|
|
cache=False,
|
|
context_for_caching=None,
|
|
):
|
|
"""Generates a unique hash key for caching.
|
|
|
|
Args:
|
|
prompt (str): The input prompt.
|
|
model_id (str): The model identifier.
|
|
instruction (str, optional): System instruction for caching. Defaults to None.
|
|
max_tokens (int, optional): Maximum tokens to generate. Defaults to None.
|
|
cache (bool, optional): Whether caching is enabled. Defaults to False.
|
|
context_for_caching (str, optional): Context text for caching. Defaults to None.
|
|
|
|
Returns:
|
|
str: SHA256 hash of the cache key components.
|
|
"""
|
|
# Build cache key from all relevant parameters
|
|
key_parts = [model_id, prompt]
|
|
if instruction is not None:
|
|
key_parts.append(str(instruction))
|
|
if max_tokens is not None:
|
|
key_parts.append(str(max_tokens))
|
|
if cache:
|
|
key_parts.append("cache_enabled")
|
|
if context_for_caching:
|
|
context_hash = hashlib.sha256(context_for_caching.encode()).hexdigest()
|
|
key_parts.append(f"context_sha256={context_hash}")
|
|
return hashlib.sha256(":".join(key_parts).encode()).hexdigest()
|
|
|
|
|
|
def invoke_claude(
|
|
prompt: str,
|
|
model_id: str,
|
|
filename: str,
|
|
max_tokens: int = 4096,
|
|
cache: bool = False,
|
|
instruction: str = None,
|
|
usage_label: str = None,
|
|
context_for_caching: str = None,
|
|
) -> str:
|
|
"""Invokes the Claude model with the given parameters.
|
|
|
|
Args:
|
|
prompt (str): The input prompt to send to the model.
|
|
model_id (str): The ID of the model to use. Supports both direct model IDs
|
|
(e.g. 'anthropic.claude-3-5-sonnet-20240620-v1:0') as well as aliases (e.g.
|
|
'sonnet_latest', 'haiku_latest', etc.).
|
|
filename (str): The name of the file being processed, for cost logging purposes.
|
|
max_tokens (int, optional): The maximum number of tokens to generate. Defaults to 4096.
|
|
cache (bool, optional): Whether to enable prompt caching. Defaults to False.
|
|
instruction (str, optional): System instruction to use with caching. Defaults to None.
|
|
context_for_caching (str, optional): Context text to cache separately before prompt.
|
|
Useful for caching large exhibit text that's reused across multiple field extractions.
|
|
When provided with cache=True, this context will be cached at the Anthropic API level.
|
|
|
|
Raises:
|
|
ValueError: If the model_id is not supported after alias resolution
|
|
|
|
Returns:
|
|
str: The model's response.
|
|
|
|
Note:
|
|
- Responses are cached using LRU eviction (limit: 20,000 entries)
|
|
- Includes automatic retry logic with exponential backoff for throttling
|
|
- Cost logging is performed automatically for all successful requests
|
|
- Model aliases are resolved via config.resolve_model_id()
|
|
- Prompt caching can be enabled to reduce costs for large repeated prompts
|
|
- Context caching allows caching exhibit text separately from field questions
|
|
"""
|
|
if usage_label is None:
|
|
try:
|
|
usage_label = inspect.stack()[1].function
|
|
except Exception:
|
|
usage_label = "UNLABELED"
|
|
|
|
cache_key = get_cache_key(
|
|
prompt,
|
|
model_id,
|
|
instruction=instruction,
|
|
max_tokens=max_tokens,
|
|
cache=cache,
|
|
context_for_caching=context_for_caching,
|
|
)
|
|
|
|
# Resolve alias up-front so cache-policy classification in tracking
|
|
# logs uses the durable id, not whatever alias the caller happened to
|
|
# pass.
|
|
resolved_model_id = config.resolve_model_id(model_id)
|
|
model_supports_cache = _supports_prompt_cache(resolved_model_id)
|
|
|
|
# Check cache first - thread-safe access
|
|
with _cache_lock:
|
|
if cache_key in claude_cache:
|
|
claude_cache.move_to_end(cache_key) # Mark as recently used
|
|
prompt_call_tracking.log_prompt_call(
|
|
contract_filename=filename,
|
|
usage_label=usage_label,
|
|
model_id=model_id,
|
|
cache_enabled=cache,
|
|
max_tokens=max_tokens,
|
|
prompt=prompt,
|
|
instruction=instruction,
|
|
context_for_caching=context_for_caching,
|
|
status="success",
|
|
from_local_memory_cache=True,
|
|
resolved_model_id=resolved_model_id,
|
|
model_supports_cache=model_supports_cache,
|
|
)
|
|
# One llm_call_inmem_hit event per in-memory cache return.
|
|
from src.utils import instrumentation, instrumentation_context
|
|
|
|
if instrumentation.is_enabled():
|
|
ctx = instrumentation_context.get_ctx()
|
|
segment = instrumentation.resolve_segment(
|
|
usage_label, ctx.get("segment")
|
|
)
|
|
instrumentation.log(
|
|
"llm_call_inmem_hit",
|
|
filename=filename,
|
|
usage_label=usage_label,
|
|
segment=segment,
|
|
model=model_id,
|
|
prompt_chars=len(prompt) if prompt else 0,
|
|
context_chars=(
|
|
len(context_for_caching) if context_for_caching else 0
|
|
),
|
|
has_context_cache=bool(context_for_caching),
|
|
cache_hit_inmem=True,
|
|
**_ctx_minus_explicit_keys(
|
|
dict(ctx),
|
|
"filename",
|
|
"usage_label",
|
|
"model",
|
|
"prompt_chars",
|
|
"context_chars",
|
|
"has_context_cache",
|
|
"cache_hit_inmem",
|
|
),
|
|
)
|
|
return claude_cache[cache_key]
|
|
|
|
# Validate model support (checks for deprecated and unsupported models)
|
|
_validate_model_support(model_id, resolved_model_id)
|
|
|
|
# Push call metadata into context so _track_usage_from_response can read it.
|
|
from src.utils.instrumentation_context import instr_scope
|
|
|
|
try:
|
|
with instr_scope(
|
|
_usage_label=usage_label,
|
|
_has_context_cache=bool(context_for_caching),
|
|
_context_chars=len(context_for_caching) if context_for_caching else 0,
|
|
_prompt_chars=len(prompt) if prompt else 0,
|
|
_cache_key=cache_key,
|
|
):
|
|
if config.RUN_MODE == "local":
|
|
response = local_claude_3_and_up(
|
|
prompt,
|
|
resolved_model_id,
|
|
filename,
|
|
max_tokens,
|
|
cache=cache,
|
|
instruction=instruction,
|
|
usage_label=usage_label,
|
|
context_for_caching=context_for_caching,
|
|
)
|
|
elif config.RUN_MODE == "ec2":
|
|
response = ec2_claude_3_and_up(
|
|
prompt,
|
|
resolved_model_id,
|
|
filename,
|
|
max_tokens,
|
|
cache=cache,
|
|
instruction=instruction,
|
|
usage_label=usage_label,
|
|
context_for_caching=context_for_caching,
|
|
)
|
|
else:
|
|
logging.error(
|
|
"Usage: python local_main.py <-local> OR python main.py <ec2>"
|
|
)
|
|
raise
|
|
except Exception as e:
|
|
prompt_call_tracking.log_prompt_call(
|
|
contract_filename=filename,
|
|
usage_label=usage_label,
|
|
model_id=resolved_model_id,
|
|
cache_enabled=cache,
|
|
max_tokens=max_tokens,
|
|
prompt=prompt,
|
|
instruction=instruction,
|
|
context_for_caching=context_for_caching,
|
|
status="error",
|
|
error=str(e),
|
|
resolved_model_id=resolved_model_id,
|
|
model_supports_cache=model_supports_cache,
|
|
)
|
|
raise
|
|
|
|
# Store response in cache with LRU eviction - thread-safe
|
|
_store_in_cache(cache_key, response)
|
|
|
|
return response
|
|
|
|
|
|
def get_cache_key_with_image_array(
|
|
prompt, model_id, image_base64_list, instruction=None, max_tokens=None, cache=False
|
|
):
|
|
"""Generates a unique hash key for caching with a list of base64 images.
|
|
|
|
Args:
|
|
prompt (str): The input prompt.
|
|
model_id (str): The model identifier.
|
|
image_base64_list (list): List of base64-encoded images.
|
|
instruction (str, optional): System instruction for caching. Defaults to None.
|
|
max_tokens (int, optional): Maximum tokens to generate. Defaults to None.
|
|
cache (bool, optional): Whether caching is enabled. Defaults to False.
|
|
|
|
Returns:
|
|
str: SHA256 hash of the cache key components.
|
|
"""
|
|
combined_images = ":".join(
|
|
image_base64_list
|
|
) # Combine all base64 images into a single string
|
|
# Build cache key from all relevant parameters
|
|
key_parts = [model_id, prompt, combined_images]
|
|
if instruction is not None:
|
|
key_parts.append(str(instruction))
|
|
if max_tokens is not None:
|
|
key_parts.append(str(max_tokens))
|
|
if cache:
|
|
key_parts.append("cache_enabled")
|
|
return hashlib.sha256(":".join(key_parts).encode()).hexdigest()
|
|
|
|
|
|
def invoke_claude_with_image_array(
|
|
prompt, model_id, filename, base64_images, max_tokens=4096
|
|
):
|
|
usage_label = inspect.stack()[1].function
|
|
cache_key = get_cache_key_with_image_array(
|
|
prompt, model_id, base64_images, max_tokens=max_tokens
|
|
)
|
|
# Check cache first - thread-safe access
|
|
with _cache_lock:
|
|
if cache_key in claude_cache:
|
|
claude_cache.move_to_end(cache_key) # Mark as recently used
|
|
return claude_cache[cache_key]
|
|
|
|
# Resolve alias to actual model ID
|
|
resolved_model_id = config.resolve_model_id(model_id)
|
|
|
|
# Validate model support (checks for deprecated and unsupported models)
|
|
_validate_model_support(model_id, resolved_model_id, context="for image array")
|
|
|
|
if config.RUN_MODE == "local":
|
|
if resolved_model_id == config.MODEL_ID_CLAUDE35_SONNET:
|
|
response = local_claude_3_and_up_with_image_array(
|
|
prompt,
|
|
resolved_model_id,
|
|
filename,
|
|
max_tokens,
|
|
base64_images,
|
|
usage_label=usage_label,
|
|
)
|
|
logging.debug(
|
|
f"Effective Date Response from local Claude 3 for {filename}: {response}"
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
f"Unsupported model_id for image array: {resolved_model_id}"
|
|
)
|
|
elif config.RUN_MODE == "ec2":
|
|
# EC2 mode with image arrays - currently only Claude 3.5 Sonnet is supported for image arrays
|
|
if resolved_model_id == config.MODEL_ID_CLAUDE35_SONNET:
|
|
# For EC2, we need to use the image array function, but it's only implemented for local mode
|
|
# For now, raise an error suggesting to use local mode or a different approach
|
|
raise ValueError(
|
|
f"Image arrays with EC2 mode are currently only supported in local mode. "
|
|
f"Please use RUN_MODE='local' or use a text-only prompt."
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
f"Unsupported model_id for image array in EC2 mode: {resolved_model_id}. "
|
|
f"Only Claude 3.5 Sonnet is supported for image arrays."
|
|
)
|
|
|
|
else:
|
|
logging.error("Usage: python local_main.py <-local> OR python main.py <ec2>")
|
|
raise
|
|
|
|
# Store response in cache with LRU eviction - thread-safe
|
|
_store_in_cache(cache_key, response)
|
|
|
|
return response
|
|
|
|
|
|
# Claude calls
|
|
# Note: local_claude_2() and ec2_claude2() have been removed as Claude 2 is no longer supported
|
|
|
|
|
|
def local_claude_3_and_up(
|
|
prompt: str,
|
|
model_id: str,
|
|
filename: str,
|
|
max_tokens: int,
|
|
max_retries: int = 5,
|
|
initial_delay: int = 1,
|
|
cache: bool = False,
|
|
instruction: str = None,
|
|
usage_label: str = None,
|
|
context_for_caching: str = None,
|
|
):
|
|
"""
|
|
Handles local invocation of Claude 3 and above models (3.x, 3.5, 4.x) via AWS Bedrock.
|
|
|
|
This function provides a unified interface for all modern Claude models using the
|
|
messages API format. Includes exponential backoff retry logic for throttling errors.
|
|
|
|
Args:
|
|
prompt (str): Input text prompt to send to the model
|
|
model_id (str): AWS Bedrock model identifier (e.g., 'anthropic.claude-3-5-sonnet-20241022-v2:0')
|
|
filename (str): Name of file being processed (for logging/tracking purposes)
|
|
max_tokens (int): Maximum number of tokens to generate in the response
|
|
max_retries (int, optional): Maximum number of retry attempts for throttling errors. Defaults to 5.
|
|
initial_delay (int, optional): Initial delay in seconds for exponential backoff. Defaults to 1.
|
|
cache (bool, optional): Whether to enable prompt caching. Defaults to False.
|
|
instruction (str, optional): System instruction to use with caching. Defaults to None.
|
|
|
|
Returns:
|
|
str: Generated text response from the Claude model
|
|
|
|
Raises:
|
|
ClientError: For AWS Bedrock API errors (non-throttling)
|
|
Exception: When max retries are exceeded or other unexpected errors occur
|
|
|
|
Note:
|
|
- Uses the modern messages API format (anthropic_version: "bedrock-2023-05-31")
|
|
- Supports answer tracking via config.ENABLE_ANSWER_TRACKING for caching/debugging
|
|
- Temperature is set to 0.0 for deterministic responses
|
|
- Automatically handles throttling with exponential backoff retry logic
|
|
- Prompt caching can be enabled to reduce costs for large repeated prompts
|
|
"""
|
|
if config.ENABLE_ANSWER_TRACKING and prompt in config.ANSWER_TRACKING_DICT.keys():
|
|
return config.ANSWER_TRACKING_DICT[prompt]
|
|
|
|
# Build request body with caching support
|
|
request_body = _build_claude_3_request_body(
|
|
prompt,
|
|
max_tokens,
|
|
instruction=instruction,
|
|
cache=cache,
|
|
model_id=model_id,
|
|
context_for_caching=context_for_caching,
|
|
)
|
|
body = json.dumps(request_body)
|
|
|
|
for attempt in range(max_retries):
|
|
try:
|
|
response = config.BEDROCK_RUNTIME.invoke_model(
|
|
body=body,
|
|
modelId=model_id,
|
|
accept="application/json",
|
|
contentType="application/json",
|
|
)
|
|
response_body = json.loads(response.get("body").read())
|
|
response_text = _extract_text_from_claude_response(response_body)
|
|
|
|
(
|
|
input_tokens,
|
|
output_tokens,
|
|
cache_creation_tokens,
|
|
cache_read_tokens,
|
|
) = usage_tracking.extract_usage_from_response(response_body)
|
|
|
|
# Track usage and costs from API response (including cache tokens)
|
|
_track_usage_from_response(response_body, filename, model_id)
|
|
|
|
prompt_call_tracking.log_prompt_call(
|
|
contract_filename=filename,
|
|
usage_label=usage_label or "UNLABELED",
|
|
model_id=model_id,
|
|
cache_enabled=cache,
|
|
max_tokens=max_tokens,
|
|
prompt=prompt,
|
|
instruction=instruction,
|
|
context_for_caching=context_for_caching,
|
|
status="success",
|
|
input_tokens=input_tokens,
|
|
output_tokens=output_tokens,
|
|
cache_creation_tokens=cache_creation_tokens,
|
|
cache_read_tokens=cache_read_tokens,
|
|
resolved_model_id=model_id,
|
|
model_supports_cache=_supports_prompt_cache(model_id),
|
|
)
|
|
|
|
if config.ENABLE_ANSWER_TRACKING:
|
|
config.ANSWER_TRACKING_DICT[prompt] = response_text
|
|
|
|
return response_text
|
|
|
|
except ValueError as e:
|
|
if attempt < max_retries - 1:
|
|
delay = (2**attempt + random.uniform(0, 1)) * initial_delay
|
|
_log_llm_retry(
|
|
filename, model_id, attempt + 1, "ValueError", str(e), delay
|
|
)
|
|
logging.warning(
|
|
f"Attempt {attempt + 1} failed due to malformed Claude response. "
|
|
f"Retrying in {delay:.2f} seconds. Error: {e}"
|
|
)
|
|
time.sleep(delay)
|
|
else:
|
|
_log_llm_error(filename, model_id, attempt + 1, "ValueError", str(e))
|
|
raise
|
|
|
|
except ClientError as e:
|
|
error_code = e.response["Error"]["Code"]
|
|
if error_code in [
|
|
"ThrottlingException",
|
|
"TooManyRequestsException",
|
|
]:
|
|
if attempt < max_retries - 1:
|
|
delay = (2**attempt + random.uniform(0, 1)) * initial_delay
|
|
_log_llm_retry(
|
|
filename, model_id, attempt + 1, error_code, str(e), delay
|
|
)
|
|
time.sleep(delay)
|
|
else:
|
|
_log_llm_error(filename, model_id, attempt + 1, error_code, str(e))
|
|
raise
|
|
else:
|
|
_log_llm_error(filename, model_id, attempt + 1, error_code, str(e))
|
|
raise
|
|
|
|
raise Exception("Max retries exceeded")
|
|
|
|
|
|
def local_claude_3_and_up_with_image_array(
|
|
prompt,
|
|
model_id,
|
|
filename,
|
|
max_tokens,
|
|
base64_images,
|
|
max_retries=5,
|
|
initial_delay=1,
|
|
image_media_type="image/png",
|
|
usage_label: str = None,
|
|
):
|
|
"""
|
|
Processes a prompt along with an image in base64 format using the Claude 3 model via AWS Bedrock.
|
|
|
|
Args:
|
|
prompt (str): Input prompt to LLM.
|
|
model_id (str): Model ID to use.
|
|
filename (str): Name of the file being operated upon.
|
|
max_tokens (int): Maximum tokens to sample.
|
|
image_base64 (str): Base64 encoded image data.
|
|
max_retries (int): Maximum number of retries for the request.
|
|
initial_delay (int): Initial delay for exponential backoff.
|
|
image_media_type (str): Media type of the image (default: "image/jpeg").
|
|
|
|
Returns:
|
|
str: Response text from the model.
|
|
"""
|
|
if config.ENABLE_ANSWER_TRACKING and prompt in config.ANSWER_TRACKING_DICT.keys():
|
|
return config.ANSWER_TRACKING_DICT[prompt]
|
|
|
|
# Create the request body with the new message format
|
|
body = json.dumps(
|
|
{
|
|
"anthropic_version": "bedrock-2023-05-31",
|
|
"max_tokens": max_tokens,
|
|
"temperature": 0.0,
|
|
"messages": [
|
|
{
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": "Images:"},
|
|
*[
|
|
{
|
|
"type": "image",
|
|
"source": {
|
|
"type": "base64",
|
|
"media_type": image_media_type,
|
|
"data": image_base64,
|
|
},
|
|
}
|
|
for image_base64 in base64_images
|
|
],
|
|
{"type": "text", "text": prompt},
|
|
],
|
|
}
|
|
],
|
|
}
|
|
)
|
|
|
|
for attempt in range(max_retries):
|
|
try:
|
|
response = config.BEDROCK_RUNTIME.invoke_model(
|
|
body=body,
|
|
modelId=model_id,
|
|
accept="application/json",
|
|
contentType="application/json",
|
|
)
|
|
response_body = json.loads(response.get("body").read())
|
|
response_text = _extract_text_from_claude_response(response_body)
|
|
|
|
# Track usage and costs from API response (including cache tokens)
|
|
_track_usage_from_response(response_body, filename, model_id)
|
|
|
|
if config.ENABLE_ANSWER_TRACKING:
|
|
config.ANSWER_TRACKING_DICT[prompt] = response_text
|
|
|
|
return response_text
|
|
|
|
except ValueError as e:
|
|
if attempt < max_retries - 1:
|
|
delay = (2**attempt + random.uniform(0, 1)) * initial_delay
|
|
logging.warning(
|
|
f"Attempt {attempt + 1} failed due to malformed Claude response. "
|
|
f"Retrying in {delay:.2f} seconds. Error: {e}"
|
|
)
|
|
time.sleep(delay)
|
|
else:
|
|
raise
|
|
|
|
except ClientError as e:
|
|
if e.response["Error"]["Code"] in [
|
|
"ThrottlingException",
|
|
"TooManyRequestsException",
|
|
]:
|
|
if attempt < max_retries - 1:
|
|
delay = (2**attempt + random.uniform(0, 1)) * initial_delay
|
|
time.sleep(delay)
|
|
else:
|
|
raise
|
|
else:
|
|
raise
|
|
|
|
raise Exception("Max retries exceeded")
|
|
|
|
|
|
def ec2_claude_3_and_up(
|
|
prompt: str,
|
|
model_id: str,
|
|
filename: str,
|
|
max_tokens: int,
|
|
max_retries: int = 5,
|
|
initial_delay: int = 1,
|
|
cache: bool = False,
|
|
instruction: str = None,
|
|
usage_label: str = None,
|
|
context_for_caching: str = None,
|
|
):
|
|
"""
|
|
Handles EC2 invocation of Claude 3 and above models.
|
|
|
|
When ENABLE_RUNTIME_ROTATION is True:
|
|
Provides multi-runtime failover (UAT -> DEV -> PROD) with exponential backoff.
|
|
|
|
When ENABLE_RUNTIME_ROTATION is False:
|
|
Uses single runtime with exponential backoff retry logic.
|
|
|
|
Both options include comprehensive usage tracking via usage_tracking.py.
|
|
|
|
Args:
|
|
prompt (str): Input text prompt to send to the model
|
|
model_id (str): AWS Bedrock model identifier (e.g., 'anthropic.claude-3-5-sonnet-20241022-v2:0')
|
|
filename (str): Name of file being processed (for usage tracking and logging)
|
|
max_tokens (int): Maximum number of tokens to generate in the response
|
|
max_retries (int, optional): Maximum retry attempts per runtime environment. Defaults to 5.
|
|
initial_delay (int, optional): Initial delay in seconds for exponential backoff. Defaults to 1.
|
|
cache (bool, optional): Whether to enable prompt caching. Defaults to False.
|
|
instruction (str, optional): System instruction to use with caching. Defaults to None.
|
|
usage_label (str, optional): Label to surface in usage diagnostics. Defaults to None.
|
|
|
|
Returns:
|
|
str: Generated text response from the Claude model
|
|
|
|
Raises:
|
|
ClientError: For non-recoverable AWS Bedrock API errors
|
|
ReadTimeoutError: For network timeout issues
|
|
Exception: When all runtimes (when applicable) and retries are exhausted
|
|
|
|
Runtime Failover Order (when ENABLE_RUNTIME_ROTATION is True):
|
|
1. UAT (primary runtime from config.BEDROCK_RUNTIME)
|
|
2. DEV (fallback via role assumption)
|
|
3. PROD (final fallback via role assumption)
|
|
|
|
Usage Tracking:
|
|
- All token usage is tracked via usage_tracking.track_usage()
|
|
- Uses exact token counts from API responses (not estimates)
|
|
- Tracks all 4 token types: input, output, cache_creation, cache_read
|
|
|
|
Note:
|
|
- Uses messages API format for all Claude 3+ models
|
|
- When rotation is enabled, runtime gets full retry cycle before switching to next
|
|
- Temperature set to 0.0 for consistent, deterministic outputs
|
|
- Prompt caching can be enabled to reduce costs for large repeated prompts
|
|
"""
|
|
# Build request body with caching support
|
|
request_body = _build_claude_3_request_body(
|
|
prompt,
|
|
max_tokens,
|
|
instruction=instruction,
|
|
cache=cache,
|
|
model_id=model_id,
|
|
context_for_caching=context_for_caching,
|
|
)
|
|
body = json.dumps(request_body)
|
|
|
|
if config.ENABLE_RUNTIME_ROTATION:
|
|
dev_runtime = config.assume_role_and_get_runtime(
|
|
config.dev_account_id, config.dev_role_name
|
|
)
|
|
prod_runtime = config.assume_role_and_get_runtime(
|
|
config.prod_account_id, config.prod_role_name
|
|
)
|
|
runtimes = {
|
|
"UAT": config.BEDROCK_RUNTIME, # This is the default runtime (UAT)
|
|
"DEV": dev_runtime,
|
|
"PROD": prod_runtime,
|
|
}
|
|
for runtime_key in runtimes.keys():
|
|
current_runtime = runtimes[runtime_key]
|
|
for attempt in range(max_retries):
|
|
try:
|
|
response = current_runtime.invoke_model(
|
|
body=body,
|
|
modelId=model_id,
|
|
accept="application/json",
|
|
contentType="application/json",
|
|
)
|
|
response_body = json.loads(response.get("body").read())
|
|
response_text = _extract_text_from_claude_response(response_body)
|
|
|
|
(
|
|
input_tokens,
|
|
output_tokens,
|
|
cache_creation_tokens,
|
|
cache_read_tokens,
|
|
) = usage_tracking.extract_usage_from_response(response_body)
|
|
|
|
# Track usage and costs from API response (including cache tokens)
|
|
_track_usage_from_response(response_body, filename, model_id)
|
|
|
|
prompt_call_tracking.log_prompt_call(
|
|
contract_filename=filename,
|
|
usage_label=usage_label or "UNLABELED",
|
|
model_id=model_id,
|
|
cache_enabled=cache,
|
|
max_tokens=max_tokens,
|
|
prompt=prompt,
|
|
instruction=instruction,
|
|
context_for_caching=context_for_caching,
|
|
status="success",
|
|
input_tokens=input_tokens,
|
|
output_tokens=output_tokens,
|
|
cache_creation_tokens=cache_creation_tokens,
|
|
cache_read_tokens=cache_read_tokens,
|
|
resolved_model_id=model_id,
|
|
model_supports_cache=_supports_prompt_cache(model_id),
|
|
)
|
|
|
|
return response_text
|
|
|
|
except ValueError as e:
|
|
if attempt < max_retries - 1:
|
|
delay = (2**attempt + random.uniform(0, 1)) * initial_delay
|
|
_log_llm_retry(
|
|
filename, model_id, attempt + 1, "ValueError", str(e), delay
|
|
)
|
|
logging.warning(
|
|
f"Attempt {attempt + 1} failed due to malformed Claude response. "
|
|
f"Retrying in {delay:.2f} seconds... Error: {e}"
|
|
)
|
|
time.sleep(delay)
|
|
else:
|
|
_log_llm_retry(
|
|
filename,
|
|
model_id,
|
|
attempt + 1,
|
|
f"runtime_switch:{runtime_key}:ValueError",
|
|
str(e),
|
|
0.0,
|
|
)
|
|
logging.warning(
|
|
f"Switching to next runtime due to malformed Claude response: {e}."
|
|
)
|
|
break # Switch to the next runtime
|
|
|
|
except (ReadTimeoutError, ClientError) as e:
|
|
error_code = e.response["Error"]["Code"]
|
|
if error_code in [
|
|
"ThrottlingException",
|
|
"TooManyRequestsException",
|
|
]:
|
|
if attempt < max_retries - 1:
|
|
delay = (2**attempt + random.uniform(0, 1)) * initial_delay
|
|
_log_llm_retry(
|
|
filename,
|
|
model_id,
|
|
attempt + 1,
|
|
error_code,
|
|
str(e),
|
|
delay,
|
|
)
|
|
logging.warning(
|
|
f"Attempt {attempt + 1} failed. Retrying in {delay:.2f} seconds..."
|
|
)
|
|
time.sleep(delay)
|
|
else:
|
|
_log_llm_retry(
|
|
filename,
|
|
model_id,
|
|
attempt + 1,
|
|
f"runtime_switch:{runtime_key}:{error_code}",
|
|
str(e),
|
|
0.0,
|
|
)
|
|
logging.warning(
|
|
f"Switching to next runtime due to {error_code}: {str(e)}."
|
|
)
|
|
break # Switch to the next runtime
|
|
else:
|
|
_log_llm_error(
|
|
filename, model_id, attempt + 1, error_code, str(e)
|
|
)
|
|
logging.error(f"Unexpected error: {str(e)}")
|
|
raise # Reraise other exceptions
|
|
|
|
except Exception as e:
|
|
_log_llm_error(
|
|
filename,
|
|
model_id,
|
|
attempt + 1,
|
|
type(e).__name__,
|
|
str(e),
|
|
)
|
|
logging.error(f"Unexpected error: {str(e)}")
|
|
raise
|
|
else: # Runtime rotation is disabled
|
|
for attempt in range(max_retries):
|
|
try:
|
|
response = config.BEDROCK_RUNTIME.invoke_model(
|
|
body=body,
|
|
modelId=model_id,
|
|
accept="application/json",
|
|
contentType="application/json",
|
|
)
|
|
response_body = json.loads(response.get("body").read())
|
|
response_text = _extract_text_from_claude_response(response_body)
|
|
|
|
(
|
|
input_tokens,
|
|
output_tokens,
|
|
cache_creation_tokens,
|
|
cache_read_tokens,
|
|
) = usage_tracking.extract_usage_from_response(response_body)
|
|
|
|
# Track usage and costs from API response (including cache tokens)
|
|
_track_usage_from_response(response_body, filename, model_id)
|
|
|
|
prompt_call_tracking.log_prompt_call(
|
|
contract_filename=filename,
|
|
usage_label=usage_label or "UNLABELED",
|
|
model_id=model_id,
|
|
cache_enabled=cache,
|
|
max_tokens=max_tokens,
|
|
prompt=prompt,
|
|
instruction=instruction,
|
|
context_for_caching=context_for_caching,
|
|
status="success",
|
|
input_tokens=input_tokens,
|
|
output_tokens=output_tokens,
|
|
cache_creation_tokens=cache_creation_tokens,
|
|
cache_read_tokens=cache_read_tokens,
|
|
resolved_model_id=model_id,
|
|
model_supports_cache=_supports_prompt_cache(model_id),
|
|
)
|
|
|
|
return response_text
|
|
|
|
except ValueError as e:
|
|
if attempt < max_retries - 1:
|
|
delay = (2**attempt + random.uniform(0, 1)) * initial_delay
|
|
_log_llm_retry(
|
|
filename, model_id, attempt + 1, "ValueError", str(e), delay
|
|
)
|
|
logging.warning(
|
|
f"Attempt {attempt + 1} failed due to malformed Claude response. "
|
|
f"Retrying in {delay:.2f} seconds... Error: {e}"
|
|
)
|
|
time.sleep(delay)
|
|
else:
|
|
_log_llm_error(
|
|
filename, model_id, attempt + 1, "ValueError", str(e)
|
|
)
|
|
raise
|
|
|
|
except ReadTimeoutError as e:
|
|
# Timeout errors should probably retry
|
|
if attempt < max_retries - 1:
|
|
delay = (2**attempt + random.uniform(0, 1)) * initial_delay
|
|
_log_llm_retry(
|
|
filename,
|
|
model_id,
|
|
attempt + 1,
|
|
"ReadTimeoutError",
|
|
str(e),
|
|
delay,
|
|
)
|
|
logging.warning(
|
|
f"Attempt {attempt + 1} failed due to timeout. Retrying in {delay:.2f} seconds..."
|
|
)
|
|
time.sleep(delay)
|
|
else:
|
|
_log_llm_error(
|
|
filename,
|
|
model_id,
|
|
attempt + 1,
|
|
"ReadTimeoutError",
|
|
str(e),
|
|
)
|
|
raise
|
|
except ClientError as e:
|
|
error_code = e.response["Error"]["Code"]
|
|
if error_code in ["ThrottlingException", "TooManyRequestsException"]:
|
|
if attempt < max_retries - 1:
|
|
delay = (2**attempt + random.uniform(0, 1)) * initial_delay
|
|
_log_llm_retry(
|
|
filename,
|
|
model_id,
|
|
attempt + 1,
|
|
error_code,
|
|
str(e),
|
|
delay,
|
|
)
|
|
logging.warning(
|
|
f"Attempt {attempt + 1} failed due to rate limiting. Retrying in {delay:.2f} seconds..."
|
|
)
|
|
time.sleep(delay)
|
|
else:
|
|
_log_llm_error(
|
|
filename, model_id, attempt + 1, error_code, str(e)
|
|
)
|
|
raise
|
|
else:
|
|
_log_llm_error(filename, model_id, attempt + 1, error_code, str(e))
|
|
logging.error(f"Unexpected error: {str(e)}")
|
|
raise # Non-throttling client errors should fail immediately
|
|
|
|
except Exception as e:
|
|
_log_llm_error(
|
|
filename, model_id, attempt + 1, type(e).__name__, str(e)
|
|
)
|
|
logging.error(f"Unexpected error: {str(e)}")
|
|
raise
|
|
raise Exception(f"Max retries ({max_retries}) exceeded for model {model_id}.")
|
|
|
|
|
|
def get_image_array_answer(prompt, filename, base64_images):
|
|
response_text = invoke_claude_with_image_array(
|
|
prompt, config.MODEL_ID_CLAUDE35_SONNET, filename, base64_images
|
|
)
|
|
response_dict = json.loads(response_text)
|
|
return response_dict
|
|
|
|
|
|
def warm_prompt_cache(
|
|
instruction: str,
|
|
cache_key: str,
|
|
model_id: str = "sonnet_latest",
|
|
max_tokens: int = 10,
|
|
):
|
|
"""
|
|
Warms the prompt cache by sending a minimal request with the instruction.
|
|
This should be called once before parallel processing to ensure the cache
|
|
is registered.
|
|
|
|
Cache entries are model-specific in Bedrock — a cache created against
|
|
`sonnet_latest` is NOT readable from a runtime call invoked against
|
|
`haiku_latest`. The warmed-state map is therefore keyed by
|
|
(cache_key, resolved_model_id) so repeated calls with different models
|
|
each get warmed exactly once.
|
|
|
|
Args:
|
|
instruction (str): The instruction text to cache
|
|
cache_key (str): A unique identifier for this cache (e.g., "REIMBURSEMENT_PRIMARY")
|
|
model_id (str): The model ID or alias to use. Defaults to "sonnet_latest"
|
|
max_tokens (int): Minimal tokens for cache warming. Defaults to 10
|
|
|
|
Returns:
|
|
bool: True if cache was warmed, False if it was already warm or
|
|
the resolved model does not support prompt caching.
|
|
"""
|
|
global _cache_warmed
|
|
|
|
# Resolve model ID up-front so the warmed-state key matches what the
|
|
# runtime call will use after its own resolve_model_id().
|
|
resolved_model_id = config.resolve_model_id(model_id)
|
|
|
|
# Hard skip on models that do not support prompt caching in this
|
|
# Bedrock setup. Without this, warming spends create-tokens against a
|
|
# cache the runtime call can never read.
|
|
if not _supports_prompt_cache(resolved_model_id):
|
|
logging.info(
|
|
"Skipping cache warming for %s on %s (resolved %s): model does "
|
|
"not support prompt caching.",
|
|
cache_key,
|
|
model_id,
|
|
resolved_model_id,
|
|
)
|
|
return False
|
|
|
|
warmed_key = (cache_key, resolved_model_id)
|
|
|
|
# Check if already warmed (thread-safe)
|
|
with _cache_warming_lock:
|
|
if _cache_warmed.get(warmed_key, False):
|
|
logging.debug(
|
|
f"Cache already warmed for {cache_key} on {resolved_model_id}"
|
|
)
|
|
return False
|
|
|
|
# Mark as being warmed
|
|
_cache_warmed[warmed_key] = True
|
|
|
|
logging.info(f"Warming prompt cache for {cache_key} on {resolved_model_id}...")
|
|
|
|
# Create minimal prompt with the instruction to register the cache
|
|
minimal_prompt = "Respond with OK."
|
|
|
|
try:
|
|
# Build request with caching enabled
|
|
request_body = _build_claude_3_request_body(
|
|
minimal_prompt,
|
|
max_tokens,
|
|
instruction=instruction,
|
|
cache=True,
|
|
model_id=resolved_model_id,
|
|
)
|
|
body = json.dumps(request_body)
|
|
|
|
# Send the request to register the cache
|
|
if config.RUN_MODE == "local":
|
|
response = config.BEDROCK_RUNTIME.invoke_model(
|
|
body=body,
|
|
modelId=resolved_model_id,
|
|
accept="application/json",
|
|
contentType="application/json",
|
|
)
|
|
elif config.RUN_MODE == "ec2":
|
|
response = config.BEDROCK_RUNTIME.invoke_model(
|
|
body=body,
|
|
modelId=resolved_model_id,
|
|
accept="application/json",
|
|
contentType="application/json",
|
|
)
|
|
else:
|
|
logging.error("Invalid RUN_MODE for cache warming")
|
|
return False
|
|
|
|
response_body = json.loads(response.get("body").read())
|
|
logging.info(f"Successfully warmed cache for {cache_key}")
|
|
|
|
# Track usage for cache warming call (use cache_key as filename identifier)
|
|
_track_usage_from_response(
|
|
response_body, f"cache_warming_{cache_key}", resolved_model_id
|
|
)
|
|
|
|
# Log cache usage info if available
|
|
if "usage" in response_body:
|
|
usage = response_body["usage"]
|
|
if "cache_creation_input_tokens" in usage:
|
|
logging.info(
|
|
f"Cache created with {usage.get('cache_creation_input_tokens', 0)} tokens for {cache_key}"
|
|
)
|
|
|
|
return True
|
|
|
|
except Exception as e:
|
|
logging.error(f"Error warming cache for {cache_key}: {str(e)}")
|
|
# Mark as not warmed on error so it can be retried
|
|
with _cache_warming_lock:
|
|
_cache_warmed[warmed_key] = False
|
|
return False
|