mirror of
https://github.com/codeflash-ai/codeflash-internal.git
synced 2026-05-04 18:25:18 +00:00
Migrate to AWS bedrock (#2430)
AWS_ACCESS_KEY_ID= AWS_SECRET_ACCESS_KEY= AWS_REGION=us-east-1 Will require these for boto3 authentication
This commit is contained in:
Sarthak Agarwal
GitHub
parent
7005156190
commit
eb5f4b460e
3 changed files with 75 additions and 20 deletions
|
|
@ -8,7 +8,7 @@ import time
|
|||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, Literal
|
||||
|
||||
from anthropic import AsyncAnthropicFoundry
|
||||
from anthropic import AsyncAnthropicBedrock
|
||||
from openai import AsyncAzureOpenAI
|
||||
from pydantic.dataclasses import dataclass as pydantic_dataclass
|
||||
|
||||
|
|
@ -66,9 +66,9 @@ class OpenAI_GPT_5_Mini(LLM):
|
|||
|
||||
@pydantic_dataclass
|
||||
class Anthropic_Claude_Sonnet_4_5(LLM):
|
||||
"""Anthropic Claude 4.5 Sonnet via Azure Foundry."""
|
||||
"""Anthropic Claude 4.5 Sonnet via AWS Bedrock."""
|
||||
|
||||
name: str = "claude-sonnet-4-5"
|
||||
name: str = "us.anthropic.claude-sonnet-4-5-20250929-v1:0"
|
||||
model_type: Literal["openai", "anthropic", "google"] = "anthropic"
|
||||
max_tokens: int = 200000
|
||||
input_cost: float = 3.00
|
||||
|
|
@ -77,9 +77,9 @@ class Anthropic_Claude_Sonnet_4_5(LLM):
|
|||
|
||||
@pydantic_dataclass
|
||||
class Anthropic_Claude_Haiku_4_5(LLM):
|
||||
"""Anthropic Claude 4.5 Haiku via Azure Foundry."""
|
||||
"""Anthropic Claude 4.5 Haiku via AWS Bedrock."""
|
||||
|
||||
name: str = "claude-haiku-4-5"
|
||||
name: str = "us.anthropic.claude-haiku-4-5-20251001-v1:0"
|
||||
model_type: Literal["openai", "anthropic", "google"] = "anthropic"
|
||||
max_tokens: int = 200000
|
||||
input_cost: float = 1.00
|
||||
|
|
@ -92,11 +92,9 @@ class Anthropic_Claude_Haiku_4_5(LLM):
|
|||
|
||||
# Read environment variables once at module load
|
||||
_AZURE_OPENAI_API_KEY = os.environ.get("AZURE_OPENAI_API_KEY")
|
||||
_ANTHROPIC_API_KEY = os.environ.get("ANTHROPIC_FOUNDRY_API_KEY")
|
||||
_ANTHROPIC_BASE_URL = os.environ.get("ANTHROPIC_FOUNDRY_BASE_URL")
|
||||
|
||||
# Clear ANTHROPIC_FOUNDRY_RESOURCE to prevent SDK conflict (mutually exclusive with base_url)
|
||||
os.environ.pop("ANTHROPIC_FOUNDRY_RESOURCE", None)
|
||||
_AWS_ACCESS_KEY = os.environ.get("AWS_ACCESS_KEY_ID")
|
||||
_AWS_SECRET_KEY = os.environ.get("AWS_SECRET_ACCESS_KEY")
|
||||
_AWS_REGION = os.environ.get("AWS_REGION", "us-east-1")
|
||||
|
||||
|
||||
def _create_openai_client() -> AsyncAzureOpenAI | None:
|
||||
|
|
@ -105,13 +103,15 @@ def _create_openai_client() -> AsyncAzureOpenAI | None:
|
|||
return None
|
||||
|
||||
|
||||
def _create_anthropic_client() -> AsyncAnthropicFoundry | None:
|
||||
if _ANTHROPIC_API_KEY and _ANTHROPIC_BASE_URL:
|
||||
return AsyncAnthropicFoundry(api_key=_ANTHROPIC_API_KEY, base_url=_ANTHROPIC_BASE_URL)
|
||||
def _create_anthropic_client() -> AsyncAnthropicBedrock | None:
|
||||
if _AWS_ACCESS_KEY and _AWS_SECRET_KEY:
|
||||
return AsyncAnthropicBedrock(
|
||||
aws_access_key=_AWS_ACCESS_KEY, aws_secret_key=_AWS_SECRET_KEY, aws_region=_AWS_REGION
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def get_llm_client(model_type: str) -> AsyncAzureOpenAI | AsyncAnthropicFoundry | None:
|
||||
def get_llm_client(model_type: str) -> AsyncAzureOpenAI | AsyncAnthropicBedrock | None:
|
||||
"""Get a fresh LLM client for the request.
|
||||
|
||||
Creates a new client for each request to avoid event loop issues
|
||||
|
|
@ -128,7 +128,7 @@ def get_llm_client(model_type: str) -> AsyncAzureOpenAI | AsyncAnthropicFoundry
|
|||
_openai_client = _create_openai_client()
|
||||
_anthropic_client = _create_anthropic_client()
|
||||
|
||||
llm_clients: dict[str, AsyncAzureOpenAI | AsyncAnthropicFoundry | None] = {
|
||||
llm_clients: dict[str, AsyncAzureOpenAI | AsyncAnthropicBedrock | None] = {
|
||||
"openai": _openai_client,
|
||||
"anthropic": _anthropic_client,
|
||||
}
|
||||
|
|
@ -186,7 +186,7 @@ async def call_llm(
|
|||
|
||||
try:
|
||||
if llm.model_type == "anthropic":
|
||||
assert isinstance(client, AsyncAnthropicFoundry)
|
||||
assert isinstance(client, AsyncAnthropicBedrock)
|
||||
system_prompt_content = next((m["content"] for m in messages if m["role"] == "system"), None)
|
||||
anthropic_messages = [
|
||||
{"role": m["role"], "content": m["content"]} for m in messages if m["role"] != "system"
|
||||
|
|
@ -289,7 +289,7 @@ def calculate_llm_cost(response: ChatCompletion | AnthropicMessage, llm: LLM) ->
|
|||
# Prefer OpenAI: use OpenAI if available, fall back to Anthropic
|
||||
OPENAI_MODEL: LLM = OpenAI_GPT_5_Mini() if _openai_client else Anthropic_Claude_Sonnet_4_5()
|
||||
|
||||
# Prefer Anthropic: use Anthropic (Azure Foundry) if available, fall back to OpenAI
|
||||
# Prefer Anthropic: use Anthropic (AWS Bedrock) if available, fall back to OpenAI
|
||||
ANTHROPIC_MODEL: LLM = Anthropic_Claude_Sonnet_4_5() if _anthropic_client else OpenAI_GPT_5_Mini()
|
||||
|
||||
# Haiku model for cost-effective tasks (testgen diversity)
|
||||
|
|
|
|||
|
|
@ -24,7 +24,7 @@ dependencies = [
|
|||
"sentry-sdk[django]>=2.35.0",
|
||||
"stamina>=25.1.0",
|
||||
"jedi>=0.19.2",
|
||||
"anthropic>=0.75.0",
|
||||
"anthropic[bedrock]>=0.75.0",
|
||||
"wcwidth>=0.2.15",
|
||||
"tree-sitter>=0.25.2",
|
||||
"tree-sitter-javascript>=0.25.0",
|
||||
|
|
|
|||
|
|
@ -119,7 +119,7 @@ name = "aiservice"
|
|||
version = "0.0.1"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "anthropic" },
|
||||
{ name = "anthropic", extra = ["bedrock"] },
|
||||
{ name = "dj-database-url" },
|
||||
{ name = "django" },
|
||||
{ name = "django-ninja" },
|
||||
|
|
@ -164,7 +164,7 @@ dev = [
|
|||
|
||||
[package.metadata]
|
||||
requires-dist = [
|
||||
{ name = "anthropic", specifier = ">=0.75.0" },
|
||||
{ name = "anthropic", extras = ["bedrock"], specifier = ">=0.75.0" },
|
||||
{ name = "dj-database-url", specifier = ">=2.2.0,<3" },
|
||||
{ name = "django", specifier = ">=5.0.6,<6" },
|
||||
{ name = "django-ninja", specifier = ">=1.3.0,<2" },
|
||||
|
|
@ -235,6 +235,12 @@ wheels = [
|
|||
{ url = "https://files.pythonhosted.org/packages/e5/70/7b0fd9c1a738f59d3babe2b4212031c34ab7d0fda4ffef15b58a55c5bcea/anthropic-0.76.0-py3-none-any.whl", hash = "sha256:81efa3113901192af2f0fe977d3ec73fdadb1e691586306c4256cd6d5ccc331c", size = 390309, upload-time = "2026-01-13T18:41:13.483Z" },
|
||||
]
|
||||
|
||||
[package.optional-dependencies]
|
||||
bedrock = [
|
||||
{ name = "boto3" },
|
||||
{ name = "botocore" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "anyio"
|
||||
version = "4.12.1"
|
||||
|
|
@ -284,6 +290,34 @@ wheels = [
|
|||
{ url = "https://files.pythonhosted.org/packages/df/73/b6e24bd22e6720ca8ee9a85a0c4a2971af8497d8f3193fa05390cbd46e09/backoff-2.2.1-py3-none-any.whl", hash = "sha256:63579f9a0628e06278f7e47b7d7d5b6ce20dc65c5e96a6f3ca99a6adca0396e8", size = 15148, upload-time = "2022-10-05T19:19:30.546Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "boto3"
|
||||
version = "1.42.50"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "botocore" },
|
||||
{ name = "jmespath" },
|
||||
{ name = "s3transfer" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/59/41/7a7280875ec000e280b0392478a5d6247bc88e7ecf2ae6ec8f4ddb35b014/boto3-1.42.50.tar.gz", hash = "sha256:38545d7e6e855fefc8a11e899ccbd6d2c9f64671d6648c2acfb1c78c1057a480", size = 112851, upload-time = "2026-02-16T20:42:09.203Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/5e/14/bf4077d843d737bec6f4176e113182a4435a1864e2a819ca07004da8a9ac/boto3-1.42.50-py3-none-any.whl", hash = "sha256:2fdf8f5349b130d62576068a6c47b3eec368a70bc28f16d8cce17c5f7e74fc2e", size = 140604, upload-time = "2026-02-16T20:42:06.652Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "botocore"
|
||||
version = "1.42.50"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "jmespath" },
|
||||
{ name = "python-dateutil" },
|
||||
{ name = "urllib3" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/93/fd/e63789133b2bf044c8550cd6766ec93628b0ac18a03f2aa0b80171f0697a/botocore-1.42.50.tar.gz", hash = "sha256:de1e128e4898f4e66877bfabbbb03c61f99366f27520442539339e8a74afe3a5", size = 14958074, upload-time = "2026-02-16T20:41:58.814Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/aa/b8/b02ad16c5198e652eafdd8bad76aa62ac094afabbe1241b4be1cd4075666/botocore-1.42.50-py3-none-any.whl", hash = "sha256:3ec7004009d1557a881b1d076d54b5768230849fa9ccdebfd409f0571490e691", size = 14631256, upload-time = "2026-02-16T20:41:55.004Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "certifi"
|
||||
version = "2026.1.4"
|
||||
|
|
@ -835,6 +869,15 @@ wheels = [
|
|||
{ url = "https://files.pythonhosted.org/packages/2f/9c/6753e6522b8d0ef07d3a3d239426669e984fb0eba15a315cdbc1253904e4/jiter-0.12.0-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c24e864cb30ab82311c6425655b0cdab0a98c5d973b065c66a3f020740c2324c", size = 346110, upload-time = "2025-11-09T20:49:21.817Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "jmespath"
|
||||
version = "1.1.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/d3/59/322338183ecda247fb5d1763a6cbe46eff7222eaeebafd9fa65d4bf5cb11/jmespath-1.1.0.tar.gz", hash = "sha256:472c87d80f36026ae83c6ddd0f1d05d4e510134ed462851fd5f754c8c3cbb88d", size = 27377, upload-time = "2026-01-22T16:35:26.279Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/14/2f/967ba146e6d58cf6a652da73885f52fc68001525b4197effc174321d70b4/jmespath-1.1.0-py3-none-any.whl", hash = "sha256:a5663118de4908c91729bea0acadca56526eb2698e83de10cd116ae0f4e97c64", size = 20419, upload-time = "2026-01-22T16:35:24.919Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "libcst"
|
||||
version = "1.8.6"
|
||||
|
|
@ -1701,6 +1744,18 @@ wheels = [
|
|||
{ url = "https://files.pythonhosted.org/packages/9e/6a/40fee331a52339926a92e17ae748827270b288a35ef4a15c9c8f2ec54715/ruff-0.14.14-py3-none-win_arm64.whl", hash = "sha256:56e6981a98b13a32236a72a8da421d7839221fa308b223b9283312312e5ac76c", size = 10920448, upload-time = "2026-01-22T22:30:15.417Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "s3transfer"
|
||||
version = "0.16.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "botocore" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/05/04/74127fc843314818edfa81b5540e26dd537353b123a4edc563109d8f17dd/s3transfer-0.16.0.tar.gz", hash = "sha256:8e990f13268025792229cd52fa10cb7163744bf56e719e0b9cb925ab79abf920", size = 153827, upload-time = "2025-12-01T02:30:59.114Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/fc/51/727abb13f44c1fcf6d145979e1535a35794db0f6e450a0cb46aa24732fe2/s3transfer-0.16.0-py3-none-any.whl", hash = "sha256:18e25d66fed509e3868dc1572b3f427ff947dd2c56f844a5bf09481ad3f3b2fe", size = 86830, upload-time = "2025-12-01T02:30:57.729Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "sentry-sdk"
|
||||
version = "2.50.0"
|
||||
|
|
|
|||
Loading…
Reference in a new issue