Skip to content

⚡️ Speed up method BedrockConverseModel._map_inference_config by 6% #19

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: try-refinement
Choose a base branch
from
Open

⚡️ Speed up method BedrockConverseModel._map_inference_config by 6% #19

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 27 additions & 14 deletions pydantic_ai_slim/pydantic_ai/models/bedrock.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,9 @@
from pydantic_ai.providers.bedrock import BedrockModelProfile
from pydantic_ai.settings import ModelSettings
from pydantic_ai.tools import ToolDefinition
from botocore.client import BaseClient
from mypy_boto3_bedrock_runtime import BedrockRuntimeClient
from mypy_boto3_bedrock_runtime.type_defs import InferenceConfigurationTypeDef

if TYPE_CHECKING:
from botocore.client import BaseClient
Expand Down Expand Up @@ -222,7 +225,13 @@ def __init__(
provider = infer_provider(provider)
self.client = cast('BedrockRuntimeClient', provider.client)

super().__init__(settings=settings, profile=profile or provider.model_profile)
model_profile = profile
if model_profile is None:
# provider.model_profile is a method; call it with model_name
mp = getattr(provider, "model_profile", None)
if callable(mp):
model_profile = mp(model_name)
super().__init__(settings=settings, profile=model_profile)

def _get_tools(self, model_request_parameters: ModelRequestParameters) -> list[ToolTypeDef]:
tools = [self._map_tool_definition(r) for r in model_request_parameters.function_tools]
Expand Down Expand Up @@ -367,19 +376,23 @@ async def _messages_create(
def _map_inference_config(
model_settings: ModelSettings | None,
) -> InferenceConfigurationTypeDef:
model_settings = model_settings or {}
inference_config: InferenceConfigurationTypeDef = {}

if max_tokens := model_settings.get('max_tokens'):
inference_config['maxTokens'] = max_tokens
if (temperature := model_settings.get('temperature')) is not None:
inference_config['temperature'] = temperature
if top_p := model_settings.get('top_p'):
inference_config['topP'] = top_p
if stop_sequences := model_settings.get('stop_sequences'):
inference_config['stopSequences'] = stop_sequences

return inference_config
"""Map ModelSettings to Bedrock InferenceConfigurationTypeDef."""
ms = model_settings or {}
# Only include keys that have valid values, minimizing work/allocations
config = {}
max_tokens = ms.get("max_tokens")
if max_tokens:
config["maxTokens"] = max_tokens
temperature = ms.get("temperature")
if temperature is not None:
config["temperature"] = temperature
top_p = ms.get("top_p")
if top_p:
config["topP"] = top_p
stop_sequences = ms.get("stop_sequences")
if stop_sequences:
config["stopSequences"] = stop_sequences
return config

def _map_tool_config(self, model_request_parameters: ModelRequestParameters) -> ToolConfigurationTypeDef | None:
tools = self._get_tools(model_request_parameters)
Expand Down
Loading