code: 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.core.llms.types import (
      ChatMessage,
      ChatResponse,
      ChatResponseGen,
      CompletionResponse,
      CompletionResponseGen,
      LLMMetadata,
  )
  from llama_index.legacy.llms.ai21_utils import ai21_model_to_context_size
  from llama_index.legacy.llms.base import llm_chat_callback, llm_completion_callback
  from llama_index.legacy.llms.custom import CustomLLM
  from llama_index.legacy.llms.generic_utils import (
      completion_to_chat_decorator,
      get_from_param_or_env,
  )
  from llama_index.legacy.types import BaseOutputParser, PydanticProgramMode
  
  
  class AI21(CustomLLM):
      """AI21 Labs LLM."""
  
      model: str = Field(description="The AI21 model to use.")
      maxTokens: int = Field(description="The maximum number of tokens to generate.")
      temperature: float = Field(description="The temperature to use for sampling.")
  
      additional_kwargs: Dict[str, Any] = Field(
          default_factory=dict, description="Additional kwargs for the anthropic API."
      )
  
      _api_key = PrivateAttr()
  
      def __init__(
          self,
          api_key: Optional[str] = None,
          model: Optional[str] = "j2-mid",
          maxTokens: Optional[int] = 512,
          temperature: Optional[float] = 0.1,
          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,
      ) -> None:
          """Initialize params."""
          try:
              import ai21 as _  # noqa
          except ImportError as e:
              raise ImportError(
                  "You must install the `ai21` package to use AI21."
                  "Please `pip install ai21`"
              ) from e
  
          additional_kwargs = additional_kwargs or {}
          callback_manager = callback_manager or CallbackManager([])
  
          api_key = get_from_param_or_env("api_key", api_key, "AI21_API_KEY")
          self._api_key = api_key
  
          super().__init__(
              model=model,
              maxTokens=maxTokens,
              temperature=temperature,
              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(self) -> str:
          """Get Class Name."""
          return "AI21_LLM"
  
      @property
      def metadata(self) -> LLMMetadata:
          return LLMMetadata(
              context_window=ai21_model_to_context_size(self.model),
              num_output=self.maxTokens,
              model_name=self.model,
          )
  
      @property
      def _model_kwargs(self) -> Dict[str, Any]:
          base_kwargs = {
              "model": self.model,
              "maxTokens": self.maxTokens,
              "temperature": self.temperature,
          }
          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:
          all_kwargs = self._get_all_kwargs(**kwargs)
  
          import ai21
  
          ai21.api_key = self._api_key
  
          response = ai21.Completion.execute(**all_kwargs, prompt=prompt)
  
          return CompletionResponse(
              text=response["completions"][0]["data"]["text"], raw=response.__dict__
          )
  
      @llm_completion_callback()
      def stream_complete(
          self, prompt: str, formatted: bool = False, **kwargs: Any
      ) -> CompletionResponseGen:
          raise NotImplementedError(
              "AI21 does not currently support streaming completion."
          )
  
      @llm_chat_callback()
      def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
          all_kwargs = self._get_all_kwargs(**kwargs)
          chat_fn = completion_to_chat_decorator(self.complete)
  
          return chat_fn(messages, **all_kwargs)
  
      @llm_chat_callback()
      def stream_chat(
          self, messages: Sequence[ChatMessage], **kwargs: Any
      ) -> ChatResponseGen:
          raise NotImplementedError("AI21 does not Currently Support Streaming Chat.")
  
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.llms.base.llm_completion_callback llama_index.legacy.bridge.pydantic.PrivateAttr llama_index.legacy.llms.generic_utils.get_from_param_or_env llama_index.legacy.llms.ai21_utils.ai21_model_to_context_size llama_index.legacy.callbacks.CallbackManager llama_index.legacy.llms.generic_utils.completion_to_chat_decorator
extract_api: [((827, 870), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'description': '"""The AI21 model to use."""'}), "(description='The AI21 model to use.')\n", (832, 870), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((892, 954), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'description': '"""The maximum number of tokens to generate."""'}), "(description='The maximum number of tokens to generate.')\n", (897, 954), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((980, 1037), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'description': '"""The temperature to use for sampling."""'}), "(description='The temperature to use for sampling.')\n", (985, 1037), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((1079, 1167), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'default_factory': 'dict', 'description': '"""Additional kwargs for the anthropic API."""'}), "(default_factory=dict, description=\n    'Additional kwargs for the anthropic API.')\n", (1084, 1167), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((1193, 1206), 'llama_index.legacy.bridge.pydantic.PrivateAttr', 'PrivateAttr', ([], {}), '()\n', (1204, 1206), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((3625, 3650), 'llama_index.legacy.llms.base.llm_completion_callback', 'llm_completion_callback', ([], {}), '()\n', (3648, 3650), False, 'from llama_index.legacy.llms.base import llm_chat_callback, llm_completion_callback\n'), ((4083, 4108), 'llama_index.legacy.llms.base.llm_completion_callback', 'llm_completion_callback', ([], {}), '()\n', (4106, 4108), False, 'from llama_index.legacy.llms.base import llm_chat_callback, llm_completion_callback\n'), ((4351, 4370), 'llama_index.legacy.llms.base.llm_chat_callback', 'llm_chat_callback', ([], {}), '()\n', (4368, 4370), False, 'from llama_index.legacy.llms.base import llm_chat_callback, llm_completion_callback\n'), ((4623, 4642), 'llama_index.legacy.llms.base.llm_chat_callback', 'llm_chat_callback', ([], {}), '()\n', (4640, 4642), False, 'from llama_index.legacy.llms.base import llm_chat_callback, llm_completion_callback\n'), ((2296, 2353), 'llama_index.legacy.llms.generic_utils.get_from_param_or_env', 'get_from_param_or_env', (['"""api_key"""', 'api_key', '"""AI21_API_KEY"""'], {}), "('api_key', api_key, 'AI21_API_KEY')\n", (2317, 2353), False, 'from llama_index.legacy.llms.generic_utils import completion_to_chat_decorator, get_from_param_or_env\n'), ((3895, 3947), 'ai21.Completion.execute', 'ai21.Completion.execute', ([], {'prompt': 'prompt'}), '(**all_kwargs, prompt=prompt)\n', (3918, 3947), False, 'import ai21\n'), ((3964, 4059), 'llama_index.legacy.core.llms.types.CompletionResponse', 'CompletionResponse', ([], {'text': "response['completions'][0]['data']['text']", 'raw': 'response.__dict__'}), "(text=response['completions'][0]['data']['text'], raw=\n    response.__dict__)\n", (3982, 4059), False, 'from llama_index.legacy.core.llms.types import ChatMessage, ChatResponse, ChatResponseGen, CompletionResponse, CompletionResponseGen, LLMMetadata\n'), ((4525, 4568), 'llama_index.legacy.llms.generic_utils.completion_to_chat_decorator', 'completion_to_chat_decorator', (['self.complete'], {}), '(self.complete)\n', (4553, 4568), False, 'from llama_index.legacy.llms.generic_utils import completion_to_chat_decorator, get_from_param_or_env\n'), ((2257, 2276), 'llama_index.legacy.callbacks.CallbackManager', 'CallbackManager', (['[]'], {}), '([])\n', (2272, 2276), False, 'from llama_index.legacy.callbacks import CallbackManager\n'), ((3075, 3113), 'llama_index.legacy.llms.ai21_utils.ai21_model_to_context_size', 'ai21_model_to_context_size', (['self.model'], {}), '(self.model)\n', (3101, 3113), False, 'from llama_index.legacy.llms.ai21_utils import ai21_model_to_context_size\n')]