diff --git a/usaspending_api/llm/tests/helper.py b/usaspending_api/llm/tests/helper.py new file mode 100644 index 0000000000..cb2cbb6ab1 --- /dev/null +++ b/usaspending_api/llm/tests/helper.py @@ -0,0 +1,38 @@ +from typing import Any + +from usaspending_api.common.elasticsearch.search_wrappers import RecipientSearch +from usaspending_api.llm.tools.lookup_recipient import RecipientLookupTool +from usaspending_api.search.v2.es_sanitization import es_sanitize + +_tool = RecipientLookupTool() + + +def build_fuzzy_recipient_query(search_text: str) -> RecipientSearch: + return _tool._build_search(es_sanitize(search_text).strip().upper(), top_k=10) + + +def fuzzy_search_recipients( + search_text: str, + limit: int = 10, +) -> list | dict[str, Any]: + response = _tool._build_search(es_sanitize(search_text).strip().upper(), top_k=limit).handle_execute() + if not response.hits: + return [] + return [ + { + "recipient_name": hit.to_dict().get("recipient_name"), + "uei": hit.to_dict().get("uei"), + "duns": hit.to_dict().get("duns"), + "recipient_level": hit.to_dict().get("recipient_level"), + "recipient_hash": hit.to_dict().get("recipient_hash"), + "score": hit.meta.score, + } + for hit in response.hits + ] + + +def retrieve_recipient_names( + search_text: str, + limit: int = 5, +) -> list[str]: + return _tool.lookup_recipient(search_text, top_k=limit) diff --git a/usaspending_api/llm/tests/integration/test_recipient_retrieval.py b/usaspending_api/llm/tests/integration/test_recipient_retrieval.py new file mode 100644 index 0000000000..8013454eb3 --- /dev/null +++ b/usaspending_api/llm/tests/integration/test_recipient_retrieval.py @@ -0,0 +1,413 @@ +"""Integration tests for recipient_retrieval module""" +import pytest +from django.conf import settings +from elasticsearch import Elasticsearch +from elasticsearch_dsl import Index, connections + +from usaspending_api.common.elasticsearch.search_wrappers import RecipientSearch +from usaspending_api.llm.tests.helper import ( + build_fuzzy_recipient_query, + fuzzy_search_recipients, + retrieve_recipient_names, +) + +# ============================================================================ +# FIXTURES +# ============================================================================ + + +@pytest.fixture(scope="module") +def elasticsearch_connection(): + """ + Fixture to provide Elasticsearch connection + """ + es = Elasticsearch( + hosts=[settings.ES_HOSTNAME], + timeout=30, + ) + + connections.add_connection('default', es) + + yield es + + connections.remove_connection('default') + es.close() + + +@pytest.fixture +def elasticsearch_recipient_index(elasticsearch_connection): + """ + Fixture to set up recipient index with test data + """ + index_name = "test-recipients" + + # Create index + index = Index(index_name) + if index.exists(using='default'): + index.delete(using='default') + index.create(using='default') + + # Add test data + test_recipients = [ + { + "recipient_name": "ACME CORPORATION", + "uei": "UEI123456789", + "duns": "123456789", + "recipient_level": "P", + "recipient_hash": "hash123", + }, + { + "recipient_name": "BETA CORP", + "uei": "UEI987654321", + "duns": "987654321", + "recipient_level": "C", + "recipient_hash": "hash456", + }, + { + "recipient_name": "ACME INDUSTRIES", + "uei": "UEI111222333", + "duns": "111222333", + "recipient_level": "P", + "recipient_hash": "hash789", + }, + ] + + # Index test documents + for recipient in test_recipients: + elasticsearch_connection.index( + index=index_name, + body=recipient, + ) + + # Refresh index to make documents searchable + elasticsearch_connection.indices.refresh(index=index_name) + + yield index_name + + # Cleanup + if index.exists(using='default'): + index.delete(using='default') + + +@pytest.fixture +def setup_test_recipients(elasticsearch_recipient_index): + """ + Convenience fixture that just ensures recipient index is set up + """ + return elasticsearch_recipient_index + + +# ============================================================================ +# TESTS +# ============================================================================ + +@pytest.mark.django_db +class TestBuildFuzzyRecipientQueryIntegration: + """Integration tests for build_fuzzy_recipient_query""" + + def test_returns_recipient_search_instance(self): + """Test that function returns a RecipientSearch instance""" + result = build_fuzzy_recipient_query("ACME Corporation") + + assert isinstance(result, RecipientSearch) + + def test_query_structure_contains_fuzzy_match(self): + """Test that generated query contains fuzzy matching logic""" + search = build_fuzzy_recipient_query("ACME") + query_dict = search.to_dict() + + # Verify query structure exists + assert "query" in query_dict + # Verify it's configured for fuzzy matching + assert query_dict.get("size") == 10 + + def test_sanitizes_special_characters(self): + """Test that special characters are properly sanitized""" + # Should not raise exception with special characters + search = build_fuzzy_recipient_query("ACME & Co. (2024)") + + assert isinstance(search, RecipientSearch) + + def test_handles_empty_string(self): + """Test handling of empty search string""" + search = build_fuzzy_recipient_query("") + + assert isinstance(search, RecipientSearch) + + def test_handles_whitespace_only(self): + """Test handling of whitespace-only string""" + search = build_fuzzy_recipient_query(" ") + + assert isinstance(search, RecipientSearch) + + +@pytest.mark.django_db +@pytest.mark.elasticsearch +class TestFuzzySearchRecipientsIntegration: + """Integration tests for fuzzy_search_recipients""" + + def test_returns_list_type(self, setup_test_recipients): + """Test that function returns a list""" + result = fuzzy_search_recipients("test") + + assert isinstance(result, list) + + def test_returns_empty_list_for_no_matches(self): + """Test that empty list is returned when no matches found""" + result = fuzzy_search_recipients("NONEXISTENT_COMPANY_XYZ_12345") + + assert result == [] + + def test_result_structure_with_matches(self, setup_test_recipients): + """Test that results have correct structure when matches exist""" + result = fuzzy_search_recipients("ACME", limit=5) + + if result: # If test data exists + assert isinstance(result, list) + for item in result: + assert "recipient_name" in item + assert "uei" in item + assert "duns" in item + assert "recipient_level" in item + assert "recipient_hash" in item + assert "score" in item + assert isinstance(item["score"], (int, float)) + + def test_respects_limit_parameter(self, setup_test_recipients): + """Test that limit parameter controls result count""" + result_5 = fuzzy_search_recipients("CORP", limit=5) + result_10 = fuzzy_search_recipients("CORP", limit=10) + + # If matches exist, verify limit is respected + if result_5: + assert len(result_5) <= 5 + if result_10: + assert len(result_10) <= 10 + + def test_results_ordered_by_relevance(self, setup_test_recipients): + """Test that results are ordered by score (descending)""" + result = fuzzy_search_recipients("ACME", limit=10) + + if len(result) > 1: + scores = [item["score"] for item in result] + assert scores == sorted(scores, reverse=True) + + def test_case_insensitive_search(self, setup_test_recipients): + """Test that search is case-insensitive""" + result_lower = fuzzy_search_recipients("acme corporation") + result_upper = fuzzy_search_recipients("ACME CORPORATION") + result_mixed = fuzzy_search_recipients("AcMe CoRpOrAtIoN") + + # All should return same results (or all empty) + assert len(result_lower) == len(result_upper) == len(result_mixed) + + def test_handles_partial_matches(self, setup_test_recipients): + """Test that partial text matches work""" + result = fuzzy_search_recipients("ACM") + + # Should find results containing "ACM" (like ACME) + assert isinstance(result, list) + + def test_handles_special_characters_in_search(self): + """Test that special characters don't cause errors""" + # Should not raise exception + result = fuzzy_search_recipients("ACME & Co. (2024)") + + assert isinstance(result, list) + + +@pytest.mark.django_db +@pytest.mark.elasticsearch +class TestRetrieveRecipientNamesIntegration: + """Integration tests for retrieve_recipient_names""" + + def test_returns_list_of_strings(self, setup_test_recipients): + """Test that function returns a list of strings""" + result = retrieve_recipient_names("ACME") + + assert isinstance(result, list) + # All items should be strings + for item in result: + assert isinstance(item, str) + + def test_handles_no_matches(self): + """Test behavior when no matches are found""" + result = retrieve_recipient_names("NONEXISTENT_XYZ_12345") + + assert isinstance(result, list) + assert result == [] + + def test_respects_limit_parameter(self, setup_test_recipients): + """Test that limit parameter controls number of results""" + result_3 = retrieve_recipient_names("CORP", limit=3) + result_10 = retrieve_recipient_names("CORP", limit=10) + + # Results should respect the limit + assert isinstance(result_3, list) + assert isinstance(result_10, list) + + def test_extracts_recipient_identifiers(self, setup_test_recipients): + """Test that recipient names, UEIs, and DUNS are extracted""" + result = retrieve_recipient_names("ACME", limit=5) + + if result: + # Should contain recipient identifiers (names, UEIs, DUNS) + assert len(result) > 0 + # All should be strings + assert all(isinstance(item, str) for item in result) + + def test_case_insensitive_search(self, setup_test_recipients): + """Test that search is case-insensitive""" + result_lower = retrieve_recipient_names("acme corp") + result_upper = retrieve_recipient_names("ACME CORP") + + # Should return same results + assert len(result_lower) == len(result_upper) + + def test_handles_special_characters(self): + """Test that special characters are handled properly""" + # Should not raise exception + result = retrieve_recipient_names("ACME & Co. (2024)") + + assert isinstance(result, list) + + def test_no_duplicate_names(self, setup_test_recipients): + """Test that result contains no duplicates""" + result = retrieve_recipient_names("CORP", limit=10) + + if result: + # Should have no duplicates + assert len(result) == len(set(result)) + + def test_returns_multiple_identifiers_per_recipient(self, setup_test_recipients): + """Test that multiple identifiers (name, UEI, DUNS) are returned for each recipient""" + result = retrieve_recipient_names("ACME CORPORATION", limit=1) + + if result: + # Should include multiple identifiers for the recipient + # (name, UEI, DUNS at minimum) + assert len(result) >= 1 + + def test_handles_partial_matches(self, setup_test_recipients): + """Test that partial text matches work""" + result = retrieve_recipient_names("ACM") + + # Should find results containing "ACM" (like ACME) + assert isinstance(result, list) + + def test_handles_whitespace_in_query(self): + """Test that whitespace in query is handled correctly""" + result = retrieve_recipient_names(" ACME ") + + assert isinstance(result, list) + + def test_performance_with_large_limit(self, setup_test_recipients): + """Test that function handles large limits without errors""" + # Should not timeout or error with large limit + result = retrieve_recipient_names("CORP", limit=100) + + assert isinstance(result, list) + + +@pytest.mark.django_db +@pytest.mark.elasticsearch +class TestRecipientRetrievalEndToEnd: + """End-to-end integration tests across multiple functions""" + + def test_fuzzy_search_to_retrieve_names_workflow(self, setup_test_recipients): + """Test workflow from fuzzy search to retrieving names""" + # Step 1: Fuzzy search to get detailed results + search_results = fuzzy_search_recipients("ACME", limit=5) + + # Step 2: Retrieve just the names for the same query + names = retrieve_recipient_names("ACME", limit=5) + + # Both should return results (or both empty) + assert isinstance(search_results, list) + assert isinstance(names, list) + + # If search results exist, names should also exist + if search_results: + assert len(names) > 0 + + def test_retrieve_names_includes_fuzzy_search_data(self, setup_test_recipients): + """Test that retrieve_recipient_names includes data from fuzzy search""" + # Get detailed results + fuzzy_results = fuzzy_search_recipients("ACME CORPORATION", limit=1) + + # Get names + names = retrieve_recipient_names("ACME CORPORATION", limit=1) + + if fuzzy_results: + # Names should include identifiers from the fuzzy search results + first_result = fuzzy_results[0] + if first_result.get("recipient_name"): + # At least one identifier should be in the names list + assert any( + identifier in names + for identifier in [ + first_result.get("recipient_name"), + first_result.get("uei"), + first_result.get("duns"), + ] + if identifier + ) + + def test_consistent_results_across_functions(self, setup_test_recipients): + """Test that different functions return consistent data for same recipient""" + search_text = "ACME CORPORATION" + + # Get results from both functions + fuzzy_results = fuzzy_search_recipients(search_text, limit=5) + names = retrieve_recipient_names(search_text, limit=5) + + # Both should return data (or both empty) + assert isinstance(fuzzy_results, list) + assert isinstance(names, list) + + # If fuzzy search has results, names should too + if fuzzy_results: + assert len(names) > 0 + + def test_build_query_to_fuzzy_search_workflow(self, setup_test_recipients): + """Test workflow from building query to executing search""" + # Step 1: Build query + search = build_fuzzy_recipient_query("ACME") + + # Step 2: Execute search (fuzzy_search_recipients does this internally) + results = fuzzy_search_recipients("ACME", limit=5) + + # Both should work without errors + assert isinstance(search, RecipientSearch) + assert isinstance(results, list) + + def test_performance_with_multiple_queries(self, setup_test_recipients): + """Test that multiple queries execute efficiently""" + queries = ["ACME", "BETA", "CORP", "INDUSTRIES"] + + for query in queries: + # Should not timeout or error + result = retrieve_recipient_names(query, limit=10) + assert isinstance(result, list) + + def test_empty_results_handled_consistently(self): + """Test that empty results are handled consistently across functions""" + nonexistent = "NONEXISTENT_XYZ_12345" + + # All functions should handle non-existent queries gracefully + fuzzy_results = fuzzy_search_recipients(nonexistent) + names = retrieve_recipient_names(nonexistent) + + assert fuzzy_results == [] + assert names == [] + + def test_special_characters_handled_consistently(self): + """Test that special characters are handled consistently""" + special_query = "ACME & Co. (2024)" + + # Should not raise exceptions + fuzzy_results = fuzzy_search_recipients(special_query) + names = retrieve_recipient_names(special_query) + + assert isinstance(fuzzy_results, list) + assert isinstance(names, list) diff --git a/usaspending_api/llm/tests/unit/test_py_models.py b/usaspending_api/llm/tests/unit/test_py_models.py new file mode 100644 index 0000000000..1159b36201 --- /dev/null +++ b/usaspending_api/llm/tests/unit/test_py_models.py @@ -0,0 +1,484 @@ +"""Unit tests for recipient-related fields in py_models""" +import pytest +from pydantic import ValidationError + +from usaspending_api.llm.models.py_models import Filters + + +class TestRecipientFields: + """Tests for recipient-related fields in Filters model""" + + def test_selected_recipients_default_empty_list(self): + """Test that selectedRecipients defaults to empty list""" + filters = Filters() + + assert filters.selectedRecipients == [] + assert isinstance(filters.selectedRecipients, list) + + def test_selected_recipients_accepts_string_list(self): + """Test that selectedRecipients accepts list of strings""" + recipient_ids = ["ACME CORP", "UEI123456789", "123456789"] + filters = Filters(selectedRecipients=recipient_ids) + + assert filters.selectedRecipients == recipient_ids + assert len(filters.selectedRecipients) == 3 + + def test_selected_recipients_accepts_empty_list(self): + """Test that selectedRecipients accepts empty list""" + filters = Filters(selectedRecipients=[]) + + assert filters.selectedRecipients == [] + + def test_selected_recipients_accepts_single_recipient(self): + """Test that selectedRecipients works with single recipient""" + filters = Filters(selectedRecipients=["ACME CORPORATION"]) + + assert filters.selectedRecipients == ["ACME CORPORATION"] + assert len(filters.selectedRecipients) == 1 + + def test_selected_recipients_preserves_order(self): + """Test that selectedRecipients preserves order of recipients""" + recipients = ["BETA CORP", "ACME CORP", "ZETA INDUSTRIES"] + filters = Filters(selectedRecipients=recipients) + + assert filters.selectedRecipients == recipients + assert filters.selectedRecipients[0] == "BETA CORP" + assert filters.selectedRecipients[2] == "ZETA INDUSTRIES" + + def test_selected_recipients_allows_duplicates(self): + """Test that selectedRecipients allows duplicate entries""" + recipients = ["ACME CORP", "ACME CORP", "BETA CORP"] + filters = Filters(selectedRecipients=recipients) + + assert len(filters.selectedRecipients) == 3 + assert filters.selectedRecipients.count("ACME CORP") == 2 + + def test_selected_recipients_rejects_non_string_values(self): + """Test that selectedRecipients rejects non-string values""" + with pytest.raises(ValidationError) as exc_info: + Filters(selectedRecipients=[123, 456]) + + assert "selectedRecipients" in str(exc_info.value) + + def test_selected_recipients_rejects_mixed_types(self): + """Test that selectedRecipients rejects mixed types""" + with pytest.raises(ValidationError) as exc_info: + Filters(selectedRecipients=["ACME CORP", 123, None]) + + assert "selectedRecipients" in str(exc_info.value) + + def test_selected_recipients_handles_special_characters(self): + """Test that selectedRecipients handles special characters in names""" + recipients = [ + "O'REILLY MEDIA", + "ACME & CO.", + "BETA CORP (2024)", + "GAMMA-DELTA LLC", + ] + filters = Filters(selectedRecipients=recipients) + + assert filters.selectedRecipients == recipients + assert "O'REILLY MEDIA" in filters.selectedRecipients + + def test_selected_recipients_handles_unicode(self): + """Test that selectedRecipients handles unicode characters""" + recipients = ["Société Générale", "日本株式会社", "Москва ООО"] + filters = Filters(selectedRecipients=recipients) + + assert filters.selectedRecipients == recipients + + def test_selected_recipients_handles_whitespace(self): + """Test that selectedRecipients preserves whitespace""" + recipients = [" ACME CORP ", "BETA CORP", "GAMMA\tCORP"] + filters = Filters(selectedRecipients=recipients) + + # Whitespace should be preserved as-is + assert filters.selectedRecipients == recipients + + +class TestRecipientType: + """Tests for recipientType field in Filters model""" + + def test_recipient_type_default_empty_list(self): + """Test that recipientType defaults to empty list""" + filters = Filters() + + assert filters.recipientType == [] + assert isinstance(filters.recipientType, list) + + def test_recipient_type_accepts_valid_business_types(self): + """Test that recipientType accepts valid business types""" + types = ["business", "small_business", "other_than_small_business"] + filters = Filters(recipientType=types) + + assert filters.recipientType == types + assert len(filters.recipientType) == 3 + + def test_recipient_type_accepts_minority_owned_types(self): + """Test that recipientType accepts minority-owned business types""" + types = [ + "minority_owned_business", + "black_american_owned_business", + "hispanic_american_owned_business", + ] + filters = Filters(recipientType=types) + + assert filters.recipientType == types + + def test_recipient_type_accepts_women_owned_types(self): + """Test that recipientType accepts women-owned business types""" + types = [ + "woman_owned_business", + "women_owned_small_business", + "economically_disadvantaged_women_owned_small_business", + ] + filters = Filters(recipientType=types) + + assert filters.recipientType == types + + def test_recipient_type_accepts_veteran_owned_types(self): + """Test that recipientType accepts veteran-owned business types""" + types = [ + "veteran_owned_business", + "service_disabled_veteran_owned_business", + ] + filters = Filters(recipientType=types) + + assert filters.recipientType == types + + def test_recipient_type_accepts_special_designations(self): + """Test that recipientType accepts special designation types""" + types = [ + "8a_program_participant", + "historically_underutilized_business_firm", + "ability_one_program", + ] + filters = Filters(recipientType=types) + + assert filters.recipientType == types + + def test_recipient_type_accepts_nonprofit_types(self): + """Test that recipientType accepts nonprofit types""" + types = ["nonprofit", "foundation", "community_development_corporations"] + filters = Filters(recipientType=types) + + assert filters.recipientType == types + + def test_recipient_type_accepts_higher_education_types(self): + """Test that recipientType accepts higher education types""" + types = [ + "higher_education", + "public_institution_of_higher_education", + "private_institution_of_higher_education", + "minority_serving_institution_of_higher_education", + ] + filters = Filters(recipientType=types) + + assert filters.recipientType == types + + def test_recipient_type_accepts_government_types(self): + """Test that recipientType accepts government types""" + types = [ + "government", + "national_government", + "local_government", + "indian_native_american_tribal_government", + ] + filters = Filters(recipientType=types) + + assert filters.recipientType == types + + def test_recipient_type_accepts_individuals(self): + """Test that recipientType accepts individuals type""" + filters = Filters(recipientType=["individuals"]) + + assert filters.recipientType == ["individuals"] + + def test_recipient_type_accepts_mixed_categories(self): + """Test that recipientType accepts types from different categories""" + types = [ + "small_business", + "woman_owned_business", + "veteran_owned_business", + "nonprofit", + "local_government", + ] + filters = Filters(recipientType=types) + + assert filters.recipientType == types + assert len(filters.recipientType) == 5 + + def test_recipient_type_rejects_invalid_type(self): + """Test that recipientType rejects invalid type values""" + with pytest.raises(ValidationError) as exc_info: + Filters(recipientType=["invalid_type"]) + + assert "recipientType" in str(exc_info.value) + + def test_recipient_type_rejects_empty_string(self): + """Test that recipientType rejects empty string""" + with pytest.raises(ValidationError) as exc_info: + Filters(recipientType=[""]) + + assert "recipientType" in str(exc_info.value) + + def test_recipient_type_rejects_non_string_values(self): + """Test that recipientType rejects non-string values""" + with pytest.raises(ValidationError) as exc_info: + Filters(recipientType=[123, 456]) + + assert "recipientType" in str(exc_info.value) + + def test_recipient_type_case_sensitive(self): + """Test that recipientType is case-sensitive""" + with pytest.raises(ValidationError) as exc_info: + Filters(recipientType=["SMALL_BUSINESS"]) # Should be lowercase + + assert "recipientType" in str(exc_info.value) + + def test_recipient_type_allows_duplicates(self): + """Test that recipientType allows duplicate entries""" + types = ["small_business", "small_business", "nonprofit"] + filters = Filters(recipientType=types) + + assert len(filters.recipientType) == 3 + assert filters.recipientType.count("small_business") == 2 + + def test_recipient_type_preserves_order(self): + """Test that recipientType preserves order""" + types = ["nonprofit", "business", "government"] + filters = Filters(recipientType=types) + + assert filters.recipientType == types + assert filters.recipientType[0] == "nonprofit" + assert filters.recipientType[2] == "government" + + def test_recipient_type_accepts_all_valid_types(self): + """Test that all documented recipient types are valid""" + all_types = [ + # Business types + "business", + "small_business", + "other_than_small_business", + "corporate_entity_tax_exempt", + "corporate_entity_not_tax_exempt", + "partnership_or_limited_liability_partnership", + "sole_proprietorship", + "manufacturer_of_goods", + "subchapter_s_corporation", + "limited_liability_corporation", + # Minority owned + "minority_owned_business", + "alaskan_native_corporation_owned_firm", + "american_indian_owned_business", + "asian_pacific_american_owned_business", + "black_american_owned_business", + "hispanic_american_owned_business", + "native_american_owned_business", + "native_hawaiian_organization_owned_firm", + "subcontinent_asian_indian_american_owned_business", + "tribally_owned_firm", + "other_minority_owned_business", + # Women owned + "woman_owned_business", + "women_owned_small_business", + "economically_disadvantaged_women_owned_small_business", + "joint_venture_women_owned_small_business", + "joint_venture_economically_disadvantaged_women_owned_small_business", + # Veteran owned + "veteran_owned_business", + "service_disabled_veteran_owned_business", + # Special designations + "special_designations", + "8a_program_participant", + "ability_one_program", + "dot_certified_disadvantaged_business_enterprise", + "emerging_small_business", + "federally_funded_research_and_development_corp", + "historically_underutilized_business_firm", + "labor_surplus_area_firm", + "sba_certified_8a_joint_venture", + "self_certified_small_disadvanted_business", + "small_agricultural_cooperative", + "community_developed_corporation_owned_firm", + "us_owned_business", + "foreign_owned_and_us_located_business", + "foreign_owned", + "foreign_government", + "international_organization", + "domestic_shelter", + "hospital", + "veterinary_hospital", + # Nonprofit + "nonprofit", + "foundation", + "community_development_corporations", + # Higher education + "higher_education", + "public_institution_of_higher_education", + "private_institution_of_higher_education", + "minority_serving_institution_of_higher_education", + "school_of_forestry", + "veterinary_college", + # Government + "government", + "national_government", + "interstate_entity", + "regional_and_state_government", + "regional_organization", + "us_territory_or_possession", + "council_of_governments", + "local_government", + "indian_native_american_tribal_government", + "authorities_and_commissions", + # Individuals + "individuals", + ] + + # Should not raise validation error + filters = Filters(recipientType=all_types) + assert len(filters.recipientType) == len(all_types) + + +class TestRecipientDomesticForeign: + """Tests for recipientDomesticForeign field in Filters model""" + + def test_recipient_domestic_foreign_default_all(self): + """Test that recipientDomesticForeign defaults to 'all'""" + filters = Filters() + + assert filters.recipientDomesticForeign == "all" + + def test_recipient_domestic_foreign_accepts_all(self): + """Test that recipientDomesticForeign accepts 'all'""" + filters = Filters(recipientDomesticForeign="all") + + assert filters.recipientDomesticForeign == "all" + + def test_recipient_domestic_foreign_accepts_foreign(self): + """Test that recipientDomesticForeign accepts 'foreign'""" + filters = Filters(recipientDomesticForeign="foreign") + + assert filters.recipientDomesticForeign == "foreign" + + def test_recipient_domestic_foreign_rejects_invalid_value(self): + """Test that recipientDomesticForeign rejects invalid values""" + with pytest.raises(ValidationError) as exc_info: + Filters(recipientDomesticForeign="domestic") + + assert "recipientDomesticForeign" in str(exc_info.value) + + def test_recipient_domestic_foreign_rejects_empty_string(self): + """Test that recipientDomesticForeign rejects empty string""" + with pytest.raises(ValidationError) as exc_info: + Filters(recipientDomesticForeign="") + + assert "recipientDomesticForeign" in str(exc_info.value) + + def test_recipient_domestic_foreign_case_sensitive(self): + """Test that recipientDomesticForeign is case-sensitive""" + with pytest.raises(ValidationError) as exc_info: + Filters(recipientDomesticForeign="ALL") + + assert "recipientDomesticForeign" in str(exc_info.value) + + def test_recipient_domestic_foreign_rejects_none(self): + """Test that recipientDomesticForeign rejects None""" + with pytest.raises(ValidationError) as exc_info: + Filters(recipientDomesticForeign=None) + + assert "recipientDomesticForeign" in str(exc_info.value) + + +class TestRecipientFieldsCombinations: + """Tests for combinations of recipient-related fields""" + + def test_all_recipient_fields_together(self): + """Test that all recipient fields can be set together""" + filters = Filters( + selectedRecipients=["ACME CORP", "BETA CORP"], + recipientType=["small_business", "woman_owned_business"], + recipientDomesticForeign="foreign", + ) + + assert filters.selectedRecipients == ["ACME CORP", "BETA CORP"] + assert filters.recipientType == ["small_business", "woman_owned_business"] + assert filters.recipientDomesticForeign == "foreign" + + def test_recipient_fields_with_empty_values(self): + """Test recipient fields with empty values""" + filters = Filters( + selectedRecipients=[], + recipientType=[], + recipientDomesticForeign="all", + ) + + assert filters.selectedRecipients == [] + assert filters.recipientType == [] + assert filters.recipientDomesticForeign == "all" + + def test_recipient_fields_independent(self): + """Test that recipient fields are independent""" + # Set only selectedRecipients + filters1 = Filters(selectedRecipients=["ACME CORP"]) + assert filters1.selectedRecipients == ["ACME CORP"] + assert filters1.recipientType == [] + assert filters1.recipientDomesticForeign == "all" + + # Set only recipientType + filters2 = Filters(recipientType=["small_business"]) + assert filters2.selectedRecipients == [] + assert filters2.recipientType == ["small_business"] + assert filters2.recipientDomesticForeign == "all" + + # Set only recipientDomesticForeign + filters3 = Filters(recipientDomesticForeign="foreign") + assert filters3.selectedRecipients == [] + assert filters3.recipientType == [] + assert filters3.recipientDomesticForeign == "foreign" + + def test_recipient_fields_serialization(self): + """Test that recipient fields serialize correctly""" + filters = Filters( + selectedRecipients=["ACME CORP"], + recipientType=["small_business", "nonprofit"], + recipientDomesticForeign="foreign", + ) + + data = filters.model_dump() + + assert data["selectedRecipients"] == ["ACME CORP"] + assert data["recipientType"] == ["small_business", "nonprofit"] + assert data["recipientDomesticForeign"] == "foreign" + + def test_recipient_fields_deserialization(self): + """Test that recipient fields deserialize correctly""" + data = { + "selectedRecipients": ["BETA CORP", "GAMMA CORP"], + "recipientType": ["veteran_owned_business", "8a_program_participant"], + "recipientDomesticForeign": "all", + } + + filters = Filters(**data) + + assert filters.selectedRecipients == ["BETA CORP", "GAMMA CORP"] + assert filters.recipientType == ["veteran_owned_business", "8a_program_participant"] + assert filters.recipientDomesticForeign == "all" + + def test_recipient_fields_json_schema(self): + """Test that recipient fields have proper JSON schema""" + schema = Filters.model_json_schema() + + # Check selectedRecipients schema + assert "selectedRecipients" in schema["properties"] + assert schema["properties"]["selectedRecipients"]["type"] == "array" + assert schema["properties"]["selectedRecipients"]["items"]["type"] == "string" + + # Check recipientType schema + assert "recipientType" in schema["properties"] + assert schema["properties"]["recipientType"]["type"] == "array" + + # Check recipientDomesticForeign schema + assert "recipientDomesticForeign" in schema["properties"] + assert "enum" in schema["properties"]["recipientDomesticForeign"] + assert set(schema["properties"]["recipientDomesticForeign"]["enum"]) == {"all", "foreign"} diff --git a/usaspending_api/llm/tests/unit/test_recipient_retrieval.py b/usaspending_api/llm/tests/unit/test_recipient_retrieval.py new file mode 100644 index 0000000000..1413d30303 --- /dev/null +++ b/usaspending_api/llm/tests/unit/test_recipient_retrieval.py @@ -0,0 +1,292 @@ +"""Unit tests for recipient_retrieval module""" +from unittest.mock import Mock, patch + +from usaspending_api.llm.tests.helper import ( + build_fuzzy_recipient_query, + fuzzy_search_recipients, + retrieve_recipient_names, +) + + +class TestBuildFuzzyRecipientQuery: + """Tests for build_fuzzy_recipient_query function""" + + @patch("usaspending_api.llm.tests.helper._tool") + @patch("usaspending_api.llm.tests.helper.es_sanitize") + def test_builds_query_with_sanitized_text(self, mock_sanitize, mock_tool): + """Test that query is built with sanitized and uppercase text""" + mock_sanitize.return_value = "acme corp" + mock_search = Mock() + mock_tool._build_search.return_value = mock_search + + result = build_fuzzy_recipient_query(" ACME Corp ") + + mock_sanitize.assert_called_once_with(" ACME Corp ") + mock_tool._build_search.assert_called_once_with("ACME CORP", top_k=10) + assert result == mock_search + + @patch("usaspending_api.llm.tests.helper._tool") + @patch("usaspending_api.llm.tests.helper.es_sanitize") + def test_strips_whitespace(self, mock_sanitize, mock_tool): + """Test that whitespace is stripped from search text""" + mock_sanitize.return_value = " test " + mock_search = Mock() + mock_tool._build_search.return_value = mock_search + + build_fuzzy_recipient_query(" test ") + + mock_tool._build_search.assert_called_once_with("TEST", top_k=10) + + @patch("usaspending_api.llm.tests.helper._tool") + @patch("usaspending_api.llm.tests.helper.es_sanitize") + def test_converts_to_uppercase(self, mock_sanitize, mock_tool): + """Test that search text is converted to uppercase""" + mock_sanitize.return_value = "lowercase text" + mock_search = Mock() + mock_tool._build_search.return_value = mock_search + + build_fuzzy_recipient_query("lowercase text") + + mock_tool._build_search.assert_called_once_with("LOWERCASE TEXT", top_k=10) + + +class TestFuzzySearchRecipients: + """Tests for fuzzy_search_recipients function""" + + @patch("usaspending_api.llm.tests.helper._tool") + @patch("usaspending_api.llm.tests.helper.es_sanitize") + def test_returns_empty_list_when_no_hits(self, mock_sanitize, mock_tool): + """Test that empty list is returned when no hits found""" + mock_sanitize.return_value = "test" + mock_response = Mock() + mock_response.hits = [] + mock_search = Mock() + mock_search.handle_execute.return_value = mock_response + mock_tool._build_search.return_value = mock_search + + result = fuzzy_search_recipients("test") + + assert result == [] + + @patch("usaspending_api.llm.tests.helper._tool") + @patch("usaspending_api.llm.tests.helper.es_sanitize") + def test_returns_formatted_results_with_hits(self, mock_sanitize, mock_tool): + """Test that results are properly formatted when hits exist""" + mock_sanitize.return_value = "acme" + + # Create mock hit + mock_hit = Mock() + mock_hit.to_dict.return_value = { + "recipient_name": "ACME CORP", + "uei": "UEI123456789", + "duns": "123456789", + "recipient_level": "P", + "recipient_hash": "hash123", + } + mock_hit.meta.score = 0.95 + + mock_response = Mock() + mock_response.hits = [mock_hit] + mock_search = Mock() + mock_search.handle_execute.return_value = mock_response + mock_tool._build_search.return_value = mock_search + + result = fuzzy_search_recipients("acme", limit=5) + + assert len(result) == 1 + assert result[0]["recipient_name"] == "ACME CORP" + assert result[0]["uei"] == "UEI123456789" + assert result[0]["duns"] == "123456789" + assert result[0]["recipient_level"] == "P" + assert result[0]["recipient_hash"] == "hash123" + assert result[0]["score"] == 0.95 + mock_tool._build_search.assert_called_once_with("ACME", top_k=5) + + @patch("usaspending_api.llm.tests.helper._tool") + @patch("usaspending_api.llm.tests.helper.es_sanitize") + def test_handles_multiple_hits(self, mock_sanitize, mock_tool): + """Test that multiple hits are properly formatted""" + mock_sanitize.return_value = "corp" + + # Create multiple mock hits + mock_hit1 = Mock() + mock_hit1.to_dict.return_value = { + "recipient_name": "ACME CORP", + "uei": "UEI111", + "duns": "111", + "recipient_level": "P", + "recipient_hash": "hash1", + } + mock_hit1.meta.score = 0.95 + + mock_hit2 = Mock() + mock_hit2.to_dict.return_value = { + "recipient_name": "BETA CORP", + "uei": "UEI222", + "duns": "222", + "recipient_level": "C", + "recipient_hash": "hash2", + } + mock_hit2.meta.score = 0.85 + + mock_response = Mock() + mock_response.hits = [mock_hit1, mock_hit2] + mock_search = Mock() + mock_search.handle_execute.return_value = mock_response + mock_tool._build_search.return_value = mock_search + + result = fuzzy_search_recipients("corp") + + assert len(result) == 2 + assert result[0]["recipient_name"] == "ACME CORP" + assert result[1]["recipient_name"] == "BETA CORP" + + @patch("usaspending_api.llm.tests.helper._tool") + @patch("usaspending_api.llm.tests.helper.es_sanitize") + def test_handles_missing_fields_in_hits(self, mock_sanitize, mock_tool): + """Test that missing fields are handled gracefully""" + mock_sanitize.return_value = "test" + + mock_hit = Mock() + mock_hit.to_dict.return_value = { + "recipient_name": "TEST CORP", + # Missing uei, duns, etc. + } + mock_hit.meta.score = 0.5 + + mock_response = Mock() + mock_response.hits = [mock_hit] + mock_search = Mock() + mock_search.handle_execute.return_value = mock_response + mock_tool._build_search.return_value = mock_search + + result = fuzzy_search_recipients("test") + + assert len(result) == 1 + assert result[0]["recipient_name"] == "TEST CORP" + assert result[0]["uei"] is None + assert result[0]["duns"] is None + assert result[0]["recipient_level"] is None + assert result[0]["recipient_hash"] is None + assert result[0]["score"] == 0.5 + + @patch("usaspending_api.llm.tests.helper._tool") + @patch("usaspending_api.llm.tests.helper.es_sanitize") + def test_respects_limit_parameter(self, mock_sanitize, mock_tool): + """Test that limit parameter is passed to _build_search""" + mock_sanitize.return_value = "test" + mock_response = Mock() + mock_response.hits = [] + mock_search = Mock() + mock_search.handle_execute.return_value = mock_response + mock_tool._build_search.return_value = mock_search + + fuzzy_search_recipients("test", limit=25) + + mock_tool._build_search.assert_called_once_with("TEST", top_k=25) + + +class TestRetrieveRecipientNames: + """Tests for retrieve_recipient_names function""" + + @patch("usaspending_api.llm.tests.helper._tool") + def test_returns_list_of_recipient_names(self, mock_tool): + """Test that function returns a list of recipient names""" + mock_tool.lookup_recipient.return_value = [ + "ACME CORP", + "UEI123456789", + "123456789", + ] + + result = retrieve_recipient_names("acme") + + mock_tool.lookup_recipient.assert_called_once_with("acme", top_k=5) + assert result is not None, "Function returned None" + assert isinstance(result, list), f"Expected list, got {type(result)}" + assert result == ["ACME CORP", "UEI123456789", "123456789"] + + @patch("usaspending_api.llm.tests.helper._tool") + def test_handles_no_results(self, mock_tool): + """Test when no results are returned""" + mock_tool.lookup_recipient.return_value = [] + + result = retrieve_recipient_names("nonexistent") + + assert isinstance(result, list) + assert result == [] + + @patch("usaspending_api.llm.tests.helper._tool") + def test_respects_limit_parameter(self, mock_tool): + """Test that limit parameter is passed to lookup_recipient""" + mock_tool.lookup_recipient.return_value = [] + + retrieve_recipient_names("test", limit=10) + + mock_tool.lookup_recipient.assert_called_once_with("test", top_k=10) + + @patch("usaspending_api.llm.tests.helper._tool") + def test_uses_default_limit(self, mock_tool): + """Test that default limit is used when not specified""" + mock_tool.lookup_recipient.return_value = [] + + retrieve_recipient_names("test") + + mock_tool.lookup_recipient.assert_called_once_with("test", top_k=5) + + @patch("usaspending_api.llm.tests.helper._tool") + def test_returns_deduplicated_list(self, mock_tool): + """Test that lookup_recipient returns deduplicated results""" + mock_tool.lookup_recipient.return_value = [ + "ACME CORP", + "UEI123", + "456", # No duplicates in return value + ] + + result = retrieve_recipient_names("acme") + + assert isinstance(result, list) + assert len(result) == 3 + assert "ACME CORP" in result + assert "UEI123" in result + assert "456" in result + + @patch("usaspending_api.llm.tests.helper._tool") + def test_handles_multiple_recipients(self, mock_tool): + """Test handling multiple recipient identifiers""" + mock_tool.lookup_recipient.return_value = [ + "ACME CORP", + "UEI111", + "111", + "BETA CORP", + "UEI222", + "222", + ] + + result = retrieve_recipient_names("corp") + + assert isinstance(result, list) + assert len(result) == 6 + assert "ACME CORP" in result + assert "BETA CORP" in result + assert "UEI111" in result + assert "UEI222" in result + + @patch("usaspending_api.llm.tests.helper._tool") + def test_passes_search_text_unchanged(self, mock_tool): + """Test that search text is passed to lookup_recipient unchanged""" + mock_tool.lookup_recipient.return_value = [] + + retrieve_recipient_names(" Test Query ", limit=10) + + # Should pass the search text as-is (lookup_recipient handles sanitization) + mock_tool.lookup_recipient.assert_called_once_with(" Test Query ", top_k=10) + + @patch("usaspending_api.llm.tests.helper._tool") + def test_returns_empty_list_on_none(self, mock_tool): + """Test handling when lookup_recipient returns None""" + mock_tool.lookup_recipient.return_value = None + + result = retrieve_recipient_names("test") + + # Should handle None gracefully + assert result is None or result == [] diff --git a/usaspending_api/llm/tools/lookup_recipient.py b/usaspending_api/llm/tools/lookup_recipient.py new file mode 100644 index 0000000000..a66a32c02e --- /dev/null +++ b/usaspending_api/llm/tools/lookup_recipient.py @@ -0,0 +1,117 @@ +import logging +from typing import Any + +from opensearchpy.helpers.query import Q as ES_Q + +from usaspending_api.common.elasticsearch.search_wrappers import RecipientSearch +from usaspending_api.llm.models.py_models import AITool, AIToolDescription +from usaspending_api.search.v2.es_sanitization import es_sanitize + +logger = logging.getLogger(__name__) + + +# take a string name, uei or duns and get list of entities and subs +# uses recipient_retrieval.py + +class RecipientLookupTool: + """Tool for looking up recipients in OpenSearch with fuzzy matching support.""" + + RECIPIENT_SOURCE_FIELDS = [ + "recipient_name", + "uei", + "duns", + "recipient_level", + "recipient_hash", + ] + + def lookup_recipient( + self, + query: str, + top_k: int = 10, + ) -> list[str]: + """ + Search for recipients by name, uei, duns, and return recipient names. + """ + if not query or not query.strip(): + return [] + + top_k = max(1, min(top_k, 100)) + query_upper = es_sanitize(query).strip().upper() + + try: + search = self._build_search(query_upper, top_k) + response = search.handle_execute() + except Exception as exception: + logger.error(f"OpenSearch query failed for query='{query}': {str(exception)}", exc_info=True) + return [] + return self._extract_recipient_names(response) + + def _build_search(self, query_upper: str, top_k: int) -> RecipientSearch: + should_queries = [] + for field in ("recipient_name", "uei", "duns"): + should_queries.extend( + [ + ES_Q("term", **{f"{field}__keyword": {"value": query_upper, "boost": 10.0, }}), + ES_Q("match", **{field: {"query": query_upper, "boost": 8.0}}), + ES_Q("match", **{field: {"query": query_upper, "fuzziness": "AUTO", "boost": 5.0}}), + ES_Q("match", **{f"{field}__contains": {"query": query_upper, "boost": 3.0}}), + ES_Q("wildcard", **{f"{field}__keyword": {"value": f"{query_upper}*", "boost": 2.0}}), + ] + ) + should_queries_dict = [q.to_dict() for q in should_queries] + + return ( + RecipientSearch() + .query("bool", should=should_queries_dict, minimum_should_match=1) + .source(list(self.RECIPIENT_SOURCE_FIELDS)) + .sort({"_score": {"order": "desc"}})[:top_k] + ) + + def _extract_recipient_names(self, response: Any) -> list[str]: + recipient_names = [] + seen_names = set() + for hit in response.hits: + recipient_name = hit.to_dict().get("recipient_name") + if not recipient_name or recipient_name in seen_names: + continue + seen_names.add(recipient_name) + recipient_names.append(recipient_name) + return recipient_names + + +lookup_recipient_tool = AITool( + description=AIToolDescription( + name="lookup_recipient", + description=""" +Search for valid recipient objects by name, UEI or DUNS using fuzzy matching. + +Returns a list of strings. + +Supported inputs: +- Recipient names (eg 'BOEING COMPANY', 'Lockheed Martin') +- UEI codes (12-character alphanumeric) +- DUNS numbers (9-digit, legacy) + +Examples: +- lookup_recipient('BOEING') -> ['BOEING COMPANY', ...] +- lookup_recipient('EWN9HP5FT8A5') -> ['BOEING COMPANY', ...] + +""".strip(), + input_schema={ + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "Recipient search (name, uai, duns)" + }, + "top_k": { + "type": "integer", + "description": "Maximum number of recipient results to return (1-100, default: 10)" + } + }, + "required": ["query"], + }, + ), + function=RecipientLookupTool().lookup_recipient, + logging=lambda tool_input: f"Searching the recipient index for '{tool_input.get('query', 'N/A')}'", +)