Merged in task/unit-tests-investment-and-client (pull request #346)

Task/unit tests investment and client

* move io_utils.py

* add __init__.py

* refactor: move string_funcs and update import paths to use utils namespace

* rename string_funcs to string_utils

* move and rename table_utils.py

* refactor: rename claude_funcs to llm_utils in multiple files

* update unit tests

* added unit tests for src/string_utils.py

* added unit tests for src/table_utils.py

* added unit tests for src/llm_utils.py

* added unit tests for src/io_utils.py

* Merge branch 'main' into task/unit-tests-investment-and-client; TODO: fix unit tests

* Unit tests to fix

* fix unit tests

* fix invalid references


Approved-by: Alex Galarce

fieldExtraction/pyproject.toml

@@ -16,6 +16,7 @@ psutil = "^6.1.0"
 rapidfuzz = "^3.10.1"
 pyxlsb = "^1.0.10"
 openpyxl = "^3.1.5"
+pytest-mock = "^3.14.0"
 
 [tool.poetry.group.dev.dependencies]
 black = "^24.10.0"

fieldExtraction/src/consolidate_output.py

@@ -1,12 +1,10 @@
 import os
 from pathlib import Path
-
 import pandas as pd
 
-import src.constants.valid as valid
-import src.tracking.costlog_utils as costlog_utils
-import src.tracking.tracking_utils as tracking_utils
 from src import config
+from src.tracking import costlog_utils, tracking_utils
+from src.constants import valid
 from src.qa_qc.qa_qc_utils import perform_qc_qa
 
 

fieldExtraction/src/utils/io_utils.py

@@ -137,7 +137,7 @@ def filter_already_processed(input_dict, input_dict_b): # io_utils.py
     """Filter out already processed files based on either local files or S3 master tracking"""
     if config.WRITE_TO_S3:
         # Get processed files from S3 master tracking for this specific batch
-        already_processed_ac, already_processed_b = tracking_utils.get_already_processed_files(config.BATCH_ID)
+        already_processed_ac, already_processed_b = tracking_utils.get_processed_files(config.BATCH_ID)
 
         # Keep files that still need either AC or B processing
         input_dict_ac = {

fieldExtraction/src/utils/string_utils.py

@@ -69,8 +69,20 @@ def extract_text_from_delimiters(
 
     return matches[match_index]
 
+def json_parsing_search(response_text: str, field_list: list[str]) -> dict[str, str]:
+    """
+    Parses a JSON-like string to extract values for fields.
 
-def json_parsing_search(response_text, field_list):
+    This function attempts to clean and format the input string to resemble a valid JSON structure.
+    It then searches for the specified fields within the string and extracts their corresponding values.
+
+    Parameters:
+        response_text (str): The JSON-like string to parse.
+        field_list (list[str]): A list of field names to search for in the response_text.
+
+    Returns:
+        dict[str, str]: A dictionary where keys are field names and values are the extracted values from the response_text.
+    """
     try:
         if response_text.split("{", 1)[1].strip()[0] == '"':
             response_text = "{" + response_text.split("{", 1)[1]
@@ -86,14 +98,18 @@ def json_parsing_search(response_text, field_list):
             response_text = response_text.rsplit("}", 1)[0]
     else:
         response_text = response_text.rstrip(",") + "}"
+    
     field_l = []
     answer_l = []
     position_dict = {}
+    
     for f in field_list:
         location = response_text.find('"' + f + '"')
         if location != -1:
             position_dict[location] = f
+    
     field_list = list(dict(sorted(position_dict.items())).values())
+    
     for f in field_list:
         if f in response_text:
             field_l.append(f)
@@ -135,7 +151,7 @@ def json_parsing_search(response_text, field_list):
     return dict(zip(field_l, answer_l))
 
 
-def secondary_string_to_dict(dict_string, filename):
+def secondary_string_to_dict(dict_string: str, filename: str) -> dict:
     """
     Converts a string representation of a dictionary into an actual dictionary object, handling potential formatting issues.
 
@@ -145,6 +161,7 @@ def secondary_string_to_dict(dict_string, filename):
 
     Parameters:
         dict_string (str): The string containing the dictionary-like content, potentially surrounded by extra text or characters.
+        filename (str): The filename associated with these entries, used to tag each resulting dictionary.
 
     Returns:
         dict: The dictionary obtained from parsing the cleaned and corrected string.

fieldExtraction/tests/io_utils_test.py

@@ -0,0 +1,265 @@
+import pytest
+from io import StringIO
+import pandas as pd
+
+from src.utils.io_utils import (
+    read_local,
+    list_s3_files,
+    read_s3,
+    read_s3_csv,
+    read_input,
+    read_input_csv,
+    filter_already_processed,
+    remove_txt_extension,
+    read_xlsb,
+)
+
+class TestIOUtils:
+    # Fixtures for reusable test data
+    @pytest.fixture
+    def mock_csv_data(self):
+        return "col1,col2\n1,2\n3,4"
+
+    @pytest.fixture
+    def mock_txt_data(self):
+        return "This is a test file."
+
+    @pytest.fixture
+    def mock_s3_response(self):
+        return {
+            "Contents": [
+                {"Key": "file1.txt"},
+                {"Key": "file2.csv"},
+                {"Key": "file3.xlsx"},
+            ]
+        }
+
+    @pytest.fixture
+    def mock_s3_csv_response(self):
+        return {
+            "Body": StringIO("col1,col2\n1,2\n3,4")
+        }
+
+    # Tests for read_local
+    def test_read_local_txt_utf8(self, mock_txt_data, mocker):
+        mocker.patch("builtins.open", mocker.mock_open(read_data=mock_txt_data))
+        mocker.patch("os.path.isfile", return_value=True)
+        result = read_local("test.txt")
+        assert result == mock_txt_data
+
+    def test_read_local_csv(self, mock_csv_data, mocker):
+        mock_df = pd.read_csv(StringIO(mock_csv_data))
+        mocker.patch("pandas.read_csv", return_value=mock_df)
+        mocker.patch("os.path.isfile", return_value=True)
+        result = read_local("test.csv")
+        pd.testing.assert_frame_equal(result, mock_df)
+
+    def test_read_local_nonexistent_file(self, mocker):
+        mocker.patch("os.path.isfile", return_value=False)
+        result = read_local("nonexistent.txt")
+        assert result is None
+
+    def test_list_s3_files(self, mocker):
+        # Test case 1: Empty S3 bucket
+        mocker.patch("src.config.S3_CLIENT.list_objects_v2", return_value={"Contents": []})
+        result = list_s3_files()
+        assert result == []
+
+        # Test case 2: S3 bucket with dummy files
+        mock_s3_response = {
+            "Contents": [
+                {"Key": "file1.txt"},
+                {"Key": "file2.csv"},
+                {"Key": "file3.xlsx"},
+            ]
+        }
+        mocker.patch("src.config.S3_CLIENT.list_objects_v2", return_value=mock_s3_response)
+        result = list_s3_files()
+        assert result == ["file1.txt", "file2.csv", "file3.xlsx"]
+
+    # Tests for read_s3
+    def test_read_s3(self, mock_txt_data, mocker):
+        mock_s3_client = mocker.MagicMock()
+        mock_s3_client.list_objects_v2.return_value = {
+            "Contents": [{"Key": "file1.txt"}],
+            "IsTruncated": False,
+        }
+        mock_s3_client.get_object.return_value = {
+            "Body": mocker.MagicMock(read=mocker.MagicMock(return_value=mock_txt_data.encode("utf-8")))
+        }
+        mocker.patch("src.config.S3_CLIENT", mock_s3_client)
+        result = read_s3()
+        assert result == {"file1.txt": mock_txt_data}
+
+    # Tests for read_s3_csv
+    def test_read_s3_csv(self, mock_csv_data, mocker):
+        mock_s3_client = mocker.MagicMock()
+        mock_s3_client.list_objects_v2.return_value = {
+            "Contents": [{"Key": "file2.csv"}],
+        }
+        # Mock get_object to return a response with a bytes object
+        mock_s3_client.get_object.return_value = {
+            "Body": mocker.MagicMock(read=mocker.MagicMock(return_value=mock_csv_data.encode("utf-8")))
+        }
+        mocker.patch("src.config.S3_CLIENT", mock_s3_client)
+        result = read_s3_csv("s3://bucket/file2.csv")
+        expected_df = pd.read_csv(StringIO(mock_csv_data))
+        pd.testing.assert_frame_equal(result, expected_df)
+
+        # Verify the S3 client was called correctly
+        mock_s3_client.list_objects_v2.assert_called_once_with(Bucket="bucket", Prefix="file2.csv")
+        mock_s3_client.get_object.assert_called_once_with(Bucket="bucket", Key="file2.csv")
+        
+    # Tests for read_input    
+    def test_read_input_local(self, mock_txt_data, mocker):
+        mocker.patch("os.listdir", return_value=["file1.txt"])
+        mocker.patch("src.utils.io_utils.read_local", return_value=mock_txt_data)
+        mocker.patch("src.config.READ_MODE", "local")
+        result = read_input()
+        assert result == {"file1.txt": mock_txt_data}
+
+    def test_read_input_s3(self, mock_txt_data, mocker):
+        mocker.patch("src.utils.io_utils.read_s3", return_value={"file1.txt": mock_txt_data})
+        mocker.patch("src.config.READ_MODE", "s3")
+        result = read_input()
+        assert result == {"file1.txt": mock_txt_data}
+
+    # Tests for read_input_csv
+    def test_read_input_csv_local(self, mock_csv_data, mocker):
+        mock_df = pd.read_csv(StringIO(mock_csv_data))
+        mock_read_local = mocker.patch("src.utils.io_utils.read_local", return_value=mock_df)
+        mocker.patch("src.config.DF_READ_MODE", "local")
+        mocker.patch("src.config.AC_DF", "ac.csv")
+        mocker.patch("src.config.B_DF", "b.csv")
+        
+        ac_df, b_df = read_input_csv()
+        
+        pd.testing.assert_frame_equal(ac_df, mock_df)
+        pd.testing.assert_frame_equal(b_df, mock_df)
+        
+        # Verify read_local was called with the correct paths
+        mock_read_local.assert_any_call("ac.csv")
+        mock_read_local.assert_any_call("b.csv")
+
+    def test_read_input_csv_s3(self, mock_csv_data, mocker):
+        # Create a mock DataFrame
+        mock_df = pd.read_csv(StringIO(mock_csv_data))
+
+        # Mock read_s3_csv to return the mock DataFrame
+        mock_read_s3_csv = mocker.patch("src.utils.io_utils.read_s3_csv", return_value=mock_df)
+
+        # Set the src.config to use S3 mode
+        mocker.patch("src.config.DF_READ_MODE", "s3")
+        mocker.patch("src.config.AC_DF", "s3://bucket/ac.csv")
+        mocker.patch("src.config.B_DF", "s3://bucket/b.csv")
+
+        # Call the function under test
+        ac_df, b_df = read_input_csv()
+
+        # Verify the results
+        pd.testing.assert_frame_equal(ac_df, mock_df)
+        pd.testing.assert_frame_equal(b_df, mock_df)
+
+        # Verify read_s3_csv was called with the correct S3 paths
+        mock_read_s3_csv.assert_any_call("s3://bucket/ac.csv")
+        mock_read_s3_csv.assert_any_call("s3://bucket/b.csv")
+
+    def test_filter_already_processed_s3(self, mocker):
+        # Case 1: WRITE_TO_S3 is True (S3 mode)
+        # Mock src.config.BATCH_ID
+        mock_batch_id = "mock_batch_123"
+        mocker.patch("src.config.BATCH_ID", mock_batch_id)
+        mock_get_processed_files = mocker.patch("src.tracking.tracking_utils.get_processed_files", return_value=(set(["file1.txt"]), set([])))
+        mock_write_to_s3 = mocker.patch("src.config.WRITE_TO_S3", True)
+        input_dict = {"file1.txt": "content1", "file2.txt": "content2"}
+        input_dict_b = {"file1.txt": "content1", "file2.txt": "content2"}
+        result_ac, result_b = filter_already_processed(input_dict, input_dict_b)
+        assert result_ac == {"file2.txt": "content2"}
+        assert result_b == {"file1.txt": "content1", "file2.txt": "content2"}
+
+        # Verify src.tracking.tracking_utils.get_processed_files was called
+        mock_get_processed_files.assert_called_once_with(mock_batch_id)
+    
+    def test_filter_already_processed_local(self, mocker):
+        # Case 2: WRITE_TO_S3 is False (local mode)
+        mocker.patch("src.config.WRITE_TO_S3", False)
+        mocker.patch("src.config.OUTPUT_DIRECTORY", "/mock/output/dir")
+        mocker.patch("src.config.AC_RESULTS_NAME", "ac_results.txt")
+        mocker.patch("src.config.B_RESULTS_NAME", "b_results.txt")
+
+        # Mock os.listdir to return a list of folders for the top-level directory
+        mocker.patch(
+            "os.listdir",
+            side_effect=lambda x: {
+                "/mock/output/dir": ["folder1", "folder2"],  # Top-level directory
+                "/mock/output/dir/folder1": ["ac_results.txt"],  # Folder 1 contents
+                "/mock/output/dir/folder2": ["b_results.txt"],  # Folder 2 contents
+            }[x],
+        )
+
+        # Mock os.path.isdir to return True for folders
+        mocker.patch("os.path.isdir", return_value=True)
+
+        # Mock os.path.join to construct paths
+        mocker.patch("os.path.join", side_effect=lambda *args: "/".join(args))
+
+        # Input dictionaries
+        input_dict = {"folder1.txt": "content1", "folder2.txt": "content2", "folder3.txt": "content3"}
+        input_dict_b = {"folder1.txt": "content1", "folder2.txt": "content2", "folder3.txt": "content3"}
+
+        # Call the function under test
+        result_ac, result_b = filter_already_processed(input_dict, input_dict_b)
+
+        # Verify the results
+        assert result_ac == {"folder2.txt": "content2", "folder3.txt": "content3"}  # folder1.txt is already processed
+        assert result_b == {"folder1.txt": "content1", "folder3.txt": "content3"}  # folder2.txt is already processed
+
+    # Tests for remove_txt_extension
+    def test_remove_txt_extension(self):
+        assert remove_txt_extension("file.txt") == "file"
+        assert remove_txt_extension("file.csv") == "file.csv"
+        assert remove_txt_extension(123) == 123
+        assert remove_txt_extension(None) == None
+
+    def test_read_xlsb(self,mocker):
+        # Mock the sheet
+        mock_sheet = mocker.MagicMock()
+        
+        # Mock the rows returned by the sheet
+        # Use a list of rows, where each row is a list of MagicMock objects with a 'v' attribute
+        mock_rows = [
+            [mocker.MagicMock(v='Header1'), mocker.MagicMock(v='Header2')],  # Header row
+            [mocker.MagicMock(v='Data1'), mocker.MagicMock(v='Data2')],      # Data row 1
+            [mocker.MagicMock(v='Data3'), mocker.MagicMock(v='Data4')],      # Data row 2
+        ]
+        
+        # Mock sheet.rows to return an iterable (e.g., a generator)
+        mock_sheet.rows.return_value = iter(mock_rows)  # Use return_value for the method
+
+        # Mock the workbook
+        mock_workbook = mocker.MagicMock()
+        # Ensure get_sheet() returns the mocked sheet
+        mock_workbook.get_sheet.return_value.__enter__.return_value = mock_sheet
+
+        # Mock open_workbook to return the mocked workbook
+        mock_open_workbook = mocker.patch('src.utils.io_utils.open_workbook')
+        mock_open_workbook.return_value.__enter__.return_value = mock_workbook
+
+        # Call the function
+        result = read_xlsb('dummy_path.xlsb')
+
+        # Expected DataFrame
+        expected_df = pd.DataFrame(
+            [
+                ['Data1', 'Data2'],  # Data row 1
+                ['Data3', 'Data4'],  # Data row 2
+            ],
+            columns=['Header1', 'Header2']  # Header row
+        )
+
+        # Assert that the result matches the expected DataFrame
+        pd.testing.assert_frame_equal(result, expected_df)
+
+        # Verify that the mock was called correctly
+        mock_open_workbook.assert_called_once_with('dummy_path.xlsb')
+        mock_workbook.get_sheet.assert_called_once()

fieldExtraction/tests/llm_utils_test.py

@@ -1,6 +1,166 @@
-from src.utils.llm_utils import count_tokens
+import pytest
+from src.utils.llm_utils import (
+    count_tokens,
+    create_cost_log_entry,
+    update_stats,
+    invoke_claude,
+)
+from src import config
 
+class TestLLMUtils:
+    # Fixture to reset global stats before each test
+    @pytest.fixture(autouse=True)
+    def reset_global_stats(self):
+        """Fixture to reset global stats before each test."""
+        config.MODEL_STATS = {
+            "file1.txt": {"requests": 0, "tokens_sent": 0, "tokens_received": 0, "total_time": 0},
+            "file2.txt": {"requests": 0, "tokens_sent": 0, "tokens_received": 0, "total_time": 0},
+            "file3.txt": {"requests": 0, "tokens_sent": 0, "tokens_received": 0, "total_time": 0},
+            "file4.txt": {"requests": 0, "tokens_sent": 0, "tokens_received": 0, "total_time": 0},
+        }
+        config.GLOBAL_STATS = {
+            "total_requests": 0,
+            "total_tokens_sent": 0,
+            "total_tokens_received": 0,
+            "total_time": 0,
+        }
 
-def test_count_tokens():
-    assert count_tokens("") == 0
-    assert count_tokens("a") == 1
+    # Test cases for count_tokens
+    @pytest.mark.parametrize("input_str, expected_tokens", [
+        ("hello world", 2),  # 11 characters / 6 = 1.83 → ceil to 2
+        ("", 0),             # Empty string
+        ("a" * 12, 2),       # 12 characters / 6 = 2
+        ("a" * 13, 3),       # 13 characters / 6 = 2.16 → ceil to 3
+        ("😊", 1),           # Unicode characters
+    ])
+    def test_count_tokens(self, input_str, expected_tokens):
+        """Test the count_tokens function."""
+        assert count_tokens(input_str) == expected_tokens
+
+    # Test cases for create_cost_log_entry
+    @pytest.mark.parametrize("filename, caller_name, prompt, response, input_cost_per_1k, output_cost_per_1k, expected_output", [
+        (
+            "file1.txt", "func1", "hello", "world", 0.1, 0.2,
+            {
+                "Filename": "file1.txt",
+                "Caller Function": "func1",
+                "Input Prompt Length": 5,
+                "Input Prompt Tokens": 1,  # 5 chars / 6 = 0.83 → ceil to 1
+                "Input Cost": 0.1 * 1 / 1000,
+                "Output Response Length": 5,
+                "Output Response Tokens": 1,  # 5 chars / 6 = 0.83 → ceil to 1
+                "Output Cost": 0.2 * 1 / 1000,
+                "Total Cost": (0.1 * 1 / 1000) + (0.2 * 1 / 1000),
+            }
+        ),
+        (
+            "file2.txt", "func2", "", "", 0.1, 0.2,
+            {
+                "Filename": "file2.txt",
+                "Caller Function": "func2",
+                "Input Prompt Length": 0,
+                "Input Prompt Tokens": 0,  # Empty string
+                "Input Cost": 0.1 * 0 / 1000,
+                "Output Response Length": 0,
+                "Output Response Tokens": 0,  # Empty string
+                "Output Cost": 0.2 * 0 / 1000,
+                "Total Cost": (0.1 * 0 / 1000) + (0.2 * 0 / 1000),
+            }
+        ),
+        (
+            "file3.txt", "func3", "a" * 12, "b" * 13, 0.1, 0.2,
+            {
+                "Filename": "file3.txt",
+                "Caller Function": "func3",
+                "Input Prompt Length": 12,
+                "Input Prompt Tokens": 2,  # 12 chars / 6 = 2
+                "Input Cost": 0.1 * 2 / 1000,
+                "Output Response Length": 13,
+                "Output Response Tokens": 3,  # 13 chars / 6 = 2.16 → ceil to 3
+                "Output Cost": 0.2 * 3 / 1000,
+                "Total Cost": (0.1 * 2 / 1000) + (0.2 * 3 / 1000),
+            }
+        ),
+        (
+            "file4.txt", "func4", "😊", "😊😊", 0.1, 0.2,
+            {
+                "Filename": "file4.txt",
+                "Caller Function": "func4",
+                "Input Prompt Length": 1,
+                "Input Prompt Tokens": 1,  # Unicode character
+                "Input Cost": 0.1 * 1 / 1000,
+                "Output Response Length": 2,
+                "Output Response Tokens": 1,  # Unicode characters
+                "Output Cost": 0.2 * 1 / 1000,
+                "Total Cost": (0.1 * 1 / 1000) + (0.2 * 1 / 1000),
+            }
+        ),
+    ])
+    
+    def test_create_cost_log_entry(self, filename, caller_name, prompt, response, input_cost_per_1k, output_cost_per_1k, expected_output):
+        """Test the create_cost_log_entry function."""
+        result = create_cost_log_entry(filename, caller_name, prompt, response, input_cost_per_1k, output_cost_per_1k)
+        assert result == expected_output
+
+    # Test cases for update_stats
+    def test_update_stats(self):
+        """Test the update_stats function."""
+        # Initialize the global state to zeros
+        config.MODEL_STATS = {"file1.txt": {"requests": 0, "tokens_sent": 0, "tokens_received": 0, "total_time": 0}}
+        config.GLOBAL_STATS = {"total_requests": 0, "total_tokens_sent": 0, "total_tokens_received": 0, "total_time": 0}
+
+        # Verify initial state
+        assert config.MODEL_STATS["file1.txt"]["requests"] == 0
+        assert config.GLOBAL_STATS["total_requests"] == 0
+
+        # Update stats
+        update_stats("file1.txt", 10, 20, 5.0, "claude2")
+
+        # Check updated state
+        assert config.MODEL_STATS["file1.txt"]["requests"] == 1
+        assert config.MODEL_STATS["file1.txt"]["tokens_sent"] == 10
+        assert config.MODEL_STATS["file1.txt"]["tokens_received"] == 20
+        assert config.MODEL_STATS["file1.txt"]["total_time"] == 5.0
+        assert config.GLOBAL_STATS["total_requests"] == 1
+        assert config.GLOBAL_STATS["total_tokens_sent"] == 10
+        assert config.GLOBAL_STATS["total_tokens_received"] == 20
+        assert config.GLOBAL_STATS["total_time"] == 5.0
+
+    # Test cases for invoke_claude
+    @pytest.mark.parametrize("run_mode, model_id, expected_response", [
+    ("local", config.MODEL_ID_CLAUDE2, "mocked response from claude 2 (local)"),
+    ("local", config.MODEL_ID_CLAUDE3_HAIKU, "mocked response from claude 3 (local)"),
+    ("local", config.MODEL_ID_CLAUDE35_SONNET, "mocked response from claude 3.5 (local)"),
+    ("ec2", config.MODEL_ID_CLAUDE2, "mocked response from claude 2 (ec2)"),
+    ("ec2", config.MODEL_ID_CLAUDE3_HAIKU, "mocked response from claude 3 (ec2)"),
+    ("ec2", config.MODEL_ID_CLAUDE35_SONNET, "mocked response from claude 3.5 (ec2)"),
+    ])
+    
+    def test_invoke_claude(self, run_mode, model_id, expected_response, mocker):
+        """Test the invoke_claude function."""
+        # Map model IDs to their corresponding mock functions
+        mock_functions = {
+            ("local", config.MODEL_ID_CLAUDE2): "src.utils.llm_utils.local_claude_2",
+            ("local", config.MODEL_ID_CLAUDE3_HAIKU): "src.utils.llm_utils.local_claude_3",
+            ("local", config.MODEL_ID_CLAUDE35_SONNET): "src.utils.llm_utils.local_claude_3",
+            ("ec2", config.MODEL_ID_CLAUDE2): "src.utils.llm_utils.ec2_claude2",
+            ("ec2", config.MODEL_ID_CLAUDE3_HAIKU): "src.utils.llm_utils.ec2_claude_3",
+            ("ec2", config.MODEL_ID_CLAUDE35_SONNET): "src.utils.llm_utils.ec2_claude_35",
+        }
+
+        # Mock the appropriate function based on run_mode and model_id
+        mock_function_path = mock_functions.get((run_mode, model_id))
+        if mock_function_path:
+            mocker.patch(mock_function_path, return_value=expected_response)
+
+        # Mock config.RUN_MODE to simulate the run mode
+        mocker.patch("src.utils.llm_utils.config.RUN_MODE", run_mode)
+
+        # Mock CSV writer to avoid writing to disk
+        mocker.patch("csv.DictWriter.writerow")
+
+        # Call the function
+        response = invoke_claude("test prompt", model_id, "file1.txt")
+
+        # Assertions
+        assert response == expected_response

fieldExtraction/tests/string_utils_test.py

@@ -1,53 +1,202 @@
 import pytest
+import src.utils as utils
+from src.utils.string_utils import extract_text_from_delimiters, json_parsing_search, secondary_string_to_dict, primary_string_to_dict, contains_reimbursement, is_empty
+import src.utils.llm_utils as llm_utils
+import numpy as np
+import pandas as pd
+
 
 from src.enums.delimiters import Delimiter
-from src.utils.string_utils import extract_text_from_delimiters
-
 
 class TestExtractTextFromDelimiters:
-    # Test case for valid inputs
+    @pytest.mark.parametrize("raw_text, delimiter, match_index, expected", [
+        ("Here is the answer |this is a pipe answer| and |more text.|", Delimiter.PIPE, 0, "this is a pipe answer"),
+        ("Here is the answer |this is a pipe answer| and |more text.|", Delimiter.PIPE, -1, "more text."),
+        ("|Here| is the answer |this is a pipe answer| and |more text.|", Delimiter.PIPE, 1, "this is a pipe answer"),
+        ("Here is the answer `this is a backtick answer` and `more text.`", Delimiter.BACKTICK, 0, "this is a backtick answer"),
+        ("Here is the answer `this is a backtick answer` and `more text.`", Delimiter.BACKTICK, -1, "more text."),
+        ("`Here` is the answer `this` is a `backtick answer` and `more text.`", Delimiter.BACKTICK, 2, "backtick answer"),
+        ("Here is the answer ```this is a triple backtick answer``` and more text.", Delimiter.TRIPLE_BACKTICK, 0, "this is a triple backtick answer"),
+        ("Here is the answer ```this``` is a ```triple``` backtick ```answer``` and more text.", Delimiter.TRIPLE_BACKTICK, -1, "answer"),
+        ("```Here``` is the answer ```this is a triple backtick answer``` and ```more text.```", Delimiter.TRIPLE_BACKTICK, 1, "this is a triple backtick answer"),
+    ])
+    def test_extract_text_from_delimiters(self, raw_text, delimiter, match_index, expected):
+        result = extract_text_from_delimiters(raw_text, delimiter, match_index)
+        assert result == expected
 
-    def test_pipe_delimiter_match_first_position(self):
-        raw_text = "Here is the answer |this is a pipe answer| and |more text.|"
-        result = extract_text_from_delimiters(raw_text,delimiter=Delimiter.PIPE,match_index=0)
-        assert result == "this is a pipe answer"
+class TestJsonParsingSearch:
+    @pytest.mark.parametrize("response_text, field_list, expected", [
+        ('{"name": "John", "age": "30", "city": "New York"}', ["name", "age", "city"], {"name": "John", "age": "30", "city": "New York"}),
+        ('{"name": "John", "city": "New York"}', ["name", "age", "city"], {"name": "John", "city": "New York"}),
+        ('', ["name", "age", "city"], {}),
+        ('{"name": "John", "age": "30", "city": "New York"}', [], {}),
+        ('{"name": "", "age": "", "city": ""}', ["name", "age", "city"], {"name": "", "age": "", "city": ""}),
+        ('{\n"name": "John",\n"age": "30",\n"city": "New York"\n}', ["name", "age", "city"], {"name": "John", "age": "30", "city": "New York"}),
+    ])
+    def test_json_parsing_search(self, response_text, field_list, expected):
+        assert json_parsing_search(response_text, field_list) == expected
 
-    def test_pipe_delimiter_match_last_position(self):
-        raw_text = "Here is the answer |this is a pipe answer| and |more text.|"
-        result = extract_text_from_delimiters(raw_text,delimiter=Delimiter.PIPE,match_index=-1)
-        assert result == "more text."
-    
-    def test_pipe_delimiter_match_any_position(self):
-        raw_text = "|Here| is the answer |this is a pipe answer| and |more text.|"
-        result = extract_text_from_delimiters(raw_text,delimiter=Delimiter.PIPE,match_index=1)
-        assert result == "this is a pipe answer"
-                
-    def test_backtick_delimiter_match_first_position(self):
-        raw_text = "Here is the answer `this is a backtick answer` and `more text.`"
-        result = extract_text_from_delimiters(raw_text, delimiter=Delimiter.BACKTICK, match_index=0)
-        assert result == "this is a backtick answer"
+class TestSecondaryStringToDict:
+    @pytest.fixture
+    def mock_invoke_claude(self, mocker):
+        return mocker.patch.object(llm_utils, 'invoke_claude')
 
-    def test_backtick_delimiter_match_last_position(self):
-        raw_text = "Here is the answer `this is a backtick answer` and `more text.`"
-        result = extract_text_from_delimiters(raw_text, delimiter=Delimiter.BACKTICK, match_index=-1)
-        assert result == "more text."
+    def test_valid_json(self, mock_invoke_claude):
+        dict_string = '{"key": "value"}'
+        filename = "test_file.txt"
+        result = secondary_string_to_dict(dict_string, filename)
+        assert result == {"key": "value"}
+        mock_invoke_claude.assert_not_called()
 
-    def test_backtick_delimiter_match_any_position(self):
-        raw_text = "`Here` is the answer `this` is a `backtick answer` and `more text.`"
-        result = extract_text_from_delimiters(raw_text, delimiter=Delimiter.BACKTICK, match_index=2)
-        assert result == "backtick answer"
+    def test_invalid_json_fixed_by_replacing_quotes(self, mock_invoke_claude):
+        dict_string = "{'key': 'value'}"
+        filename = "test_file.txt"
+        result = secondary_string_to_dict(dict_string, filename)
+        assert result == {"key": "value"}
+        mock_invoke_claude.assert_not_called()
 
-    def test_triple_backtick_delimiter_match_first_position(self):
-        raw_text = "Here is the answer ```this is a triple backtick answer``` and more text."
-        result = extract_text_from_delimiters(raw_text, delimiter=Delimiter.TRIPLE_BACKTICK, match_index=0)
-        assert result == "this is a triple backtick answer"
+    def test_invalid_json_fixed_by_llm(self, mock_invoke_claude):
+        mock_invoke_claude.return_value = '{"key": "value"}'
+        dict_string = "{key: value}"
+        filename = "test_file.txt"
+        result = secondary_string_to_dict(dict_string, filename)
+        assert result == {"key": "value"}
+        mock_invoke_claude.assert_called_once()
 
-    def test_triple_backtick_delimiter_match_last_position(self):
-        raw_text = "Here is the answer ```this``` is a ```triple``` backtick ```answer``` and more text."
-        result = extract_text_from_delimiters(raw_text, delimiter=Delimiter.TRIPLE_BACKTICK, match_index=-1)
-        assert result == "answer"
+    def test_invalid_json_failed_by_llm(self, mock_invoke_claude):
+        mock_invoke_claude.side_effect = Exception("LLM call failed")
+        dict_string = "{key: value}"
+        filename = "test_file.txt"
+        result = secondary_string_to_dict(dict_string, filename)
+        assert result == {}
+        mock_invoke_claude.assert_called_once()
 
-    def test_triple_backtick_delimiter_match_any_position(self):
-        raw_text = "```Here``` is the answer ```this is a triple backtick answer``` and ```more text.```"
-        result = extract_text_from_delimiters(raw_text, delimiter=Delimiter.TRIPLE_BACKTICK, match_index=1)
-        assert result == "this is a triple backtick answer"
+    def test_no_dict_found(self, mock_invoke_claude):
+        mock_invoke_claude.return_value = '{}'
+        dict_string = "This is a plain text without any dictionary."
+        filename = "test_file.txt"
+        result = secondary_string_to_dict(dict_string, filename)
+        assert result == {}
+        mock_invoke_claude.assert_called_once()
+
+class TestPrimaryStringToDict:
+    @pytest.fixture
+    def mock_secondary_string_to_dict(self, mocker):
+        return mocker.patch.object(utils.string_utils, "secondary_string_to_dict")
+
+    def test_valid_input(self, mock_secondary_string_to_dict):
+        mock_secondary_string_to_dict.side_effect = [
+            {"key1": "value1"},
+            {"key2": "value2"},
+            {"key3": "value3"},
+        ]
+        string_dict = {
+            "1": "[{'key1': 'value1'}, {'key2': 'value2'}]",
+            "2": "[{'key3': 'value3'}]",
+        }
+        filename = "test_file.txt"
+        expected_output = [
+            {"key1": "value1", "page_num": "1", "Filename": "test_file.txt"},
+            {"key2": "value2", "page_num": "1", "Filename": "test_file.txt"},
+            {"key3": "value3", "page_num": "2", "Filename": "test_file.txt"},
+        ]
+        result = primary_string_to_dict(string_dict, filename)
+        assert result == expected_output
+        assert mock_secondary_string_to_dict.call_count == len(expected_output)
+
+    def test_invalid_input_no_dicts(self, mock_secondary_string_to_dict):
+        mock_secondary_string_to_dict.return_value = {}
+        string_dict = {
+            "1": "[This is not a dictionary]",
+            "2": "[Still not a dictionary]",
+        }
+        filename = "test_file.txt"
+        expected_output = []
+        result = primary_string_to_dict(string_dict, filename)
+        assert result == expected_output
+        assert mock_secondary_string_to_dict.call_count == len(expected_output)
+
+    def test_empty_input(self, mock_secondary_string_to_dict):
+        string_dict = {}
+        filename = "test_file.txt"
+        expected_output = []
+        result = primary_string_to_dict(string_dict, filename)
+        assert result == expected_output
+        mock_secondary_string_to_dict.assert_not_called()
+
+    def test_malformed_input(self, mock_secondary_string_to_dict):
+        mock_secondary_string_to_dict.side_effect = [
+            {"key1": "value1"},
+            {"key2": "value2"},
+        ]
+        string_dict = {
+            "1": "[{'key1': 'value1'}, malformed]",
+            "2": "[{'key2': 'value2'}, <<<invalid>>>]",
+        }
+        filename = "test_file.txt"
+        expected_output = [
+            {"key1": "value1", "page_num": "1", "Filename": "test_file.txt"},
+            {"key2": "value2", "page_num": "2", "Filename": "test_file.txt"},
+        ]
+        result = primary_string_to_dict(string_dict, filename)
+        assert result == expected_output
+        assert mock_secondary_string_to_dict.call_count == 2
+
+    def test_multiple_pages(self, mock_secondary_string_to_dict):
+        mock_secondary_string_to_dict.side_effect = [
+            {"key1": "value1"},
+            {"key2": "value2"},
+            {"key3": "value3"},
+            {"key4": "value4"},
+        ]
+        string_dict = {
+            "1": "[{'key1': 'value1'}, {'key2': 'value2'}]",
+            "2": "[{'key3': 'value3'}, {'key4': 'value4'}]",
+        }
+        filename = "test_file.txt"
+        expected_output = [
+            {"key1": "value1", "page_num": "1", "Filename": "test_file.txt"},
+            {"key2": "value2", "page_num": "1", "Filename": "test_file.txt"},
+            {"key3": "value3", "page_num": "2", "Filename": "test_file.txt"},
+            {"key4": "value4", "page_num": "2", "Filename": "test_file.txt"},
+        ]
+        result = primary_string_to_dict(string_dict, filename)
+        assert result == expected_output
+        assert mock_secondary_string_to_dict.call_count == 4
+
+class TestContainsReimbursement:
+    @pytest.mark.parametrize("text, page, expected", [
+        ({"1": "This page contains a 10% reimbursement schedule.", "2": "No relevant content here."}, "1", True),
+        ({"1": "This page has no relevant content.", "2": "Still no relevant content here."}, "1", False),
+        ({"1": "This page contains a 10% reimbursement schedule.", "2": "No relevant content here."}, "invalid_page", False),
+        ("This text contains a $100 reimbursement schedule.", "1", True),
+        ("This text has no relevant content.", "1", False),
+    ])
+    def test_contains_reimbursement(self, text, page, expected):
+        result = contains_reimbursement(text, page)
+        assert result == expected
+
+    def test_invalid_input_type(self, capsys):
+        text = ["This is a list, not a dictionary or string."]
+        page = "1"
+        result = contains_reimbursement(text, page)
+        assert result is None
+        captured = capsys.readouterr()
+        assert "contains_reimbursement - Invalid data type" in captured.out
+
+class TestIsEmpty:
+    @pytest.mark.parametrize("value, expected", [
+        (None, True),
+        ("", True),
+        ("N/A", True),
+        ("null", True),
+        ("none", True),
+        ("NaN", True),
+        (np.nan, True),
+        (pd.NA, True),
+        ("This is not empty.", False),
+        (42, False),
+    ])
+    def test_is_empty(self, value, expected):
+        result = is_empty(value)
+        assert result == expected

fieldExtraction/tests/table_utils_test.py

@@ -0,0 +1,114 @@
+import pytest
+import src.utils as utils
+from src.utils.table_utils import align_and_format_tables
+from src import (prompts, config)
+import copy
+
+class TestAlignAndFormatTables:
+    @pytest.fixture
+    def mock_utils(self, mocker):
+        mock_contains_reimbursement = mocker.patch.object(
+            utils.string_utils,
+            "contains_reimbursement",
+        )
+        mock_invoke_claude = mocker.patch.object(
+            utils.llm_utils,
+            "invoke_claude",
+        )
+        return mock_contains_reimbursement, mock_invoke_claude
+
+    def test_align_and_format_tables_with_reimbursement(self, mock_utils):
+        """Test aligning and formatting tables when reimbursement keywords are present."""
+        mock_contains_reimbursement, mock_invoke_claude = mock_utils
+        mock_contains_reimbursement.return_value = True
+        mock_invoke_claude.return_value = (
+            "Table Header\nColumn 1 : Value 1% | Column 2 : Value 3\nColumn 1 : Value 2 | Column 2 : Value 4"
+        )
+
+        text_dict = {
+            "1": "-------Table Start-------- Table Header {'Column 1': ['Value 1%', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}-------Table End--------",
+        }
+        filename = "test_file.txt"
+
+        expected_prompt = prompts.ALIGN_AND_FORMAT_TABLES(text_dict["1"])
+        input_text_dict = copy.deepcopy(text_dict)
+
+        result = align_and_format_tables(text_dict, filename)
+
+        assert result == {
+            "1": "Table Header\nColumn 1 : Value 1% | Column 2 : Value 3\nColumn 1 : Value 2 | Column 2 : Value 4",
+        }
+        mock_contains_reimbursement.assert_called_once_with(input_text_dict["1"])
+        mock_invoke_claude.assert_called_once_with(
+            expected_prompt,
+            config.MODEL_ID_CLAUDE35_SONNET,
+            filename,
+            8192,
+        )
+
+    def test_align_and_format_tables_without_reimbursement(self, mock_utils):
+        """Test aligning and formatting tables when no reimbursement keywords are present."""
+        mock_contains_reimbursement, mock_invoke_claude = mock_utils
+        mock_contains_reimbursement.return_value = False
+
+        text_dict = {
+            "1": "-------Table Start-------- Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3', 'Value 4']}-------Table End--------",
+        }
+        filename = "test_file.txt"
+
+        result = align_and_format_tables(text_dict, filename)
+
+        assert result == text_dict  # No changes should be made
+        mock_contains_reimbursement.assert_called_once_with(text_dict["1"])
+        mock_invoke_claude.assert_not_called()
+
+    def test_align_and_format_tables_no_table(self, mock_utils):
+        """Test aligning and formatting tables when no table is present."""
+        mock_contains_reimbursement, mock_invoke_claude = mock_utils
+
+        text_dict = {
+            "1": "This page does not contain a table.",
+        }
+        filename = "test_file.txt"
+
+        result = align_and_format_tables(text_dict, filename)
+
+        assert result == text_dict  # No changes should be made
+        mock_contains_reimbursement.assert_not_called()
+        mock_invoke_claude.assert_not_called()
+
+    def test_align_and_format_tables_multiple_pages(self, mock_utils):
+        """Test aligning and formatting tables with multiple pages."""
+        mock_contains_reimbursement, mock_invoke_claude = mock_utils
+        mock_contains_reimbursement.side_effect = [True, False]
+        mock_invoke_claude.return_value = (
+            "Table Header\nColumn 1 : Value 1 | Column 2 : Value 3%\nColumn 1 : Value 2 | Column 2 : Value 4"
+        )
+
+        text_dict = {
+            "1": "-------Table Start-------- Table Header {'Column 1': ['Value 1', 'Value 2'], 'Column 2': ['Value 3%', 'Value 4']}-------Table End--------",
+            "2": "This page has 'Table Start' but does not contain reimbursement.",
+        }
+        filename = "test_file.txt"
+
+        expected_prompt = prompts.ALIGN_AND_FORMAT_TABLES(text_dict["1"])
+        input_text_dict = copy.deepcopy(text_dict)
+
+        result = align_and_format_tables(text_dict, filename)
+
+        assert result == {
+            "1": "Table Header\nColumn 1 : Value 1 | Column 2 : Value 3%\nColumn 1 : Value 2 | Column 2 : Value 4",
+            "2": "This page has 'Table Start' but does not contain reimbursement.",
+        }
+        assert mock_contains_reimbursement.call_count == 2
+        mock_contains_reimbursement.assert_any_call(input_text_dict["1"])
+        mock_contains_reimbursement.assert_any_call(input_text_dict["2"])
+        mock_invoke_claude.assert_called_once_with(
+            expected_prompt,
+            config.MODEL_ID_CLAUDE35_SONNET,
+            filename,
+            8192,
+        )
+        for call in mock_invoke_claude.call_args_list:
+            prompt, _, _, _ = call[0]  # Unpack the arguments
+            assert input_text_dict["2"] not in prompt