Spaces:
Running
Running
| import unittest | |
| import pandas as pd | |
| import numpy as np | |
| from data_assistant import ( | |
| DataExtractionError, | |
| load_csv_dataframe, | |
| generate_dataset_profile, | |
| search_dataframe, | |
| prepare_dataframe_llm_context, | |
| ) | |
| class DataAssistantTests(unittest.TestCase): | |
| def setUp(self): | |
| self.sample_csv_text = """Issue_ID,Issue_Topic,Category,Functional_Owner,Status,Count_Affected | |
| 101,Form 13 Transfer Request stuck at employer,Transfers,Field Office Delhi,Open,15 | |
| 102,DSC registration error on Employer Portal,DSC,Tech Backend,In Progress,8 | |
| 103,Form 13 Member Passbook not updating after transfer,Transfers,Tech Backend,Open,24 | |
| 104,Joint Declaration name mismatch,Member Profile,Field Office Mumbai,Resolved,5 | |
| 105,Form 19 Claim rejection due to KYC,Claims,Settlement Team,Open,12 | |
| 106,Form 13 Annexure K generation failed,Transfers,Tech Backend,Resolved,3 | |
| """ | |
| self.sample_csv_bytes = self.sample_csv_text.encode("utf-8") | |
| def test_load_valid_csv(self): | |
| df = load_csv_dataframe(self.sample_csv_bytes, "cites_sample.csv") | |
| self.assertEqual(len(df), 6) | |
| self.assertIn("Issue_Topic", df.columns) | |
| self.assertIn("Functional_Owner", df.columns) | |
| def test_arbitrary_schema_and_double_csv_suffix_are_supported(self): | |
| location_csv = ( | |
| b"place,latitude,longitude,population\n" | |
| b"North Point,28.61,77.21,1200\n" | |
| b"South Point,19.08,72.88,900\n" | |
| ) | |
| df = load_csv_dataframe( | |
| location_csv, | |
| "~/Downloads/location.csv - location.csv.csv", | |
| ) | |
| self.assertEqual(list(df.columns), ["place", "latitude", "longitude", "population"]) | |
| self.assertEqual(len(df), 2) | |
| def test_domain_words_are_not_discarded_from_deterministic_search(self): | |
| df = load_csv_dataframe(b"category\nissues\nother\n", "generic.csv") | |
| context = prepare_dataframe_llm_context(df, "How many issues are there?") | |
| self.assertIn("Exactly **1** matching row(s) found", context) | |
| def test_load_empty_csv_raises_error(self): | |
| with self.assertRaises(DataExtractionError): | |
| load_csv_dataframe(b"", "empty.csv") | |
| with self.assertRaises(DataExtractionError): | |
| load_csv_dataframe(b" \n ", "whitespace.csv") | |
| def test_load_non_csv_and_malformed_rows_raise_error(self): | |
| with self.assertRaises(DataExtractionError): | |
| load_csv_dataframe(b"not actually csv", "plain.txt") | |
| with self.assertRaises(DataExtractionError): | |
| load_csv_dataframe(b"a,b\n1,2\n3,4,5\n6,7\n", "broken.csv") | |
| def test_duplicate_trimmed_columns_are_made_unique(self): | |
| df = load_csv_dataframe(b"Name, Name\nAlice,Bob\n", "duplicates.csv") | |
| self.assertEqual(list(df.columns), ["Name", "Name_2"]) | |
| profile = generate_dataset_profile(df) | |
| self.assertEqual(profile.column_count, 2) | |
| def test_load_semicolon_delimited_csv(self): | |
| semicolon_csv = b"ColA;ColB;ColC\n1;2;3\n4;5;6\n" | |
| df = load_csv_dataframe(semicolon_csv, "semi.csv") | |
| self.assertEqual(len(df), 2) | |
| self.assertEqual(len(df.columns), 3) | |
| self.assertIn("ColA", df.columns) | |
| def test_dataset_profile_generation(self): | |
| df = load_csv_dataframe(self.sample_csv_bytes) | |
| profile = generate_dataset_profile(df) | |
| self.assertEqual(profile.row_count, 6) | |
| self.assertEqual(profile.column_count, 6) | |
| self.assertIn("Count_Affected", profile.numeric_columns) | |
| self.assertIn("Issue_Topic", profile.categorical_columns) | |
| self.assertIn("Count_Affected", profile.summary_stats) | |
| self.assertEqual(profile.summary_stats["Count_Affected"]["min"], 3.0) | |
| self.assertEqual(profile.summary_stats["Count_Affected"]["max"], 24.0) | |
| def test_search_dataframe_form_13_exact_count(self): | |
| df = load_csv_dataframe(self.sample_csv_bytes) | |
| matched_df, total_matches, breakdown = search_dataframe(df, "Form 13") | |
| # In our sample CSV: rows 101, 103, 106 have "Form 13" in Issue_Topic | |
| self.assertEqual(total_matches, 3) | |
| self.assertEqual(len(matched_df), 3) | |
| self.assertEqual(breakdown.get("Issue_Topic"), 3) | |
| def test_search_dataframe_case_insensitive(self): | |
| df = load_csv_dataframe(self.sample_csv_bytes) | |
| matched_df, total_matches, _ = search_dataframe(df, "form 13", case_sensitive=False) | |
| self.assertEqual(total_matches, 3) | |
| def test_search_dataframe_owner_filter(self): | |
| df = load_csv_dataframe(self.sample_csv_bytes) | |
| matched_df, total_matches, _ = search_dataframe(df, "Tech Backend") | |
| self.assertEqual(total_matches, 3) | |
| def test_prepare_dataframe_llm_context_includes_deterministic_count(self): | |
| df = load_csv_dataframe(self.sample_csv_bytes) | |
| user_query = "How many Form 13 issues are there in the dataset?" | |
| context = prepare_dataframe_llm_context(df, user_query) | |
| self.assertIn("Dataset Structure", context) | |
| self.assertIn("Form 13", context) | |
| self.assertIn("Exactly **3** matching row(s) found", context) | |
| def test_prepare_context_includes_combined_filter_count(self): | |
| df = load_csv_dataframe(self.sample_csv_bytes) | |
| context = prepare_dataframe_llm_context( | |
| df, | |
| "How many Form 13 issues are Open?", | |
| ) | |
| self.assertIn("Combined filter", context) | |
| self.assertIn("Exactly **2** row(s) match all identified terms", context) | |
| def test_prepare_context_respects_character_limit(self): | |
| df = pd.DataFrame( | |
| {f"Column_{index}": ["x" * 200 for _ in range(20)] for index in range(40)} | |
| ) | |
| context = prepare_dataframe_llm_context( | |
| df, | |
| "summarize this dataset", | |
| max_context_chars=2_500, | |
| ) | |
| self.assertLessEqual(len(context), 2_500) | |
| self.assertIn("dataset context truncated", context) | |
| if __name__ == "__main__": | |
| unittest.main() | |