Skip to content

Commit

Permalink
Alleviate throttling error (#761)
Browse files Browse the repository at this point in the history 
* Reduce max results from 10 to 5

* Implement retry mechanism for Bedrock API throttling
  • Loading branch information
@statefb
statefb authored Mar 6, 2025
1 parent dc18f0d commit dbf945e
Show file tree
Hide file tree
Showing 5 changed files with 138 additions and 44 deletions.
24 changes: 22 additions & 2 deletions backend/app/bedrock.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@
from app.repositories.models.custom_bot_guardrails import BedrockGuardrailsModel
from app.routes.schemas.conversation import type_model_name
from app.utils import get_bedrock_runtime_client
from botocore.exceptions import ClientError
from retry import retry

if TYPE_CHECKING:
from app.agents.tools.agent_tool import AgentTool
Expand All @@ -26,6 +28,7 @@
SystemContentBlockTypeDef,
)


logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

Expand All @@ -43,6 +46,9 @@
client = get_bedrock_runtime_client()


class BedrockThrottlingException(Exception): ...


def _is_conversation_role(role: str) -> TypeGuard[ConversationRoleType]:
return role in ["user", "assistant"]

Expand Down Expand Up @@ -272,12 +278,26 @@ def process_content(c: ContentModel, role: str) -> list[ContentBlockTypeDef]:
return args


@retry(
exceptions=(BedrockThrottlingException,),
tries=3,
delay=60,
backoff=2,
jitter=(0, 2),
logger=logger,
)
def call_converse_api(
args: ConverseStreamRequestTypeDef,
) -> ConverseResponseTypeDef:
client = get_bedrock_runtime_client()

return client.converse(**args)
try:
return client.converse(**args)
except ClientError as e:
if e.response["Error"]["Code"] == "ThrottlingException":
raise BedrockThrottlingException(
"Bedrock API is throttling requests"
) from e
raise


