From 80a132e08cf4eaccce3f55b6a062b6587a29ff72 Mon Sep 17 00:00:00 2001 From: PoJu Chen Date: Thu, 14 Dec 2023 21:05:43 -0600 Subject: [PATCH] Create-endpoints-for-inference --- src/llm_vm/server/routes.py | 76 +++++++++++++++++++++++-------------- 1 file changed, 47 insertions(+), 29 deletions(-) diff --git a/src/llm_vm/server/routes.py b/src/llm_vm/server/routes.py index ef36bd14..140b9b4c 100644 --- a/src/llm_vm/server/routes.py +++ b/src/llm_vm/server/routes.py @@ -6,19 +6,36 @@ from llm_vm.agents.REBEL import agent from llm_vm.client import Client from llm_vm.config import settings + # load optimizer for endpoint use # optimizer = LocalOptimizer(MIN_TRAIN_EXS=2,openai_key=None) -client = Client( big_model=settings.big_model, small_model=settings.small_model) +client = Client(big_model=settings.big_model, small_model=settings.small_model) print('optimizer loaded', file=sys.stderr) -bp = Blueprint('bp',__name__) +bp = Blueprint('bp', __name__) + @bp.route('/', methods=['GET']) def home(): return '''home''' + +@bp.route('/v1/inference', methods=['GET']) +def inference(): + data = json.loads(request.data) + try: + big_model = data["big_model"] + openai_key = data["openai_key"] + prompt = data["prompt"] + except KeyError as e: + return {"status": 0, "resp": f"Key {str(e)} not found in the request data"} + client = Client(big_model=big_model) + response = client.complete(prompt=prompt, context='', openai_key=openai_key) + return response + + @bp.route('/v1/complete', methods=['POST']) def optimizing_complete(): rebel_agent = agent.Agent("", [], verbose=1) @@ -30,77 +47,78 @@ def optimizing_complete(): use_rebel_agent = False kwargs = {} if "openai_key" not in data.keys(): - return {"status":0, "resp":"No OpenAI key provided"} + return {"status": 0, "resp": "No OpenAI key provided"} if "temperature" in data.keys(): if type(data["temperature"]) != float and type(data["temperature"]) != int: - return {"status":0, "resp":"Wrong Data Type for temperature"} + return {"status": 0, "resp": "Wrong Data Type for temperature"} else: - kwargs.update({"temperature":data["temperature"]}) + kwargs.update({"temperature": data["temperature"]}) if "stoptoken" in data.keys(): - if type(data["stoptoken"]) != str and type(data["stoptoken"]) != list: # stop can either be a string or array of strings - return {"status":0, "resp":"Wrong Data Type for stop"} + return {"status": 0, "resp": "Wrong Data Type for stop"} elif type(data["stoptoken"]) == list: if len(data["stoptoken"]) > 4: - return {"status":0, "resp":"Too many stop tokens in array limit to 4 or less"} + return {"status": 0, "resp": "Too many stop tokens in array limit to 4 or less"} # check that every element in the list is a string for j in data["stoptoken"]: if type(j) != str: - return {"status":0, "resp":"Wrong Data Type for stop"} + return {"status": 0, "resp": "Wrong Data Type for stop"} kwargs.update({"stop": data["stoptoken"]}) else: - kwargs.update({"stop":data["stoptoken"]}) + kwargs.update({"stop": data["stoptoken"]}) if "data_synthesis" in data.keys(): - if type(data["data_synthesis"])==bool: + if type(data["data_synthesis"]) == bool: data_synthesis = data["data_synthesis"] else: - return {"status":0, "resp":"Wrong Data Type for data_synthesis"} + return {"status": 0, "resp": "Wrong Data Type for data_synthesis"} if "finetune" in data.keys(): - if type(data["finetune"])==bool: + if type(data["finetune"]) == bool: finetune = data["finetune"] else: - return {"status":0, "resp":"Wrong Data Type for finetune"} + return {"status": 0, "resp": "Wrong Data Type for finetune"} if "tools" in data.keys(): if type(data["tools"]) != list: - return {"status":0, "resp":"Wrong data type for tools list"} + return {"status": 0, "resp": "Wrong data type for tools list"} else: - tools=[] + tools = [] for i in data["tools"]: temp_tool_dict = {} temp_args_dict = {} - temp_tool_dict.update({"description":i["description"]}) - temp_tool_dict.update({"dynamic_params":i["dynamic_params"]}) - temp_tool_dict.update({"method":i["method"]}) - temp_args_dict.update({"url":i["url"]}) - temp_args_dict.update({"params":{}}) + temp_tool_dict.update({"description": i["description"]}) + temp_tool_dict.update({"dynamic_params": i["dynamic_params"]}) + temp_tool_dict.update({"method": i["method"]}) + temp_args_dict.update({"url": i["url"]}) + temp_args_dict.update({"params": {}}) for j in i["static_params"].keys(): - temp_args_dict["params"].update({j:i["static_params"][j]}) + temp_args_dict["params"].update({j: i["static_params"][j]}) for k in i["dynamic_params"].keys(): - temp_args_dict["params"].update({k:"{"+k+"}"}) - temp_tool_dict.update({"args":temp_args_dict}) + temp_args_dict["params"].update({k: "{" + k + "}"}) + temp_tool_dict.update({"args": temp_args_dict}) tools.append(temp_tool_dict) rebel_agent.set_tools(tools) use_rebel_agent = True try: openai.api_key = data["openai_key"] except: - return {"status":0, "resp":"Issue with OpenAI key"} + return {"status": 0, "resp": "Issue with OpenAI key"} # optimizer.openai_key = openai.api_key - agent.set_api_key(openai.api_key,"OPENAI_API_KEY") + agent.set_api_key(openai.api_key, "OPENAI_API_KEY") try: if not use_rebel_agent: - completion = client.complete(static_context,dynamic_prompt,openai_key=openai.api_key, data_synthesis=data_synthesis,finetune = finetune, **kwargs) + completion = client.complete( + static_context, dynamic_prompt, openai_key=openai.api_key, data_synthesis=data_synthesis, finetune=finetune, **kwargs + ) else: - completion = rebel_agent.run(static_context+dynamic_prompt,[])[0] + completion = rebel_agent.run(static_context + dynamic_prompt, [])[0] except Exception as e: - return {"status":0, "resp": str(e)} + return {"status": 0, "resp": str(e)} # return {"completion":completion, "status": 200} return completion