code: import json
  from typing import Any, Callable, Dict, Optional, Sequence
  
  from llama_index.legacy.bridge.pydantic import Field, PrivateAttr
  from llama_index.legacy.callbacks import CallbackManager
  from llama_index.legacy.constants import (
      DEFAULT_TEMPERATURE,
  )
  from llama_index.legacy.core.llms.types import (
      ChatMessage,
      ChatResponse,
      ChatResponseAsyncGen,
      ChatResponseGen,
      CompletionResponse,
      CompletionResponseAsyncGen,
      CompletionResponseGen,
      LLMMetadata,
  )
  from llama_index.legacy.llms.base import (
      llm_chat_callback,
      llm_completion_callback,
  )
  from llama_index.legacy.llms.bedrock_utils import (
      BEDROCK_FOUNDATION_LLMS,
      CHAT_ONLY_MODELS,
      STREAMING_MODELS,
      Provider,
      completion_with_retry,
      get_provider,
  )
  from llama_index.legacy.llms.generic_utils import (
      completion_response_to_chat_response,
      stream_completion_response_to_chat_response,
  )
  from llama_index.legacy.llms.llm import LLM
  from llama_index.legacy.types import BaseOutputParser, PydanticProgramMode
  
  
  class Bedrock(LLM):
      model: str = Field(description="The modelId of the Bedrock model to use.")
      temperature: float = Field(description="The temperature to use for sampling.")
      max_tokens: int = Field(description="The maximum number of tokens to generate.")
      context_size: int = Field("The maximum number of tokens available for input.")
      profile_name: Optional[str] = Field(
          description="The name of aws profile to use. If not given, then the default profile is used."
      )
      aws_access_key_id: Optional[str] = Field(
          description="AWS Access Key ID to use", exclude=True
      )
      aws_secret_access_key: Optional[str] = Field(
          description="AWS Secret Access Key to use", exclude=True
      )
      aws_session_token: Optional[str] = Field(
          description="AWS Session Token to use", exclude=True
      )
      region_name: Optional[str] = Field(
          description="AWS region name to use. Uses region configured in AWS CLI if not passed",
          exclude=True,
      )
      botocore_session: Optional[Any] = Field(
          description="Use this Botocore session instead of creating a new default one.",
          exclude=True,
      )
      botocore_config: Optional[Any] = Field(
          description="Custom configuration object to use instead of the default generated one.",
          exclude=True,
      )
      max_retries: int = Field(
          default=10, description="The maximum number of API retries.", gt=0
      )
      timeout: float = Field(
          default=60.0,
          description="The timeout for the Bedrock API request in seconds. It will be used for both connect and read timeouts.",
      )
      additional_kwargs: Dict[str, Any] = Field(
          default_factory=dict,
          description="Additional kwargs for the bedrock invokeModel request.",
      )
  
      _client: Any = PrivateAttr()
      _aclient: Any = PrivateAttr()
      _provider: Provider = PrivateAttr()
  
      def __init__(
          self,
          model: str,
          temperature: Optional[float] = DEFAULT_TEMPERATURE,
          max_tokens: Optional[int] = 512,
          context_size: Optional[int] = None,
          profile_name: Optional[str] = None,
          aws_access_key_id: Optional[str] = None,
          aws_secret_access_key: Optional[str] = None,
          aws_session_token: Optional[str] = None,
          region_name: Optional[str] = None,
          botocore_session: Optional[Any] = None,
          client: Optional[Any] = None,
          timeout: Optional[float] = 60.0,
          max_retries: Optional[int] = 10,
          botocore_config: Optional[Any] = None,
          additional_kwargs: Optional[Dict[str, Any]] = None,
          callback_manager: Optional[CallbackManager] = None,
          system_prompt: Optional[str] = None,
          messages_to_prompt: Optional[Callable[[Sequence[ChatMessage]], str]] = None,
          completion_to_prompt: Optional[Callable[[str], str]] = None,
          pydantic_program_mode: PydanticProgramMode = PydanticProgramMode.DEFAULT,
          output_parser: Optional[BaseOutputParser] = None,
          **kwargs: Any,
      ) -> None:
          if context_size is None and model not in BEDROCK_FOUNDATION_LLMS:
              raise ValueError(
                  "`context_size` argument not provided and"
                  "model provided refers to a non-foundation model."
                  " Please specify the context_size"
              )
  
          session_kwargs = {
              "profile_name": profile_name,
              "region_name": region_name,
              "aws_access_key_id": aws_access_key_id,
              "aws_secret_access_key": aws_secret_access_key,
              "aws_session_token": aws_session_token,
              "botocore_session": botocore_session,
          }
          config = None
          try:
              import boto3
              from botocore.config import Config
  
              config = (
                  Config(
                      retries={"max_attempts": max_retries, "mode": "standard"},
                      connect_timeout=timeout,
                      read_timeout=timeout,
                  )
                  if botocore_config is None
                  else botocore_config
              )
              session = boto3.Session(**session_kwargs)
          except ImportError:
              raise ImportError(
                  "boto3 package not found, install with" "'pip install boto3'"
              )
  
          # Prior to general availability, custom boto3 wheel files were
          # distributed that used the bedrock service to invokeModel.
          # This check prevents any services still using those wheel files
          # from breaking
          if client is not None:
              self._client = client
          elif "bedrock-runtime" in session.get_available_services():
              self._client = session.client("bedrock-runtime", config=config)
          else:
              self._client = session.client("bedrock", config=config)
  
          additional_kwargs = additional_kwargs or {}
          callback_manager = callback_manager or CallbackManager([])
          context_size = context_size or BEDROCK_FOUNDATION_LLMS[model]
          self._provider = get_provider(model)
          messages_to_prompt = messages_to_prompt or self._provider.messages_to_prompt
          completion_to_prompt = (
              completion_to_prompt or self._provider.completion_to_prompt
          )
          super().__init__(
              model=model,
              temperature=temperature,
              max_tokens=max_tokens,
              context_size=context_size,
              profile_name=profile_name,
              timeout=timeout,
              max_retries=max_retries,
              botocore_config=config,
              additional_kwargs=additional_kwargs,
              callback_manager=callback_manager,
              system_prompt=system_prompt,
              messages_to_prompt=messages_to_prompt,
              completion_to_prompt=completion_to_prompt,
              pydantic_program_mode=pydantic_program_mode,
              output_parser=output_parser,
          )
  
      @classmethod
      def class_name(cls) -> str:
          """Get class name."""
          return "Bedrock_LLM"
  
      @property
      def metadata(self) -> LLMMetadata:
          return LLMMetadata(
              context_window=self.context_size,
              num_output=self.max_tokens,
              is_chat_model=self.model in CHAT_ONLY_MODELS,
              model_name=self.model,
          )
  
      @property
      def _model_kwargs(self) -> Dict[str, Any]:
          base_kwargs = {
              "temperature": self.temperature,
              self._provider.max_tokens_key: self.max_tokens,
          }
          return {
              **base_kwargs,
              **self.additional_kwargs,
          }
  
      def _get_all_kwargs(self, **kwargs: Any) -> Dict[str, Any]:
          return {
              **self._model_kwargs,
              **kwargs,
          }
  
      @llm_completion_callback()
      def complete(
          self, prompt: str, formatted: bool = False, **kwargs: Any
      ) -> CompletionResponse:
          if not formatted:
              prompt = self.completion_to_prompt(prompt)
          all_kwargs = self._get_all_kwargs(**kwargs)
          request_body = self._provider.get_request_body(prompt, all_kwargs)
          request_body_str = json.dumps(request_body)
          response = completion_with_retry(
              client=self._client,
              model=self.model,
              request_body=request_body_str,
              max_retries=self.max_retries,
              **all_kwargs,
          )["body"].read()
          response = json.loads(response)
          return CompletionResponse(
              text=self._provider.get_text_from_response(response), raw=response
          )
  
      @llm_completion_callback()
      def stream_complete(
          self, prompt: str, formatted: bool = False, **kwargs: Any
      ) -> CompletionResponseGen:
          if self.model in BEDROCK_FOUNDATION_LLMS and self.model not in STREAMING_MODELS:
              raise ValueError(f"Model {self.model} does not support streaming")
  
          if not formatted:
              prompt = self.completion_to_prompt(prompt)
  
          all_kwargs = self._get_all_kwargs(**kwargs)
          request_body = self._provider.get_request_body(prompt, all_kwargs)
          request_body_str = json.dumps(request_body)
          response = completion_with_retry(
              client=self._client,
              model=self.model,
              request_body=request_body_str,
              max_retries=self.max_retries,
              stream=True,
              **all_kwargs,
          )["body"]
  
          def gen() -> CompletionResponseGen:
              content = ""
              for r in response:
                  r = json.loads(r["chunk"]["bytes"])
                  content_delta = self._provider.get_text_from_stream_response(r)
                  content += content_delta
                  yield CompletionResponse(text=content, delta=content_delta, raw=r)
  
          return gen()
  
      @llm_chat_callback()
      def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
          prompt = self.messages_to_prompt(messages)
          completion_response = self.complete(prompt, formatted=True, **kwargs)
          return completion_response_to_chat_response(completion_response)
  
      def stream_chat(
          self, messages: Sequence[ChatMessage], **kwargs: Any
      ) -> ChatResponseGen:
          prompt = self.messages_to_prompt(messages)
          completion_response = self.stream_complete(prompt, formatted=True, **kwargs)
          return stream_completion_response_to_chat_response(completion_response)
  
      async def achat(
          self, messages: Sequence[ChatMessage], **kwargs: Any
      ) -> ChatResponse:
          """Chat asynchronously."""
          # TODO: do synchronous chat for now
          return self.chat(messages, **kwargs)
  
      async def acomplete(
          self, prompt: str, formatted: bool = False, **kwargs: Any
      ) -> CompletionResponse:
          raise NotImplementedError
  
      async def astream_chat(
          self, messages: Sequence[ChatMessage], **kwargs: Any
      ) -> ChatResponseAsyncGen:
          raise NotImplementedError
  
      async def astream_complete(
          self, prompt: str, formatted: bool = False, **kwargs: Any
      ) -> CompletionResponseAsyncGen:
          raise NotImplementedError
  