def calculate_price(
Expand Down
25 changes: 23 additions & 2 deletions backend/app/stream.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,11 @@
from typing import Callable, TypedDict, TypeGuard

from app.agents.tools.agent_tool import AgentTool
from app.bedrock import calculate_price, compose_args_for_converse_api
from app.bedrock import (
BedrockThrottlingException,
calculate_price,
compose_args_for_converse_api,
)
from app.repositories.models.conversation import (
ContentModel,
MessageModel,
Expand All @@ -17,9 +21,11 @@
from app.repositories.models.custom_bot_guardrails import BedrockGuardrailsModel
from app.routes.schemas.conversation import type_model_name
from app.utils import get_bedrock_runtime_client, get_current_time
from botocore.exceptions import ClientError
from mypy_boto3_bedrock_runtime.literals import ConversationRoleType, StopReasonType
from mypy_boto3_bedrock_runtime.type_defs import GuardrailConverseContentBlockTypeDef
from pydantic import JsonValue
from retry import retry

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
Expand Down Expand Up @@ -174,6 +180,14 @@ def __init__(
self.on_thinking = on_thinking
self.on_reasoning = on_reasoning

@retry(
exceptions=(BedrockThrottlingException,),
tries=3,
delay=60,
backoff=2,
jitter=(0, 2),
logger=logger,
)
def run(
self,
messages: list[SimpleMessageModel],
Expand All @@ -196,7 +210,14 @@ def run(
logger.info(f"args for converse_stream: {args}")

client = get_bedrock_runtime_client()
response = client.converse_stream(**args)
try:
response = client.converse_stream(**args)
except ClientError as e:
if e.response["Error"]["Code"] == "ThrottlingException":
raise BedrockThrottlingException(
"Bedrock API is throttling requests"
) from e
raise

current_message = _PartialMessage(
role="assistant",
Expand Down
103 changes: 69 additions & 34 deletions frontend/src/features/agent/components/AvailableTools.tsx
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,11 @@
import { useTranslation } from 'react-i18next';
import { AgentTool, FirecrawlConfig, InternetAgentTool, SearchEngine, ToolType } from '../types';
import {
AgentTool,
FirecrawlConfig,
InternetAgentTool,
SearchEngine,
ToolType,
} from '../types';
import { isInternetTool } from '../utils/typeGuards';
import Toggle from '../../../components/Toggle';
import { Dispatch, useCallback, useState, useEffect } from 'react';
Expand All @@ -10,6 +16,7 @@ import { TooltipDirection } from '../../../constants';
import { FirecrawlConfig as FirecrawlConfigComponent } from './FirecrawlConfig';
import ExpandableDrawerGroup from '../../../components/ExpandableDrawerGroup';
import RadioButton from '../../../components/RadioButton';
import { DEFAULT_FIRECRAWL_CONFIG } from '../constants';

type Props = {
availableTools: AgentTool[] | undefined;
Expand All @@ -23,19 +30,23 @@ export const AvailableTools = ({ availableTools, tools, setTools }: Props) => {

const handleChangeTool = useCallback(
(tool: AgentTool) => () => {

if (tool.name === 'internet_search') {
setTools((preTools) => {
const isEnabled = preTools?.map(({ name }) => name).includes(tool.name);

const isEnabled = preTools
?.map(({ name }) => name)
.includes(tool.name);

const newTools = isEnabled
? [...preTools.filter(({ name }) => name != tool.name)]
: [...preTools, {
...tool,
toolType: "internet" as ToolType,
name: 'internet_search',
searchEngine: searchEngine || 'duckduckgo'
} as AgentTool];
: [
...preTools,
{
...tool,
toolType: 'internet' as ToolType,
name: 'internet_search',
searchEngine: searchEngine || 'duckduckgo',
} as AgentTool,
];

return newTools;
});
Expand All @@ -57,10 +68,12 @@ export const AvailableTools = ({ availableTools, tools, setTools }: Props) => {
if (tool.name === 'internet_search') {
return {
...tool,
toolType: "internet" as ToolType,
toolType: 'internet' as ToolType,
name: 'internet_search',
searchEngine: isInternetTool(tool) ? tool.searchEngine : 'duckduckgo',
firecrawlConfig: config
searchEngine: isInternetTool(tool)
? tool.searchEngine
: 'duckduckgo',
firecrawlConfig: config,
} as AgentTool;
}
return tool;
Expand All @@ -74,10 +87,12 @@ export const AvailableTools = ({ availableTools, tools, setTools }: Props) => {
(value: string) => {
const newEngine = value as SearchEngine;
setSearchEngine(newEngine);

// Update existing internet_search tool if it exists
setTools((prevTools) => {
const internetSearchTool = prevTools.find((t) => t.name === 'internet_search');
const internetSearchTool = prevTools.find(
(t) => t.name === 'internet_search'
);
if (!internetSearchTool) {
return prevTools;
}
Expand All @@ -86,11 +101,14 @@ export const AvailableTools = ({ availableTools, tools, setTools }: Props) => {
tool.name === 'internet_search'
? {
...tool,
toolType: "internet" as ToolType,
toolType: 'internet' as ToolType,
name: 'internet_search',
searchEngine: newEngine as SearchEngine,
// Reset firecrawlConfig when switching away from firecrawl
firecrawlConfig: newEngine === 'firecrawl' && isInternetTool(tool) ? tool.firecrawlConfig : undefined
firecrawlConfig:
newEngine === 'firecrawl' && isInternetTool(tool)
? tool.firecrawlConfig
: undefined,
}
: tool
);
Expand All @@ -103,15 +121,21 @@ export const AvailableTools = ({ availableTools, tools, setTools }: Props) => {
// Initialize searchEngine from existing tool if present
useEffect(() => {
const internetSearchTool = tools.find((t) => t.name === 'internet_search');
if (internetSearchTool && isInternetTool(internetSearchTool) && internetSearchTool.searchEngine) {
if (
internetSearchTool &&
isInternetTool(internetSearchTool) &&
internetSearchTool.searchEngine
) {
setSearchEngine(internetSearchTool.searchEngine);
}
}, [tools]);

return (
<>
<div className="flex items-center gap-1">
<div className="text-lg font-bold dark:text-aws-font-color-dark">{t('agent.label')}</div>
<div className="text-lg font-bold dark:text-aws-font-color-dark">
{t('agent.label')}
</div>
<Help direction={TooltipDirection.RIGHT} message={t('agent.hint')} />
</div>

Expand All @@ -136,45 +160,56 @@ export const AvailableTools = ({ availableTools, tools, setTools }: Props) => {
<ExpandableDrawerGroup
className="ml-8 mt-2"
isDefaultShow={false}
label={t('agent.tools.internet_search.settings')}
>
label={t('agent.tools.internet_search.settings')}>
<div className="space-y-4">
<div className="space-y-4">
<div className="flex flex-col gap-2">
<RadioButton
name="searchEngine"
value="duckduckgo"
checked={searchEngine === 'duckduckgo'}
label={t('agent.tools.internet_search.engines.duckduckgo.label')}
label={t(
'agent.tools.internet_search.engines.duckduckgo.label'
)}
onChange={handleSearchEngineChange}
/>
<div className="ml-6 text-sm text-aws-font-color-light/50 dark:text-aws-font-color-dark">
{t('agent.tools.internet_search.engines.duckduckgo.hint')}
{t(
'agent.tools.internet_search.engines.duckduckgo.hint'
)}
</div>
</div>
<div className="flex flex-col gap-2">
<RadioButton
name="searchEngine"
value="firecrawl"
checked={searchEngine === 'firecrawl'}
label={t('agent.tools.internet_search.engines.firecrawl.label')}
label={t(
'agent.tools.internet_search.engines.firecrawl.label'
)}
onChange={handleSearchEngineChange}
/>
<div className="ml-6 text-sm text-aws-font-color-light/50 dark:text-aws-font-color-dark">
{t('agent.tools.internet_search.engines.firecrawl.hint')}
{t(
'agent.tools.internet_search.engines.firecrawl.hint'
)}
</div>
<div className="ml-6 text-sm">
{searchEngine === 'firecrawl' && (
<FirecrawlConfigComponent
config={
tools.find((t): t is InternetAgentTool => t.name === 'internet_search' && isInternetTool(t))?.firecrawlConfig || {
apiKey: '',
maxResults: 10,
<FirecrawlConfigComponent
config={
tools.find(
(t): t is InternetAgentTool =>
t.name === 'internet_search' &&
isInternetTool(t)
)?.firecrawlConfig || {
apiKey: DEFAULT_FIRECRAWL_CONFIG.apiKey,
maxResults: DEFAULT_FIRECRAWL_CONFIG.maxResults,
}
}
}
onChange={handleFirecrawlConfigChange}
/>
)}
onChange={handleFirecrawlConfigChange}
/>
)}
</div>
</div>
</div>
Expand Down
16 changes: 10 additions & 6 deletions frontend/src/features/agent/components/FirecrawlConfig.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import { PiEyeLight, PiEyeSlashLight } from 'react-icons/pi';
import InputText from '../../../components/InputText';
import { Slider } from '../../../components/Slider';
import { FirecrawlConfig as FirecrawlConfigType } from '../types';
import { EDGE_FIRECRAWL_CONFIG } from '../constants';

type Props = {
config: FirecrawlConfigType;
Expand Down Expand Up @@ -47,19 +48,22 @@ export const FirecrawlConfig = ({ config, onChange }: Props) => {
<button
type="button"
className="size-9 rounded border border-aws-font-color-light/50 p-2 text-sm dark:border-aws-font-color-dark/50"
onClick={() => setShowPassword(!showPassword)}
>
{showPassword ? <PiEyeLight className="size-5" /> : <PiEyeSlashLight className="size-5" />}
onClick={() => setShowPassword(!showPassword)}>
{showPassword ? (
<PiEyeLight className="size-5" />
) : (
<PiEyeSlashLight className="size-5" />
)}
</button>
</div>
<Slider
label={t('agent.tools.firecrawl.maxResults')}
value={config.maxResults}
onChange={handleMaxResultsChange}
range={{
min: 1,
max: 50,
step: 1,
min: EDGE_FIRECRAWL_CONFIG.maxResults.MIN,
max: EDGE_FIRECRAWL_CONFIG.maxResults.MAX,
step: EDGE_FIRECRAWL_CONFIG.maxResults.STEP,
}}
/>
</div>
Expand Down
14 changes: 14 additions & 0 deletions frontend/src/features/agent/constants/index.ts
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,14 @@
import { FirecrawlConfig } from '../types';

export const DEFAULT_FIRECRAWL_CONFIG: FirecrawlConfig = {
apiKey: '',
maxResults: 5,
};

export const EDGE_FIRECRAWL_CONFIG = {
maxResults: {
MIN: 1,
MAX: 50,
STEP: 1,
},
};

0 comments on commit dbf945e

Please sign in to comment.