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.constants import DEFAULT_TEMPERATURE
  
  # from mistralai.models.chat_completion import ChatMessage
  from llama_index.legacy.core.llms.types import (
      ChatMessage,
      ChatResponse,
      ChatResponseAsyncGen,
      ChatResponseGen,
      CompletionResponse,
      CompletionResponseAsyncGen,
      CompletionResponseGen,
      LLMMetadata,
      MessageRole,
  )
  from llama_index.legacy.llms.base import (
      llm_chat_callback,
      llm_completion_callback,
  )
  from llama_index.legacy.llms.generic_utils import (
      achat_to_completion_decorator,
      astream_chat_to_completion_decorator,
      chat_to_completion_decorator,
      get_from_param_or_env,
      stream_chat_to_completion_decorator,
  )
  from llama_index.legacy.llms.llm import LLM
  from llama_index.legacy.llms.mistralai_utils import (
      mistralai_modelname_to_contextsize,
  )
  from llama_index.legacy.types import BaseOutputParser, PydanticProgramMode
  
  DEFAULT_MISTRALAI_MODEL = "mistral-tiny"
  DEFAULT_MISTRALAI_ENDPOINT = "https://api.mistral.ai"
  DEFAULT_MISTRALAI_MAX_TOKENS = 512
  
  
  class MistralAI(LLM):
      model: str = Field(
          default=DEFAULT_MISTRALAI_MODEL, description="The mistralai model to use."
      )
      temperature: float = Field(
          default=DEFAULT_TEMPERATURE,
          description="The temperature to use for sampling.",
          gte=0.0,
          lte=1.0,
      )
      max_tokens: int = Field(
          default=DEFAULT_MISTRALAI_MAX_TOKENS,
          description="The maximum number of tokens to generate.",
          gt=0,
      )
  
      timeout: float = Field(
          default=120, description="The timeout to use in seconds.", gte=0
      )
      max_retries: int = Field(
          default=5, description="The maximum number of API retries.", gte=0
      )
      safe_mode: bool = Field(
          default=False,
          description="The parameter to enforce guardrails in chat generations.",
      )
      random_seed: str = Field(
          default=None, description="The random seed to use for sampling."
      )
      additional_kwargs: Dict[str, Any] = Field(
          default_factory=dict, description="Additional kwargs for the MistralAI API."
      )
  
      _client: Any = PrivateAttr()
      _aclient: Any = PrivateAttr()
  
      def __init__(
          self,
          model: str = DEFAULT_MISTRALAI_MODEL,
          temperature: float = DEFAULT_TEMPERATURE,
          max_tokens: int = DEFAULT_MISTRALAI_MAX_TOKENS,
          timeout: int = 120,
          max_retries: int = 5,
          safe_mode: bool = False,
          random_seed: Optional[int] = None,
          api_key: Optional[str] = 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,
      ) -> None:
          try:
              from mistralai.async_client import MistralAsyncClient
              from mistralai.client import MistralClient
          except ImportError as e:
              raise ImportError(
                  "You must install the `mistralai` package to use mistralai."
                  "Please `pip install mistralai`"
              ) from e
  
          additional_kwargs = additional_kwargs or {}
          callback_manager = callback_manager or CallbackManager([])
  
          api_key = get_from_param_or_env("api_key", api_key, "MISTRAL_API_KEY", "")
  
          if not api_key:
              raise ValueError(
                  "You must provide an API key to use mistralai. "
                  "You can either pass it in as an argument or set it `MISTRAL_API_KEY`."
              )
  
          self._client = MistralClient(
              api_key=api_key,
              endpoint=DEFAULT_MISTRALAI_ENDPOINT,
              timeout=timeout,
              max_retries=max_retries,
          )
          self._aclient = MistralAsyncClient(
              api_key=api_key,
              endpoint=DEFAULT_MISTRALAI_ENDPOINT,
              timeout=timeout,
              max_retries=max_retries,
          )
  
          super().__init__(
              temperature=temperature,
              max_tokens=max_tokens,
              additional_kwargs=additional_kwargs,
              timeout=timeout,
              max_retries=max_retries,
              safe_mode=safe_mode,
              random_seed=random_seed,
              model=model,
              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:
          return "MistralAI_LLM"
  
      @property
      def metadata(self) -> LLMMetadata:
          return LLMMetadata(
              context_window=mistralai_modelname_to_contextsize(self.model),
              num_output=self.max_tokens,
              is_chat_model=True,
              model_name=self.model,
              safe_mode=self.safe_mode,
              random_seed=self.random_seed,
          )
  
      @property
      def _model_kwargs(self) -> Dict[str, Any]:
          base_kwargs = {
              "model": self.model,
              "temperature": self.temperature,
              "max_tokens": self.max_tokens,
              "random_seed": self.random_seed,
              "safe_mode": self.safe_mode,
          }
          return {
              **base_kwargs,
              **self.additional_kwargs,
          }
  
      def _get_all_kwargs(self, **kwargs: Any) -> Dict[str, Any]:
          return {
              **self._model_kwargs,
              **kwargs,
          }
  
      @llm_chat_callback()
      def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
          # convert messages to mistral ChatMessage
          from mistralai.client import ChatMessage as mistral_chatmessage
  
          messages = [
              mistral_chatmessage(role=x.role, content=x.content) for x in messages
          ]
          all_kwargs = self._get_all_kwargs(**kwargs)
          response = self._client.chat(messages=messages, **all_kwargs)
          return ChatResponse(
              message=ChatMessage(
                  role=MessageRole.ASSISTANT, content=response.choices[0].message.content
              ),
              raw=dict(response),
          )
  
      @llm_completion_callback()
      def complete(
          self, prompt: str, formatted: bool = False, **kwargs: Any
      ) -> CompletionResponse:
          complete_fn = chat_to_completion_decorator(self.chat)
          return complete_fn(prompt, **kwargs)
  
      @llm_chat_callback()
      def stream_chat(
          self, messages: Sequence[ChatMessage], **kwargs: Any
      ) -> ChatResponseGen:
          # convert messages to mistral ChatMessage
          from mistralai.client import ChatMessage as mistral_chatmessage
  
          messages = [
              mistral_chatmessage(role=message.role, content=message.content)
              for message in messages
          ]
          all_kwargs = self._get_all_kwargs(**kwargs)
  
          response = self._client.chat_stream(messages=messages, **all_kwargs)
  
          def gen() -> ChatResponseGen:
              content = ""
              role = MessageRole.ASSISTANT
              for chunk in response:
                  content_delta = chunk.choices[0].delta.content
                  if content_delta is None:
                      continue
                  content += content_delta
                  yield ChatResponse(
                      message=ChatMessage(role=role, content=content),
                      delta=content_delta,
                      raw=chunk,
                  )
  
          return gen()
  
      @llm_completion_callback()
      def stream_complete(
          self, prompt: str, formatted: bool = False, **kwargs: Any
      ) -> CompletionResponseGen:
          stream_complete_fn = stream_chat_to_completion_decorator(self.stream_chat)
          return stream_complete_fn(prompt, **kwargs)
  
      @llm_chat_callback()
      async def achat(
          self, messages: Sequence[ChatMessage], **kwargs: Any
      ) -> ChatResponse:
          # convert messages to mistral ChatMessage
          from mistralai.client import ChatMessage as mistral_chatmessage
  
          messages = [
              mistral_chatmessage(role=message.role, content=message.content)
              for message in messages
          ]
          all_kwargs = self._get_all_kwargs(**kwargs)
          response = await self._aclient.chat(messages=messages, **all_kwargs)
          return ChatResponse(
              message=ChatMessage(
                  role=MessageRole.ASSISTANT, content=response.choices[0].message.content
              ),
              raw=dict(response),
          )
  
      @llm_completion_callback()
      async def acomplete(
          self, prompt: str, formatted: bool = False, **kwargs: Any
      ) -> CompletionResponse:
          acomplete_fn = achat_to_completion_decorator(self.achat)
          return await acomplete_fn(prompt, **kwargs)
  
      @llm_chat_callback()
      async def astream_chat(
          self, messages: Sequence[ChatMessage], **kwargs: Any
      ) -> ChatResponseAsyncGen:
          # convert messages to mistral ChatMessage
          from mistralai.client import ChatMessage as mistral_chatmessage
  
          messages = [
              mistral_chatmessage(role=x.role, content=x.content) for x in messages
          ]
          all_kwargs = self._get_all_kwargs(**kwargs)
  
          response = await self._aclient.chat_stream(messages=messages, **all_kwargs)
  
          async def gen() -> ChatResponseAsyncGen:
              content = ""
              role = MessageRole.ASSISTANT
              async for chunk in response:
                  content_delta = chunk.choices[0].delta.content
                  if content_delta is None:
                      continue
                  content += content_delta
                  yield ChatResponse(
                      message=ChatMessage(role=role, content=content),
                      delta=content_delta,
                      raw=chunk,
                  )
  
          return gen()
  
      @llm_completion_callback()
      async def astream_complete(
          self, prompt: str, formatted: bool = False, **kwargs: Any
      ) -> CompletionResponseAsyncGen:
          astream_complete_fn = astream_chat_to_completion_decorator(self.astream_chat)
          return await astream_complete_fn(prompt, **kwargs)
  
apis: llama_index.legacy.llms.generic_utils.astream_chat_to_completion_decorator llama_index.legacy.llms.base.llm_chat_callback llama_index.legacy.bridge.pydantic.Field llama_index.legacy.llms.generic_utils.stream_chat_to_completion_decorator llama_index.legacy.llms.base.llm_completion_callback llama_index.legacy.core.llms.types.ChatMessage llama_index.legacy.bridge.pydantic.PrivateAttr llama_index.legacy.llms.generic_utils.get_from_param_or_env llama_index.legacy.callbacks.CallbackManager llama_index.legacy.llms.generic_utils.achat_to_completion_decorator llama_index.legacy.llms.generic_utils.chat_to_completion_decorator llama_index.legacy.llms.mistralai_utils.mistralai_modelname_to_contextsize
extract_api: [((1271, 1357), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'default': 'DEFAULT_MISTRALAI_MODEL', 'description': '"""The mistralai model to use."""'}), "(default=DEFAULT_MISTRALAI_MODEL, description=\n    'The mistralai model to use.')\n", (1276, 1357), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((1392, 1501), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'default': 'DEFAULT_TEMPERATURE', 'description': '"""The temperature to use for sampling."""', 'gte': '(0.0)', 'lte': '(1.0)'}), "(default=DEFAULT_TEMPERATURE, description=\n    'The temperature to use for sampling.', gte=0.0, lte=1.0)\n", (1397, 1501), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((1558, 1669), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'default': 'DEFAULT_MISTRALAI_MAX_TOKENS', 'description': '"""The maximum number of tokens to generate."""', 'gt': '(0)'}), "(default=DEFAULT_MISTRALAI_MAX_TOKENS, description=\n    'The maximum number of tokens to generate.', gt=0)\n", (1563, 1669), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((1718, 1789), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'default': '(120)', 'description': '"""The timeout to use in seconds."""', 'gte': '(0)'}), "(default=120, description='The timeout to use in seconds.', gte=0)\n", (1723, 1789), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((1827, 1900), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'default': '(5)', 'description': '"""The maximum number of API retries."""', 'gte': '(0)'}), "(default=5, description='The maximum number of API retries.', gte=0)\n", (1832, 1900), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((1937, 2034), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'default': '(False)', 'description': '"""The parameter to enforce guardrails in chat generations."""'}), "(default=False, description=\n    'The parameter to enforce guardrails in chat generations.')\n", (1942, 2034), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((2076, 2147), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'default': 'None', 'description': '"""The random seed to use for sampling."""'}), "(default=None, description='The random seed to use for sampling.')\n", (2081, 2147), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((2202, 2290), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'default_factory': 'dict', 'description': '"""Additional kwargs for the MistralAI API."""'}), "(default_factory=dict, description=\n    'Additional kwargs for the MistralAI API.')\n", (2207, 2290), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((2320, 2333), 'llama_index.legacy.bridge.pydantic.PrivateAttr', 'PrivateAttr', ([], {}), '()\n', (2331, 2333), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((2354, 2367), 'llama_index.legacy.bridge.pydantic.PrivateAttr', 'PrivateAttr', ([], {}), '()\n', (2365, 2367), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((5957, 5976), 'llama_index.legacy.llms.base.llm_chat_callback', 'llm_chat_callback', ([], {}), '()\n', (5974, 5976), False, 'from llama_index.legacy.llms.base import llm_chat_callback, llm_completion_callback\n'), ((6632, 6657), 'llama_index.legacy.llms.base.llm_completion_callback', 'llm_completion_callback', ([], {}), '()\n', (6655, 6657), False, 'from llama_index.legacy.llms.base import llm_chat_callback, llm_completion_callback\n'), ((6884, 6903), 'llama_index.legacy.llms.base.llm_chat_callback', 'llm_chat_callback', ([], {}), '()\n', (6901, 6903), False, 'from llama_index.legacy.llms.base import llm_chat_callback, llm_completion_callback\n'), ((7946, 7971), 'llama_index.legacy.llms.base.llm_completion_callback', 'llm_completion_callback', ([], {}), '()\n', (7969, 7971), False, 'from llama_index.legacy.llms.base import llm_chat_callback, llm_completion_callback\n'), ((8236, 8255), 'llama_index.legacy.llms.base.llm_chat_callback', 'llm_chat_callback', ([], {}), '()\n', (8253, 8255), False, 'from llama_index.legacy.llms.base import llm_chat_callback, llm_completion_callback\n'), ((8969, 8994), 'llama_index.legacy.llms.base.llm_completion_callback', 'llm_completion_callback', ([], {}), '()\n', (8992, 8994), False, 'from llama_index.legacy.llms.base import llm_chat_callback, llm_completion_callback\n'), ((9238, 9257), 'llama_index.legacy.llms.base.llm_chat_callback', 'llm_chat_callback', ([], {}), '()\n', (9255, 9257), False, 'from llama_index.legacy.llms.base import llm_chat_callback, llm_completion_callback\n'), ((10306, 10331), 'llama_index.legacy.llms.base.llm_completion_callback', 'llm_completion_callback', ([], {}), '()\n', (10329, 10331), False, 'from llama_index.legacy.llms.base import llm_chat_callback, llm_completion_callback\n'), ((3684, 3748), 'llama_index.legacy.llms.generic_utils.get_from_param_or_env', 'get_from_param_or_env', (['"""api_key"""', 'api_key', '"""MISTRAL_API_KEY"""', '""""""'], {}), "('api_key', api_key, 'MISTRAL_API_KEY', '')\n", (3705, 3748), False, 'from llama_index.legacy.llms.generic_utils import achat_to_completion_decorator, astream_chat_to_completion_decorator, chat_to_completion_decorator, get_from_param_or_env, stream_chat_to_completion_decorator\n'), ((3995, 4109), 'mistralai.client.MistralClient', 'MistralClient', ([], {'api_key': 'api_key', 'endpoint': 'DEFAULT_MISTRALAI_ENDPOINT', 'timeout': 'timeout', 'max_retries': 'max_retries'}), '(api_key=api_key, endpoint=DEFAULT_MISTRALAI_ENDPOINT, timeout\n    =timeout, max_retries=max_retries)\n', (4008, 4109), False, 'from mistralai.client import MistralClient\n'), ((4188, 4306), 'mistralai.async_client.MistralAsyncClient', 'MistralAsyncClient', ([], {'api_key': 'api_key', 'endpoint': 'DEFAULT_MISTRALAI_ENDPOINT', 'timeout': 'timeout', 'max_retries': 'max_retries'}), '(api_key=api_key, endpoint=DEFAULT_MISTRALAI_ENDPOINT,\n    timeout=timeout, max_retries=max_retries)\n', (4206, 4306), False, 'from mistralai.async_client import MistralAsyncClient\n'), ((6793, 6832), 'llama_index.legacy.llms.generic_utils.chat_to_completion_decorator', 'chat_to_completion_decorator', (['self.chat'], {}), '(self.chat)\n', (6821, 6832), False, 'from llama_index.legacy.llms.generic_utils import achat_to_completion_decorator, astream_chat_to_completion_decorator, chat_to_completion_decorator, get_from_param_or_env, stream_chat_to_completion_decorator\n'), ((8124, 8177), 'llama_index.legacy.llms.generic_utils.stream_chat_to_completion_decorator', 'stream_chat_to_completion_decorator', (['self.stream_chat'], {}), '(self.stream_chat)\n', (8159, 8177), False, 'from llama_index.legacy.llms.generic_utils import achat_to_completion_decorator, astream_chat_to_completion_decorator, chat_to_completion_decorator, get_from_param_or_env, stream_chat_to_completion_decorator\n'), ((9138, 9179), 'llama_index.legacy.llms.generic_utils.achat_to_completion_decorator', 'achat_to_completion_decorator', (['self.achat'], {}), '(self.achat)\n', (9167, 9179), False, 'from llama_index.legacy.llms.generic_utils import achat_to_completion_decorator, astream_chat_to_completion_decorator, chat_to_completion_decorator, get_from_param_or_env, stream_chat_to_completion_decorator\n'), ((10497, 10552), 'llama_index.legacy.llms.generic_utils.astream_chat_to_completion_decorator', 'astream_chat_to_completion_decorator', (['self.astream_chat'], {}), '(self.astream_chat)\n', (10533, 10552), False, 'from llama_index.legacy.llms.generic_utils import achat_to_completion_decorator, astream_chat_to_completion_decorator, chat_to_completion_decorator, get_from_param_or_env, stream_chat_to_completion_decorator\n'), ((3645, 3664), 'llama_index.legacy.callbacks.CallbackManager', 'CallbackManager', (['[]'], {}), '([])\n', (3660, 3664), False, 'from llama_index.legacy.callbacks import CallbackManager\n'), ((6217, 6268), 'mistralai.client.ChatMessage', 'mistral_chatmessage', ([], {'role': 'x.role', 'content': 'x.content'}), '(role=x.role, content=x.content)\n', (6236, 6268), True, 'from mistralai.client import ChatMessage as mistral_chatmessage\n'), ((7168, 7231), 'mistralai.client.ChatMessage', 'mistral_chatmessage', ([], {'role': 'message.role', 'content': 'message.content'}), '(role=message.role, content=message.content)\n', (7187, 7231), True, 'from mistralai.client import ChatMessage as mistral_chatmessage\n'), ((8517, 8580), 'mistralai.client.ChatMessage', 'mistral_chatmessage', ([], {'role': 'message.role', 'content': 'message.content'}), '(role=message.role, content=message.content)\n', (8536, 8580), True, 'from mistralai.client import ChatMessage as mistral_chatmessage\n'), ((9534, 9585), 'mistralai.client.ChatMessage', 'mistral_chatmessage', ([], {'role': 'x.role', 'content': 'x.content'}), '(role=x.role, content=x.content)\n', (9553, 9585), True, 'from mistralai.client import ChatMessage as mistral_chatmessage\n'), ((5163, 5209), 'llama_index.legacy.llms.mistralai_utils.mistralai_modelname_to_contextsize', 'mistralai_modelname_to_contextsize', (['self.model'], {}), '(self.model)\n', (5197, 5209), False, 'from llama_index.legacy.llms.mistralai_utils import mistralai_modelname_to_contextsize\n'), ((6468, 6557), 'llama_index.legacy.core.llms.types.ChatMessage', 'ChatMessage', ([], {'role': 'MessageRole.ASSISTANT', 'content': 'response.choices[0].message.content'}), '(role=MessageRole.ASSISTANT, content=response.choices[0].message\n    .content)\n', (6479, 6557), False, 'from llama_index.legacy.core.llms.types import ChatMessage, ChatResponse, ChatResponseAsyncGen, ChatResponseGen, CompletionResponse, CompletionResponseAsyncGen, CompletionResponseGen, LLMMetadata, MessageRole\n'), ((8805, 8894), 'llama_index.legacy.core.llms.types.ChatMessage', 'ChatMessage', ([], {'role': 'MessageRole.ASSISTANT', 'content': 'response.choices[0].message.content'}), '(role=MessageRole.ASSISTANT, content=response.choices[0].message\n    .content)\n', (8816, 8894), False, 'from llama_index.legacy.core.llms.types import ChatMessage, ChatResponse, ChatResponseAsyncGen, ChatResponseGen, CompletionResponse, CompletionResponseAsyncGen, CompletionResponseGen, LLMMetadata, MessageRole\n'), ((7787, 7826), 'llama_index.legacy.core.llms.types.ChatMessage', 'ChatMessage', ([], {'role': 'role', 'content': 'content'}), '(role=role, content=content)\n', (7798, 7826), False, 'from llama_index.legacy.core.llms.types import ChatMessage, ChatResponse, ChatResponseAsyncGen, ChatResponseGen, CompletionResponse, CompletionResponseAsyncGen, CompletionResponseGen, LLMMetadata, MessageRole\n'), ((10147, 10186), 'llama_index.legacy.core.llms.types.ChatMessage', 'ChatMessage', ([], {'role': 'role', 'content': 'content'}), '(role=role, content=content)\n', (10158, 10186), False, 'from llama_index.legacy.core.llms.types import ChatMessage, ChatResponse, ChatResponseAsyncGen, ChatResponseGen, CompletionResponse, CompletionResponseAsyncGen, CompletionResponseGen, LLMMetadata, MessageRole\n')]