llama-index

1
"""
2
Portkey integration with Llama_index for enhanced monitoring.
3
"""
4

5
from typing import TYPE_CHECKING, Any, Callable, List, Optional, Sequence, Union, cast
6

7
from llama_index.legacy.bridge.pydantic import Field, PrivateAttr
8
from llama_index.legacy.core.llms.types import (
9
    ChatMessage,
10
    ChatResponse,
11
    ChatResponseGen,
12
    CompletionResponse,
13
    CompletionResponseGen,
14
    LLMMetadata,
15
)
16
from llama_index.legacy.llms.base import llm_chat_callback, llm_completion_callback
17
from llama_index.legacy.llms.custom import CustomLLM
18
from llama_index.legacy.llms.generic_utils import (
19
    chat_to_completion_decorator,
20
    completion_to_chat_decorator,
21
    stream_chat_to_completion_decorator,
22
    stream_completion_to_chat_decorator,
23
)
24
from llama_index.legacy.llms.portkey_utils import (
25
    IMPORT_ERROR_MESSAGE,
26
    generate_llm_metadata,
27
    get_llm,
28
    is_chat_model,
29
)
30
from llama_index.legacy.types import BaseOutputParser, PydanticProgramMode
31

32
if TYPE_CHECKING:
33
    from portkey import (
34
        LLMOptions,
35
        Modes,
36
        ModesLiteral,
37
        PortkeyResponse,
38
    )
39

40
DEFAULT_PORTKEY_MODEL = "gpt-3.5-turbo"
41

42

43
class Portkey(CustomLLM):
44
    """_summary_.
45

46
    Args:
47
        LLM (_type_): _description_
48
    """
49

50
    mode: Optional[Union["Modes", "ModesLiteral"]] = Field(
51
        description="The mode for using the Portkey integration"
52
    )
53

54
    model: Optional[str] = Field(default=DEFAULT_PORTKEY_MODEL)
55
    llm: "LLMOptions" = Field(description="LLM parameter", default_factory=dict)
56

57
    llms: List["LLMOptions"] = Field(description="LLM parameters", default_factory=list)
58

59
    _client: Any = PrivateAttr()
60

61
    def __init__(
62
        self,
63
        *,
64
        mode: Union["Modes", "ModesLiteral"],
65
        api_key: Optional[str] = None,
66
        base_url: Optional[str] = None,
67
        system_prompt: Optional[str] = None,
68
        messages_to_prompt: Optional[Callable[[Sequence[ChatMessage]], str]] = None,
69
        completion_to_prompt: Optional[Callable[[str], str]] = None,
70
        pydantic_program_mode: PydanticProgramMode = PydanticProgramMode.DEFAULT,
71
        output_parser: Optional[BaseOutputParser] = None,
72
    ) -> None:
73
        """
74
        Initialize a Portkey instance.
75

76
        Args:
77
            mode (Optional[Modes]): The mode for using the Portkey integration
78
            (default: Modes.SINGLE).
79
            api_key (Optional[str]): The API key to authenticate with Portkey.
80
            base_url (Optional[str]): The Base url to the self hosted rubeus \
81
                (the opensource version of portkey) or any other self hosted server.
82
        """
83
        try:
84
            import portkey
85
        except ImportError as exc:
86
            raise ImportError(IMPORT_ERROR_MESSAGE) from exc
87

88
        super().__init__(
89
            base_url=base_url,
90
            api_key=api_key,
91
            system_prompt=system_prompt,
92
            messages_to_prompt=messages_to_prompt,
93
            completion_to_prompt=completion_to_prompt,
94
            pydantic_program_mode=pydantic_program_mode,
95
            output_parser=output_parser,
96
        )
97
        if api_key is not None:
98
            portkey.api_key = api_key
99

100
        if base_url is not None:
101
            portkey.base_url = base_url
102

103
        portkey.mode = mode
104

105
        self._client = portkey
106
        self.model = None
107
        self.mode = mode
108

109
    @property
110
    def metadata(self) -> LLMMetadata:
111
        """LLM metadata."""
112
        return generate_llm_metadata(self.llms[0])
113

114
    def add_llms(
115
        self, llm_params: Union["LLMOptions", List["LLMOptions"]]
116
    ) -> "Portkey":
117
        """
118
        Adds the specified LLM parameters to the list of LLMs. This may be used for
119
        fallbacks or load-balancing as specified in the mode.
120

121
        Args:
122
            llm_params (Union[LLMOptions, List[LLMOptions]]): A single LLM parameter \
123
            set or a list of LLM parameter sets. Each set should be an instance of \
124
            LLMOptions with
125
            the specified attributes.
126
                > provider: Optional[ProviderTypes]
127
                > model: str
128
                > temperature: float
129
                > max_tokens: Optional[int]
130
                > max_retries: int
131
                > trace_id: Optional[str]
132
                > cache_status: Optional[CacheType]
133
                > cache: Optional[bool]
134
                > metadata: Dict[str, Any]
135
                > weight: Optional[float]
136
                > **kwargs : Other additional parameters that are supported by \
137
                    LLMOptions in portkey-ai
138

139
            NOTE: User may choose to pass additional params as well.
140

141
        Returns:
142
            self
143
        """
144
        try:
145
            from portkey import LLMOptions
146
        except ImportError as exc:
147
            raise ImportError(IMPORT_ERROR_MESSAGE) from exc
148
        if isinstance(llm_params, LLMOptions):
149
            llm_params = [llm_params]
150
        self.llms.extend(llm_params)
151
        if self.model is None:
