Michael McGuinness
60 lines
2.3 KiB
Python
60 lines
2.3 KiB
Python
import prompts
|
|
import claude_funcs
|
|
|
|
|
|
def prompt_select_multiple(key, value, bu_dict, page):
|
|
prompt = prompts.MERGE_SELECT_PROMPT(key, value, bu_dict, page)
|
|
answer = claude_funcs.invoke_claude_3(prompt, max_tokens=4000)
|
|
return answer
|
|
|
|
|
|
def merge_results(td_results, bu_results, text_dict):
|
|
def find_td_dict(page_num):
|
|
"""Helper function to find dictionary in td_results with matching page_num"""
|
|
for td in td_results:
|
|
if td["page_num"] == page_num:
|
|
return td
|
|
return None
|
|
|
|
def get_value(td_dict, key, page_num, text_dict):
|
|
"""Helper function to handle value selection based on the rules specified."""
|
|
if td_dict is None:
|
|
# If there's no matching page, use recursion to look one page earlier
|
|
if int(page_num) > 1:
|
|
return get_value(
|
|
find_td_dict(int(page_num) - 1), key, int(page_num) - 1, text_dict
|
|
), str(int(page_num) - 1)
|
|
else:
|
|
return "N/A", "N/A" # Base case if no previous pages exist
|
|
else:
|
|
value = td_dict.get(key, "N/A")
|
|
if value == "N/A" or value == "":
|
|
# If value is 'N/A' or empty, use recursion to look one page earlier
|
|
return get_value(
|
|
find_td_dict(int(page_num) - 1), key, int(page_num) - 1, text_dict
|
|
), str(int(page_num) - 1)
|
|
elif "," in value and "DATE" not in key:
|
|
# Randomly select one of the comma-separated values
|
|
return prompt_select_multiple(
|
|
key, value, bu_dict, text_dict[page_num]
|
|
), str(int(page_num))
|
|
else:
|
|
return value, page_num
|
|
|
|
merged_results = []
|
|
for bu_dict in bu_results:
|
|
page_num = bu_dict["page_num"]
|
|
td_dict = find_td_dict(page_num)
|
|
|
|
new_dict = bu_dict.copy() # Start with the bu_dict's data
|
|
# Add or overwrite keys from td_dict
|
|
if td_dict:
|
|
for key, value in td_dict.items():
|
|
if key not in ["Filename", "page_num"]:
|
|
new_dict[key], new_dict[f"{key}_PG"] = get_value(
|
|
td_dict, key, page_num, text_dict
|
|
)
|
|
|
|
merged_results.append(new_dict)
|
|
return merged_results
|