/ 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)