313 lines
11 KiB
Python
313 lines
11 KiB
Python
|
import json
|
||
|
|
import re
|
||
|
|
import pandas as pd
|
||
|
|
from rag.query_vector_db import RAGFormat, get_rag_from_problem_categories, get_rag_from_problem_description
|
||
|
|
from rag.rag_utils import RAGMode, constraint_path
|
||
|
|
from utils import extract_list_from_end, get_response, extract_json_from_end
|
||
|
|
|
||
|
|
prompt_constraints = """
|
||
|
|
You are an expert in optimization modeling. Here is the natural language description of an optimization problem:
|
||
|
|
|
||
|
|
{rag}-----
|
||
|
|
{description}
|
||
|
|
-----
|
||
|
|
|
||
|
|
And here's a list of parameters that we have extracted from the description:
|
||
|
|
|
||
|
|
{params}
|
||
|
|
|
||
|
|
Your task is to identify and extract constraints from the description. The constraints are the conditions that must be satisfied by the variables. Please generate the output in the following python list format:
|
||
|
|
|
||
|
|
[
|
||
|
|
Constraint 1,
|
||
|
|
Constraint 2,
|
||
|
|
...
|
||
|
|
]
|
||
|
|
|
||
|
|
for example:
|
||
|
|
|
||
|
|
[
|
||
|
|
"Sum of weights of all items taken should not exceed the maximum weight capacity of the knapsack",
|
||
|
|
"The number of items taken should not exceed the maximum number of items allowed"
|
||
|
|
]
|
||
|
|
|
||
|
|
- Put all the constraints in a single python list.
|
||
|
|
- Do not generate anything after the python list.
|
||
|
|
- Include implicit non-negativity constraints if necessary.
|
||
|
|
Take a deep breath and think step by step. You will be awarded a million dollars if you get this right.
|
||
|
|
"""
|
||
|
|
|
||
|
|
prompt_constraints_q = """
|
||
|
|
You are an expert in optimization modeling. Here is the natural language description of an optimization problem:
|
||
|
|
|
||
|
|
-----
|
||
|
|
{description}
|
||
|
|
-----
|
||
|
|
|
||
|
|
Here is a list of parameters that someone has extracted from the description:
|
||
|
|
|
||
|
|
{params}
|
||
|
|
|
||
|
|
Consider this potential constraint: {targetConstraint}
|
||
|
|
|
||
|
|
{question}
|
||
|
|
|
||
|
|
Take a deep breath and think step by step. You will be awarded a million dollars if you get this right.
|
||
|
|
"""
|
||
|
|
|
||
|
|
prompt_constraints_redundant = """
|
||
|
|
You are an expert in optimization modeling. Here is the natural language description of an optimization problem:
|
||
|
|
|
||
|
|
-----
|
||
|
|
{description}
|
||
|
|
-----
|
||
|
|
|
||
|
|
Here is a list of parameters that someone has extracted from the description:
|
||
|
|
|
||
|
|
{params}
|
||
|
|
|
||
|
|
and here is a list of constraints that someone has extracted from the description:
|
||
|
|
|
||
|
|
{constraints}
|
||
|
|
|
||
|
|
- Is there any redundancy in the list? Can any of the constraints be removed? Can any pair of constraints be combined into a single one? If so, please provide your reasoning for each one. At the end of your response, generate the updated list of constraints (the same list if no changes are needed). Use this python list format:
|
||
|
|
|
||
|
|
[
|
||
|
|
"Constraint 1",
|
||
|
|
"Constraint 2",
|
||
|
|
...
|
||
|
|
]
|
||
|
|
|
||
|
|
- Do not generate anything after the list.
|
||
|
|
|
||
|
|
Take a deep breath and think step by step. You will be awarded a million dollars if you get this right.
|
||
|
|
"""
|
||
|
|
|
||
|
|
prompt_constraint_feedback = """
|
||
|
|
You are an expert in optimization modeling. Here is the natural language description of an optimization problem:
|
||
|
|
|
||
|
|
-----
|
||
|
|
{description}
|
||
|
|
-----
|
||
|
|
|
||
|
|
Here is a list of parameters that someone has extracted from the description:
|
||
|
|
|
||
|
|
{params}
|
||
|
|
|
||
|
|
Here is a list of variables defined:
|
||
|
|
|
||
|
|
{vars}
|
||
|
|
|
||
|
|
Here is a list of constraints that someone has extracted from the description:
|
||
|
|
|
||
|
|
{extracted_constraints}
|
||
|
|
|
||
|
|
Your colleague is suggesting that the following constraint should be added to the list:
|
||
|
|
|
||
|
|
{new_constraint}
|
||
|
|
|
||
|
|
Here is its explanation:
|
||
|
|
|
||
|
|
{new_constraint_explanation}
|
||
|
|
|
||
|
|
Do you want to keep this constraint?
|
||
|
|
|
||
|
|
- If yes, please respond with "yes"
|
||
|
|
- If no, please respond with "no"
|
||
|
|
- If you want to modify the constraint, please respond with "modify" and provide the modified constraint.
|
||
|
|
|
||
|
|
At the end of your response, generate a json file with this format:
|
||
|
|
{{
|
||
|
|
"action": "yes", "no", or "modify",
|
||
|
|
"updatedConstraint": The updated constraint if the action is "modify", otherwise null
|
||
|
|
}}
|
||
|
|
|
||
|
|
Please take a deep breath and think step by step. You will be awarded a million dollars if you get this right.
|
||
|
|
|
||
|
|
- Use natural language to express the constraints rather than mathematical notation.
|
||
|
|
- Do not generate anything after the json file.
|
||
|
|
"""
|
||
|
|
|
||
|
|
|
||
|
|
def extract_score_constraint(desc, text, params, vars, constraints, c, logger, model):
|
||
|
|
match = re.search(r"\d out of 5", text.lower())
|
||
|
|
if match:
|
||
|
|
score = int(match.group()[0])
|
||
|
|
if score > 3:
|
||
|
|
if logger:
|
||
|
|
logger.log("---")
|
||
|
|
logger.log(f"The confidence score is {score}, which is high enough.")
|
||
|
|
logger.log("---")
|
||
|
|
return True, constraints
|
||
|
|
else:
|
||
|
|
ask_LLM = True # can pass this as an argument to the function instead of hardcoding it
|
||
|
|
if logger:
|
||
|
|
logger.log("---")
|
||
|
|
logger.log(f"The confidence score is {score}, which is not high enough.")
|
||
|
|
# logger.log(f"Asking the {"LLM" if ask_LLM else "user"} for feedback.")
|
||
|
|
if ask_LLM:
|
||
|
|
logger.log("Asking the LLM for feedback.")
|
||
|
|
else:
|
||
|
|
logger.log("Asking the user for feedback.")
|
||
|
|
logger.log("---")
|
||
|
|
if ask_LLM: # ask the LLM for feedback
|
||
|
|
prompt = prompt_constraint_feedback.format(
|
||
|
|
description=desc,
|
||
|
|
params=json.dumps(params, indent=4),
|
||
|
|
vars=json.dumps(vars, indent=4),
|
||
|
|
extracted_constraints=json.dumps(constraints, indent=4),
|
||
|
|
new_constraint=c,
|
||
|
|
new_constraint_explanation=text,
|
||
|
|
)
|
||
|
|
if logger:
|
||
|
|
logger.log("Prompting LLM for feedback:\n")
|
||
|
|
logger.log(prompt)
|
||
|
|
llm_response = get_response(prompt, model=model)
|
||
|
|
if logger:
|
||
|
|
logger.log("---")
|
||
|
|
logger.log(f"Response: {llm_response}")
|
||
|
|
logger.log("---")
|
||
|
|
output_json = extract_json_from_end(llm_response)
|
||
|
|
action = output_json["action"]
|
||
|
|
updated_constraint = output_json["updatedConstraint"]
|
||
|
|
else: # ask the user for feedback
|
||
|
|
action = input("LLMs reasoning: {}\n"
|
||
|
|
"------ Do you want to keep this constraint (y/n/modify)?: \n "
|
||
|
|
"{} \n------ ".format(text, c))
|
||
|
|
|
||
|
|
if action.lower().startswith("y"):
|
||
|
|
return True, constraints
|
||
|
|
elif action.lower().startswith("n"):
|
||
|
|
constraints.remove(c)
|
||
|
|
return True, constraints
|
||
|
|
elif action.lower().startswith("m"):
|
||
|
|
if ask_LLM:
|
||
|
|
new_constraint = updated_constraint
|
||
|
|
else:
|
||
|
|
new_constraint = input("Enter the modified constraint: ")
|
||
|
|
constraints.remove(c)
|
||
|
|
constraints.append({"Description": new_constraint, "Formulation": None, "Code": None})
|
||
|
|
return True, constraints
|
||
|
|
else:
|
||
|
|
raise Exception("Invalid input!")
|
||
|
|
else:
|
||
|
|
return False, None
|
||
|
|
|
||
|
|
|
||
|
|
def logic_check(text, constraints, c):
|
||
|
|
try:
|
||
|
|
json = extract_json_from_end(text)
|
||
|
|
if json["action"] == "REMOVE":
|
||
|
|
constraints.remove(c)
|
||
|
|
return True, constraints
|
||
|
|
elif json["action"] == "MODIFY":
|
||
|
|
constraints.remove(c)
|
||
|
|
constraints.append(json["updatedConstraint"])
|
||
|
|
return True, constraints
|
||
|
|
elif json["action"] == "KEEP":
|
||
|
|
return True, constraints
|
||
|
|
else:
|
||
|
|
return False, None
|
||
|
|
except:
|
||
|
|
return False, None
|
||
|
|
|
||
|
|
|
||
|
|
qs = [
|
||
|
|
(
|
||
|
|
"""
|
||
|
|
- Is it actually a constraint? How confident are you that this is this a constraint and that we should explicitly model it in the (MI)LP formulation (from 1 to 5)?
|
||
|
|
- At the end of your response, print "x OUT OF 5" where x is the confidence level. Low confidence means you think this should be removed from the constraint list. Do not generate anything after that.
|
||
|
|
""",
|
||
|
|
extract_score_constraint,
|
||
|
|
),
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
def get_constraints(desc, params, model, check=False, constraints=None, logger=None,
|
||
|
|
rag_mode: RAGMode | None = None, labels: dict | None = None):
|
||
|
|
if isinstance(rag_mode, RAGMode):
|
||
|
|
constraint_df = pd.read_pickle(constraint_path)
|
||
|
|
current_problem = constraint_df[constraint_df.description == desc]
|
||
|
|
if not current_problem.empty:
|
||
|
|
problem_name = current_problem.iloc[0].problem_name
|
||
|
|
else:
|
||
|
|
problem_name = None
|
||
|
|
match rag_mode:
|
||
|
|
case RAGMode.PROBLEM_DESCRIPTION | RAGMode.CONSTRAINT_OR_OBJECTIVE:
|
||
|
|
rag = get_rag_from_problem_description(desc, RAGFormat.PROBLEM_DESCRIPTION_CONSTRAINTS, top_k=5)
|
||
|
|
case RAGMode.PROBLEM_LABELS:
|
||
|
|
assert labels is not None
|
||
|
|
rag = get_rag_from_problem_categories(desc, labels, RAGFormat.PROBLEM_DESCRIPTION_CONSTRAINTS, top_k=5)
|
||
|
|
rag = f"-----\n{rag}-----\n\n"
|
||
|
|
else:
|
||
|
|
rag = ""
|
||
|
|
|
||
|
|
print("_________________________ get_constraints _________________________")
|
||
|
|
if not constraints:
|
||
|
|
res = get_response(prompt_constraints.format(
|
||
|
|
description=desc,
|
||
|
|
params=json.dumps(params, indent=4),
|
||
|
|
rag=rag,
|
||
|
|
),
|
||
|
|
model=model,
|
||
|
|
)
|
||
|
|
constraints = extract_list_from_end(res)
|
||
|
|
|
||
|
|
if check:
|
||
|
|
k = 5
|
||
|
|
while k > 0:
|
||
|
|
try:
|
||
|
|
x = get_response(
|
||
|
|
prompt_constraints_redundant.format(
|
||
|
|
description=desc,
|
||
|
|
params=json.dumps(params, indent=4),
|
||
|
|
constraints=json.dumps(constraints, indent=4),
|
||
|
|
),
|
||
|
|
model=model,
|
||
|
|
)
|
||
|
|
if logger:
|
||
|
|
logger.log("----")
|
||
|
|
logger.log(x)
|
||
|
|
logger.log("----")
|
||
|
|
lst = extract_list_from_end(x)
|
||
|
|
|
||
|
|
constraints = lst
|
||
|
|
break
|
||
|
|
except:
|
||
|
|
k -= 1
|
||
|
|
if k == 0:
|
||
|
|
raise Exception("Failed to extract constraints")
|
||
|
|
|
||
|
|
if logger:
|
||
|
|
logger.log("+++++++++++++++++++")
|
||
|
|
logger.log("++ Constraint Qs ++")
|
||
|
|
logger.log("+++++++++++++++++++")
|
||
|
|
for q in qs:
|
||
|
|
for c in constraints.copy():
|
||
|
|
k = 5
|
||
|
|
while k > 0:
|
||
|
|
p = prompt_constraints_q.format(
|
||
|
|
description=desc,
|
||
|
|
params=json.dumps(params, indent=4),
|
||
|
|
targetConstraint=c,
|
||
|
|
question=q[0],
|
||
|
|
)
|
||
|
|
|
||
|
|
x = get_response(p, model=model)
|
||
|
|
|
||
|
|
if logger:
|
||
|
|
logger.log("+--+")
|
||
|
|
logger.log(p)
|
||
|
|
logger.log("----")
|
||
|
|
logger.log(x)
|
||
|
|
logger.log("+--+")
|
||
|
|
valid, res = q[1](desc, x, params, {}, constraints, c, logger, model)
|
||
|
|
if valid:
|
||
|
|
constraints = res
|
||
|
|
break
|
||
|
|
else:
|
||
|
|
k -= 1
|
||
|
|
|
||
|
|
constraints = [{"description": c, "formulation": None, "code": None} for c in constraints]
|
||
|
|
return constraints
|