AIML-84 Respond to PR feedback and rename extension files to smartfix_* & reduce filesystem MCP server timeouts

Author: Jacob Mages-Haskins

src/agent_handler.py

@@ -52,8 +52,8 @@
 try:
     from google.adk.agents import Agent
     from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
-    from src.extensions.extended_litellm import ExtendedLiteLlm
-    from src.extensions.extended_llm_agent import ExtendedLlmAgent
+    from src.extensions.smartfix_litellm import SmartFixLiteLlm
+    from src.extensions.smartfix_llm_agent import SmartFixLlmAgent
     from google.adk.runners import Runner
     from google.adk.sessions import InMemorySessionService
     from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset, StdioServerParameters, StdioConnectionParams
@@ -80,10 +80,10 @@
 async def _create_mcp_toolset(target_folder_str: str) -> MCPToolset:
     """Create MCP toolset with platform-specific configuration."""
     if platform.system() == 'Windows':
-        connection_timeout = 300
+        connection_timeout = 180
         debug_log("Using Windows-specific MCP connection settings")
     else:
-        connection_timeout = 180
+        connection_timeout = 120
 
     return MCPToolset(
         connection_params=StdioConnectionParams(
@@ -208,13 +208,13 @@ async def create_agent(target_folder: Path, remediation_id: str, agent_type: str
     agent_name = f"contrast_{agent_type}_agent"
 
     try:
-        model_instance = ExtendedLiteLlm(
+        model_instance = SmartFixLiteLlm(
             model=config.AGENT_MODEL,
             temperature=0.2,  # Set low temperature for more deterministic output
             # seed=42, # The random seed for reproducibility (not supported by bedrock/anthropic atm call throws error)
             stream_options={"include_usage": True}
         )
-        root_agent = ExtendedLlmAgent(
+        root_agent = SmartFixLlmAgent(
             model=model_instance,
             name=agent_name,
             instruction=agent_instruction,

src/extensions/extended_litellm.py renamed to src/extensions/smartfix_litellm.py

@@ -1,6 +1,6 @@
 # -
 # #%L
-# Extended LiteLLM with Prompt Caching
+# SmartFix LiteLLM with Prompt Caching
 # %%
 # Copyright (C) 2025 Contrast Security, Inc.
 # %%
@@ -58,7 +58,7 @@ def reset(self):
 
     def add_usage(self, input_tokens: int, output_tokens: int, cache_read_tokens: int,
                   cache_write_tokens: int, new_input_cost: float, cache_read_cost: float,
-                  cache_write_cost: float, output_cost: float):
+                  cache_write_cost: float, output_cost: float) -> None:
         """Add usage statistics from a single LLM call."""
         # Accumulate tokens
         self.total_new_input_tokens += input_tokens
@@ -114,8 +114,8 @@ def cache_savings_percentage(self):
         return 0.0
 
 
-class ExtendedLiteLlm(LiteLlm):
-    """Extended LiteLlm with automatic prompt caching and comprehensive cost analysis.
+class SmartFixLiteLlm(LiteLlm):
+    """SmartFix LiteLlm with automatic prompt caching and comprehensive cost analysis.
 
     This class extends the base LiteLlm to automatically apply prompt caching
     and provide detailed cost analysis for all LLM interactions:
@@ -131,16 +131,16 @@ class ExtendedLiteLlm(LiteLlm):
     Example usage:
     ```python
     # Bedrock Claude - Works with all features and cost tracking
-    model = ExtendedLiteLlm(model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0")
+    model = SmartFixLiteLlm(model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0")
 
     # Direct Anthropic - Works with all features and cost tracking
-    model = ExtendedLiteLlm(model="anthropic/claude-3-7-sonnet-20250219")
+    model = SmartFixLiteLlm(model="anthropic/claude-3-7-sonnet-20250219")
 
     # OpenAI - works with cost tracking (no caching applied)
-    model = ExtendedLiteLlm(model="openai/gpt-4o")
+    model = SmartFixLiteLlm(model="openai/gpt-4o")
 
     # Other models - work normally with cost tracking
-    model = ExtendedLiteLlm(model="gemini/gemini-1.5-pro")
+    model = SmartFixLiteLlm(model="gemini/gemini-1.5-pro")
     ```
     """
 

src/extensions/extended_llm_agent.py renamed to src/extensions/smartfix_llm_agent.py

@@ -1,6 +1,6 @@
 # -
 # #%L
-# Extended LLM Agent
+# SmartFix LLM Agent
 # %%
 # Copyright (C) 2025 Contrast Security, Inc.
 # %%
@@ -31,32 +31,32 @@
 from google.adk.agents import LlmAgent
 from src.utils import debug_log
 
-# Import ExtendedLiteLlm at module level so it's available for model_rebuild()
-ExtendedLiteLlm = None
+# Import SmartFixLiteLlm at module level so it's available for model_rebuild()
+SmartFixLiteLlm = None
 try:
-    from .extended_litellm import ExtendedLiteLlm
+    from .smartfix_litellm import SmartFixLiteLlm
 except ImportError:
     pass
 
 
-class ExtendedLlmAgent(LlmAgent):
-    """Extended LLM Agent that preserves ExtendedLiteLlm statistics across calls.
+class SmartFixLlmAgent(LlmAgent):
+    """SmartFix LLM Agent that preserves SmartFixLiteLlm statistics across calls.
 
-    This class solves the issue where ExtendedLiteLlm accumulated statistics
+    This class solves the issue where SmartFixLiteLlm accumulated statistics
     aren't preserved when using the model within a Google ADK Agent. It ensures
-    that the same ExtendedLiteLlm instance is used for all LLM calls and provides
+    that the same SmartFixLiteLlm instance is used for all LLM calls and provides
     convenient methods to access accumulated statistics.
 
     Example usage:
     ```python
-    from src.extensions.extended_litellm import ExtendedLiteLlm
-    from src.extensions.extended_llm_agent import ExtendedLlmAgent
+    from src.extensions.smartfix_litellm import SmartFixLiteLlm
+    from src.extensions.smartfix_llm_agent import SmartFixLlmAgent
 
     # Create the extended model
-    model = ExtendedLiteLlm(model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0")
+    model = SmartFixLiteLlm(model="bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0")
 
     # Create the extended agent
-    agent = ExtendedLlmAgent(
+    agent = SmartFixLlmAgent(
         name="my-agent",
         model=model,
         instruction="You are a helpful assistant."
@@ -77,48 +77,48 @@ class ExtendedLlmAgent(LlmAgent):
     original_extended_model: Optional[Any] = Field(
         default=None,
         exclude=True,
-        description="Reference to the original ExtendedLiteLlm instance for stats access"
+        description="Reference to the original SmartFixLiteLlm instance for stats access"
     )
 
     @model_validator(mode='after')
     def _preserve_extended_model_reference(self):
-        """Preserve a reference to the ExtendedLiteLlm instance if provided."""
+        """Preserve a reference to the SmartFixLiteLlm instance if provided."""
         # Import here to avoid circular imports
         try:
-            from .extended_litellm import ExtendedLiteLlm
+            from .smartfix_litellm import SmartFixLiteLlm
 
-            if isinstance(self.model, ExtendedLiteLlm):
+            if isinstance(self.model, SmartFixLiteLlm):
                 # Store reference to the original instance
                 self.original_extended_model = self.model
-                debug_log(f"[EXTENDED_AGENT] Preserved reference to ExtendedLiteLlm instance for agent: {self.name}")
+                debug_log(f"[SMARTFIX_AGENT] Preserved reference to SmartFixLiteLlm instance for agent: {self.name}")
 
         except ImportError:
-            # ExtendedLiteLlm not available, ignore
+            # SmartFixLiteLlm not available, ignore
             pass
 
         return self
 
     @override
     @property
     def canonical_model(self):
-        """Override to ensure we return the original ExtendedLiteLlm instance."""
-        # If we have a preserved ExtendedLiteLlm instance, return it
+        """Override to ensure we return the original SmartFixLiteLlm instance."""
+        # If we have a preserved SmartFixLiteLlm instance, return it
         if self.original_extended_model is not None:
             return self.original_extended_model
 
         # Otherwise, use the parent's canonical_model
         return super().canonical_model
 
     def has_extended_model(self) -> bool:
-        """Check if this agent is using an ExtendedLiteLlm model."""
+        """Check if this agent is using an SmartFixLiteLlm model."""
         try:
-            from .extended_litellm import ExtendedLiteLlm
-            return isinstance(self.canonical_model, ExtendedLiteLlm)
+            from .smartfix_litellm import SmartFixLiteLlm
+            return isinstance(self.canonical_model, SmartFixLiteLlm)
         except ImportError:
             return False
 
     def get_extended_model(self) -> Optional[Any]:
-        """Get the ExtendedLiteLlm instance if available."""
+        """Get the SmartFixLiteLlm instance if available."""
         if self.has_extended_model():
             return self.canonical_model
         return None
@@ -127,18 +127,18 @@ def gather_accumulated_stats_dict(self) -> dict:
         """Get accumulated token usage and cost statistics as dictionary.
 
         This method provides programmatic access to the accumulated statistics
-        from the ExtendedLiteLlm instance being used by this agent.
+        from the SmartFixLiteLlm instance being used by this agent.
 
         Returns:
             dict: Dictionary containing accumulated statistics
 
         Raises:
-            ValueError: If the agent is not using an ExtendedLiteLlm model.
+            ValueError: If the agent is not using an SmartFixLiteLlm model.
         """
         extended_model = self.get_extended_model()
         if extended_model is None:
             raise ValueError(
-                f"Agent '{self.name}' is not using an ExtendedLiteLlm model. "
+                f"Agent '{self.name}' is not using an SmartFixLiteLlm model. "
                 "Cannot access accumulated statistics. "
                 f"Current model type: {type(self.canonical_model).__name__}"
             )
@@ -149,18 +149,18 @@ def gather_accumulated_stats(self) -> str:
         """Get accumulated token usage and cost statistics as JSON string.
 
         This method provides programmatic access to the accumulated statistics
-        from the ExtendedLiteLlm instance being used by this agent.
+        from the SmartFixLiteLlm instance being used by this agent.
 
         Returns:
             str: JSON formatted string containing accumulated statistics
 
         Raises:
-            ValueError: If the agent is not using an ExtendedLiteLlm model.
+            ValueError: If the agent is not using an SmartFixLiteLlm model.
         """
         extended_model = self.get_extended_model()
         if extended_model is None:
             raise ValueError(
-                f"Agent '{self.name}' is not using an ExtendedLiteLlm model. "
+                f"Agent '{self.name}' is not using an SmartFixLiteLlm model. "
                 "Cannot access accumulated statistics. "
                 f"Current model type: {type(self.canonical_model).__name__}"
             )
@@ -171,12 +171,12 @@ def reset_accumulated_stats(self) -> None:
         """Reset accumulated statistics to start fresh.
 
         Raises:
-            ValueError: If the agent is not using an ExtendedLiteLlm model.
+            ValueError: If the agent is not using an SmartFixLiteLlm model.
         """
         extended_model = self.get_extended_model()
         if extended_model is None:
             raise ValueError(
-                f"Agent '{self.name}' is not using an ExtendedLiteLlm model. "
+                f"Agent '{self.name}' is not using an SmartFixLiteLlm model. "
                 "Cannot reset accumulated statistics. "
                 f"Current model type: {type(self.canonical_model).__name__}"
             )
@@ -197,12 +197,12 @@ def get_accumulated_stats_summary(self) -> dict:
                 - And more detailed breakdowns
 
         Raises:
-            ValueError: If the agent is not using an ExtendedLiteLlm model.
+            ValueError: If the agent is not using an SmartFixLiteLlm model.
         """
         extended_model = self.get_extended_model()
         if extended_model is None:
             raise ValueError(
-                f"Agent '{self.name}' is not using an ExtendedLiteLlm model. "
+                f"Agent '{self.name}' is not using an SmartFixLiteLlm model. "
                 "Cannot access accumulated statistics. "
                 f"Current model type: {type(self.canonical_model).__name__}"
             )
@@ -236,7 +236,7 @@ def get_model_info(self) -> dict:
             dict: Dictionary containing model information including:
                 - model_name: The model name/identifier
                 - model_type: The class name of the model
-                - is_extended: Whether it's an ExtendedLiteLlm instance
+                - is_extended: Whether it's an SmartFixLiteLlm instance
                 - has_stats: Whether accumulated statistics are available
         """
         model = self.canonical_model
@@ -251,11 +251,11 @@ def get_model_info(self) -> dict:
         }
 
 
-# Rebuild the model schema after ExtendedLiteLlm is available
+# Rebuild the model schema after SmartFixLiteLlm is available
 # This resolves forward references and ensures Pydantic can fully validate the model
-if ExtendedLiteLlm is not None:
+if SmartFixLiteLlm is not None:
     try:
-        ExtendedLlmAgent.model_rebuild()
+        SmartFixLlmAgent.model_rebuild()
     except Exception:
         # If rebuild fails for any reason, just continue
         # The class will still work, just without perfect type validation

test/test_extended_litellm.py renamed to test/test_smartfix_litellm.py

@@ -19,7 +19,7 @@
 #
 
 """
-Unit tests for ExtendedLiteLlm and TokenCostAccumulator classes.
+Unit tests for SmartFixLiteLlm and TokenCostAccumulator classes.
 
 This module tests the extended LiteLLM functionality including:
 - Token cost accumulation and statistics gathering
@@ -44,7 +44,7 @@
 _ = get_config(testing=True)
 
 # Import the classes under test AFTER config initialization
-from src.extensions.extended_litellm import ExtendedLiteLlm, TokenCostAccumulator  # noqa: E402
+from src.extensions.smartfix_litellm import SmartFixLiteLlm, TokenCostAccumulator  # noqa: E402
 
 
 class TestTokenCostAccumulator(unittest.TestCase):
@@ -262,21 +262,21 @@ def test_reset(self):
         self.assertEqual(self.accumulator.call_count, 0)
 
 
-class TestExtendedLiteLlm(unittest.TestCase):
-    """Test cases for ExtendedLiteLlm class."""
+class TestSmartFixLiteLlm(unittest.TestCase):
+    """Test cases for SmartFixLiteLlm class."""
 
     def setUp(self):
         """Set up test fixtures before each test method."""
         # Mock LiteLlm initialization to avoid dependencies
         with patch('litellm.completion'):
-            self.extended_model = ExtendedLiteLlm(model="test-model")
+            self.extended_model = SmartFixLiteLlm(model="test-model")
 
     def test_initialization(self):
-        """Test that ExtendedLiteLlm initializes correctly."""
+        """Test that SmartFixLiteLlm initializes correctly."""
         self.assertEqual(self.extended_model.model, "test-model")
         self.assertIsInstance(self.extended_model.cost_accumulator, TokenCostAccumulator)
 
-    @patch('src.extensions.extended_litellm.debug_log')
+    @patch('src.extensions.smartfix_litellm.debug_log')
     def test_gather_accumulated_stats_dict(self, mock_debug_log):
         """Test statistics dictionary generation."""
         # Add some usage to the accumulator
@@ -299,7 +299,7 @@ def test_gather_accumulated_stats_dict(self, mock_debug_log):
         self.assertIn('cost_analysis', stats)
         self.assertIn('averages', stats)
 
-    @patch('src.extensions.extended_litellm.debug_log')
+    @patch('src.extensions.smartfix_litellm.debug_log')
     def test_gather_accumulated_stats_json(self, mock_debug_log):
         """Test JSON statistics generation."""
         # Add some usage to the accumulator
@@ -321,7 +321,7 @@ def test_gather_accumulated_stats_json(self, mock_debug_log):
         self.assertEqual(stats_dict['call_count'], 1)
         self.assertIn('token_usage', stats_dict)
 
-    @patch('src.extensions.extended_litellm.debug_log')
+    @patch('src.extensions.smartfix_litellm.debug_log')
     def test_reset_accumulated_stats(self, mock_debug_log):
         """Test that reset clears accumulated statistics."""
         # Add some usage first
@@ -349,15 +349,15 @@ def test_reset_accumulated_stats(self, mock_debug_log):
         mock_debug_log.assert_called_with("Accumulated statistics have been reset.")
 
 
-class TestExtendedLiteLlmIntegration(unittest.TestCase):
-    """Integration tests for ExtendedLiteLlm functionality."""
+class TestSmartFixLiteLlmIntegration(unittest.TestCase):
+    """Integration tests for SmartFixLiteLlm functionality."""
 
     @patch('litellm.completion')
-    @patch('src.extensions.extended_litellm.debug_log')
+    @patch('src.extensions.smartfix_litellm.debug_log')
     def test_cost_accumulator_integration(self, mock_debug_log, mock_completion):
-        """Test that cost accumulator integrates properly with ExtendedLiteLlm."""
-        # Create a real ExtendedLiteLlm instance
-        model = ExtendedLiteLlm(model="test-integration-model")
+        """Test that cost accumulator integrates properly with SmartFixLiteLlm."""
+        # Create a real SmartFixLiteLlm instance
+        model = SmartFixLiteLlm(model="test-integration-model")
 
         # Verify it has a cost accumulator
         self.assertIsInstance(model.cost_accumulator, TokenCostAccumulator)