code: from typing import Any, Callable, Optional, Sequence
  
  from typing_extensions import override
  
  from llama_index.legacy.bridge.pydantic import Field, PrivateAttr
  from llama_index.legacy.callbacks import CallbackManager
  from llama_index.legacy.constants import DEFAULT_NUM_OUTPUTS
  from llama_index.legacy.core.llms.types import (
      ChatMessage,
      CompletionResponse,
      CompletionResponseGen,
      LLMMetadata,
  )
  from llama_index.legacy.llms.base import llm_completion_callback
  from llama_index.legacy.llms.custom import CustomLLM
  from llama_index.legacy.types import BaseOutputParser, PydanticProgramMode
  
  
  class _BaseGradientLLM(CustomLLM):
      _gradient = PrivateAttr()
      _model = PrivateAttr()
  
      # Config
      max_tokens: Optional[int] = Field(
          default=DEFAULT_NUM_OUTPUTS,
          description="The number of tokens to generate.",
          gt=0,
          lt=512,
      )
  
      # Gradient client config
      access_token: Optional[str] = Field(
          description="The Gradient access token to use.",
      )
      host: Optional[str] = Field(
          description="The url of the Gradient service to access."
      )
      workspace_id: Optional[str] = Field(
          description="The Gradient workspace id to use.",
      )
      is_chat_model: bool = Field(
          default=False, description="Whether the model is a chat model."
      )
  
      def __init__(
          self,
          *,
          access_token: Optional[str] = None,
          host: Optional[str] = None,
          max_tokens: Optional[int] = None,
          workspace_id: Optional[str] = None,
          callback_manager: Optional[CallbackManager] = None,
          is_chat_model: bool = False,
          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:
          super().__init__(
              max_tokens=max_tokens,
              access_token=access_token,
              host=host,
              workspace_id=workspace_id,
              callback_manager=callback_manager,
              is_chat_model=is_chat_model,
              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,
              **kwargs,
          )
          try:
              from gradientai import Gradient
  
              self._gradient = Gradient(
                  access_token=access_token, host=host, workspace_id=workspace_id
              )
          except ImportError as e:
              raise ImportError(
                  "Could not import Gradient Python package. "
                  "Please install it with `pip install gradientai`."
              ) from e
  
      def close(self) -> None:
          self._gradient.close()
  
      @llm_completion_callback()
      @override
      def complete(
          self, prompt: str, formatted: bool = False, **kwargs: Any
      ) -> CompletionResponse:
          return CompletionResponse(
              text=self._model.complete(
                  query=prompt,
                  max_generated_token_count=self.max_tokens,
                  **kwargs,
              ).generated_output
          )
  
      @llm_completion_callback()
      @override
      async def acomplete(
          self, prompt: str, formatted: bool = False, **kwargs: Any
      ) -> CompletionResponse:
          grdt_reponse = await self._model.acomplete(
              query=prompt,
              max_generated_token_count=self.max_tokens,
              **kwargs,
          )
  
          return CompletionResponse(text=grdt_reponse.generated_output)
  
      @override
      def stream_complete(
          self,
          prompt: str,
          formatted: bool = False,
          **kwargs: Any,
      ) -> CompletionResponseGen:
          raise NotImplementedError
  
      @property
      @override
      def metadata(self) -> LLMMetadata:
          return LLMMetadata(
              context_window=1024,
              num_output=self.max_tokens or 20,
              is_chat_model=self.is_chat_model,
              is_function_calling_model=False,
              model_name=self._model.id,
          )
  
  
  class GradientBaseModelLLM(_BaseGradientLLM):
      base_model_slug: str = Field(
          description="The slug of the base model to use.",
      )
  
      def __init__(
          self,
          *,
          access_token: Optional[str] = None,
          base_model_slug: str,
          host: Optional[str] = None,
          max_tokens: Optional[int] = None,
          workspace_id: Optional[str] = None,
          callback_manager: Optional[CallbackManager] = None,
          is_chat_model: bool = False,
      ) -> None:
          super().__init__(
              access_token=access_token,
              base_model_slug=base_model_slug,
              host=host,
              max_tokens=max_tokens,
              workspace_id=workspace_id,
              callback_manager=callback_manager,
              is_chat_model=is_chat_model,
          )
  
          self._model = self._gradient.get_base_model(
              base_model_slug=base_model_slug,
          )
  
  
  class GradientModelAdapterLLM(_BaseGradientLLM):
      model_adapter_id: str = Field(
          description="The id of the model adapter to use.",
      )
  
      def __init__(
          self,
          *,
          access_token: Optional[str] = None,
          host: Optional[str] = None,
          max_tokens: Optional[int] = None,
          model_adapter_id: str,
          workspace_id: Optional[str] = None,
          callback_manager: Optional[CallbackManager] = None,
          is_chat_model: bool = False,
      ) -> None:
          super().__init__(
              access_token=access_token,
              host=host,
              max_tokens=max_tokens,
              model_adapter_id=model_adapter_id,
              workspace_id=workspace_id,
              callback_manager=callback_manager,
              is_chat_model=is_chat_model,
          )
          self._model = self._gradient.get_model_adapter(
              model_adapter_id=model_adapter_id
          )
  
apis: llama_index.legacy.core.llms.types.CompletionResponse llama_index.legacy.core.llms.types.LLMMetadata llama_index.legacy.bridge.pydantic.Field llama_index.legacy.llms.base.llm_completion_callback llama_index.legacy.bridge.pydantic.PrivateAttr
extract_api: [((660, 673), 'llama_index.legacy.bridge.pydantic.PrivateAttr', 'PrivateAttr', ([], {}), '()\n', (671, 673), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((687, 700), 'llama_index.legacy.bridge.pydantic.PrivateAttr', 'PrivateAttr', ([], {}), '()\n', (698, 700), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((747, 849), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'default': 'DEFAULT_NUM_OUTPUTS', 'description': '"""The number of tokens to generate."""', 'gt': '(0)', 'lt': '(512)'}), "(default=DEFAULT_NUM_OUTPUTS, description=\n    'The number of tokens to generate.', gt=0, lt=512)\n", (752, 849), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((948, 1002), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'description': '"""The Gradient access token to use."""'}), "(description='The Gradient access token to use.')\n", (953, 1002), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((1044, 1107), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'description': '"""The url of the Gradient service to access."""'}), "(description='The url of the Gradient service to access.')\n", (1049, 1107), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((1156, 1210), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'description': '"""The Gradient workspace id to use."""'}), "(description='The Gradient workspace id to use.')\n", (1161, 1210), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((1252, 1322), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'default': '(False)', 'description': '"""Whether the model is a chat model."""'}), "(default=False, description='Whether the model is a chat model.')\n", (1257, 1322), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((3019, 3044), 'llama_index.legacy.llms.base.llm_completion_callback', 'llm_completion_callback', ([], {}), '()\n', (3042, 3044), False, 'from llama_index.legacy.llms.base import llm_completion_callback\n'), ((3408, 3433), 'llama_index.legacy.llms.base.llm_completion_callback', 'llm_completion_callback', ([], {}), '()\n', (3431, 3433), False, 'from llama_index.legacy.llms.base import llm_completion_callback\n'), ((4391, 4446), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'description': '"""The slug of the base model to use."""'}), "(description='The slug of the base model to use.')\n", (4396, 4446), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((5307, 5363), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'description': '"""The id of the model adapter to use."""'}), "(description='The id of the model adapter to use.')\n", (5312, 5363), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((3749, 3803), 'llama_index.legacy.core.llms.types.CompletionResponse', 'CompletionResponse', ([], {'text': 'grdt_reponse.generated_output'}), '(text=grdt_reponse.generated_output)\n', (3767, 3803), False, 'from llama_index.legacy.core.llms.types import ChatMessage, CompletionResponse, CompletionResponseGen, LLMMetadata\n'), ((4084, 4252), 'llama_index.legacy.core.llms.types.LLMMetadata', 'LLMMetadata', ([], {'context_window': '(1024)', 'num_output': '(self.max_tokens or 20)', 'is_chat_model': 'self.is_chat_model', 'is_function_calling_model': '(False)', 'model_name': 'self._model.id'}), '(context_window=1024, num_output=self.max_tokens or 20,\n    is_chat_model=self.is_chat_model, is_function_calling_model=False,\n    model_name=self._model.id)\n', (4095, 4252), False, 'from llama_index.legacy.core.llms.types import ChatMessage, CompletionResponse, CompletionResponseGen, LLMMetadata\n'), ((2635, 2708), 'gradientai.Gradient', 'Gradient', ([], {'access_token': 'access_token', 'host': 'host', 'workspace_id': 'workspace_id'}), '(access_token=access_token, host=host, workspace_id=workspace_id)\n', (2643, 2708), False, 'from gradientai import Gradient\n')]