code: from portkey_ai import Message, Portkey
  from typing import Optional, Union, List, Any, Mapping, cast, Sequence
  from portkey_ai.api_resources.utils import PortkeyResponse
  
  from portkey_ai.llms.llama_index.utils import (
      IMPORT_ERROR_MESSAGE,
      is_chat_model,
      modelname_to_contextsize,
  )
  
  try:
      from llama_index.llms.custom import CustomLLM
      from llama_index.bridge.pydantic import PrivateAttr
      from llama_index.llms.base import (
          ChatMessage,
          ChatResponse,
          ChatResponseGen,
          CompletionResponse,
          CompletionResponseGen,
          LLMMetadata,
          llm_chat_callback,
          llm_completion_callback,
      )
  except ImportError as exc:
      raise ImportError(IMPORT_ERROR_MESSAGE) from exc
  
  
  class PortkeyLLM(CustomLLM):
      """_summary_.
  
      Args:
          LLM (_type_): _description_
      """
  
      _client: Any = PrivateAttr()
      model: str = ""
  
      api_key: Optional[str] = None
      base_url: Optional[str] = None
      virtual_key: Optional[str] = None
      config: Optional[Union[Mapping, str]] = None
      provider: Optional[str] = None
      trace_id: Optional[str] = None
      custom_metadata: Optional[str] = None
  
      def __init__(
          self,
          *,
          api_key: Optional[str] = None,
          base_url: Optional[str] = None,
          virtual_key: Optional[str] = None,
          config: Optional[Union[Mapping, str]] = None,
          provider: Optional[str] = None,
          trace_id: Optional[str] = None,
          custom_metadata: Optional[str] = None,
          **kwargs,
      ) -> None:
          """
          Initialize a Portkey instance.
  
          Args:
              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.
          """
          super().__init__(
              base_url=base_url,
              api_key=api_key,
          )
          self._client = Portkey(
              api_key=api_key,
              base_url=base_url,
              virtual_key=virtual_key,
              config=config,
              provider=provider,
              trace_id=trace_id,
              metadata=custom_metadata,
              **kwargs,
          )
          self.model = ""
  
      @property
      def metadata(self) -> LLMMetadata:
          """LLM metadata."""
          try:
              from llama_index.llms.base import (
                  LLMMetadata,
              )
          except ImportError as exc:
              raise ImportError(IMPORT_ERROR_MESSAGE) from exc
          return LLMMetadata(
              _context_window=modelname_to_contextsize(self.model) if self.model else 0,
              is_chat_model=is_chat_model(self.model),
              model_name=self.model,
          )
  
      @llm_completion_callback()
      def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
          """Completion endpoint for LLM."""
          complete_fn = self._complete
          return complete_fn(prompt, **kwargs)
  
      @llm_chat_callback()
      def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
          chat_fn = self._chat
          return chat_fn(messages, **kwargs)
  
      @llm_completion_callback()
      def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
          """Completion endpoint for LLM."""
          complete_fn = self._stream_complete
          return complete_fn(prompt, **kwargs)
  
      @llm_chat_callback()
      def stream_chat(
          self, messages: Sequence[ChatMessage], **kwargs: Any
      ) -> ChatResponseGen:
          stream_chat_fn = self._stream_chat
          return stream_chat_fn(messages, **kwargs)
  
      def _chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
          _messages = cast(
              List[Message],
              [{"role": i.role.value, "content": i.content} for i in messages],
          )
          response = self._client.chat.completions.create(messages=_messages, **kwargs)
          self.model = self._get_model(response)
  
          message = response.choices[0].message
          return ChatResponse(message=message, raw=response)
  
      def _complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
          response = self._client.completions.create(prompt=prompt, **kwargs)
          text = response.choices[0].text
          return CompletionResponse(text=text, raw=response)
  
      def _stream_chat(
          self, messages: Sequence[ChatMessage], **kwargs: Any
      ) -> ChatResponseGen:
          _messages = cast(
              List[Message],
              [{"role": i.role.value, "content": i.content} for i in messages],
          )
          response = self._client.chat.completions.create(
              messages=_messages, 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:
          response = self._client.completions.create(prompt=prompt, 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_model(self, response: PortkeyResponse) -> str:
          return response.model
  
apis: llama_index.llms.base.ChatResponse llama_index.llms.base.ChatMessage llama_index.llms.base.llm_completion_callback llama_index.llms.base.CompletionResponse llama_index.bridge.pydantic.PrivateAttr llama_index.llms.base.llm_chat_callback
extract_api: [((866, 879), 'llama_index.bridge.pydantic.PrivateAttr', 'PrivateAttr', ([], {}), '()\n', (877, 879), False, 'from llama_index.bridge.pydantic import PrivateAttr\n'), ((2779, 2804), 'llama_index.llms.base.llm_completion_callback', 'llm_completion_callback', ([], {}), '()\n', (2802, 2804), False, 'from llama_index.llms.base import ChatMessage, ChatResponse, ChatResponseGen, CompletionResponse, CompletionResponseGen, LLMMetadata, llm_chat_callback, llm_completion_callback\n'), ((3010, 3029), 'llama_index.llms.base.llm_chat_callback', 'llm_chat_callback', ([], {}), '()\n', (3027, 3029), False, 'from llama_index.llms.base import ChatMessage, ChatResponse, ChatResponseGen, CompletionResponse, CompletionResponseGen, LLMMetadata, llm_chat_callback, llm_completion_callback\n'), ((3192, 3217), 'llama_index.llms.base.llm_completion_callback', 'llm_completion_callback', ([], {}), '()\n', (3215, 3217), False, 'from llama_index.llms.base import ChatMessage, ChatResponse, ChatResponseGen, CompletionResponse, CompletionResponseGen, LLMMetadata, llm_chat_callback, llm_completion_callback\n'), ((3440, 3459), 'llama_index.llms.base.llm_chat_callback', 'llm_chat_callback', ([], {}), '()\n', (3457, 3459), False, 'from llama_index.llms.base import ChatMessage, ChatResponse, ChatResponseGen, CompletionResponse, CompletionResponseGen, LLMMetadata, llm_chat_callback, llm_completion_callback\n'), ((1989, 2155), 'portkey_ai.Portkey', 'Portkey', ([], {'api_key': 'api_key', 'base_url': 'base_url', 'virtual_key': 'virtual_key', 'config': 'config', 'provider': 'provider', 'trace_id': 'trace_id', 'metadata': 'custom_metadata'}), '(api_key=api_key, base_url=base_url, virtual_key=virtual_key, config\n    =config, provider=provider, trace_id=trace_id, metadata=custom_metadata,\n    **kwargs)\n', (1996, 2155), False, 'from portkey_ai import Message, Portkey\n'), ((3767, 3856), '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", (3771, 3856), False, 'from typing import Optional, Union, List, Any, Mapping, cast, Sequence\n'), ((4083, 4126), 'llama_index.llms.base.ChatResponse', 'ChatResponse', ([], {'message': 'message', 'raw': 'response'}), '(message=message, raw=response)\n', (4095, 4126), False, 'from llama_index.llms.base import ChatMessage, ChatResponse, ChatResponseGen, CompletionResponse, CompletionResponseGen, LLMMetadata, llm_chat_callback, llm_completion_callback\n'), ((4334, 4377), 'llama_index.llms.base.CompletionResponse', 'CompletionResponse', ([], {'text': 'text', 'raw': 'response'}), '(text=text, raw=response)\n', (4352, 4377), False, 'from llama_index.llms.base import ChatMessage, ChatResponse, ChatResponseGen, CompletionResponse, CompletionResponseGen, LLMMetadata, llm_chat_callback, llm_completion_callback\n'), ((4508, 4597), '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", (4512, 4597), False, 'from typing import Optional, Union, List, Any, Mapping, cast, Sequence\n'), ((7130, 7161), 'portkey_ai.llms.llama_index.utils.is_chat_model', 'is_chat_model', (["(self.model or '')"], {}), "(self.model or '')\n", (7143, 7161), False, 'from portkey_ai.llms.llama_index.utils import IMPORT_ERROR_MESSAGE, is_chat_model, modelname_to_contextsize\n'), ((2701, 2726), 'portkey_ai.llms.llama_index.utils.is_chat_model', 'is_chat_model', (['self.model'], {}), '(self.model)\n', (2714, 2726), False, 'from portkey_ai.llms.llama_index.utils import IMPORT_ERROR_MESSAGE, is_chat_model, modelname_to_contextsize\n'), ((2616, 2652), 'portkey_ai.llms.llama_index.utils.modelname_to_contextsize', 'modelname_to_contextsize', (['self.model'], {}), '(self.model)\n', (2640, 2652), False, 'from portkey_ai.llms.llama_index.utils import IMPORT_ERROR_MESSAGE, is_chat_model, modelname_to_contextsize\n'), ((6706, 6758), 'llama_index.llms.base.CompletionResponse', 'CompletionResponse', ([], {'delta': 'delta', 'text': 'text', 'raw': 'resp'}), '(delta=delta, text=text, raw=resp)\n', (6724, 6758), False, 'from llama_index.llms.base import ChatMessage, ChatResponse, ChatResponseGen, CompletionResponse, CompletionResponseGen, LLMMetadata, llm_chat_callback, llm_completion_callback\n'), ((6043, 6119), 'llama_index.llms.base.ChatMessage', 'ChatMessage', ([], {'role': 'role', 'content': 'content', 'additional_kwargs': 'additional_kwargs'}), '(role=role, content=content, additional_kwargs=additional_kwargs)\n', (6054, 6119), False, 'from llama_index.llms.base import ChatMessage, ChatResponse, ChatResponseGen, CompletionResponse, CompletionResponseGen, LLMMetadata, llm_chat_callback, llm_completion_callback\n')]