Re-apply a couple of recent changes to AWSNovaSonicLLMService that were lost in a rebase

This commit is contained in:
Paul Kompfner
2025-10-16 14:51:27 -04:00
parent 4820f1c059
commit 87d9e8c9cd

View File

@@ -65,7 +65,7 @@ try:
BedrockRuntimeClient,
InvokeModelWithBidirectionalStreamOperationInput,
)
from aws_sdk_bedrock_runtime.config import Config, HTTPAuthSchemeResolver, SigV4AuthScheme
from aws_sdk_bedrock_runtime.config import Config
from aws_sdk_bedrock_runtime.models import (
BidirectionalInputPayloadPart,
InvokeModelWithBidirectionalStreamInput,
@@ -73,8 +73,8 @@ try:
InvokeModelWithBidirectionalStreamOperationOutput,
InvokeModelWithBidirectionalStreamOutput,
)
from smithy_aws_core.credentials_resolvers.static import StaticCredentialsResolver
from smithy_aws_core.identity import AWSCredentialsIdentity
from smithy_aws_core.auth.sigv4 import SigV4AuthScheme
from smithy_aws_core.identity.static import StaticCredentialsResolver
from smithy_core.aio.eventstream import DuplexEventStream
except ModuleNotFoundError as e:
logger.error(f"Exception: {e}")
@@ -453,7 +453,7 @@ class AWSNovaSonicLLMService(LLMService):
await self._finish_connecting_if_context_available()
except Exception as e:
logger.error(f"{self} initialization error: {e}")
self._disconnect()
await self._disconnect()
async def _process_completed_function_calls(self, send_new_results: bool):
# Check for set of completed function calls in the context
@@ -582,15 +582,11 @@ class AWSNovaSonicLLMService(LLMService):
config = Config(
endpoint_uri=f"https://bedrock-runtime.{self._region}.amazonaws.com",
region=self._region,
aws_credentials_identity_resolver=StaticCredentialsResolver(
credentials=AWSCredentialsIdentity(
access_key_id=self._access_key_id,
secret_access_key=self._secret_access_key,
session_token=self._session_token,
)
),
http_auth_scheme_resolver=HTTPAuthSchemeResolver(),
http_auth_schemes={"aws.auth#sigv4": SigV4AuthScheme()},
aws_access_key_id=self._access_key_id,
aws_secret_access_key=self._secret_access_key,
aws_session_token=self._session_token,
aws_credentials_identity_resolver=StaticCredentialsResolver(),
auth_schemes={"aws.auth#sigv4": SigV4AuthScheme(service="bedrock")},
)
return BedrockRuntimeClient(config=config)