Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-llm/src/lexigram/ai/llm/clients/aws_bedrock.py: 20%
181 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""AWS Bedrock LLM client for the Lexigram LLM routing system.
3Implements the :class:`~lexigram.contracts.ai.protocols.LLMClientProtocol`
4protocol against the AWS Bedrock ``Converse`` API. The ``Converse`` API
5provides a unified interface for all Bedrock models and natively supports
6multi-turn conversations and tool (function) calling.
8Authentication and request signing use the ``boto3``/``botocore`` AWS SDK,
9which honours the standard AWS credential chain (environment variables, shared
10credentials file, EC2 instance profile, ECS task role, etc.).
12Configuration is sourced from ``ClientConfig.extra``:
14* ``aws_region`` — AWS region, e.g. ``us-east-1`` (required)
15* ``aws_access_key_id`` — AWS access key ID (optional; from credential chain)
16* ``aws_secret_access_key`` — AWS secret access key (optional; from credential chain)
17* ``aws_session_token`` — AWS session token for temporary credentials (optional)
18* ``aws_profile`` — Named profile from ``~/.aws/credentials`` (optional)
20Notes:
21 ``boto3`` and ``botocore`` are optional dependencies. An
22 :class:`ImportError` is raised at construction time if they are absent.
24 The Bedrock ``Converse`` API does not natively expose a server-sent event
25 (SSE) streaming response over HTTP like OpenAI. Streaming uses
26 ``ConverseStream``, which returns an event-stream that ``botocore``
27 processes asynchronously via ``aiobotocore`` if available, or
28 synchronously via a thread-pool executor otherwise.
29"""
31from __future__ import annotations
33import asyncio
34import atexit
35from concurrent.futures import ThreadPoolExecutor
36from concurrent.futures.thread import _threads_queues, _worker
37import threading
38from typing import TYPE_CHECKING, Any
39import weakref
41from lexigram.ai.llm.clients._bedrock_mappers import (
42 bedrock_stream_chunks,
43 extract_system,
44 parse_bedrock_response,
45 tool_to_bedrock,
46)
47from lexigram.ai.llm.clients._tools_utils import (
48 _tool_schema_fields,
49 parse_json_arguments,
50)
51from lexigram.ai.llm.clients.base import AbstractLLMClient
52from lexigram.ai.llm.exceptions import (
53 LLMAuthenticationError,
54 LLMContentFilterError,
55 LLMError,
56 LLMModelNotFoundError,
57 LLMRateLimitError,
58)
59from lexigram.ai.llm.multimodal.fetcher import fetch_image_as_base64
60from lexigram.ai.llm.types import AIError, Completion, StreamChunk
61from lexigram.contracts.ai.multimodal import (
62 ImageBase64Part,
63 ImageUrlPart,
64 MessageContent,
65 TextPart,
66)
67from lexigram.contracts.core import HealthCheckResult, HealthStatus
68from lexigram.logging import (
69 get_logger,
70)
71from lexigram.result import Err, Ok, Result
73if TYPE_CHECKING:
74 from collections.abc import AsyncIterator
76 from lexigram.ai.llm.config import ClientConfig
78logger = get_logger(__name__)
80__all__ = ["BedrockClient"]
82_thread_pool: ThreadPoolExecutor | None = None
85class _DaemonThreadPoolExecutor(ThreadPoolExecutor):
86 """ThreadPoolExecutor whose workers are daemon threads.
88 Keeps the pool off the non-daemon thread list so it never blocks
89 interpreter shutdown or trips test teardown assertions. Mirrors the
90 CPython 3.13 worker-spawn logic with ``daemon=True``.
91 """
93 def _adjust_thread_count(self) -> None:
94 if self._idle_semaphore.acquire(timeout=0):
95 return
97 def weakref_cb(_: Any, q: Any = self._work_queue) -> None:
98 q.put(None)
100 num_threads = len(self._threads)
101 if num_threads < self._max_workers:
102 thread_name = f"{self._thread_name_prefix or self}_{num_threads}"
103 t = threading.Thread(
104 name=thread_name,
105 target=_worker,
106 args=(
107 weakref.ref(self, weakref_cb),
108 self._work_queue,
109 self._initializer,
110 self._initargs,
111 ),
112 daemon=True,
113 )
114 t.start()
115 self._threads.add(t) # type: ignore[attr-defined]
116 _threads_queues[t] = self._work_queue # type: ignore[index]
119def _get_thread_pool() -> ThreadPoolExecutor:
120 """Get the shared executor, creating it lazily on first use."""
121 global _thread_pool
122 if _thread_pool is None:
123 _thread_pool = _DaemonThreadPoolExecutor(
124 max_workers=4, thread_name_prefix="bedrock-sync"
125 )
126 atexit.register(_thread_pool.shutdown)
127 return _thread_pool
130def _content_to_text(content: Any) -> str:
131 """Extract plain text from message content for tool-result payloads.
133 Args:
134 content: Message content (string or list of content parts).
136 Returns:
137 Plain text string.
138 """
139 if isinstance(content, str):
140 return content
141 if isinstance(content, list):
142 return " ".join(
143 str(p.get("text", ""))
144 for p in content
145 if isinstance(p, dict) and p.get("type") == "text"
146 )
147 return str(content)
150class BedrockClient(AbstractLLMClient):
151 """AWS Bedrock client using the ``Converse`` API.
153 Routes requests to the Bedrock ``converse`` (non-streaming) and
154 ``converse_stream`` (streaming) endpoints. The ``Converse`` API
155 normalises request/response format across all Bedrock models, so
156 no per-model message translation is needed.
158 Args:
159 config: LLM configuration. ``config.extra`` must contain
160 ``aws_region``.
161 """
163 def __init__(self, config: ClientConfig) -> None:
164 """Initialise the Bedrock client.
166 Args:
167 config: LLM configuration with AWS-specific ``extra`` keys.
169 Raises:
170 ImportError: If ``boto3`` is not installed.
171 ValueError: If ``aws_region`` is missing from ``config.extra``.
172 """
173 super().__init__(config=config)
175 try:
176 import boto3
177 except ImportError as exc:
178 raise ImportError(
179 "BedrockClient requires 'boto3'. Install with: pip install boto3"
180 ) from exc
182 extra: dict[str, Any] = config.extra or {}
183 region = extra.get("aws_region", "")
184 if not region:
185 raise ValueError(
186 "BedrockClient requires 'aws_region' in ClientConfig.extra"
187 )
189 self._region = region
190 boto_kwargs: dict[str, Any] = {"region_name": region}
191 if extra.get("aws_access_key_id"):
192 boto_kwargs["aws_access_key_id"] = extra["aws_access_key_id"]
193 if extra.get("aws_secret_access_key"):
194 boto_kwargs["aws_secret_access_key"] = extra["aws_secret_access_key"]
195 if extra.get("aws_session_token"):
196 boto_kwargs["aws_session_token"] = extra["aws_session_token"]
197 if extra.get("aws_profile"):
198 import boto3
200 session = boto3.Session(profile_name=extra["aws_profile"], **boto_kwargs)
201 self._client = session.client("bedrock-runtime")
202 else:
203 import boto3
205 self._client = boto3.client("bedrock-runtime", **boto_kwargs)
207 # ──────────────────────────────────────────────────────────────────
208 # LLMClientProtocol implementation
209 # ──────────────────────────────────────────────────────────────────
211 async def _do_complete(
212 self,
213 messages: list[Any],
214 *,
215 model: str | None = None,
216 temperature: float = 0.2,
217 max_tokens: int | None = None,
218 **kwargs: Any,
219 ) -> Result[Completion, LLMError]:
220 """Generate completion via Bedrock ``Converse``.
222 Args:
223 messages: OpenAI-compatible message list.
224 model: Model ID override (Bedrock model ARN or ID).
225 temperature: Sampling temperature.
226 max_tokens: Maximum output tokens.
227 **kwargs: Ignored for protocol compatibility.
229 Returns:
230 ``Ok(Completion)`` on success. ``Err(LLMError)`` for recoverable
231 failures.
233 Raises:
234 LLMAuthenticationError: On credential or authorisation failure.
235 AIError: For unexpected infrastructure failures.
236 """
237 active_model = model or self.config.model
238 bedrock_messages = await self._to_bedrock_messages_async(messages)
239 request: dict[str, Any] = {
240 "modelId": active_model,
241 "messages": bedrock_messages,
242 "inferenceConfig": {},
243 }
244 # Extended thinking is incompatible with temperature on Bedrock Claude
245 self._apply_thinking(request)
246 if "additionalModelRequestFields" not in request:
247 request["inferenceConfig"]["temperature"] = temperature
249 if max_tokens is not None:
250 request["inferenceConfig"]["maxTokens"] = max_tokens
252 system = extract_system(messages)
253 if system:
254 request["system"] = [{"text": system}]
256 tools = kwargs.pop("tools", None)
257 if tools:
258 converted_tools = [
259 tool_to_bedrock(t) for t in tools if _tool_schema_fields(t)[0]
260 ]
261 if converted_tools:
262 request["toolConfig"] = {
263 "tools": converted_tools,
264 "toolChoice": {"auto": {}},
265 }
267 try:
268 raw = await asyncio.get_event_loop().run_in_executor(
269 _get_thread_pool(),
270 lambda: self._client.converse(**request),
271 )
272 except Exception as exc: # noqa: BLE001 - botocore raises dynamic provider exceptions
273 return self._handle_error_as_result(exc)
275 return Ok(parse_bedrock_response(raw, active_model))
277 async def _do_stream_chat(
278 self,
279 messages: list[Any],
280 *,
281 model: str | None = None,
282 temperature: float = 0.2,
283 max_tokens: int | None = None,
284 **kwargs: Any,
285 ) -> Result[AsyncIterator[StreamChunk], LLMError]:
286 """Stream completion tokens from Bedrock ``ConverseStream``.
288 Args:
289 messages: OpenAI-compatible message list.
290 model: Model ID override.
291 temperature: Sampling temperature.
292 max_tokens: Maximum output tokens.
293 **kwargs: Ignored for protocol compatibility.
295 Returns:
296 ``Ok(AsyncIterator[StreamChunk])`` on success.
297 ``Err(LLMError)`` for recoverable failures.
299 Raises:
300 LLMAuthenticationError: On credential or authorisation failure.
301 AIError: For unexpected infrastructure failures.
302 """
303 active_model = model or self.config.model
304 bedrock_messages = await self._to_bedrock_messages_async(messages)
305 request: dict[str, Any] = {
306 "modelId": active_model,
307 "messages": bedrock_messages,
308 "inferenceConfig": {},
309 }
310 # Extended thinking is incompatible with temperature on Bedrock Claude
311 self._apply_thinking(request)
312 if "additionalModelRequestFields" not in request:
313 request["inferenceConfig"]["temperature"] = temperature
315 if max_tokens is not None:
316 request["inferenceConfig"]["maxTokens"] = max_tokens
318 system = extract_system(messages)
319 if system:
320 request["system"] = [{"text": system}]
322 try:
323 raw_stream = await asyncio.get_event_loop().run_in_executor(
324 _get_thread_pool(),
325 lambda: self._client.converse_stream(**request),
326 )
327 except Exception as exc: # noqa: BLE001 - botocore raises dynamic provider exceptions
328 return self._handle_error_as_result(exc)
330 return Ok(bedrock_stream_chunks(raw_stream, active_model, _get_thread_pool()))
332 async def _do_chat(
333 self,
334 messages: list[Any],
335 tools: list[Any] | None = None,
336 *,
337 model: str | None = None,
338 temperature: float = 0.2,
339 max_tokens: int | None = None,
340 **kwargs: Any,
341 ) -> Result[Completion, LLMError]:
342 """Generate completion with optional tool calling on Bedrock.
344 Tool calling is handled by :meth:`_do_complete` via
345 ``complete(..., tools=...)``; this method forwards the tool
346 descriptors to keep the ``chat`` code path consistent.
348 Args:
349 messages: OpenAI-compatible message list.
350 tools: Optional tool descriptors.
351 model: Model ID override.
352 temperature: Sampling temperature.
353 max_tokens: Maximum output tokens.
354 **kwargs: Ignored for protocol compatibility.
356 Returns:
357 ``Ok(Completion)`` on success. ``Err(LLMError)`` for recoverable
358 failures.
359 """
360 return await self._do_complete(
361 messages,
362 model=model,
363 temperature=temperature,
364 max_tokens=max_tokens,
365 tools=tools,
366 **kwargs,
367 )
369 async def health_check(self, timeout: float = 5.0) -> HealthCheckResult:
370 """Probe Bedrock by listing foundation models.
372 Args:
373 timeout: Informational only.
375 Returns:
376 Structured :class:`~lexigram.contracts.core.health.HealthCheckResult`.
377 """
378 try:
379 import boto3
381 mgmt = boto3.client("bedrock", region_name=self._region)
382 await asyncio.get_event_loop().run_in_executor(
383 _get_thread_pool(),
384 lambda: mgmt.list_foundation_models(maxResults=1),
385 )
386 except Exception as exc: # noqa: BLE001 - boto3 exposes provider-specific exceptions
387 return HealthCheckResult(
388 component="llm.bedrock",
389 status=HealthStatus.UNHEALTHY,
390 error=str(exc),
391 details={"region": self._region, "model": self.config.model},
392 )
394 return HealthCheckResult(
395 component="llm.bedrock",
396 status=HealthStatus.HEALTHY,
397 details={"region": self._region, "model": self.config.model},
398 )
400 async def close(self) -> None:
401 """Release the Bedrock boto3 client."""
402 if self._client is not None:
403 self._client.close()
404 self._client = None
405 await super().close()
407 def _apply_thinking(self, request: dict[str, Any]) -> None:
408 """Inject Bedrock extended-thinking parameters into the request payload.
410 Sets ``additionalModelRequestFields.thinking`` when ``config.thinking``
411 is configured. Also removes ``temperature`` from ``inferenceConfig``
412 because Bedrock Claude rejects that combination.
414 Args:
415 request: Mutable Bedrock request dict.
416 """
417 if self.config.thinking is None:
418 return
419 request["additionalModelRequestFields"] = {
420 "thinking": {
421 "type": "enabled",
422 "budget_tokens": self.config.thinking.budget_tokens,
423 }
424 }
425 request.get("inferenceConfig", {}).pop("temperature", None)
427 async def _build_content_blocks(
428 self, content: MessageContent
429 ) -> list[dict[str, Any]]:
430 """Convert MessageContent to Bedrock ``Converse`` content blocks.
432 Handles plain strings, TextPart, ImageBase64Part, and ImageUrlPart.
433 ImageUrlPart entries are fetched and converted to base64 via
434 ``fetch_image_as_base64``.
436 Args:
437 content: Message content (string or list of ContentPart).
439 Returns:
440 List of Bedrock content block dicts with wire format:
441 - ``{"type": "text", "text": "..."}``.
442 - ``{"type": "image", "source": {"type": "base64", "mediaType": "...", "data": "..."}}``
443 (camelCase ``mediaType`` as per Bedrock API).
445 Raises:
446 LLMError: If image fetching fails (infrastructure error).
447 """
448 if isinstance(content, str):
449 return [{"type": "text", "text": content}]
451 blocks: list[dict[str, Any]] = []
452 for part in content:
453 if isinstance(part, TextPart):
454 blocks.append({"type": "text", "text": part.text})
455 elif isinstance(part, ImageBase64Part):
456 blocks.append(
457 {
458 "type": "image",
459 "source": {
460 "type": "base64",
461 "mediaType": part.media_type,
462 "data": part.data,
463 },
464 }
465 )
466 elif isinstance(part, ImageUrlPart):
467 # Fetch and convert to base64
468 fetched = await fetch_image_as_base64(part.url)
469 blocks.append(
470 {
471 "type": "image",
472 "source": {
473 "type": "base64",
474 "mediaType": fetched.media_type,
475 "data": fetched.data,
476 },
477 }
478 )
479 return blocks
481 async def _to_bedrock_messages_async(
482 self, messages: list[Any]
483 ) -> list[dict[str, Any]]:
484 """Async conversion of OpenAI-compatible messages to Bedrock format.
486 System messages are excluded (handled separately). Multimodal content
487 is converted using ``_build_content_blocks``.
489 Args:
490 messages: OpenAI-compatible message list (dicts or ChatMessage).
492 Returns:
493 List of Bedrock ``ConversationTurn`` dicts.
494 """
495 result: list[dict[str, Any]] = []
496 for msg in messages:
497 role = (
498 msg.get("role", "user")
499 if isinstance(msg, dict)
500 else getattr(msg, "role", "user")
501 )
502 role_str = role.value if hasattr(role, "value") else str(role)
503 if role_str == "system":
504 continue
506 content_raw = (
507 msg.get("content", "")
508 if isinstance(msg, dict)
509 else getattr(msg, "content", "")
510 )
511 tool_calls = (
512 msg.get("tool_calls")
513 if isinstance(msg, dict)
514 else getattr(msg, "tool_calls", None)
515 )
516 tool_call_id = (
517 msg.get("tool_call_id", "")
518 if isinstance(msg, dict)
519 else getattr(msg, "tool_call_id", "")
520 )
522 # Tool results become user turns with a toolResult block
523 if role_str == "tool":
524 result.append(
525 {
526 "role": "user",
527 "content": [
528 {
529 "toolResult": {
530 "toolUseId": tool_call_id,
531 "content": [
532 {"text": _content_to_text(content_raw)}
533 ],
534 }
535 }
536 ],
537 }
538 )
539 continue
541 # Use _build_content_blocks for content serialization
542 bedrock_content = await self._build_content_blocks(content_raw)
544 # Assistant turns that requested tools gain toolUse blocks
545 if role_str == "assistant" and tool_calls:
546 for call in tool_calls:
547 fn = getattr(call, "function", None)
548 if fn is None or not getattr(fn, "name", None):
549 continue
550 bedrock_content.append(
551 {
552 "toolUse": {
553 "toolUseId": getattr(call, "id", ""),
554 "name": fn.name,
555 "input": parse_json_arguments(fn.arguments),
556 }
557 }
558 )
560 # Bedrock uses "user" and "assistant" roles only
561 bedrock_role = "assistant" if role_str == "assistant" else "user"
562 result.append({"role": bedrock_role, "content": bedrock_content})
563 return result
565 def _handle_error_as_result(self, error: Exception) -> Result[Any, LLMError]:
566 """Map a caught exception to ``Err`` or re-raise for infrastructure failures."""
567 err_str = str(error)
568 err_code = getattr(
569 getattr(error, "response", {}).get("Error", {}), "Code", None
570 )
571 if err_code is None and hasattr(error, "response"):
572 resp = error.response
573 if isinstance(resp, dict):
574 err_code = resp.get("Error", {}).get("Code", "")
576 if err_code in (
577 "AccessDeniedException",
578 "AuthorizationException",
579 "UnauthorizedException",
580 ):
581 raise LLMAuthenticationError(f"bedrock: auth failed: {error}") from error
582 if err_code in ("ThrottlingException", "TooManyRequestsException"):
583 return Err(LLMRateLimitError(f"bedrock: rate limit: {error}"))
584 if err_code in (
585 "ModelNotReadyException",
586 "ModelNotFoundException",
587 "ResourceNotFoundException",
588 ):
589 return Err(LLMModelNotFoundError(f"bedrock: model not found: {error}"))
590 if err_code == "ValidationException" and "content filter" in err_str.lower():
591 return Err(LLMContentFilterError(f"bedrock: content filtered: {error}"))
592 raise AIError(f"bedrock: infrastructure error: {error}") from error