Karan Desai
04a190d5b7
Bugfix/hmk issues * fix for missing AARETE_DERIVED_LOB * fix for AARETE_DERIVED_SID * bugfix in universal_json_load * bugfix for normalize_state_field * reimb eff date assignmnet prompt update * prompt update for missing reimb terms * Merged main into bugfix/hmk-issues * test case for universal_json_loads * reveresed the prompt update for missing reimb * prompt fix for reimb dates Approved-by: Katon Minhas
234 lines
12 KiB
Python
234 lines
12 KiB
Python
import unittest
|
|
import json
|
|
import pandas as pd
|
|
import numpy as np
|
|
from src.utils import string_utils
|
|
from constants.delimiters import Delimiter
|
|
|
|
class TestStringUtils(unittest.TestCase):
|
|
def test_extract_text_from_delimiters_pipe(self):
|
|
s = "foo |bar| baz |qux|"
|
|
self.assertEqual(string_utils.extract_text_from_delimiters(s, Delimiter.PIPE), "qux")
|
|
self.assertEqual(string_utils.extract_text_from_delimiters(s, Delimiter.PIPE, 0), "bar")
|
|
self.assertEqual(string_utils.extract_text_from_delimiters("no delimiters", Delimiter.PIPE), "N/A")
|
|
with self.assertRaises(TypeError):
|
|
string_utils.extract_text_from_delimiters(123, Delimiter.PIPE)
|
|
with self.assertRaises(TypeError):
|
|
string_utils.extract_text_from_delimiters(s, "notadelimiter")
|
|
|
|
def test_extract_text_from_delimiters_backtick(self):
|
|
s = "foo `bar` baz `qux`"
|
|
self.assertEqual(string_utils.extract_text_from_delimiters(s, Delimiter.BACKTICK), "qux")
|
|
|
|
def test_extract_text_from_delimiters_triple_backtick(self):
|
|
s = "foo ```bar``` baz ```qux```"
|
|
self.assertEqual(string_utils.extract_text_from_delimiters(s, Delimiter.TRIPLE_BACKTICK), "qux")
|
|
|
|
def test_page_key_sort(self):
|
|
keys = ["10", "2", "A", "1", "B", "1.5"]
|
|
sorted_keys = sorted(keys, key=string_utils.page_key_sort)
|
|
self.assertEqual(sorted_keys, ["1", "1.5", "2", "10", "A", "B"])
|
|
|
|
def test_json_parsing_search(self):
|
|
s = '{"FIELD1": "value1", "FIELD2": "value2"}'
|
|
fields = ["FIELD1", "FIELD2"]
|
|
result = string_utils.json_parsing_search(s, fields)
|
|
self.assertEqual(result, {"FIELD1": "value1", "FIELD2": "value2"})
|
|
|
|
def test_universal_json_load_valid(self):
|
|
# Simple dict
|
|
s = '{"a": 1, "b": 2}'
|
|
self.assertEqual(string_utils.universal_json_load(s), {"a": 1, "b": 2})
|
|
# Simple list
|
|
s = '[1, 2, 3]'
|
|
self.assertEqual(string_utils.universal_json_load(s), [1, 2, 3])
|
|
# Nested dict
|
|
s = '{"a": {"b": [1, 2, {"c": "d"}]}}'
|
|
self.assertEqual(string_utils.universal_json_load(s), {"a": {"b": [1, 2, {"c": "d"}]}})
|
|
# List of dicts
|
|
s = '[{"x": 1}, {"y": 2}]'
|
|
self.assertEqual(string_utils.universal_json_load(s), [{"x": 1}, {"y": 2}])
|
|
# Dict of lists
|
|
s = '{"a": [1, 2], "b": [3, 4]}'
|
|
self.assertEqual(string_utils.universal_json_load(s), {"a": [1, 2], "b": [3, 4]})
|
|
# List of strings
|
|
s = '["foo", "bar", "baz"]'
|
|
self.assertEqual(string_utils.universal_json_load(s), ["foo", "bar", "baz"])
|
|
# JSON with text before/after
|
|
s = 'start {"field": "value"} end'
|
|
self.assertEqual(string_utils.universal_json_load(s), {"field": "value"})
|
|
s = 'before [1, 2, 3] after'
|
|
self.assertEqual(string_utils.universal_json_load(s), [1, 2, 3])
|
|
# Multiple objects, last wins
|
|
s = 'first {"a": 1} middle {"b": 2} last {"c": 3}'
|
|
self.assertEqual(string_utils.universal_json_load(s), {"c": 3})
|
|
# Multiple arrays, last wins
|
|
s = 'first [1] middle [2] last [3]'
|
|
self.assertEqual(string_utils.universal_json_load(s), [3])
|
|
# Nested JSON with text 1
|
|
s = 'junk {"outer": {"inner": [1, 2, {"deep": true}]}} more junk'
|
|
self.assertEqual(string_utils.universal_json_load(s), {"outer": {"inner": [1, 2, {"deep": True}]}})
|
|
# Nested JSON with text 2
|
|
s = 'junk [{"dict1_key": "dict1_val"}, {"dict2_key" : "dict2_val"}] more junk'
|
|
self.assertEqual(string_utils.universal_json_load(s), [{"dict1_key": "dict1_val"}, {"dict2_key" : "dict2_val"}])
|
|
# List of dicts
|
|
s = 'junk```json[{"a1": "value_a1","b1": "value_b1"},{"a2": "value_a2","b2": "value_b2"}]```'
|
|
self.assertEqual(string_utils.universal_json_load(s), [{"a1": "value_a1","b1": "value_b1"},{"a2": "value_a2","b2": "value_b2"}])
|
|
# Boolean and null values
|
|
s = '{"ok": true, "val": null}'
|
|
self.assertEqual(string_utils.universal_json_load(s), {"ok": True, "val": None})
|
|
|
|
|
|
def test_universal_json_load_invalid(self):
|
|
# No valid JSON
|
|
s = 'no json here'
|
|
with self.assertRaises(ValueError):
|
|
string_utils.universal_json_load(s)
|
|
# Broken JSON
|
|
s = '{broken: "json"}'
|
|
with self.assertRaises(ValueError):
|
|
string_utils.universal_json_load(s)
|
|
# Non-string input
|
|
with self.assertRaises(TypeError):
|
|
string_utils.universal_json_load(123)
|
|
# Unbalanced brackets
|
|
s = '{"a": [1, 2, 3}'
|
|
with self.assertRaises(ValueError):
|
|
string_utils.universal_json_load(s)
|
|
# Empty string
|
|
s = ''
|
|
with self.assertRaises(ValueError):
|
|
string_utils.universal_json_load(s)
|
|
|
|
def test_universal_json_load_edge_cases(self):
|
|
# Empty object
|
|
s = 'text {} more text'
|
|
self.assertEqual(string_utils.universal_json_load(s), {})
|
|
# Empty array
|
|
s = 'text [] more text'
|
|
self.assertEqual(string_utils.universal_json_load(s), [])
|
|
# Large object
|
|
large = {f"k{i}": i for i in range(100)}
|
|
s = f'before {json.dumps(large)} after'
|
|
self.assertEqual(string_utils.universal_json_load(s), large)
|
|
|
|
def test_contains_reimbursement(self):
|
|
s = "90% reimbursement of billed charges"
|
|
self.assertTrue(string_utils.contains_reimbursement(s, method="keyword"))
|
|
s = "Payment of $500 per visit"
|
|
self.assertTrue(string_utils.contains_reimbursement(s, method="regex"))
|
|
s = "No reimbursement here"
|
|
self.assertTrue(string_utils.contains_reimbursement(s, method="keyword"))
|
|
with self.assertRaises(ValueError):
|
|
string_utils.contains_reimbursement(s, method="invalid")
|
|
|
|
def test_is_empty_string(self):
|
|
self.assertTrue(string_utils.is_empty(""))
|
|
self.assertTrue(string_utils.is_empty("N/A"))
|
|
self.assertFalse(string_utils.is_empty("valid"))
|
|
|
|
def test_is_empty_list(self):
|
|
self.assertTrue(string_utils.is_empty([]))
|
|
self.assertTrue(string_utils.is_empty([None, "", "N/A"]))
|
|
self.assertFalse(string_utils.is_empty(["valid"]))
|
|
|
|
def test_is_empty_pandas_series(self):
|
|
s = pd.Series(["valid", "", None, "N/A"])
|
|
mask = string_utils.is_empty(s, pd_mask=True)
|
|
self.assertTrue(mask.iloc[1])
|
|
self.assertTrue(mask.iloc[2])
|
|
self.assertTrue(mask.iloc[3])
|
|
self.assertFalse(mask.iloc[0])
|
|
self.assertFalse(string_utils.is_empty(pd.Series(["valid", None]), pd_mask=False))
|
|
self.assertTrue(string_utils.is_empty(pd.Series([None, "", "N/A"]), pd_mask=False))
|
|
|
|
def test_datetime_str(self):
|
|
result = string_utils.datetime_str()
|
|
self.assertTrue(result.startswith("["))
|
|
self.assertTrue(result.endswith("]"))
|
|
|
|
def test_extract_signature_page(self):
|
|
text_dict = {"1": "Contract page 1", "2": "This page has 2 signatures at the bottom", "3": "Final page"}
|
|
result = string_utils.extract_signature_page(text_dict, "test.pdf")
|
|
self.assertIn("2", result)
|
|
self.assertEqual(result["2"], "This page has 2 signatures at the bottom")
|
|
text_dict = {"1": "Contract page 1", "2": "In Witness Whereof", "3": "Final page"}
|
|
result = string_utils.extract_signature_page(text_dict, "test.pdf")
|
|
self.assertIn("2", result)
|
|
self.assertEqual(result["2"], "In Witness Whereof")
|
|
text_dict = {"1": "Contract page 1", "2": "Signature page follows this section", "3": "signature page"}
|
|
result = string_utils.extract_signature_page(text_dict, "test.pdf")
|
|
self.assertIn("3", result)
|
|
|
|
def test_extract_effective_date_pages(self):
|
|
text_dict = {"1": "Contract effective date: January 1, 2024", "2": "General terms and conditions", "3": "This page has 1 signature"}
|
|
result = string_utils.extract_effective_date_pages(text_dict, "test.pdf")
|
|
self.assertIn("1", result)
|
|
self.assertIn("3", result)
|
|
self.assertEqual(len(result), 2)
|
|
|
|
def test_normalize_state_to_abbreviation(self):
|
|
self.assertEqual(string_utils.normalize_state_to_abbreviation("California"), "CA")
|
|
self.assertEqual(string_utils.normalize_state_to_abbreviation("ca"), "CA")
|
|
self.assertEqual(string_utils.normalize_state_to_abbreviation("N.Y."), "NY")
|
|
self.assertEqual(string_utils.normalize_state_to_abbreviation("Unknown State"), "Unknown State")
|
|
self.assertEqual(string_utils.normalize_state_to_abbreviation("N/A"), "N/A")
|
|
|
|
def test_flatten_to_strings(self):
|
|
self.assertEqual(string_utils.flatten_to_strings(["a", ["b", "c"], "d"]), ["a", "b", "c", "d"])
|
|
self.assertEqual(string_utils.flatten_to_strings("single"), ["single"])
|
|
self.assertEqual(string_utils.flatten_to_strings([1, [2, 3], "text"]), ["1", "2", "3", "text"])
|
|
self.assertEqual(string_utils.flatten_to_strings([" text ", [" nested "]]), ["text", "nested"])
|
|
|
|
def test_universal_json_load_multiple_invalid(self):
|
|
# Multiple invalid JSON objects
|
|
s = '{broken: true} some text {not: valid} more text'
|
|
with self.assertRaises(ValueError):
|
|
string_utils.universal_json_load(s)
|
|
|
|
def test_universal_json_load_one_invalid_multiple_valid(self):
|
|
# One invalid, multiple valid (should take the last valid)
|
|
s = '{broken: true} before {"a": 1} middle {"b": 2} end'
|
|
self.assertEqual(string_utils.universal_json_load(s), {"b": 2})
|
|
|
|
def test_universal_json_load_first_nested(self):
|
|
# Multiple, with the first one being nested
|
|
s = '{"outer": {"inner": [1, 2]}} middle {"simple": true} end'
|
|
self.assertEqual(string_utils.universal_json_load(s), {"simple": True})
|
|
|
|
def test_universal_json_load_last_nested(self):
|
|
# Multiple, with the last one being nested
|
|
s = '{"simple": true} middle {"outer": {"inner": [1, 2]}} end'
|
|
self.assertEqual(string_utils.universal_json_load(s), {"outer": {"inner": [1, 2]}})
|
|
|
|
def test_universal_json_load_brackets_in_strings(self):
|
|
# Brackets inside string values should be ignored during bracket matching
|
|
s = '{"text": "array [1] and {2}"}'
|
|
self.assertEqual(string_utils.universal_json_load(s), {"text": "array [1] and {2}"})
|
|
# Double brackets in strings (the original bug case)
|
|
s = '{"field": "Exhibit [[ to Attachment ]"}'
|
|
self.assertEqual(string_utils.universal_json_load(s), {"field": "Exhibit [[ to Attachment ]"})
|
|
# Unbalanced brackets in strings
|
|
s = '{"text": "missing close [ bracket"}'
|
|
self.assertEqual(string_utils.universal_json_load(s), {"text": "missing close [ bracket"})
|
|
# Curly braces in strings
|
|
s = '{"text": "object {key: value} here"}'
|
|
self.assertEqual(string_utils.universal_json_load(s), {"text": "object {key: value} here"})
|
|
# Escaped quotes should not break string detection
|
|
s = '{"text": "say \\"hello\\""}'
|
|
self.assertEqual(string_utils.universal_json_load(s), {"text": 'say "hello"'})
|
|
# Escaped quotes with brackets
|
|
s = '{"text": "say \\"hello [world]\\"" }'
|
|
self.assertEqual(string_utils.universal_json_load(s), {"text": 'say "hello [world]"'})
|
|
# Real-world LLM response case with markdown and brackets in JSON strings
|
|
raw_input = 'Here is the analysis:\n\n```json\n[{"SERVICE_TERM": "Inpatient services", "REIMB_TERM": "100% of rates in Exhibit [[ to Attachment ]"}, {"SERVICE_TERM": "Outpatient services", "REIMB_TERM": "63% of rates in Exhibit [[ to Attachment ]"}]\n```'
|
|
result = string_utils.universal_json_load(raw_input)
|
|
self.assertEqual(len(result), 2)
|
|
self.assertEqual(result[0]["SERVICE_TERM"], "Inpatient services")
|
|
self.assertIn("Exhibit [[", result[0]["REIMB_TERM"])
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|
|
|
|
|