apis: llama_index.legacy.core.llms.types.CompletionResponse llama_index.legacy.llms.base.llm_chat_callback llama_index.legacy.bridge.pydantic.Field llama_index.legacy.core.llms.types.LLMMetadata llama_index.legacy.llms.generic_utils.stream_completion_response_to_chat_response llama_index.legacy.llms.generic_utils.completion_response_to_chat_response llama_index.legacy.llms.base.llm_completion_callback llama_index.legacy.bridge.pydantic.PrivateAttr llama_index.legacy.llms.bedrock_utils.get_provider llama_index.legacy.callbacks.CallbackManager llama_index.legacy.llms.bedrock_utils.completion_with_retry
extract_api: [((1084, 1145), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'description': '"""The modelId of the Bedrock model to use."""'}), "(description='The modelId of the Bedrock model to use.')\n", (1089, 1145), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((1171, 1228), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'description': '"""The temperature to use for sampling."""'}), "(description='The temperature to use for sampling.')\n", (1176, 1228), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((1251, 1313), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'description': '"""The maximum number of tokens to generate."""'}), "(description='The maximum number of tokens to generate.')\n", (1256, 1313), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((1338, 1396), 'llama_index.legacy.bridge.pydantic.Field', 'Field', (['"""The maximum number of tokens available for input."""'], {}), "('The maximum number of tokens available for input.')\n", (1343, 1396), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((1431, 1541), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'description': '"""The name of aws profile to use. If not given, then the default profile is used."""'}), "(description=\n    'The name of aws profile to use. If not given, then the default profile is used.'\n    )\n", (1436, 1541), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((1585, 1644), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'description': '"""AWS Access Key ID to use"""', 'exclude': '(True)'}), "(description='AWS Access Key ID to use', exclude=True)\n", (1590, 1644), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((1702, 1765), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'description': '"""AWS Secret Access Key to use"""', 'exclude': '(True)'}), "(description='AWS Secret Access Key to use', exclude=True)\n", (1707, 1765), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((1819, 1878), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'description': '"""AWS Session Token to use"""', 'exclude': '(True)'}), "(description='AWS Session Token to use', exclude=True)\n", (1824, 1878), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((1926, 2041), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'description': '"""AWS region name to use. Uses region configured in AWS CLI if not passed"""', 'exclude': '(True)'}), "(description=\n    'AWS region name to use. Uses region configured in AWS CLI if not passed',\n    exclude=True)\n", (1931, 2041), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((2094, 2202), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'description': '"""Use this Botocore session instead of creating a new default one."""', 'exclude': '(True)'}), "(description=\n    'Use this Botocore session instead of creating a new default one.',\n    exclude=True)\n", (2099, 2202), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((2254, 2370), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'description': '"""Custom configuration object to use instead of the default generated one."""', 'exclude': '(True)'}), "(description=\n    'Custom configuration object to use instead of the default generated one.',\n    exclude=True)\n", (2259, 2370), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((2408, 2481), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'default': '(10)', 'description': '"""The maximum number of API retries."""', 'gt': '(0)'}), "(default=10, description='The maximum number of API retries.', gt=0)\n", (2413, 2481), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((2517, 2665), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'default': '(60.0)', 'description': '"""The timeout for the Bedrock API request in seconds. It will be used for both connect and read timeouts."""'}), "(default=60.0, description=\n    'The timeout for the Bedrock API request in seconds. It will be used for both connect and read timeouts.'\n    )\n", (2522, 2665), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((2719, 2821), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'default_factory': 'dict', 'description': '"""Additional kwargs for the bedrock invokeModel request."""'}), "(default_factory=dict, description=\n    'Additional kwargs for the bedrock invokeModel request.')\n", (2724, 2821), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((2860, 2873), 'llama_index.legacy.bridge.pydantic.PrivateAttr', 'PrivateAttr', ([], {}), '()\n', (2871, 2873), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((2894, 2907), 'llama_index.legacy.bridge.pydantic.PrivateAttr', 'PrivateAttr', ([], {}), '()\n', (2905, 2907), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((2934, 2947), 'llama_index.legacy.bridge.pydantic.PrivateAttr', 'PrivateAttr', ([], {}), '()\n', (2945, 2947), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((7807, 7832), 'llama_index.legacy.llms.base.llm_completion_callback', 'llm_completion_callback', ([], {}), '()\n', (7830, 7832), False, 'from llama_index.legacy.llms.base import llm_chat_callback, llm_completion_callback\n'), ((8617, 8642), 'llama_index.legacy.llms.base.llm_completion_callback', 'llm_completion_callback', ([], {}), '()\n', (8640, 8642), False, 'from llama_index.legacy.llms.base import llm_chat_callback, llm_completion_callback\n'), ((9840, 9859), 'llama_index.legacy.llms.base.llm_chat_callback', 'llm_chat_callback', ([], {}), '()\n', (9857, 9859), False, 'from llama_index.legacy.llms.base import llm_chat_callback, llm_completion_callback\n'), ((6106, 6125), 'llama_index.legacy.llms.bedrock_utils.get_provider', 'get_provider', (['model'], {}), '(model)\n', (6118, 6125), False, 'from llama_index.legacy.llms.bedrock_utils import BEDROCK_FOUNDATION_LLMS, CHAT_ONLY_MODELS, STREAMING_MODELS, Provider, completion_with_retry, get_provider\n'), ((7158, 7304), 'llama_index.legacy.core.llms.types.LLMMetadata', 'LLMMetadata', ([], {'context_window': 'self.context_size', 'num_output': 'self.max_tokens', 'is_chat_model': '(self.model in CHAT_ONLY_MODELS)', 'model_name': 'self.model'}), '(context_window=self.context_size, num_output=self.max_tokens,\n    is_chat_model=self.model in CHAT_ONLY_MODELS, model_name=self.model)\n', (7169, 7304), False, 'from llama_index.legacy.core.llms.types import ChatMessage, ChatResponse, ChatResponseAsyncGen, ChatResponseGen, CompletionResponse, CompletionResponseAsyncGen, CompletionResponseGen, LLMMetadata\n'), ((8181, 8205), 'json.dumps', 'json.dumps', (['request_body'], {}), '(request_body)\n', (8191, 8205), False, 'import json\n'), ((8466, 8486), 'json.loads', 'json.loads', (['response'], {}), '(response)\n', (8476, 8486), False, 'import json\n'), ((9171, 9195), 'json.dumps', 'json.dumps', (['request_body'], {}), '(request_body)\n', (9181, 9195), False, 'import json\n'), ((10088, 10145), 'llama_index.legacy.llms.generic_utils.completion_response_to_chat_response', 'completion_response_to_chat_response', (['completion_response'], {}), '(completion_response)\n', (10124, 10145), False, 'from llama_index.legacy.llms.generic_utils import completion_response_to_chat_response, stream_completion_response_to_chat_response\n'), ((10406, 10470), 'llama_index.legacy.llms.generic_utils.stream_completion_response_to_chat_response', 'stream_completion_response_to_chat_response', (['completion_response'], {}), '(completion_response)\n', (10449, 10470), False, 'from llama_index.legacy.llms.generic_utils import completion_response_to_chat_response, stream_completion_response_to_chat_response\n'), ((5180, 5211), 'boto3.Session', 'boto3.Session', ([], {}), '(**session_kwargs)\n', (5193, 5211), False, 'import boto3\n'), ((5991, 6010), 'llama_index.legacy.callbacks.CallbackManager', 'CallbackManager', (['[]'], {}), '([])\n', (6006, 6010), False, 'from llama_index.legacy.callbacks import CallbackManager\n'), ((9215, 9368), 'llama_index.legacy.llms.bedrock_utils.completion_with_retry', 'completion_with_retry', ([], {'client': 'self._client', 'model': 'self.model', 'request_body': 'request_body_str', 'max_retries': 'self.max_retries', 'stream': '(True)'}), '(client=self._client, model=self.model, request_body=\n    request_body_str, max_retries=self.max_retries, stream=True, **all_kwargs)\n', (9236, 9368), False, 'from llama_index.legacy.llms.bedrock_utils import BEDROCK_FOUNDATION_LLMS, CHAT_ONLY_MODELS, STREAMING_MODELS, Provider, completion_with_retry, get_provider\n'), ((4872, 4988), 'botocore.config.Config', 'Config', ([], {'retries': "{'max_attempts': max_retries, 'mode': 'standard'}", 'connect_timeout': 'timeout', 'read_timeout': 'timeout'}), "(retries={'max_attempts': max_retries, 'mode': 'standard'},\n    connect_timeout=timeout, read_timeout=timeout)\n", (4878, 4988), False, 'from botocore.config import Config\n'), ((9576, 9607), 'json.loads', 'json.loads', (["r['chunk']['bytes']"], {}), "(r['chunk']['bytes'])\n", (9586, 9607), False, 'import json\n'), ((8225, 8365), 'llama_index.legacy.llms.bedrock_utils.completion_with_retry', 'completion_with_retry', ([], {'client': 'self._client', 'model': 'self.model', 'request_body': 'request_body_str', 'max_retries': 'self.max_retries'}), '(client=self._client, model=self.model, request_body=\n    request_body_str, max_retries=self.max_retries, **all_kwargs)\n', (8246, 8365), False, 'from llama_index.legacy.llms.bedrock_utils import BEDROCK_FOUNDATION_LLMS, CHAT_ONLY_MODELS, STREAMING_MODELS, Provider, completion_with_retry, get_provider\n'), ((9751, 9811), 'llama_index.legacy.core.llms.types.CompletionResponse', 'CompletionResponse', ([], {'text': 'content', 'delta': 'content_delta', 'raw': 'r'}), '(text=content, delta=content_delta, raw=r)\n', (9769, 9811), False, 'from llama_index.legacy.core.llms.types import ChatMessage, ChatResponse, ChatResponseAsyncGen, ChatResponseGen, CompletionResponse, CompletionResponseAsyncGen, CompletionResponseGen, LLMMetadata\n')]