diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 8ed65aa3dc..c7f703f99b 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -26,8 +26,8 @@ repos: hooks: - id: mypy name: mypy - entry: poetry run mypy - args: ["--scripts-are-modules", "--config-file=pyproject.toml"] + entry: poetry run mypy . + args: ["--config-file=pyproject.toml"] language: system - types: [python] + pass_filenames: false exclude: tests/llm/fixtures/.* diff --git a/holmes/common/env_vars.py b/holmes/common/env_vars.py index 79d6145920..6e9c72c91c 100644 --- a/holmes/common/env_vars.py +++ b/holmes/common/env_vars.py @@ -81,3 +81,7 @@ def load_bool(env_var, default: Optional[bool]) -> Optional[bool]: TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT = float( os.environ.get("TOOL_MAX_ALLOCATED_CONTEXT_WINDOW_PCT", 15) ) + +MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION = int( + os.environ.get("MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION", 3000) +) diff --git a/holmes/core/supabase_dal.py b/holmes/core/supabase_dal.py index ef027d27ce..82db84b18c 100644 --- a/holmes/core/supabase_dal.py +++ b/holmes/core/supabase_dal.py @@ -30,6 +30,9 @@ ResourceInstructionDocument, ResourceInstructions, ) +from holmes.core.truncation.dal_truncation_utils import ( + truncate_evidences_entities_if_necessary, +) from holmes.utils.definitions import RobustaConfig from holmes.utils.env import get_env_replacement from holmes.utils.global_instructions import Instructions @@ -46,6 +49,9 @@ SCANS_META_TABLE = "ScansMeta" SCANS_RESULTS_TABLE = "ScansResults" +ENRICHMENT_BLACKLIST = ["text_file", "graph", "ai_analysis", "holmes"] +ENRICHMENT_BLACKLIST_SET = set(ENRICHMENT_BLACKLIST) + class RobustaToken(BaseModel): store_url: str @@ -262,11 +268,14 @@ def get_configuration_changes( .select("*") .eq("account_id", self.account_id) .in_("issue_id", changes_ids) + .not_.in_("enrichment_type", ENRICHMENT_BLACKLIST) .execute() ) if not len(change_data_response.data): return None + truncate_evidences_entities_if_necessary(change_data_response.data) + except Exception: logging.exception("Supabase error while retrieving change content") return None @@ -323,11 +332,10 @@ def unzip_evidence_file(self, data): return data def extract_relevant_issues(self, evidence): - enrichment_blacklist = {"text_file", "graph", "ai_analysis", "holmes"} data = [ enrich for enrich in evidence.data - if enrich.get("enrichment_type") not in enrichment_blacklist + if enrich.get("enrichment_type") not in ENRICHMENT_BLACKLIST_SET ] unzipped_files = [ @@ -370,12 +378,14 @@ def get_issue_data(self, issue_id: Optional[str]) -> Optional[Dict]: evidence = ( self.client.table(EVIDENCE_TABLE) .select("*") - .filter("issue_id", "eq", issue_id) + .eq("issue_id", issue_id) + .not_.in_("enrichment_type", ENRICHMENT_BLACKLIST) .execute() ) - data = self.extract_relevant_issues(evidence) + relevant_evidence = self.extract_relevant_issues(evidence) + truncate_evidences_entities_if_necessary(relevant_evidence) - issue_data["evidence"] = data + issue_data["evidence"] = relevant_evidence # build issue investigation dates started_at = issue_data.get("starts_at") @@ -518,10 +528,13 @@ def get_workload_issues(self, resource: dict, since_hours: float) -> List[str]: self.client.table(EVIDENCE_TABLE) .select("data, enrichment_type") .in_("issue_id", unique_issues) + .not_.in_("enrichment_type", ENRICHMENT_BLACKLIST) .execute() ) - return self.extract_relevant_issues(res) + relevant_issues = self.extract_relevant_issues(res) + truncate_evidences_entities_if_necessary(relevant_issues) + return relevant_issues except Exception: logging.exception("failed to fetch workload issues data", exc_info=True) diff --git a/holmes/core/truncation/dal_truncation_utils.py b/holmes/core/truncation/dal_truncation_utils.py new file mode 100644 index 0000000000..560f0aac22 --- /dev/null +++ b/holmes/core/truncation/dal_truncation_utils.py @@ -0,0 +1,23 @@ +from holmes.common.env_vars import MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION + + +def truncate_string(data_str: str) -> str: + if data_str and len(data_str) > MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION: + return ( + data_str[:MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION] + + "-- DATA TRUNCATED TO AVOID HITTING CONTEXT WINDOW LIMITS" + ) + return data_str + + +def truncate_evidences_entities_if_necessary(evidence_list: list[dict]): + if ( + not MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION + or MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION <= 0 + ): + return + + for evidence in evidence_list: + data = evidence.get("data") + if data: + evidence["data"] = truncate_string(str(data)) diff --git a/pyproject.toml b/pyproject.toml index cfb915dba4..85194c3702 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -83,6 +83,7 @@ ignore_missing_imports = true scripts_are_modules = true exclude = [ "tests/llm/fixtures/.*", + "dist/.*", ] [tool.pytest.ini_options] diff --git a/tests/core/truncation/__init__.py b/tests/core/truncation/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/core/truncation/test_dal_truncation_utils.py b/tests/core/truncation/test_dal_truncation_utils.py new file mode 100644 index 0000000000..b74b36dda1 --- /dev/null +++ b/tests/core/truncation/test_dal_truncation_utils.py @@ -0,0 +1,282 @@ +from unittest.mock import patch +from holmes.core.truncation.dal_truncation_utils import ( + truncate_evidences_entities_if_necessary, +) + + +class TestTruncateEvidencesEntitiesIfNecessary: + """Test cases for the truncate_evidences_entities_if_necessary function.""" + + @patch( + "holmes.core.truncation.dal_truncation_utils.MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION", + 100, + ) + def test_truncate_long_evidence_data(self): + """Test that evidence data longer than the limit gets truncated.""" + long_data = "a" * 150 # 150 characters, exceeds limit of 100 + evidence_list = [ + {"data": long_data, "id": "test-1"}, + {"data": "short", "id": "test-2"}, + ] + + truncate_evidences_entities_if_necessary(evidence_list) + + # First evidence should be truncated + expected_truncated = ( + "a" * 100 + "-- DATA TRUNCATED TO AVOID HITTING CONTEXT WINDOW LIMITS" + ) + assert evidence_list[0]["data"] == expected_truncated + # Second evidence should remain unchanged + assert evidence_list[1]["data"] == "short" + + @patch( + "holmes.core.truncation.dal_truncation_utils.MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION", + 100, + ) + def test_no_truncation_when_data_within_limit(self): + """Test that evidence data within the limit remains unchanged.""" + short_data = "a" * 50 # 50 characters, within limit of 100 + evidence_list = [ + {"data": short_data, "id": "test-1"}, + {"data": "very short", "id": "test-2"}, + ] + + original_data_0 = evidence_list[0]["data"] + original_data_1 = evidence_list[1]["data"] + + truncate_evidences_entities_if_necessary(evidence_list) + + # Both should remain unchanged + assert evidence_list[0]["data"] == original_data_0 + assert evidence_list[1]["data"] == original_data_1 + + @patch( + "holmes.core.truncation.dal_truncation_utils.MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION", + 100, + ) + def test_truncation_at_exact_limit(self): + """Test behavior when data is exactly at the limit.""" + exact_limit_data = "a" * 100 # Exactly 100 characters + evidence_list = [{"data": exact_limit_data, "id": "test-1"}] + + original_data = evidence_list[0]["data"] + + truncate_evidences_entities_if_necessary(evidence_list) + + # Should remain unchanged (not greater than limit) + assert evidence_list[0]["data"] == original_data + + @patch( + "holmes.core.truncation.dal_truncation_utils.MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION", + 100, + ) + def test_truncation_one_character_over_limit(self): + """Test behavior when data is one character over the limit.""" + over_limit_data = "a" * 101 # 101 characters, one over limit of 100 + evidence_list = [{"data": over_limit_data, "id": "test-1"}] + + truncate_evidences_entities_if_necessary(evidence_list) + + # Should be truncated + expected_truncated = ( + "a" * 100 + "-- DATA TRUNCATED TO AVOID HITTING CONTEXT WINDOW LIMITS" + ) + assert evidence_list[0]["data"] == expected_truncated + + @patch( + "holmes.core.truncation.dal_truncation_utils.MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION", + None, + ) + def test_no_truncation_when_limit_is_none(self): + """Test that no truncation occurs when MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION is None.""" + long_data = "a" * 10000 + evidence_list = [{"data": long_data, "id": "test-1"}] + + original_data = evidence_list[0]["data"] + + truncate_evidences_entities_if_necessary(evidence_list) + + # Should remain unchanged + assert evidence_list[0]["data"] == original_data + + @patch( + "holmes.core.truncation.dal_truncation_utils.MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION", + 0, + ) + def test_no_truncation_when_limit_is_zero(self): + """Test that no truncation occurs when MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION is 0.""" + long_data = "a" * 1000 + evidence_list = [{"data": long_data, "id": "test-1"}] + + original_data = evidence_list[0]["data"] + + truncate_evidences_entities_if_necessary(evidence_list) + + # Should remain unchanged + assert evidence_list[0]["data"] == original_data + + @patch( + "holmes.core.truncation.dal_truncation_utils.MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION", + -1, + ) + def test_no_truncation_when_limit_is_negative(self): + """Test that no truncation occurs when MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION is negative.""" + long_data = "a" * 1000 + evidence_list = [{"data": long_data, "id": "test-1"}] + + original_data = evidence_list[0]["data"] + + truncate_evidences_entities_if_necessary(evidence_list) + + # Should remain unchanged + assert evidence_list[0]["data"] == original_data + + @patch( + "holmes.core.truncation.dal_truncation_utils.MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION", + 100, + ) + def test_empty_evidence_list(self): + """Test that function handles empty evidence list without errors.""" + evidence_list = [] + + truncate_evidences_entities_if_necessary(evidence_list) + + # Should remain empty + assert evidence_list == [] + + @patch( + "holmes.core.truncation.dal_truncation_utils.MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION", + 100, + ) + def test_evidence_without_data_field(self): + """Test that evidence without 'data' field is handled gracefully.""" + evidence_list = [ + {"id": "test-1", "type": "log"}, + {"data": "valid_data", "id": "test-2"}, + ] + + truncate_evidences_entities_if_necessary(evidence_list) + + # First evidence should remain unchanged (no data field) + assert evidence_list[0] == {"id": "test-1", "type": "log"} + # Second evidence should remain unchanged (data is short) + assert evidence_list[1]["data"] == "valid_data" + + @patch( + "holmes.core.truncation.dal_truncation_utils.MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION", + 100, + ) + def test_evidence_with_none_data(self): + """Test that evidence with None data is handled gracefully.""" + evidence_list = [ + {"data": None, "id": "test-1"}, + {"data": "valid_data", "id": "test-2"}, + ] + + truncate_evidences_entities_if_necessary(evidence_list) + + # First evidence should remain unchanged (data is None) + assert evidence_list[0]["data"] is None + # Second evidence should remain unchanged (data is short) + assert evidence_list[1]["data"] == "valid_data" + + @patch( + "holmes.core.truncation.dal_truncation_utils.MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION", + 100, + ) + def test_evidence_with_non_string_data(self): + """Test that evidence with non-string data is converted to string and truncated if needed.""" + large_dict = { + "key" + str(i): "value" + str(i) for i in range(20) + } # Creates a long string representation + evidence_list = [ + {"data": large_dict, "id": "test-1"}, + {"data": 12345, "id": "test-2"}, + ] + + truncate_evidences_entities_if_necessary(evidence_list) + + # Dict data should be converted to string and potentially truncated + dict_str = str(large_dict) + if len(dict_str) > 100: + expected_truncated = ( + dict_str[:100] + + "-- DATA TRUNCATED TO AVOID HITTING CONTEXT WINDOW LIMITS" + ) + assert evidence_list[0]["data"] == expected_truncated + else: + assert evidence_list[0]["data"] == dict_str + + # Integer data should be converted to string and remain unchanged (short) + assert evidence_list[1]["data"] == "12345" + + @patch( + "holmes.core.truncation.dal_truncation_utils.MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION", + 50, + ) + def test_multiple_evidences_with_mixed_lengths(self): + """Test truncation with multiple evidences of varying lengths.""" + evidence_list = [ + {"data": "a" * 25, "id": "short"}, # Within limit + {"data": "b" * 75, "id": "long"}, # Over limit + {"data": "c" * 50, "id": "exact"}, # Exactly at limit + {"data": "d" * 100, "id": "very_long"}, # Way over limit + ] + + truncate_evidences_entities_if_necessary(evidence_list) + + # Short data should remain unchanged + assert evidence_list[0]["data"] == "a" * 25 + + # Long data should be truncated + expected_long = ( + "b" * 50 + "-- DATA TRUNCATED TO AVOID HITTING CONTEXT WINDOW LIMITS" + ) + assert evidence_list[1]["data"] == expected_long + + # Exact limit should remain unchanged + assert evidence_list[2]["data"] == "c" * 50 + + # Very long data should be truncated + expected_very_long = ( + "d" * 50 + "-- DATA TRUNCATED TO AVOID HITTING CONTEXT WINDOW LIMITS" + ) + assert evidence_list[3]["data"] == expected_very_long + + @patch( + "holmes.core.truncation.dal_truncation_utils.MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION", + 100, + ) + def test_function_modifies_original_list(self): + """Test that the function modifies the original list in-place.""" + long_data = "x" * 150 + evidence_list = [{"data": long_data, "id": "test-1"}] + original_list_id = id(evidence_list) + original_dict_id = id(evidence_list[0]) + + truncate_evidences_entities_if_necessary(evidence_list) + + # The list and dictionary objects should be the same (modified in-place) + assert id(evidence_list) == original_list_id + assert id(evidence_list[0]) == original_dict_id + + # But the data should be different + assert evidence_list[0]["data"] != long_data + assert evidence_list[0]["data"].endswith( + "-- DATA TRUNCATED TO AVOID HITTING CONTEXT WINDOW LIMITS" + ) + + @patch( + "holmes.core.truncation.dal_truncation_utils.MAX_EVIDENCE_DATA_CHARACTERS_BEFORE_TRUNCATION", + 20, + ) + def test_truncation_message_consistency(self): + """Test that the truncation message is consistent.""" + long_data = "a" * 100 + evidence_list = [{"data": long_data, "id": "test-1"}] + + truncate_evidences_entities_if_necessary(evidence_list) + + expected_suffix = "-- DATA TRUNCATED TO AVOID HITTING CONTEXT WINDOW LIMITS" + assert evidence_list[0]["data"].endswith(expected_suffix) + assert evidence_list[0]["data"].startswith("a" * 20)