55"""
66
77import asyncio
8+ import base64
89import json
910import logging
1011import os
11- from typing import Any , AsyncGenerator , Callable , Optional , Type , TypeVar , Union , cast
12+ from collections .abc import AsyncGenerator , Callable
13+ from typing import Any , TypeVar , cast
1214
1315import boto3
1416from botocore .config import Config as BotocoreConfig
@@ -54,23 +56,23 @@ class BedrockModelInvoke(Model):
5456 class BedrockInvokeConfig (TypedDict , total = False ):
5557 """Configuration options for Bedrock InvokeModel."""
5658
57- guardrail_id : Optional [ str ]
58- guardrail_version : Optional [ str ]
59- max_tokens : Optional [ int ]
59+ guardrail_id : str | None
60+ guardrail_version : str | None
61+ max_tokens : int | None
6062 model_id : str
61- streaming : Optional [ bool ]
62- temperature : Optional [ float ]
63- top_p : Optional [ float ]
64- top_k : Optional [ int ]
65- stop_sequences : Optional [ list [str ]]
63+ streaming : bool | None
64+ temperature : float | None
65+ top_p : float | None
66+ top_k : int | None
67+ stop_sequences : list [str ] | None
6668
6769 def __init__ (
6870 self ,
6971 * ,
70- boto_session : Optional [ boto3 .Session ] = None ,
71- boto_client_config : Optional [ BotocoreConfig ] = None ,
72- region_name : Optional [ str ] = None ,
73- endpoint_url : Optional [ str ] = None ,
72+ boto_session : boto3 .Session | None = None ,
73+ boto_client_config : BotocoreConfig | None = None ,
74+ region_name : str | None = None ,
75+ endpoint_url : str | None = None ,
7476 ** model_config : Unpack [BedrockInvokeConfig ],
7577 ):
7678 """Initialize provider instance."""
@@ -108,7 +110,7 @@ def __init__(
108110 logger .debug ("region=<%s> | bedrock client created" , self .client .meta .region_name )
109111
110112 @override
111- def update_config (self , ** model_config : Unpack [BedrockInvokeConfig ]) -> None :
113+ def update_config (self , ** model_config : Unpack [BedrockInvokeConfig ]) -> None : # type: ignore[override]
112114 """Update the Bedrock Model configuration."""
113115 validate_config_keys (model_config , self .BedrockInvokeConfig )
114116 self .config .update (model_config )
@@ -138,12 +140,12 @@ def _get_model_family(self) -> str:
138140 def _format_anthropic_request (
139141 self ,
140142 messages : Messages ,
141- tool_specs : Optional [ list [ToolSpec ]] = None ,
142- system_prompt_content : Optional [ list [SystemContentBlock ]] = None ,
143+ tool_specs : list [ToolSpec ] | None = None ,
144+ system_prompt_content : list [SystemContentBlock ] | None = None ,
143145 tool_choice : ToolChoice | None = None ,
144146 ) -> dict [str , Any ]:
145147 """Format request for Anthropic Claude models."""
146- request = {
148+ request : dict [ str , Any ] = {
147149 "anthropic_version" : "bedrock-2023-05-31" ,
148150 "max_tokens" : self .config .get ("max_tokens" , 4096 ),
149151 "messages" : [],
@@ -158,7 +160,7 @@ def _format_anthropic_request(
158160 # Convert messages
159161 for msg in messages :
160162 role = msg ["role" ]
161- content = []
163+ content : list [ dict [ str , Any ]] = []
162164
163165 for block in msg ["content" ]:
164166 if "text" in block :
@@ -171,7 +173,7 @@ def _format_anthropic_request(
171173 "source" : {
172174 "type" : "base64" ,
173175 "media_type" : image_data ["format" ],
174- "data" : image_data ["source" ]["bytes" ],
176+ "data" : base64 . b64encode ( image_data ["source" ]["bytes" ]). decode ( "utf-8" ) ,
175177 },
176178 }
177179 )
@@ -226,12 +228,12 @@ def _format_anthropic_request(
226228 def _format_openai_request (
227229 self ,
228230 messages : Messages ,
229- tool_specs : Optional [ list [ToolSpec ]] = None ,
230- system_prompt_content : Optional [ list [SystemContentBlock ]] = None ,
231+ tool_specs : list [ToolSpec ] | None = None ,
232+ system_prompt_content : list [SystemContentBlock ] | None = None ,
231233 tool_choice : ToolChoice | None = None ,
232234 ) -> dict [str , Any ]:
233235 """Format request for OpenAI-compatible models."""
234- request = {
236+ request : dict [ str , Any ] = {
235237 "model" : self .config ["model_id" ],
236238 "messages" : [],
237239 "max_tokens" : self .config .get ("max_tokens" , 4096 ),
@@ -284,8 +286,8 @@ def _format_openai_request(
284286 def _format_request (
285287 self ,
286288 messages : Messages ,
287- tool_specs : Optional [ list [ToolSpec ]] = None ,
288- system_prompt_content : Optional [ list [SystemContentBlock ]] = None ,
289+ tool_specs : list [ToolSpec ] | None = None ,
290+ system_prompt_content : list [SystemContentBlock ] | None = None ,
289291 tool_choice : ToolChoice | None = None ,
290292 ) -> dict [str , Any ]:
291293 """Format request based on model family."""
@@ -295,7 +297,7 @@ def _format_request(
295297
296298 def _parse_anthropic_response (self , response : dict [str , Any ]) -> list [StreamEvent ]:
297299 """Parse response into StreamEvent format."""
298- events = []
300+ events : list [ StreamEvent ] = []
299301
300302 # Start message
301303 events .append ({"messageStart" : {"role" : "assistant" }})
@@ -318,7 +320,7 @@ def _extract_text_from_response(self, response: dict[str, Any]) -> str:
318320 # Try common text fields
319321 for field in ["completion" , "outputText" , "text" , "response" , "output" ]:
320322 if field in response and isinstance (response [field ], str ):
321- return response [field ]
323+ return str ( response [field ])
322324
323325 # Try Anthropic content format
324326 if "content" in response and isinstance (response ["content" ], list ):
@@ -334,7 +336,7 @@ def _extract_text_from_response(self, response: dict[str, Any]) -> str:
334336 # Fallback: convert entire response to string if no text found
335337 return json .dumps (response )
336338
337- def _parse_anthropic_streaming_chunk (self , chunk : dict [str , Any ]) -> Optional [ StreamEvent ] :
339+ def _parse_anthropic_streaming_chunk (self , chunk : dict [str , Any ]) -> StreamEvent | None :
338340 """Parse a single streaming chunk."""
339341 # Handle OpenAI-style chat completion chunks
340342 if "choices" in chunk and chunk ["choices" ]:
@@ -354,8 +356,8 @@ def _parse_anthropic_streaming_chunk(self, chunk: dict[str, Any]) -> Optional[St
354356 return None
355357
356358 def _extract_usage_from_response (
357- self , response : dict [str , Any ], response_body : dict [str , Any ] = None
358- ) -> Optional [ dict [str , Any ]] :
359+ self , response : dict [str , Any ], response_body : dict [str , Any ] | None = None
360+ ) -> dict [str , Any ] | None :
359361 """Extract usage information from response body."""
360362 usage = {}
361363
@@ -385,20 +387,20 @@ def _extract_usage_from_response(
385387 async def stream (
386388 self ,
387389 messages : Messages ,
388- tool_specs : Optional [ list [ToolSpec ]] = None ,
389- system_prompt : Optional [ str ] = None ,
390+ tool_specs : list [ToolSpec ] | None = None ,
391+ system_prompt : str | None = None ,
390392 * ,
391393 tool_choice : ToolChoice | None = None ,
392- system_prompt_content : Optional [ list [SystemContentBlock ]] = None ,
394+ system_prompt_content : list [SystemContentBlock ] | None = None ,
393395 ** kwargs : Any ,
394396 ) -> AsyncGenerator [StreamEvent , None ]:
395397 """Stream conversation with the Bedrock model using InvokeModel APIs."""
396398
397- def callback (event : Optional [ StreamEvent ] = None ) -> None :
399+ def callback (event : StreamEvent | None = None ) -> None :
398400 loop .call_soon_threadsafe (queue .put_nowait , event )
399401
400402 loop = asyncio .get_event_loop ()
401- queue : asyncio .Queue [Optional [ StreamEvent ] ] = asyncio .Queue ()
403+ queue : asyncio .Queue [StreamEvent | None ] = asyncio .Queue ()
402404
403405 # Handle backward compatibility
404406 if system_prompt and system_prompt_content is None :
@@ -419,8 +421,8 @@ def _stream(
419421 self ,
420422 callback : Callable [..., None ],
421423 messages : Messages ,
422- tool_specs : Optional [ list [ToolSpec ]] = None ,
423- system_prompt_content : Optional [ list [SystemContentBlock ]] = None ,
424+ tool_specs : list [ToolSpec ] | None = None ,
425+ system_prompt_content : list [SystemContentBlock ] | None = None ,
424426 tool_choice : ToolChoice | None = None ,
425427 ) -> None :
426428 """Stream conversation in separate thread."""
@@ -512,11 +514,11 @@ def _stream(
512514 @override
513515 async def structured_output (
514516 self ,
515- output_model : Type [T ],
517+ output_model : type [T ],
516518 prompt : Messages ,
517- system_prompt : Optional [ str ] = None ,
519+ system_prompt : str | None = None ,
518520 ** kwargs : Any ,
519- ) -> AsyncGenerator [dict [str , Union [ T , Any ] ], None ]:
521+ ) -> AsyncGenerator [dict [str , T | Any ], None ]:
520522 """Get structured output from the model."""
521523 tool_spec = convert_pydantic_to_tool_spec (output_model )
522524
0 commit comments