Coverage for mindsdb / integrations / handlers / portkey_handler / portkey_handler.py: 0%

1from typing import Dict, Optional 

2 

3import pandas as pd 

4from portkey_ai import Portkey 

5 

6from mindsdb.integrations.libs.base import BaseMLEngine 

7from mindsdb.utilities import log 

8 

9from mindsdb.integrations.utilities.handler_utils import get_api_key 

10 

11logger = log.getLogger(__name__) 

12 

13DEFAULT_METADATA = { 

14 "_source": "portkey-mindsdb-integration", 

15} 

16 

17 

18class PortkeyHandler(BaseMLEngine): 

19 """ 

20 Integration with the Portkey LLM Python Library 

21 """ 

22 

23 name = "portkey" 

24 

25 def __init__(self, *args, **kwargs): 

26 super().__init__(*args, **kwargs) 

27 self.generative = True 

28 

29 def create( 

30 self, 

31 target: str, 

32 df: Optional[pd.DataFrame] = None, 

33 args: Optional[Dict] = None, 

34 ) -> None: 

35 

36 if "using" not in args: 

37 raise Exception( 

38 "Portkey engine requires a USING clause! Refer to its documentation for more details." 

39 ) 

40 

41 self.model_storage.json_set("args", args) 

42 

43 def predict( 

44 self, df: Optional[pd.DataFrame] = None, args: Optional[Dict] = None 

45 ) -> None: 

46 

47 args = self.model_storage.json_get("args") 

48 api_key = get_api_key('portkey', args["using"], self.engine_storage, strict=False) 

49 

50 self.client = Portkey( 

51 **self.engine_storage.get_connection_args(), 

52 api_key=api_key, 

53 metadata=DEFAULT_METADATA 

54 ) 

55 

56 result_df = pd.DataFrame() 

57 

58 result_df["predictions"] = df["question"].apply(self._predict_answer) 

59 

60 result_df = result_df.rename(columns={"predictions": args["target"]}) 

61 

62 return result_df 

63 

64 def _predict_answer(self, text): 

65 """ 

66 connects with portkey messages api to predict the answer for the particular question 

67 

68 """ 

69 

70 model_args = self.model_storage.json_get("args") 

71 

72 message = self.client.chat.completions.create( 

73 **model_args, 

74 messages=[ 

75 {"role": "user", "content": text} 

76 ] 

77 ) 

78 

79 return message.choices[0].message.content