import json import requests import time import re from consts.consts import CodeBlock, Claude3 from config import config from utils.logging import logging # ⭐ Inference profile for Bearer Token auth CLAUDE_PROFILE_ARN = "arn:aws:bedrock:us-west-2:841162690310:application-inference-profile/pe6vvdc9orkr" BEDROCK_REGION = "us-west-2" BEDROCK_INVOKE_URL = "https://bedrock-runtime.{region}.amazonaws.com/model/{model_id}/invoke" REQUEST_TIMEOUT = 600 MAX_RETRIES = 3 BEARER_TOKEN = config.AWS_BEARER_TOKEN_BEDROCK if BEARER_TOKEN: logging.info("AWS_BEARER_TOKEN_BEDROCK loaded, value = %s...%s", BEARER_TOKEN[:8], BEARER_TOKEN[-4:]) else: logging.error("AWS_BEARER_TOKEN_BEDROCK is empty, bearer token auth will not work") def _invoke_bedrock(model_id, body, bearer_token): from urllib.parse import quote url = BEDROCK_INVOKE_URL.format(region=BEDROCK_REGION, model_id=quote(model_id, safe="")) headers = { "Authorization": f"Bearer {bearer_token}", "Content-Type": "application/json", } logging.info("_invoke_bedrock: url = %s, bearer_token = %s...%s", url, bearer_token[:8] if bearer_token else "EMPTY", bearer_token[-4:] if bearer_token else "") for attempt in range(1, MAX_RETRIES + 1): logging.info("_invoke_bedrock: attempt %d/%d, sending request...", attempt, MAX_RETRIES) try: resp = requests.post(url, headers=headers, data=json.dumps(body), timeout=REQUEST_TIMEOUT) except requests.exceptions.RequestException as e: logging.error("_invoke_bedrock: attempt %d/%d, request exception: %s", attempt, MAX_RETRIES, str(e)) if attempt < MAX_RETRIES: time.sleep(min(2 ** attempt, 8)) continue raise logging.info("_invoke_bedrock: attempt %d/%d, status_code = %d, response_length = %d", attempt, MAX_RETRIES, resp.status_code, len(resp.content)) if resp.status_code == 200: return resp.json() logging.error("_invoke_bedrock: attempt %d/%d, request failed, status = %d, response = %s", attempt, MAX_RETRIES, resp.status_code, resp.text[:500]) if resp.status_code in (429, 500, 503) and attempt < MAX_RETRIES: time.sleep(min(2 ** attempt, 8)) continue resp.raise_for_status() return resp.json() def invoke_claude_3_with_text(prompt, output_format, message_list = [], dry_run=False, ai_model=Claude3.DEFAULT, extra_info={}): model_id = CLAUDE_PROFILE_ARN if not model_id: logging.info("session_id = %s, query_id = %s, module = %s, %s is not supported", extra_info.get("session_id"), extra_info.get("query_id"), extra_info.get("module"), ai_model) time_before = int(time.time()) logging.info("session_id = %s, query_id = %s, module = %s, start to request %s. use model_id = %s, bearer_token = %s...%s, prompt = %s, dry = %s", extra_info.get("session_id"), extra_info.get("query_id"), extra_info.get("module"), ai_model, model_id, BEARER_TOKEN[:8] if BEARER_TOKEN else "EMPTY", BEARER_TOKEN[-4:] if BEARER_TOKEN else "", prompt, str(dry_run)) if dry_run: return True messages = message_list + [{"role": "user", "content": prompt}] try: body = { "anthropic_version": "bedrock-2023-05-31", "max_tokens": 4096, "temperature": 0.2, "messages": messages } result = _invoke_bedrock(model_id, body, BEARER_TOKEN) input_tokens = result["usage"]["input_tokens"] output_tokens = result["usage"]["output_tokens"] output_list = result.get("content", []) if len(output_list) == 0: raise Exception("No response received.") resp_str = output_list[0]["text"] logging.info("session_id = %s, query_id = %s, module = %s, end to request %s. model_id = %s, resp = %s, input_tokens = %s, output_tokens = %s, time_cost = %s", extra_info.get("session_id"), extra_info.get("query_id"), extra_info.get("module"), ai_model, model_id, resp_str, str(input_tokens), str(output_tokens), str(int(time.time()) - time_before)) resp_json = extract_text(resp_str) logging.info("session_id = %s, query_id = %s, module = %s, extract result finished, resp = %s", extra_info.get("session_id"), extra_info.get("query_id"), extra_info.get("module"), resp_json) messages.append({"role": "assistant", "content": resp_json}) if output_format == "json": return json.loads(resp_json) return resp_json except Exception as e: logging.error("session_id = %s, query_id = %s, module = %s, failed to request %s. model_id = %s, error = %s, message_list = %s", extra_info.get("session_id"), extra_info.get("query_id"), extra_info.get("module"), ai_model, model_id, str(e), str(message_list)) if isinstance(e, json.JSONDecodeError) and len(messages) == 2: logging.error("session_id = %s, query_id = %s, module = %s, json decode error, regenerate, message_list = %s", extra_info.get("session_id"), extra_info.get("query_id"), extra_info.get("module"), str(messages)) resp_json = invoke_claude_3_with_text("The result is not in a json format, please regenerate it", "json", message_list=messages, extra_info=extra_info) logging.error("session_id = %s, query_id = %s, module = %s, regenerate finished, resp_json= %s", extra_info.get("session_id"), extra_info.get("query_id"), extra_info.get("module"), str(resp_json)) return resp_json raise def extract_text(resp_str): # extract json if CodeBlock.JSON_CODE_BLOCK in resp_str: pattern = r'```json(.*?)```' match = re.search(pattern, resp_str, re.DOTALL) if match: resp_str = match.group(1).strip().replace("\\\n", " ").replace("\n", " ").replace("\t", " ") # extract sql elif CodeBlock.SQL_CODE_BLOCK in resp_str: pattern = r'```sql(.*?)```' match = re.search(pattern, resp_str, re.DOTALL) if match: resp_str = match.group(1).strip().replace("\\\n", " ").replace("\n", " ").replace("\t", " ") else: resp_str = resp_str.replace("\\\n", " ").replace("\n", " ").replace("\t", " ").strip() # if operator '+' in json str if '+' in resp_str: pattern = re.compile(r'"\s*\+\s*"') resp_str = pattern.sub(' ', resp_str) return resp_str def test_claude3_api_key(): model_id = CLAUDE_PROFILE_ARN body = { "anthropic_version": "bedrock-2023-05-31", "max_tokens": 4096, "temperature": 0, "messages": [{"role": "user", "content": "test"}] } _invoke_bedrock(model_id, body, BEARER_TOKEN) def test_claude_inference_profile_connectivity(prompt="Reply exactly: pong"): model_id = CLAUDE_PROFILE_ARN.strip() if not model_id: raise ValueError("CLAUDE_PROFILE_ARN is empty") body = { "anthropic_version": "bedrock-2023-05-31", "max_tokens": 64, "temperature": 0, "messages": [{"role": "user", "content": prompt}], } result = _invoke_bedrock(model_id, body, BEARER_TOKEN) content = result.get("content", []) if not content: raise Exception(f"No content in response: {result}") return { "model_id": model_id, "text": content[0].get("text", ""), "input_tokens": result.get("usage", {}).get("input_tokens", 0), "output_tokens": result.get("usage", {}).get("output_tokens", 0), }