Alex Galarce
7ca723037a
Bugfix/too many files * Implement caching for field data loading to optimize performance and prevent "too many open files" errors * Documentation * isort * fix breaking test to clear cache and ensure mock data is actually being used * Merge remote-tracking branch 'origin/main' into bugfix/too-many-files * Implement thread-safe vectorstore caching and cleanup in field context extraction * Refactor imports in smart_chunking_funcs.py for clarity and organization * Update module docstring in smart_chunking_funcs.py * Fix mypy errors * Fix variable naming and error handling in run_smart_chunked_fields function * Reorganize imports and enhance module documentation for clarity * Add logging for error handling in field_context function Approved-by: Katon Minhas
449 lines
17 KiB
Python
449 lines
17 KiB
Python
import json
|
|
import os
|
|
import unittest
|
|
from unittest.mock import patch, mock_open
|
|
|
|
from src.prompts.fieldset import Field, FieldSet, _field_data_cache, _cache_lock
|
|
from constants.constants import Constants
|
|
|
|
|
|
class TestFieldWithRealConstants(unittest.TestCase):
|
|
"""Tests for the Field class using real Constants values."""
|
|
|
|
def setUp(self):
|
|
"""Set up test fixtures before each test method."""
|
|
# Create test field with a real VALID_BILL_TYPE placeholder
|
|
self.test_field_dict = {
|
|
"field_name": "TEST_FIELD",
|
|
"relationship": "one-to-one",
|
|
"field_type": "string",
|
|
"prompt": "What is the test field? Valid values: {valid_values}",
|
|
"valid_values": "VALID_BILL_TYPE",
|
|
"base_field": "base_test_field",
|
|
"crosswalk": "test_crosswalk",
|
|
"keywords": ["test", "field"],
|
|
"methodology": "LLM",
|
|
"case_sensitive": True,
|
|
"regex_pattern": r"\d+",
|
|
"format": "text",
|
|
"allow_na": False,
|
|
}
|
|
self.field = Field(self.test_field_dict)
|
|
|
|
# Create a real Constants object - this will load actual mappings
|
|
self.constants = Constants()
|
|
|
|
def test_init(self):
|
|
"""Test Field initialization with all parameters."""
|
|
# Test all fields are correctly set
|
|
self.assertEqual(self.field.field_name, "TEST_FIELD")
|
|
self.assertEqual(self.field.relationship, "one-to-one")
|
|
self.assertEqual(self.field.field_type, "string")
|
|
self.assertEqual(
|
|
self.field.prompt, "What is the test field? Valid values: {valid_values}"
|
|
)
|
|
self.assertEqual(self.field.valid_values, "VALID_BILL_TYPE")
|
|
self.assertEqual(self.field.base_field, "base_test_field")
|
|
self.assertEqual(self.field.crosswalk, "test_crosswalk")
|
|
self.assertEqual(self.field.keywords, ["test", "field"])
|
|
self.assertEqual(self.field.methodology, "LLM")
|
|
self.assertTrue(self.field.case_sensitive)
|
|
self.assertEqual(self.field.regex_pattern, r"\d+")
|
|
self.assertEqual(self.field.format, "text")
|
|
self.assertFalse(self.field.allow_na)
|
|
self.assertIsNone(self.field.retrieved_text)
|
|
|
|
def test_init_minimal(self):
|
|
"""Test Field initialization with minimal parameters."""
|
|
minimal_field_dict = {
|
|
"field_name": "MINIMAL_FIELD",
|
|
"relationship": "one-to-one",
|
|
"field_type": "string",
|
|
}
|
|
field = Field(minimal_field_dict)
|
|
|
|
# Check required fields
|
|
self.assertEqual(field.field_name, "MINIMAL_FIELD")
|
|
self.assertEqual(field.relationship, "one-to-one")
|
|
self.assertEqual(field.field_type, "string")
|
|
|
|
# Check defaults for optional fields
|
|
self.assertIsNone(field.prompt)
|
|
self.assertIsNone(field.valid_values)
|
|
self.assertIsNone(field.base_field)
|
|
self.assertIsNone(field.crosswalk)
|
|
self.assertIsNone(field.keywords)
|
|
self.assertIsNone(field.methodology)
|
|
self.assertFalse(field.case_sensitive)
|
|
self.assertIsNone(field.regex_pattern)
|
|
self.assertIsNone(field.format)
|
|
self.assertTrue(field.allow_na)
|
|
|
|
def test_repr(self):
|
|
"""Test the string representation of a Field object."""
|
|
self.assertEqual(repr(self.field), "TEST_FIELD")
|
|
|
|
@patch(
|
|
"builtins.open",
|
|
new_callable=mock_open,
|
|
read_data=json.dumps(
|
|
[
|
|
{
|
|
"field_name": "FIELD1",
|
|
"relationship": "one-to-one",
|
|
"field_type": "string",
|
|
},
|
|
{
|
|
"field_name": "FIELD2",
|
|
"relationship": "one-to-many",
|
|
"field_type": "number",
|
|
},
|
|
]
|
|
),
|
|
)
|
|
def test_load_from_file(self, mock_file):
|
|
"""Test loading a Field object from a file."""
|
|
field = Field.load_from_file("dummy_path.json", "FIELD2")
|
|
self.assertEqual(field.field_name, "FIELD2")
|
|
self.assertEqual(field.relationship, "one-to-many")
|
|
self.assertEqual(field.field_type, "number")
|
|
|
|
def test_from_values(self):
|
|
"""Test creating a Field using the from_values class method."""
|
|
field = Field.from_values(
|
|
field_name="TEST_FIELD2",
|
|
relationship="one-to-many",
|
|
field_type="number",
|
|
prompt="Test prompt",
|
|
)
|
|
self.assertEqual(field.field_name, "TEST_FIELD2")
|
|
self.assertEqual(field.relationship, "one-to-many")
|
|
self.assertEqual(field.field_type, "number")
|
|
self.assertEqual(field.prompt, "Test prompt")
|
|
|
|
def test_to_dict(self):
|
|
"""Test converting a Field to a dictionary."""
|
|
field_dict = self.field.to_dict()
|
|
self.assertEqual(field_dict["field_name"], "TEST_FIELD")
|
|
self.assertEqual(field_dict["relationship"], "one-to-one")
|
|
self.assertEqual(field_dict["field_type"], "string")
|
|
self.assertEqual(
|
|
field_dict["prompt"], "What is the test field? Valid values: {valid_values}"
|
|
)
|
|
self.assertEqual(field_dict["valid_values"], "VALID_BILL_TYPE")
|
|
self.assertEqual(field_dict["base_field"], "base_test_field")
|
|
self.assertEqual(field_dict["keywords"], ["test", "field"])
|
|
self.assertEqual(field_dict["methodology"], "LLM")
|
|
self.assertTrue(field_dict["case_sensitive"])
|
|
self.assertEqual(field_dict["regex_pattern"], r"\d+")
|
|
self.assertEqual(field_dict["format"], "text")
|
|
|
|
@patch("builtins.print")
|
|
def test_print(self, mock_print):
|
|
"""Test printing a Field object."""
|
|
self.field.print()
|
|
# Verify print was called the expected number of times
|
|
self.assertEqual(mock_print.call_count, 9)
|
|
|
|
def test_get_prompt_dict_with_real_constants(self):
|
|
"""Test getting a prompt dictionary from a Field with real Constants."""
|
|
prompt_dict = self.field.get_prompt_dict(self.constants)
|
|
|
|
# Get the actual list of bill types from Constants
|
|
bill_types = self.constants.VALID_BILL_TYPE
|
|
|
|
# Verify the prompt contains the actual bill types from Constants
|
|
prompt = prompt_dict["TEST_FIELD"]
|
|
self.assertTrue(prompt.startswith("What is the test field? Valid values: "))
|
|
|
|
# Check if the output contains the real values
|
|
for bill_type in bill_types[:3]: # Check first few to avoid overly long test
|
|
self.assertIn(bill_type, prompt)
|
|
|
|
def test_get_prompt_no_valid_values(self):
|
|
"""Test getting a prompt with no valid_values to resolve."""
|
|
field = Field(
|
|
{
|
|
"field_name": "SIMPLE_FIELD",
|
|
"relationship": "one-to-one",
|
|
"field_type": "string",
|
|
"prompt": "What is the simple field?",
|
|
}
|
|
)
|
|
prompt = field.get_prompt()
|
|
self.assertEqual(prompt, "What is the simple field?")
|
|
|
|
def test_get_prompt_with_real_constants(self):
|
|
"""Test resolving valid values in a prompt using real Constants."""
|
|
prompt = self.field.get_prompt(self.constants)
|
|
|
|
# Check that the prompt includes the actual bill types
|
|
self.assertTrue(prompt.startswith("What is the test field? Valid values: "))
|
|
|
|
# Verify that the real constants values are used
|
|
bill_types = self.constants.VALID_BILL_TYPE
|
|
bill_types_str = str(bill_types)
|
|
self.assertIn(bill_types_str, prompt)
|
|
|
|
def test_update_valid_values(self):
|
|
"""Test updating valid values for a Field."""
|
|
self.field.update_valid_values("NEW_VALID_VALUES")
|
|
self.assertEqual(self.field.valid_values, "NEW_VALID_VALUES")
|
|
|
|
|
|
class TestFieldSetWithRealConstants(unittest.TestCase):
|
|
"""Tests for the FieldSet class using real Constants values."""
|
|
|
|
def setUp(self):
|
|
"""Set up test fixtures before each test method."""
|
|
# Create fields with real Constants placeholders
|
|
self.test_fields = [
|
|
Field(
|
|
{
|
|
"field_name": "FIELD1",
|
|
"relationship": "one-to-one",
|
|
"field_type": "string",
|
|
"prompt": "What is field 1?",
|
|
}
|
|
),
|
|
Field(
|
|
{
|
|
"field_name": "FIELD2",
|
|
"relationship": "one-to-many",
|
|
"field_type": "number",
|
|
"prompt": "What is field 2? Valid values: {valid_values}",
|
|
"valid_values": "VALID_BILL_TYPE",
|
|
}
|
|
),
|
|
Field(
|
|
{
|
|
"field_name": "FIELD3",
|
|
"relationship": "one-to-one",
|
|
"field_type": "string",
|
|
"prompt": "What is field 3? Valid values: {valid_values}",
|
|
"valid_values": "VALID_LOBS",
|
|
}
|
|
),
|
|
]
|
|
self.fieldset = FieldSet(fields=self.test_fields)
|
|
|
|
# Create a real Constants object
|
|
self.constants = Constants()
|
|
|
|
@patch(
|
|
"builtins.open",
|
|
new_callable=mock_open,
|
|
read_data=json.dumps(
|
|
[
|
|
{
|
|
"field_name": "FIELD1",
|
|
"relationship": "one-to-one",
|
|
"field_type": "string",
|
|
},
|
|
{
|
|
"field_name": "FIELD2",
|
|
"relationship": "one-to-many",
|
|
"field_type": "number",
|
|
},
|
|
{
|
|
"field_name": "FIELD3",
|
|
"relationship": "one-to-one",
|
|
"field_type": "boolean",
|
|
},
|
|
]
|
|
),
|
|
)
|
|
def test_init_with_file(self, mock_file):
|
|
"""Test FieldSet initialization with a file."""
|
|
|
|
# Clear the cache so that our mock will be used
|
|
with _cache_lock:
|
|
_field_data_cache.clear()
|
|
fieldset = FieldSet(file_path="dummy_path.json")
|
|
self.assertEqual(len(fieldset.fields), 3)
|
|
self.assertEqual(fieldset.fields[0].field_name, "FIELD1")
|
|
self.assertEqual(fieldset.fields[1].field_name, "FIELD2")
|
|
self.assertEqual(fieldset.fields[2].field_name, "FIELD3")
|
|
|
|
def test_add_field(self):
|
|
"""Test adding a field to a FieldSet."""
|
|
new_field = Field(
|
|
{
|
|
"field_name": "FIELD4",
|
|
"relationship": "one-to-one",
|
|
"field_type": "string",
|
|
}
|
|
)
|
|
self.fieldset.add_field(new_field)
|
|
self.assertEqual(len(self.fieldset.fields), 4)
|
|
self.assertEqual(self.fieldset.fields[3].field_name, "FIELD4")
|
|
|
|
# Test adding a duplicate field (should not add)
|
|
duplicate_field = Field(
|
|
{
|
|
"field_name": "FIELD1",
|
|
"relationship": "one-to-one",
|
|
"field_type": "string",
|
|
}
|
|
)
|
|
self.fieldset.add_field(duplicate_field)
|
|
self.assertEqual(len(self.fieldset.fields), 4) # No change in length
|
|
|
|
def test_get_field(self):
|
|
"""Test getting a field by name."""
|
|
field = self.fieldset.get_field("FIELD2")
|
|
self.assertEqual(field.field_name, "FIELD2")
|
|
self.assertEqual(field.relationship, "one-to-many")
|
|
|
|
# Test getting a field that doesn't exist
|
|
with self.assertRaises(ValueError):
|
|
self.fieldset.get_field("NONEXISTENT_FIELD")
|
|
|
|
def test_list_fields(self):
|
|
"""Test listing all field names."""
|
|
field_names = self.fieldset.list_fields()
|
|
self.assertEqual(field_names, ["FIELD1", "FIELD2", "FIELD3"])
|
|
|
|
def test_remove_field(self):
|
|
"""Test removing a field."""
|
|
self.fieldset.remove_field("FIELD2")
|
|
self.assertEqual(len(self.fieldset.fields), 2)
|
|
self.assertEqual(self.fieldset.fields[0].field_name, "FIELD1")
|
|
self.assertEqual(self.fieldset.fields[1].field_name, "FIELD3")
|
|
|
|
def test_update_field_attr(self):
|
|
"""Test updating a field attribute."""
|
|
self.fieldset.update_field_attr("FIELD1", "prompt", "Updated prompt")
|
|
self.assertEqual(self.fieldset.get_field("FIELD1").prompt, "Updated prompt")
|
|
|
|
def test_get_prompt_dict_with_real_constants(self):
|
|
"""Test getting prompt dictionaries with real Constants values."""
|
|
prompt_dict = self.fieldset.get_prompt_dict(self.constants)
|
|
self.assertEqual(len(prompt_dict), 3)
|
|
self.assertEqual(prompt_dict["FIELD1"], "What is field 1?")
|
|
|
|
# Check that FIELD2 has the actual bill types
|
|
self.assertTrue(
|
|
prompt_dict["FIELD2"].startswith("What is field 2? Valid values: ")
|
|
)
|
|
for bill_type in self.constants.VALID_BILL_TYPE[:3]: # Check first few
|
|
self.assertIn(bill_type, str(prompt_dict["FIELD2"]))
|
|
|
|
# Check that FIELD3 has the actual LOBs
|
|
self.assertTrue(
|
|
prompt_dict["FIELD3"].startswith("What is field 3? Valid values: ")
|
|
)
|
|
for lob in self.constants.VALID_LOBS[:3]: # Check first few
|
|
self.assertIn(lob, str(prompt_dict["FIELD3"]))
|
|
|
|
# Test filtering by field_type
|
|
prompt_dict = self.fieldset.get_prompt_dict(self.constants, field_type="string")
|
|
self.assertEqual(len(prompt_dict), 2)
|
|
self.assertEqual(prompt_dict["FIELD1"], "What is field 1?")
|
|
self.assertIn("FIELD3", prompt_dict)
|
|
|
|
# Test filtering by relationship
|
|
prompt_dict = self.fieldset.get_prompt_dict(
|
|
self.constants, relationship="one-to-many"
|
|
)
|
|
self.assertEqual(len(prompt_dict), 1)
|
|
self.assertIn("FIELD2", prompt_dict)
|
|
|
|
def test_print_prompt_dict_with_real_constants(self):
|
|
"""Test printing all prompts as a string with real Constants."""
|
|
prompt_str = self.fieldset.print_prompt_dict(self.constants)
|
|
|
|
# Verify each field is in the output
|
|
self.assertIn("FIELD1 : What is field 1?", prompt_str)
|
|
self.assertIn("FIELD2 : What is field 2? Valid values: ", prompt_str)
|
|
self.assertIn("FIELD3 : What is field 3? Valid values: ", prompt_str)
|
|
|
|
# Verify real constants values are used
|
|
for bill_type in self.constants.VALID_BILL_TYPE[:2]: # Check first couple
|
|
self.assertIn(bill_type, prompt_str)
|
|
|
|
for lob in self.constants.VALID_LOBS[:2]: # Check first couple
|
|
self.assertIn(lob, prompt_str)
|
|
|
|
@patch("builtins.print")
|
|
def test_print(self, mock_print):
|
|
"""Test printing all fields."""
|
|
self.fieldset.print()
|
|
# Each field calls print multiple times
|
|
self.assertTrue(mock_print.call_count > 20)
|
|
|
|
def test_repr(self):
|
|
"""Test the string representation of a FieldSet."""
|
|
self.assertEqual(repr(self.fieldset), "['FIELD1', 'FIELD2', 'FIELD3']")
|
|
|
|
# Test empty FieldSet
|
|
empty_fieldset = FieldSet()
|
|
self.assertEqual(repr(empty_fieldset), "FieldSet: No fields loaded")
|
|
|
|
@patch("builtins.print")
|
|
def test_print_fields(self, mock_print):
|
|
"""Test printing field names."""
|
|
self.fieldset.print_fields()
|
|
mock_print.assert_called_with(["FIELD1", "FIELD2", "FIELD3"])
|
|
|
|
def test_filter(self):
|
|
"""Test filtering fields based on criteria."""
|
|
# Filter by relationship
|
|
filtered = self.fieldset.filter(relationship="one-to-one")
|
|
self.assertEqual(len(filtered.fields), 2)
|
|
self.assertEqual(filtered.fields[0].field_name, "FIELD1")
|
|
self.assertEqual(filtered.fields[1].field_name, "FIELD3")
|
|
|
|
# Filter by field_type
|
|
filtered = self.fieldset.filter(field_type="number")
|
|
self.assertEqual(len(filtered.fields), 1)
|
|
self.assertEqual(filtered.fields[0].field_name, "FIELD2")
|
|
|
|
def test_contains_fields(self):
|
|
"""Test checking if FieldSet contains fields."""
|
|
self.assertTrue(self.fieldset.contains_fields())
|
|
|
|
# Test empty FieldSet
|
|
empty_fieldset = FieldSet()
|
|
self.assertFalse(empty_fieldset.contains_fields())
|
|
|
|
def test_combine(self):
|
|
"""Test combining FieldSets."""
|
|
other_fields = [
|
|
Field(
|
|
{
|
|
"field_name": "FIELD4",
|
|
"relationship": "one-to-one",
|
|
"field_type": "string",
|
|
}
|
|
),
|
|
Field(
|
|
{
|
|
"field_name": "FIELD1", # Duplicate, should not be added
|
|
"relationship": "one-to-one",
|
|
"field_type": "string",
|
|
}
|
|
),
|
|
]
|
|
other_fieldset = FieldSet(fields=other_fields)
|
|
|
|
# Test combine with new FieldSet (not in-place)
|
|
combined = self.fieldset.combine(other_fieldset)
|
|
self.assertEqual(len(combined.fields), 4)
|
|
self.assertEqual(combined.fields[0].field_name, "FIELD1")
|
|
self.assertEqual(combined.fields[1].field_name, "FIELD2")
|
|
self.assertEqual(combined.fields[2].field_name, "FIELD3")
|
|
self.assertEqual(combined.fields[3].field_name, "FIELD4")
|
|
|
|
# Original should be unchanged
|
|
self.assertEqual(len(self.fieldset.fields), 3)
|
|
|
|
# Test combine in-place
|
|
self.fieldset.combine(other_fieldset, inplace=True)
|
|
self.assertEqual(len(self.fieldset.fields), 4)
|
|
self.assertEqual(self.fieldset.fields[3].field_name, "FIELD4")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|