HarshSanghavi's picture
Upload 5 files
17a7095 verified
from fastapi import FastAPI, status, Response
from agent import agent, system_prompt
from functions import process_input_json, evalute_final_json, get_variables_and_boundaries
import traceback
app = FastAPI()
@app.post("/generate-step")
async def generate_step(input_json: dict, response: Response = None):
try:
rl_input_json = input_json['rl_input_json']
usecase_details = input_json['step_json']
if "title" not in usecase_details or not usecase_details['title']:
response.status_code = status.HTTP_422_UNPROCESSABLE_ENTITY
return {"code": 422, "data": "Usecase details is required.", "success": False}
if "observableFactor" not in rl_input_json or not rl_input_json['observableFactor']:
response.status_code = status.HTTP_422_UNPROCESSABLE_ENTITY
return {"code": 422, "data": "Observable factor is required.", "success": False}
if "actions" not in rl_input_json or not rl_input_json['actions']:
response.status_code = status.HTTP_422_UNPROCESSABLE_ENTITY
return {"code": 422, "data": "Actions is required.", "success": False}
if "boundaries" not in rl_input_json or not rl_input_json['boundaries']:
response.status_code = status.HTTP_422_UNPROCESSABLE_ENTITY
return {"code": 422, "data": "Boundaries is required.", "success": False}
processed_json = process_input_json(rl_input_json['observableFactor'], rl_input_json.get(
'constantFactor'), rl_input_json.get('modelSelection'), rl_input_json['actions'], rl_input_json['boundaries'])
print("rl_step | processed_json:", processed_json)
try:
evalute_final_json(processed_json)
except Exception:
trace = traceback.format_exc()
if "KeyError:" in trace:
response.status_code = status.HTTP_422_UNPROCESSABLE_ENTITY
return {"code": 422, "data": trace, "success": False}
details = get_variables_and_boundaries(processed_json)
last_trace = None
current_errors = ""
current_prompt = system_prompt
existing_reward = ""
if "reward" in rl_input_json and rl_input_json['reward']["selectedSection"] == "ai" and rl_input_json['reward']['generatedCode']:
current_prompt += "\n###Existing reward:\n```" + \
rl_input_json['reward']['generatedCode'].replace(
"{", "{{").replace("}", "}}") + "```"
existing_reward = rl_input_json['reward']['generatedCode']
elif "reward" in rl_input_json and rl_input_json['reward']["selectedSection"] != "ai" and rl_input_json['reward']['customCode']:
current_prompt += "\n###Existing reward:\n```" + \
rl_input_json['reward']['customCode'].replace(
"{", "{{").replace("}", "}}") + "```"
existing_reward = rl_input_json['reward']['customCode']
for _ in range(3):
final_prompt = current_prompt + "\n" + current_errors
try:
agent_response = agent.invoke({"input": final_prompt.format(
json_data=processed_json, use_case_name=usecase_details['title'], variable_information=details, use_case_description=usecase_details['description'])})
step = agent_response['output'].get("step")
reward = agent_response['output'].get("reward")
processed_json['environment']['step'] = ["\n".join(step)]
processed_json['environment']['reward'] = [
"\n".join(reward)] if not existing_reward else [existing_reward]
try:
print(
"##################### STARTED EVALUATION ##############################")
evalute_final_json(processed_json)
except Exception:
print(
"######################## Error in evalution #########################################")
last_trace = traceback.format_exc()
current_errors = "In last executation I got below error so make sure you I don't get same error again. \nError: \n" + last_trace
print(last_trace)
print(
"######################## Error in evalution #########################################")
continue
response.status_code = status.HTTP_200_OK
return {"code": 200, "data": "\n".join(agent_response['output']['step']), "success": True}
except Exception:
last_trace = traceback.format_exc()
print("rl_step | Error trace:", last_trace)
response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
return {"code": 500, "data": last_trace, "success": False}
except Exception:
trace = traceback.format_exc()
print("rl_step | Error trace:", trace)
if "KeyError:" in trace:
response.status_code = status.HTTP_422_UNPROCESSABLE_ENTITY
return {"code": 422, "data": trace, "success": False}
@app.post("/generate-reward")
async def generate_reward(input_json: dict, response: Response = None):
try:
rl_input_json = input_json['rl_input_json']
usecase_details = input_json['reward_json']
if "title" not in usecase_details or not usecase_details['title']:
response.status_code = status.HTTP_422_UNPROCESSABLE_ENTITY
return {"code": 422, "data": "Usecase details is required.", "success": False}
if "observableFactor" not in rl_input_json or not rl_input_json['observableFactor']:
response.status_code = status.HTTP_422_UNPROCESSABLE_ENTITY
return {"code": 422, "data": "Observable factor is required.", "success": False}
if "actions" not in rl_input_json or not rl_input_json['actions']:
response.status_code = status.HTTP_422_UNPROCESSABLE_ENTITY
return {"code": 422, "data": "Actions is required.", "success": False}
if "boundaries" not in rl_input_json or not rl_input_json['boundaries']:
response.status_code = status.HTTP_422_UNPROCESSABLE_ENTITY
return {"code": 422, "data": "Boundaries is required.", "success": False}
processed_json = process_input_json(rl_input_json['observableFactor'], rl_input_json.get(
'constantFactor'), rl_input_json.get('modelSelection'), rl_input_json['actions'], rl_input_json['boundaries'])
print("rl_reward | processed_json:", processed_json)
try:
evalute_final_json(processed_json)
except Exception:
trace = traceback.format_exc()
print("rl_reward | Error trace:", trace)
if "KeyError:" in trace:
response.status_code = status.HTTP_422_UNPROCESSABLE_ENTITY
return {"code": 422, "data": trace, "success": False}
details = get_variables_and_boundaries(processed_json)
last_trace = None
current_errors = ""
current_prompt = system_prompt
existing_step = ""
if "step" in rl_input_json and rl_input_json['step']["selectedSection"] == "ai" and rl_input_json['step']['generatedCode']:
current_prompt += "\n###Existing step:\n```" + \
rl_input_json['step']['generatedCode'].replace(
"{", "{{").replace("}", "}}") + "```"
existing_step = rl_input_json['step']['generatedCode']
elif "step" in rl_input_json and rl_input_json['step']["selectedSection"] != "ai" and rl_input_json['step']['customCode']:
current_prompt += "\n###Existing step:\n```" + \
rl_input_json['step']['customCode'].replace(
"{", "{{").replace("}", "}}") + "```"
existing_step = rl_input_json['step']['customCode']
for _ in range(3):
final_prompt = current_prompt + "\n" + current_errors
try:
agent_response = agent.invoke({"input": final_prompt.format(
json_data=processed_json, use_case_name=usecase_details['title'], variable_information=details, use_case_description=usecase_details['description'])})
step = agent_response['output'].get("step")
reward = agent_response['output'].get("reward")
processed_json['environment']['step'] = [
"\n".join(step)] if not existing_step else [existing_step]
processed_json['environment']['reward'] = ["\n".join(reward)]
try:
print(
"##################### STARTED EVALUATION ##############################")
evalute_final_json(processed_json)
except Exception:
print(
"######################## Error in evalution #########################################")
last_trace = traceback.format_exc()
current_errors = "In last executation I got below error so make sure you I don't get same error again. \nError: \n" + last_trace
print(last_trace)
print(
"######################## Error in evalution #########################################")
continue
response.status_code = status.HTTP_200_OK
return {"code": 200, "data": "\n".join(agent_response['output']['reward']), "success": True}
except Exception:
last_trace = traceback.format_exc()
print("rl_reward | Error trace:", last_trace)
response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
return {"code": 500, "data": last_trace, "success": False}
except Exception:
trace = traceback.format_exc()
print("rl_reward | Error trace:", trace)
if "KeyError:" in trace:
response.status_code = status.HTTP_422_UNPROCESSABLE_ENTITY
return {"code": 422, "data": trace, "success": False}