code: """
  Portkey integration with Llama_index for enhanced monitoring.
  """
  
  from typing import TYPE_CHECKING, Any, Callable, List, Optional, Sequence, Union, cast
  
  from llama_index.legacy.bridge.pydantic import Field, PrivateAttr
  from llama_index.legacy.core.llms.types import (
      ChatMessage,
      ChatResponse,
      ChatResponseGen,
      CompletionResponse,
      CompletionResponseGen,
      LLMMetadata,
  )
  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 (
      chat_to_completion_decorator,
      completion_to_chat_decorator,
      stream_chat_to_completion_decorator,
      stream_completion_to_chat_decorator,
  )
  from llama_index.legacy.llms.portkey_utils import (
      IMPORT_ERROR_MESSAGE,
      generate_llm_metadata,
      get_llm,
      is_chat_model,
  )
  from llama_index.legacy.types import BaseOutputParser, PydanticProgramMode
  
  if TYPE_CHECKING:
      from portkey import (
          LLMOptions,
          Modes,
          ModesLiteral,
          PortkeyResponse,
      )
  
  DEFAULT_PORTKEY_MODEL = "gpt-3.5-turbo"
  
  
  class Portkey(CustomLLM):
      """_summary_.
  
      Args:
          LLM (_type_): _description_
      """
  
      mode: Optional[Union["Modes", "ModesLiteral"]] = Field(
          description="The mode for using the Portkey integration"
      )
  
      model: Optional[str] = Field(default=DEFAULT_PORTKEY_MODEL)
      llm: "LLMOptions" = Field(description="LLM parameter", default_factory=dict)
  
      llms: List["LLMOptions"] = Field(description="LLM parameters", default_factory=list)
  
      _client: Any = PrivateAttr()
  
      def __init__(
          self,
          *,
          mode: Union["Modes", "ModesLiteral"],
          api_key: Optional[str] = None,
          base_url: Optional[str] = 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 a Portkey instance.
  
          Args:
              mode (Optional[Modes]): The mode for using the Portkey integration
              (default: Modes.SINGLE).
              api_key (Optional[str]): The API key to authenticate with Portkey.
              base_url (Optional[str]): The Base url to the self hosted rubeus \
                  (the opensource version of portkey) or any other self hosted server.
          """
          try:
              import portkey
          except ImportError as exc:
              raise ImportError(IMPORT_ERROR_MESSAGE) from exc
  
          super().__init__(
              base_url=base_url,
              api_key=api_key,
              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,
          )
          if api_key is not None:
              portkey.api_key = api_key
  
          if base_url is not None:
              portkey.base_url = base_url
  
          portkey.mode = mode
  
          self._client = portkey
          self.model = None
          self.mode = mode
  
      @property
      def metadata(self) -> LLMMetadata:
          """LLM metadata."""
          return generate_llm_metadata(self.llms[0])
  
      def add_llms(
          self, llm_params: Union["LLMOptions", List["LLMOptions"]]
      ) -> "Portkey":
          """
          Adds the specified LLM parameters to the list of LLMs. This may be used for
          fallbacks or load-balancing as specified in the mode.
  
          Args:
              llm_params (Union[LLMOptions, List[LLMOptions]]): A single LLM parameter \
              set or a list of LLM parameter sets. Each set should be an instance of \
              LLMOptions with
              the specified attributes.
                  > provider: Optional[ProviderTypes]
                  > model: str
                  > temperature: float
                  > max_tokens: Optional[int]
                  > max_retries: int
                  > trace_id: Optional[str]
                  > cache_status: Optional[CacheType]
                  > cache: Optional[bool]
                  > metadata: Dict[str, Any]
                  > weight: Optional[float]
                  > **kwargs : Other additional parameters that are supported by \
                      LLMOptions in portkey-ai
  
              NOTE: User may choose to pass additional params as well.
  
          Returns:
              self
          """
          try:
              from portkey import LLMOptions
          except ImportError as exc:
              raise ImportError(IMPORT_ERROR_MESSAGE) from exc
          if isinstance(llm_params, LLMOptions):
              llm_params = [llm_params]
          self.llms.extend(llm_params)
          if self.model is None:
              self.model = self.llms[0].model
          return self
  
      @llm_completion_callback()
      def complete(
          self, prompt: str, formatted: bool = False, **kwargs: Any
      ) -> CompletionResponse:
          """Completion endpoint for LLM."""
          if self._is_chat_model:
              complete_fn = chat_to_completion_decorator(self._chat)
          else:
              complete_fn = self._complete
          return complete_fn(prompt, **kwargs)
  
      @llm_chat_callback()
      def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
          if self._is_chat_model:
              chat_fn = self._chat
          else:
              chat_fn = completion_to_chat_decorator(self._complete)
          return chat_fn(messages, **kwargs)
  
      @llm_completion_callback()
      def stream_complete(
          self, prompt: str, formatted: bool = False, **kwargs: Any
      ) -> CompletionResponseGen:
          """Completion endpoint for LLM."""
          if self._is_chat_model:
              complete_fn = stream_chat_to_completion_decorator(self._stream_chat)
          else:
              complete_fn = self._stream_complete
          return complete_fn(prompt, **kwargs)
  
      @llm_chat_callback()
      def stream_chat(
          self, messages: Sequence[ChatMessage], **kwargs: Any
      ) -> ChatResponseGen:
          if self._is_chat_model:
              stream_chat_fn = self._stream_chat
          else:
              stream_chat_fn = stream_completion_to_chat_decorator(self._stream_complete)
          return stream_chat_fn(messages, **kwargs)
  
      def _chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
          try:
              from portkey import Config, Message
          except ImportError as exc:
              raise ImportError(IMPORT_ERROR_MESSAGE) from exc
          _messages = cast(
              List[Message],
              [{"role": i.role.value, "content": i.content} for i in messages],
          )
          config = Config(llms=self.llms)
          response = self._client.ChatCompletions.create(
              messages=_messages, config=config
          )
          self.llm = self._get_llm(response)
  
          message = response.choices[0].message
          return ChatResponse(message=message, raw=response)
  
      def _complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
          try:
              from portkey import Config
          except ImportError as exc:
              raise ImportError(IMPORT_ERROR_MESSAGE) from exc
  
          config = Config(llms=self.llms)
          response = self._client.Completions.create(prompt=prompt, config=config)
          text = response.choices[0].text
          return CompletionResponse(text=text, raw=response)
  
      def _stream_chat(
          self, messages: Sequence[ChatMessage], **kwargs: Any
      ) -> ChatResponseGen:
          try:
              from portkey import Config, Message
          except ImportError as exc:
              raise ImportError(IMPORT_ERROR_MESSAGE) from exc
          _messages = cast(
              List[Message],
              [{"role": i.role.value, "content": i.content} for i in messages],
          )
          config = Config(llms=self.llms)
          response = self._client.ChatCompletions.create(
              messages=_messages, config=config, stream=True, **kwargs
          )
  
          def gen() -> ChatResponseGen:
              content = ""
              function_call: Optional[dict] = {}
              for resp in response:
                  if resp.choices is None:
                      continue
                  delta = resp.choices[0].delta
                  role = delta.get("role", "assistant")
                  content_delta = delta.get("content", "") or ""
                  content += content_delta
  
                  function_call_delta = delta.get("function_call", None)
                  if function_call_delta is not None:
                      if function_call is None:
                          function_call = function_call_delta
                          # ensure we do not add a blank function call
                          if (
                              function_call
                              and function_call.get("function_name", "") is None
                          ):
                              del function_call["function_name"]
                      else:
                          function_call["arguments"] += function_call_delta["arguments"]
  
                  additional_kwargs = {}
                  if function_call is not None:
                      additional_kwargs["function_call"] = function_call
  
                  yield ChatResponse(
                      message=ChatMessage(
                          role=role,
                          content=content,
                          additional_kwargs=additional_kwargs,
                      ),
                      delta=content_delta,
                      raw=resp,
                  )
  
          return gen()
  
      def _stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
          try:
              from portkey import Config
          except ImportError as exc:
              raise ImportError(IMPORT_ERROR_MESSAGE) from exc
  
          config = Config(llms=self.llms)
          response = self._client.Completions.create(
              prompt=prompt, config=config, stream=True, **kwargs
          )
  
          def gen() -> CompletionResponseGen:
              text = ""
              for resp in response:
                  delta = resp.choices[0].text or ""
                  text += delta
                  yield CompletionResponse(
                      delta=delta,
                      text=text,
                      raw=resp,
                  )
  
          return gen()
  
      @property
      def _is_chat_model(self) -> bool:
          """Check if a given model is a chat-based language model.
  
          Returns:
              bool: True if the provided model is a chat-based language model,
              False otherwise.
          """
          return is_chat_model(self.model or "")
  
      def _get_llm(self, response: "PortkeyResponse") -> "LLMOptions":
          return get_llm(response, self.llms)
  
apis: llama_index.legacy.core.llms.types.CompletionResponse llama_index.legacy.llms.generic_utils.stream_completion_to_chat_decorator llama_index.legacy.llms.portkey_utils.is_chat_model 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.portkey_utils.get_llm llama_index.legacy.llms.portkey_utils.generate_llm_metadata llama_index.legacy.llms.generic_utils.completion_to_chat_decorator llama_index.legacy.llms.generic_utils.chat_to_completion_decorator llama_index.legacy.core.llms.types.ChatResponse
extract_api: [((1284, 1347), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'description': '"""The mode for using the Portkey integration"""'}), "(description='The mode for using the Portkey integration')\n", (1289, 1347), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((1390, 1426), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'default': 'DEFAULT_PORTKEY_MODEL'}), '(default=DEFAULT_PORTKEY_MODEL)\n', (1395, 1426), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((1451, 1507), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'description': '"""LLM parameter"""', 'default_factory': 'dict'}), "(description='LLM parameter', default_factory=dict)\n", (1456, 1507), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((1540, 1597), 'llama_index.legacy.bridge.pydantic.Field', 'Field', ([], {'description': '"""LLM parameters"""', 'default_factory': 'list'}), "(description='LLM parameters', default_factory=list)\n", (1545, 1597), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((1618, 1631), 'llama_index.legacy.bridge.pydantic.PrivateAttr', 'PrivateAttr', ([], {}), '()\n', (1629, 1631), False, 'from llama_index.legacy.bridge.pydantic import Field, PrivateAttr\n'), ((5009, 5034), 'llama_index.legacy.llms.base.llm_completion_callback', 'llm_completion_callback', ([], {}), '()\n', (5032, 5034), False, 'from llama_index.legacy.llms.base import llm_chat_callback, llm_completion_callback\n'), ((5396, 5415), 'llama_index.legacy.llms.base.llm_chat_callback', 'llm_chat_callback', ([], {}), '()\n', (5413, 5415), False, 'from llama_index.legacy.llms.base import llm_chat_callback, llm_completion_callback\n'), ((5695, 5720), 'llama_index.legacy.llms.base.llm_completion_callback', 'llm_completion_callback', ([], {}), '()\n', (5718, 5720), False, 'from llama_index.legacy.llms.base import llm_chat_callback, llm_completion_callback\n'), ((6113, 6132), 'llama_index.legacy.llms.base.llm_chat_callback', 'llm_chat_callback', ([], {}), '()\n', (6130, 6132), False, 'from llama_index.legacy.llms.base import llm_chat_callback, llm_completion_callback\n'), ((3423, 3458), 'llama_index.legacy.llms.portkey_utils.generate_llm_metadata', 'generate_llm_metadata', (['self.llms[0]'], {}), '(self.llms[0])\n', (3444, 3458), False, 'from llama_index.legacy.llms.portkey_utils import IMPORT_ERROR_MESSAGE, generate_llm_metadata, get_llm, is_chat_model\n'), ((6735, 6824), 'typing.cast', 'cast', (['List[Message]', "[{'role': i.role.value, 'content': i.content} for i in messages]"], {}), "(List[Message], [{'role': i.role.value, 'content': i.content} for i in\n    messages])\n", (6739, 6824), False, 'from typing import TYPE_CHECKING, Any, Callable, List, Optional, Sequence, Union, cast\n'), ((6873, 6895), 'portkey.Config', 'Config', ([], {'llms': 'self.llms'}), '(llms=self.llms)\n', (6879, 6895), False, 'from portkey import Config\n'), ((7113, 7156), 'llama_index.legacy.core.llms.types.ChatResponse', 'ChatResponse', ([], {'message': 'message', 'raw': 'response'}), '(message=message, raw=response)\n', (7125, 7156), False, 'from llama_index.legacy.core.llms.types import ChatMessage, ChatResponse, ChatResponseGen, CompletionResponse, CompletionResponseGen, LLMMetadata\n'), ((7399, 7421), 'portkey.Config', 'Config', ([], {'llms': 'self.llms'}), '(llms=self.llms)\n', (7405, 7421), False, 'from portkey import Config\n'), ((7558, 7601), 'llama_index.legacy.core.llms.types.CompletionResponse', 'CompletionResponse', ([], {'text': 'text', 'raw': 'response'}), '(text=text, raw=response)\n', (7576, 7601), False, 'from llama_index.legacy.core.llms.types import ChatMessage, ChatResponse, ChatResponseGen, CompletionResponse, CompletionResponseGen, LLMMetadata\n'), ((7889, 7978), 'typing.cast', 'cast', (['List[Message]', "[{'role': i.role.value, 'content': i.content} for i in messages]"], {}), "(List[Message], [{'role': i.role.value, 'content': i.content} for i in\n    messages])\n", (7893, 7978), False, 'from typing import TYPE_CHECKING, Any, Callable, List, Optional, Sequence, Union, cast\n'), ((8027, 8049), 'portkey.Config', 'Config', ([], {'llms': 'self.llms'}), '(llms=self.llms)\n', (8033, 8049), False, 'from portkey import Config\n'), ((10014, 10036), 'portkey.Config', 'Config', ([], {'llms': 'self.llms'}), '(llms=self.llms)\n', (10020, 10036), False, 'from portkey import Config\n'), ((10791, 10822), 'llama_index.legacy.llms.portkey_utils.is_chat_model', 'is_chat_model', (["(self.model or '')"], {}), "(self.model or '')\n", (10804, 10822), False, 'from llama_index.legacy.llms.portkey_utils import IMPORT_ERROR_MESSAGE, generate_llm_metadata, get_llm, is_chat_model\n'), ((10908, 10936), 'llama_index.legacy.llms.portkey_utils.get_llm', 'get_llm', (['response', 'self.llms'], {}), '(response, self.llms)\n', (10915, 10936), False, 'from llama_index.legacy.llms.portkey_utils import IMPORT_ERROR_MESSAGE, generate_llm_metadata, get_llm, is_chat_model\n'), ((5249, 5289), 'llama_index.legacy.llms.generic_utils.chat_to_completion_decorator', 'chat_to_completion_decorator', (['self._chat'], {}), '(self._chat)\n', (5277, 5289), False, 'from llama_index.legacy.llms.generic_utils import chat_to_completion_decorator, completion_to_chat_decorator, stream_chat_to_completion_decorator, stream_completion_to_chat_decorator\n'), ((5601, 5645), 'llama_index.legacy.llms.generic_utils.completion_to_chat_decorator', 'completion_to_chat_decorator', (['self._complete'], {}), '(self._complete)\n', (5629, 5645), False, 'from llama_index.legacy.llms.generic_utils import chat_to_completion_decorator, completion_to_chat_decorator, stream_chat_to_completion_decorator, stream_completion_to_chat_decorator\n'), ((5945, 5999), 'llama_index.legacy.llms.generic_utils.stream_chat_to_completion_decorator', 'stream_chat_to_completion_decorator', (['self._stream_chat'], {}), '(self._stream_chat)\n', (5980, 5999), False, 'from llama_index.legacy.llms.generic_utils import chat_to_completion_decorator, completion_to_chat_decorator, stream_chat_to_completion_decorator, stream_completion_to_chat_decorator\n'), ((6363, 6421), 'llama_index.legacy.llms.generic_utils.stream_completion_to_chat_decorator', 'stream_completion_to_chat_decorator', (['self._stream_complete'], {}), '(self._stream_complete)\n', (6398, 6421), False, 'from llama_index.legacy.llms.generic_utils import chat_to_completion_decorator, completion_to_chat_decorator, stream_chat_to_completion_decorator, stream_completion_to_chat_decorator\n'), ((10367, 10419), 'llama_index.legacy.core.llms.types.CompletionResponse', 'CompletionResponse', ([], {'delta': 'delta', 'text': 'text', 'raw': 'resp'}), '(delta=delta, text=text, raw=resp)\n', (10385, 10419), False, 'from llama_index.legacy.core.llms.types import ChatMessage, ChatResponse, ChatResponseGen, CompletionResponse, CompletionResponseGen, LLMMetadata\n'), ((9478, 9554), 'llama_index.legacy.core.llms.types.ChatMessage', 'ChatMessage', ([], {'role': 'role', 'content': 'content', 'additional_kwargs': 'additional_kwargs'}), '(role=role, content=content, additional_kwargs=additional_kwargs)\n', (9489, 9554), False, 'from llama_index.legacy.core.llms.types import ChatMessage, ChatResponse, ChatResponseGen, CompletionResponse, CompletionResponseGen, LLMMetadata\n')]