/ ragaai_catalyst / internal_api_completion.py
internal_api_completion.py
 1  import requests
 2  import json
 3  import subprocess
 4  import logging
 5  import traceback
 6  import pandas as pd
 7  
 8  logger = logging.getLogger(__name__)
 9  
10  def api_completion(messages, model_config, kwargs):
11      attempts = 0
12      while attempts < 3:
13  
14          user_id = kwargs.get('user_id', '1')
15          internal_llm_proxy = kwargs.get('internal_llm_proxy', -1)
16              
17              
18          job_id = model_config.get('job_id',-1)
19          converted_message = convert_input(messages,model_config, user_id)
20          payload = json.dumps(converted_message)
21          headers = {
22              'Content-Type': 'application/json',
23              # 'Wd-PCA-Feature-Key':f'your_feature_key, $(whoami)'
24          }
25          try:
26              response = requests.request("POST", internal_llm_proxy, headers=headers, data=payload)
27              if model_config.get('log_level','')=='debug':
28                  logger.info(f'Model response Job ID {job_id} {response.text}')
29              if response.status_code!=200:
30                  # logger.error(f'Error in model response Job ID {job_id}:',str(response.text))
31                  raise ValueError(str(response.text))
32              
33              if response.status_code==200:
34                  response = response.json()                
35                  if "error" in response:
36                      raise ValueError(response["error"]["message"])
37                  else:
38                      result=  response["choices"][0]["message"]["content"]
39                      response1 = result.replace('\n', '').replace('```json','').replace('```', '').strip()
40                      try:
41                          json_data = json.loads(response1)
42                          df = pd.DataFrame(json_data)
43                          return(df)
44                      except json.JSONDecodeError:
45                          attempts += 1  # Increment attempts if JSON parsing fails
46                          if attempts == 3:
47                              raise Exception("Failed to generate a valid response after multiple attempts.")
48  
49          except Exception as e:
50              raise ValueError(f"{e}")
51  
52  
53  def get_username():
54      result = subprocess.run(['whoami'], capture_output=True, text=True)
55      result = result.stdout
56      return result
57  
58  
59  def convert_input(messages, model_config, user_id):
60      doc_input = {
61        "model": model_config.get('model'),
62        **model_config,
63        "messages": messages,
64        "user_id": user_id
65      }
66      return doc_input
67  
68  
69  if __name__=='__main__':
70      messages = [
71          {
72              "role": "system",
73              "content": "you are a poet well versed in shakespeare literature"
74          },
75          {
76            "role": "user",
77            "content": "write a poem on pirates and penguins"
78          }
79        ]
80      kwargs = {"internal_llm_proxy": "http://13.200.11.66:4000/chat/completions", "user_id": 1}
81      model_config = {"model": "workday_gateway", "provider":"openai", "max_tokens": 10}
82      answer = api_completion(messages, model_config, kwargs)
83      print(answer)