Katon Minhas
1050287397
Refactor/one to n * remove print * Update logging for fsbreakout and grouper breakout * Additional type hints * updated special case primary * Remove deprecated special case check * OUTLIER_TERMS-->OUTLIER_TERM standardization * Applied suggestion to fieldExtraction/src/prompts/investment_prompts.json * update unit test * Update Carveout IND and Default IND logic * Split out exhibit-level lesser * Bugfix * update lesser of exhibit prompt * Bugfix * exhibit lesser of varible name bugfix * Update exhibit lesser of passing * Update Exhibit-level prompt template * Shorten full_context prompt length if the questions are long * Special case assignment index * Resolve unit tests * updated UC IND prompt * Remove print * update one-to-n unit test * added dot based date example * updated date fix prompt * Merged in bugfix/json_error_and_unit_of_measure (pull request #708) Bugfix/json error and unit of measure * json parser fixes and prompt changes * universal json load fixed * json error fixed * json error fixed 2 * Merged refactor/one-to-n into bugfix/json_error_and_unit_of_measure Approved-by: Katon Minhas Approved-by: Alex Galarce
326 lines
14 KiB
Python
326 lines
14 KiB
Python
import unittest
|
|
import json
|
|
import os
|
|
from unittest.mock import patch, MagicMock
|
|
|
|
import pandas as pd
|
|
|
|
from src.investment import dynamic_funcs, prompt_calls
|
|
from src.prompts.fieldset import Field, FieldSet
|
|
from constants.constants import Constants
|
|
|
|
|
|
class TestDynamicFuncsWithRealConstants(unittest.TestCase):
|
|
"""Unit tests for functions in src/investment/dynamic_funcs.py using real Constants."""
|
|
|
|
def setUp(self):
|
|
"""Set up test fixtures before each test method."""
|
|
# Create real Constants object
|
|
self.constants = Constants()
|
|
|
|
# Create test fields
|
|
self.test_field = Field.from_values(
|
|
field_name="TEST_FIELD",
|
|
relationship="one-to-one",
|
|
field_type="exhibit_level",
|
|
prompt="What is the test field?",
|
|
base_field="TEST_FIELD_BASE"
|
|
)
|
|
|
|
self.test_field2 = Field.from_values(
|
|
field_name="TEST_FIELD2",
|
|
relationship="one-to-many",
|
|
field_type="exhibit_level",
|
|
prompt="What is the second test field?",
|
|
base_field="TEST_FIELD2_BASE"
|
|
)
|
|
|
|
self.escalator_field = Field.from_values(
|
|
field_name="RATE_ESCALATOR_DESC",
|
|
relationship="one-to-many",
|
|
field_type="rate_escalator",
|
|
prompt="What is the rate escalator?",
|
|
base_field="RATE_ESCALATOR_STATEMENT"
|
|
)
|
|
|
|
self.outlier_field = Field.from_values(
|
|
field_name="OUTLIER_TERMS",
|
|
relationship="one-to-many",
|
|
field_type="exhibit_level",
|
|
prompt="What are the outlier terms?",
|
|
base_field="OUTLIER_TERMS"
|
|
)
|
|
|
|
# Create test FieldSets
|
|
self.test_fieldset = FieldSet()
|
|
self.test_fieldset.add_field(self.test_field)
|
|
self.test_fieldset.add_field(self.test_field2)
|
|
|
|
# Mock filename
|
|
self.filename = "test_file.pdf"
|
|
|
|
# Sample text
|
|
self.exhibit_text = "This is a sample exhibit text with test data."
|
|
self.exhibit_header = "This is a sample exhibit header with test data."
|
|
|
|
# Sample answer dictionaries
|
|
self.answer_dicts = [
|
|
{"TEST_FIELD": "Value 1", "TEST_FIELD2": "", "RATE_ESCALATOR_DESC": "", "AARETE_DERIVED_LOB": "Commercial"},
|
|
{"TEST_FIELD": "Value 2", "TEST_FIELD2": "", "RATE_ESCALATOR_DESC": "", "AARETE_DERIVED_LOB": "Medicare"},
|
|
{"TEST_FIELD": "Value 3", "TEST_FIELD2": "", "OUTLIER_TERMS": "", "AARETE_DERIVED_LOB": ""}
|
|
]
|
|
|
|
@patch('src.utils.llm_utils.invoke_claude')
|
|
@patch('src.utils.string_utils.universal_json_load')
|
|
def test_prompt_dynamic(self, mock_json_load, mock_invoke_claude):
|
|
"""Test the prompt_dynamic function."""
|
|
# Setup
|
|
field_prompts = {
|
|
"TEST_FIELD": "What is the test field?",
|
|
"TEST_FIELD2": "What is the second test field?"
|
|
}
|
|
|
|
# Configure mocks
|
|
mock_invoke_claude.return_value = '{"TEST_FIELD": ["Value 1"], "TEST_FIELD2": ["Value 2"]}'
|
|
mock_json_load.return_value = {"TEST_FIELD": ["Value 1"], "TEST_FIELD2": ["Value 2"]}
|
|
|
|
# Execute
|
|
result = prompt_calls.prompt_dynamic(self.exhibit_text, field_prompts, self.filename)
|
|
|
|
# Assert
|
|
mock_invoke_claude.assert_called_once()
|
|
mock_json_load.assert_called_once_with('{"TEST_FIELD": ["Value 1"], "TEST_FIELD2": ["Value 2"]}')
|
|
self.assertEqual(result, {"TEST_FIELD": ["Value 1"], "TEST_FIELD2": ["Value 2"]})
|
|
|
|
# Verify that the prompt contains our text
|
|
prompt_call = mock_invoke_claude.call_args[0][0]
|
|
self.assertIn(self.exhibit_text, prompt_call)
|
|
self.assertIn("TEST_FIELD", prompt_call)
|
|
self.assertIn("TEST_FIELD2", prompt_call)
|
|
|
|
def test_add_full_context_field_normal(self):
|
|
"""Test adding a normal field to one_to_one_fields."""
|
|
# Setup
|
|
one_to_one_fields = FieldSet()
|
|
field_to_add = self.test_field
|
|
|
|
# Execute
|
|
result = dynamic_funcs.add_full_context_field(one_to_one_fields, field_to_add, self.answer_dicts)
|
|
|
|
# Assert
|
|
self.assertEqual(len(result.fields), 1)
|
|
self.assertEqual(result.fields[0].field_name, "TEST_FIELD")
|
|
self.assertEqual(result.fields[0].field_type, "full_context")
|
|
self.assertEqual(result.fields[0].relationship, "one_to_one")
|
|
|
|
def test_add_full_context_field_program_with_multiple_lobs(self):
|
|
"""Test that program fields are not added when multiple LOBs exist."""
|
|
# Setup
|
|
one_to_one_fields = FieldSet()
|
|
program_field = Field.from_values(
|
|
field_name="PROGRAM",
|
|
relationship="one-to-one",
|
|
field_type="exhibit_level",
|
|
prompt="What is the program?",
|
|
base_field="PROGRAM"
|
|
)
|
|
|
|
# Execute - use our test answer_dicts which has multiple non-empty LOBs
|
|
result = dynamic_funcs.add_full_context_field(one_to_one_fields, program_field, self.answer_dicts)
|
|
|
|
# Assert
|
|
self.assertEqual(len(result.fields), 0) # Field should not be added
|
|
|
|
def test_add_full_context_field_program_with_single_lob(self):
|
|
"""Test that program fields are added when only one LOB exists."""
|
|
# Setup
|
|
one_to_one_fields = FieldSet()
|
|
program_field = Field.from_values(
|
|
field_name="PROGRAM",
|
|
relationship="one-to-one",
|
|
field_type="exhibit_level",
|
|
prompt="What is the program?",
|
|
base_field="PROGRAM"
|
|
)
|
|
|
|
# Use modified answer_dicts with only one non-empty LOB
|
|
single_lob_answer_dicts = [
|
|
{"TEST_FIELD": "Value 1", "TEST_FIELD2": "", "RATE_ESCALATOR_DESC": "", "AARETE_DERIVED_LOB": "Commercial"},
|
|
{"TEST_FIELD": "Value 2", "TEST_FIELD2": "", "RATE_ESCALATOR_DESC": "", "AARETE_DERIVED_LOB": ""},
|
|
{"TEST_FIELD": "Value 3", "TEST_FIELD2": "", "OUTLIER_TERMS": "", "AARETE_DERIVED_LOB": ""}
|
|
]
|
|
|
|
# Execute
|
|
result = dynamic_funcs.add_full_context_field(one_to_one_fields, program_field, single_lob_answer_dicts)
|
|
|
|
# Assert
|
|
self.assertEqual(len(result.fields), 1) # Field should be added
|
|
self.assertEqual(result.fields[0].field_name, "PROGRAM")
|
|
|
|
@patch('src.investment.dynamic_funcs.Field.load_from_file')
|
|
@patch('src.investment.dynamic_funcs.FieldSet')
|
|
def test_get_dynamic_one_to_one_fields_any_empty(self, mock_fieldset_class, mock_load_field):
|
|
"""Test getting dynamic one-to-one fields with any empty criterion."""
|
|
# Setup
|
|
# Configure mock for fieldset class
|
|
mock_all_empty_fields = MagicMock()
|
|
mock_all_empty_fields.fields = []
|
|
|
|
mock_any_empty_fields = MagicMock()
|
|
mock_field1 = MagicMock()
|
|
mock_field1.field_name = "RATE_ESCALATOR_DESC"
|
|
mock_field1.base_field = "RATE_ESCALATOR_STATEMENT"
|
|
mock_any_empty_fields.fields = [mock_field1]
|
|
|
|
mock_one_to_one_fields = MagicMock()
|
|
mock_one_to_one_fields.fields = []
|
|
|
|
mock_fieldset_class.side_effect = [
|
|
mock_all_empty_fields,
|
|
mock_any_empty_fields,
|
|
mock_one_to_one_fields
|
|
]
|
|
|
|
# Configure mock for loading fields
|
|
escalator_statement_field = Field.from_values(
|
|
field_name="RATE_ESCALATOR_STATEMENT",
|
|
relationship="one-to-one",
|
|
field_type="full_context",
|
|
prompt="What is the rate escalator statement?"
|
|
)
|
|
mock_load_field.return_value = escalator_statement_field
|
|
|
|
# Execute
|
|
result = dynamic_funcs.get_dynamic_one_to_one_fields(self.answer_dicts)
|
|
|
|
@patch('src.investment.dynamic_funcs.prompt_calls.prompt_dynamic_primary')
|
|
def test_dynamic_primary_with_header_answers(self, mock_prompt_dynamic_primary):
|
|
"""Test dynamic_primary function with answers in the header."""
|
|
# Setup
|
|
dynamic_fields = self.test_fieldset
|
|
reimbursement_level_fields = FieldSet()
|
|
exhibit_level_answers = {}
|
|
|
|
# Configure mocks
|
|
mock_prompt_dynamic_primary.side_effect = ["Single Answer", "Another Answer"]
|
|
|
|
# Execute
|
|
exhibit_level_answers, updated_reimbursement_fields = dynamic_funcs.dynamic_primary(
|
|
self.exhibit_text, self.exhibit_header, exhibit_level_answers,
|
|
self.constants, self.filename, dynamic_fields, reimbursement_level_fields
|
|
)
|
|
|
|
# Assert
|
|
self.assertEqual(len(exhibit_level_answers), 2) # Both fields should be added
|
|
self.assertEqual(exhibit_level_answers["TEST_FIELD"], "Single Answer")
|
|
self.assertEqual(exhibit_level_answers["TEST_FIELD2"], "Another Answer")
|
|
mock_prompt_dynamic_primary.assert_called()
|
|
self.assertEqual(mock_prompt_dynamic_primary.call_count, 2)
|
|
|
|
@patch('src.investment.dynamic_funcs.prompt_calls.prompt_dynamic_primary')
|
|
def test_dynamic_primary_with_multiple_header_answers(self, mock_prompt_dynamic_primary):
|
|
"""Test dynamic_primary function with multiple answers in the header."""
|
|
# Setup
|
|
dynamic_fields = self.test_fieldset
|
|
reimbursement_level_fields = FieldSet()
|
|
exhibit_level_answers = {}
|
|
|
|
# Configure mocks to return multiple values for first field
|
|
mock_prompt_dynamic_primary.side_effect = ["Value 1|Value 2|Value 3", "Another Answer"]
|
|
|
|
# Execute
|
|
exhibit_level_answers, updated_reimbursement_fields = dynamic_funcs.dynamic_primary(
|
|
self.exhibit_text, self.exhibit_header, exhibit_level_answers,
|
|
self.constants, self.filename, dynamic_fields, reimbursement_level_fields
|
|
)
|
|
|
|
# Assert
|
|
self.assertEqual(len(exhibit_level_answers), 1) # One field should be added to exhibit level
|
|
self.assertEqual(exhibit_level_answers["TEST_FIELD2"], "Another Answer")
|
|
self.assertEqual(len(updated_reimbursement_fields.fields), 1) # One field should be added to reimbursement level
|
|
self.assertEqual(updated_reimbursement_fields.fields[0].field_name, "TEST_FIELD")
|
|
self.assertEqual(updated_reimbursement_fields.fields[0].valid_values, "Value 1|Value 2|Value 3")
|
|
self.assertTrue("Ensure ALL values that apply" in updated_reimbursement_fields.fields[0].prompt)
|
|
|
|
@patch('src.investment.prompt_calls.prompt_dynamic')
|
|
def test_dynamic_with_real_constants(self, mock_prompt_dynamic):
|
|
"""Test dynamic function with real Constants object."""
|
|
# Setup
|
|
dynamic_fields = self.test_fieldset
|
|
reimbursement_level_fields = FieldSet()
|
|
exhibit_level_answers = {}
|
|
|
|
# Configure mocks
|
|
mock_prompt_dynamic.side_effect = [
|
|
{"TEST_FIELD": "Header Answer", "TEST_FIELD2": ""}, # Header results
|
|
{"TEST_FIELD2": "Text Answer"} # Text results (only for remaining field)
|
|
]
|
|
|
|
# Execute
|
|
exhibit_level_answers, updated_reimbursement_fields = dynamic_funcs.dynamic(
|
|
self.exhibit_text, self.exhibit_header, exhibit_level_answers,
|
|
self.constants, self.filename, dynamic_fields, reimbursement_level_fields
|
|
)
|
|
|
|
# Assert
|
|
self.assertEqual(len(updated_reimbursement_fields.fields), 2) # Both fields should be added to reimbursement level
|
|
|
|
# Check that the first field was processed from header results
|
|
self.assertEqual(updated_reimbursement_fields.fields[0].field_name, "TEST_FIELD")
|
|
self.assertEqual(updated_reimbursement_fields.fields[0].valid_values, "Header Answer")
|
|
|
|
# Check that the second field was processed from text results
|
|
self.assertEqual(updated_reimbursement_fields.fields[1].field_name, "TEST_FIELD2")
|
|
self.assertEqual(updated_reimbursement_fields.fields[1].valid_values, "Text Answer")
|
|
|
|
# Verify that prompt_dynamic was called twice
|
|
self.assertEqual(mock_prompt_dynamic.call_count, 2)
|
|
|
|
|
|
@patch('src.investment.prompt_calls.prompt_dynamic')
|
|
def test_dynamic_with_no_fields(self, mock_prompt_dynamic):
|
|
"""Test dynamic function with an empty fieldset."""
|
|
# Setup
|
|
empty_fieldset = FieldSet()
|
|
reimbursement_level_fields = FieldSet()
|
|
exhibit_level_answers = {}
|
|
|
|
# Execute
|
|
exhibit_level_answers, updated_reimbursement_fields = dynamic_funcs.dynamic(
|
|
self.exhibit_text, self.exhibit_header, exhibit_level_answers,
|
|
self.constants, self.filename, empty_fieldset, reimbursement_level_fields
|
|
)
|
|
|
|
# Assert
|
|
self.assertEqual(len(exhibit_level_answers), 0)
|
|
self.assertEqual(len(updated_reimbursement_fields.fields), 0)
|
|
mock_prompt_dynamic.assert_not_called()
|
|
|
|
@patch('src.investment.prompt_calls.prompt_dynamic')
|
|
def test_dynamic_with_no_answers(self, mock_prompt_dynamic):
|
|
"""Test dynamic function when neither header nor text contains answers."""
|
|
# Setup
|
|
dynamic_fields = self.test_fieldset
|
|
reimbursement_level_fields = FieldSet()
|
|
exhibit_level_answers = {}
|
|
|
|
# Configure mocks - no answers anywhere
|
|
mock_prompt_dynamic.side_effect = [
|
|
{"TEST_FIELD": "", "TEST_FIELD2": ""}, # Header results
|
|
{"TEST_FIELD": "", "TEST_FIELD2": ""} # Text results
|
|
]
|
|
|
|
# Execute
|
|
exhibit_level_answers, updated_reimbursement_fields = dynamic_funcs.dynamic(
|
|
self.exhibit_text, self.exhibit_header, exhibit_level_answers,
|
|
self.constants, self.filename, dynamic_fields, reimbursement_level_fields
|
|
)
|
|
|
|
# Assert
|
|
self.assertEqual(len(updated_reimbursement_fields.fields), 0) # No fields should be added to reimbursement level
|
|
self.assertEqual(exhibit_level_answers, {"TEST_FIELD": "", "TEST_FIELD2": ""}) # Both fields should be added to exhibit level
|
|
self.assertEqual(mock_prompt_dynamic.call_count, 2)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main() |