ROB-2034: Limit evidence size

.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/.*

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)
+)

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)

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))

pyproject.toml

@@ -83,6 +83,7 @@ ignore_missing_imports = true
 scripts_are_modules = true
 exclude = [
     "tests/llm/fixtures/.*",
+    "dist/.*",
 ]
 
 [tool.pytest.ini_options]