152
            self.model = self.llms[0].model
153
        return self
154

155
    @llm_completion_callback()
156
    def complete(
157
        self, prompt: str, formatted: bool = False, **kwargs: Any
158
    ) -> CompletionResponse:
159
        """Completion endpoint for LLM."""
160
        if self._is_chat_model:
161
            complete_fn = chat_to_completion_decorator(self._chat)
162
        else:
163
            complete_fn = self._complete
164
        return complete_fn(prompt, **kwargs)
165

166
    @llm_chat_callback()
167
    def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
168
        if self._is_chat_model:
169
            chat_fn = self._chat
170
        else:
171
            chat_fn = completion_to_chat_decorator(self._complete)
172
        return chat_fn(messages, **kwargs)
173

174
    @llm_completion_callback()
175
    def stream_complete(
176
        self, prompt: str, formatted: bool = False, **kwargs: Any
177
    ) -> CompletionResponseGen:
178
        """Completion endpoint for LLM."""
179
        if self._is_chat_model:
180
            complete_fn = stream_chat_to_completion_decorator(self._stream_chat)
181
        else:
182
            complete_fn = self._stream_complete
183
        return complete_fn(prompt, **kwargs)
184

185
    @llm_chat_callback()
186
    def stream_chat(
187
        self, messages: Sequence[ChatMessage], **kwargs: Any
188
    ) -> ChatResponseGen:
189
        if self._is_chat_model:
190
            stream_chat_fn = self._stream_chat
191
        else:
192
            stream_chat_fn = stream_completion_to_chat_decorator(self._stream_complete)
193
        return stream_chat_fn(messages, **kwargs)
194

195
    def _chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
196
        try:
197
            from portkey import Config, Message
198
        except ImportError as exc:
199
            raise ImportError(IMPORT_ERROR_MESSAGE) from exc
200
        _messages = cast(
201
            List[Message],
202
            [{"role": i.role.value, "content": i.content} for i in messages],
203
        )
204
        config = Config(llms=self.llms)
205
        response = self._client.ChatCompletions.create(
206
            messages=_messages, config=config
207
        )
208
        self.llm = self._get_llm(response)
209

210
        message = response.choices[0].message
211
        return ChatResponse(message=message, raw=response)
212

213
    def _complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
214
        try:
215
            from portkey import Config
216
        except ImportError as exc:
217
            raise ImportError(IMPORT_ERROR_MESSAGE) from exc
218

219
        config = Config(llms=self.llms)
220
        response = self._client.Completions.create(prompt=prompt, config=config)
221
        text = response.choices[0].text
222
        return CompletionResponse(text=text, raw=response)
223

224
    def _stream_chat(
225
        self, messages: Sequence[ChatMessage], **kwargs: Any
226
    ) -> ChatResponseGen:
227
        try:
228
            from portkey import Config, Message
229
        except ImportError as exc:
230
            raise ImportError(IMPORT_ERROR_MESSAGE) from exc
231
        _messages = cast(
232
            List[Message],
233
            [{"role": i.role.value, "content": i.content} for i in messages],
234
        )
235
        config = Config(llms=self.llms)
236
        response = self._client.ChatCompletions.create(
237
            messages=_messages, config=config, stream=True, **kwargs
238
        )
239

240
        def gen() -> ChatResponseGen:
241
            content = ""
242
            function_call: Optional[dict] = {}
243
            for resp in response:
244
                if resp.choices is None:
245
                    continue
246
                delta = resp.choices[0].delta
247
                role = delta.get("role", "assistant")
248
                content_delta = delta.get("content", "") or ""
249
                content += content_delta
250

251
                function_call_delta = delta.get("function_call", None)
252
                if function_call_delta is not None:
253
                    if function_call is None:
254
                        function_call = function_call_delta
255
                        # ensure we do not add a blank function call
256
                        if (
257
                            function_call
258
                            and function_call.get("function_name", "") is None
259
                        ):
260
                            del function_call["function_name"]
261
                    else:
262
                        function_call["arguments"] += function_call_delta["arguments"]
263

264
                additional_kwargs = {}
265
                if function_call is not None:
266
                    additional_kwargs["function_call"] = function_call
267

268
                yield ChatResponse(
269
                    message=ChatMessage(
270
                        role=role,
271
                        content=content,
272
                        additional_kwargs=additional_kwargs,
273
                    ),
274
                    delta=content_delta,
275
                    raw=resp,
276
                )
277

278
        return gen()
279

280
    def _stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
281
        try:
282
            from portkey import Config
283
        except ImportError as exc:
284
            raise ImportError(IMPORT_ERROR_MESSAGE) from exc
285

286
        config = Config(llms=self.llms)
287
        response = self._client.Completions.create(
288
            prompt=prompt, config=config, stream=True, **kwargs
289
        )
290

291
        def gen() -> CompletionResponseGen:
292
            text = ""
293
            for resp in response:
294
                delta = resp.choices[0].text or ""
295
                text += delta
296
                yield CompletionResponse(
297
                    delta=delta,
298
                    text=text,
299
                    raw=resp,
300
                )
301

302
        return gen()
303

304
    @property
305
    def _is_chat_model(self) -> bool:
306
        """Check if a given model is a chat-based language model.
307

308
        Returns:
309
            bool: True if the provided model is a chat-based language model,
310
            False otherwise.
311
        """
312
        return is_chat_model(self.model or "")
313

314
    def _get_llm(self, response: "PortkeyResponse") -> "LLMOptions":
315
        return get_llm(response, self.llms)
316