Do not add security_risk unless security analyzer is enabled (#341)

Co-authored-by: openhands <[email protected]>

Author: xingyaoww

examples/17_llm_security_analyzer.py

@@ -5,10 +5,11 @@
 """
 
 import os
+import uuid
 
 from pydantic import SecretStr
 
-from openhands.sdk import LLM, Agent, Conversation, Message, TextContent
+from openhands.sdk import LLM, Agent, Conversation, LocalFileStore, Message, TextContent
 from openhands.sdk.conversation.state import AgentExecutionStatus
 from openhands.sdk.event.utils import get_unmatched_actions
 from openhands.sdk.security.llm_analyzer import LLMSecurityAnalyzer
@@ -36,7 +37,12 @@
 # Create agent with security analyzer
 security_analyzer = LLMSecurityAnalyzer()
 agent = Agent(llm=llm, tools=tools, security_analyzer=security_analyzer)
-conversation = Conversation(agent=agent)
+
+conversation_id = uuid.uuid4()
+file_store = LocalFileStore(f"./.conversations/{conversation_id}")
+conversation = Conversation(
+    agent=agent, conversation_id=conversation_id, persist_filestore=file_store
+)
 
 print("\n1) Safe command (LOW risk - should execute automatically)...")
 conversation.send_message(

openhands/sdk/agent/agent.py

@@ -29,6 +29,8 @@
     get_llm_metadata,
 )
 from openhands.sdk.logger import get_logger
+from openhands.sdk.security import risk
+from openhands.sdk.security.llm_analyzer import LLMSecurityAnalyzer
 from openhands.sdk.tool import (
     BUILT_IN_TOOLS,
     ActionBase,
@@ -217,7 +219,16 @@ def step(
             f"{json.dumps([m.model_dump() for m in _messages], indent=2)}"
         )
         assert isinstance(self.tools, dict)
-        tools = [tool.to_openai_tool() for tool in self.tools.values()]
+
+        tools = [
+            # add llm security risk prediction if analyzer is present
+            tool.to_openai_tool(
+                add_security_risk_prediction=isinstance(
+                    self.security_analyzer, LLMSecurityAnalyzer
+                )
+            )
+            for tool in self.tools.values()
+        ]
         response = self.llm.completion(
             messages=_messages,
             tools=tools,
@@ -368,10 +379,28 @@ def _get_action_events(
             return
 
         # Validate arguments
+        security_risk: risk.SecurityRisk = risk.SecurityRisk.UNKNOWN
         try:
-            action: ActionBase = tool.action_type.model_validate(
-                json.loads(tool_call.function.arguments)
-            )
+            arguments = json.loads(tool_call.function.arguments)
+
+            # if the tool has a security_risk field (when security analyzer = LLM),
+            # pop it out as it's not part of the tool's action schema
+            if (_predicted_risk := arguments.pop("security_risk", None)) is not None:
+                if not isinstance(self.security_analyzer, LLMSecurityAnalyzer):
+                    raise RuntimeError(
+                        "LLM provided a security_risk but no security analyzer is "
+                        "configured - THIS SHOULD NOT HAPPEN!"
+                    )
+                try:
+                    security_risk = risk.SecurityRisk(_predicted_risk)
+                except ValueError:
+                    logger.warning(
+                        f"Invalid security_risk value from LLM: {_predicted_risk}"
+                    )
+
+            # Arguments we passed in should not contains `security_risk`
+            # as a field
+            action: ActionBase = tool.action_type.model_validate(arguments)
         except (json.JSONDecodeError, ValidationError) as e:
             err = (
                 f"Error validating args {tool_call.function.arguments} for tool "
@@ -394,6 +423,7 @@ def _get_action_events(
             tool_call=tool_call,
             llm_response_id=llm_response_id,
             metrics=metrics,
+            security_risk=security_risk,
         )
         on_event(action_event)
         return action_event

openhands/sdk/event/llm_convertible.py

@@ -10,6 +10,7 @@
 from openhands.sdk.event.types import EventID, SourceType, ToolCallID
 from openhands.sdk.llm import ImageContent, Message, TextContent, content_to_str
 from openhands.sdk.llm.utils.metrics import MetricsSnapshot
+from openhands.sdk.security import risk
 from openhands.sdk.tool.schema import Action, Observation
 
 
@@ -108,6 +109,10 @@ class ActionEvent(LLMConvertibleEvent):
             "to the last action when multiple actions share the same LLM response."
         ),
     )
+    security_risk: risk.SecurityRisk = Field(
+        default=risk.SecurityRisk.UNKNOWN,
+        description="The LLM's assessment of the safety risk of this action.",
+    )
 
     @property
     def visualize(self) -> Text:

openhands/sdk/security/analyzer.py

@@ -5,7 +5,6 @@
 from openhands.sdk.event.llm_convertible import ActionEvent
 from openhands.sdk.logger import get_logger
 from openhands.sdk.security.risk import SecurityRisk
-from openhands.sdk.tool.schema import Action
 from openhands.sdk.utils.discriminated_union import (
     DiscriminatedUnionMixin,
     DiscriminatedUnionType,
@@ -26,15 +25,15 @@ class SecurityAnalyzerBase(DiscriminatedUnionMixin, ABC):
     """
 
     @abstractmethod
-    def security_risk(self, action: Action) -> SecurityRisk:
-        """Evaluate the security risk of an action.
+    def security_risk(self, action: ActionEvent) -> SecurityRisk:
+        """Evaluate the security risk of an ActionEvent.
 
-        This is the core method that analyzes an action and returns its risk level.
+        This is the core method that analyzes an ActionEvent and returns its risk level.
         Implementations should examine the action's content, context, and potential
         impact to determine the appropriate risk level.
 
         Args:
-            action: The action to analyze for security risks
+            action: The ActionEvent to analyze for security risks
 
         Returns:
             ActionSecurityRisk enum indicating the risk level
@@ -54,7 +53,7 @@ def analyze_event(self, event: Event) -> SecurityRisk | None:
             ActionSecurityRisk if event is an action, None otherwise
         """
         if isinstance(event, ActionEvent):
-            return self.security_risk(event.action)
+            return self.security_risk(event)
         return None
 
     def should_require_confirmation(
@@ -103,7 +102,7 @@ def analyze_pending_actions(
 
         for action_event in pending_actions:
             try:
-                risk = self.security_risk(action_event.action)
+                risk = self.security_risk(action_event)
                 analyzed_actions.append((action_event, risk))
                 logger.debug(f"Action {action_event} analyzed with risk level: {risk}")
             except Exception as e:

openhands/sdk/security/llm_analyzer.py

@@ -1,7 +1,7 @@
+from openhands.sdk.event import ActionEvent
 from openhands.sdk.logger import get_logger
 from openhands.sdk.security.analyzer import SecurityAnalyzer
 from openhands.sdk.security.risk import SecurityRisk
-from openhands.sdk.tool.schema import Action
 
 
 logger = get_logger(__name__)
@@ -17,7 +17,7 @@ class LLMSecurityAnalyzer(SecurityAnalyzer):
     understanding of action context and potential risks.
     """
 
-    def security_risk(self, action: Action) -> SecurityRisk:
+    def security_risk(self, action: ActionEvent) -> SecurityRisk:
         """Evaluate security risk based on LLM-provided assessment.
 
         This method checks if the action has a security_risk attribute set by the LLM

openhands/sdk/tool/schema.py

@@ -4,7 +4,6 @@
 from pydantic import BaseModel, ConfigDict, Field, create_model
 from rich.text import Text
 
-import openhands.sdk.security.risk as risk
 from openhands.sdk.llm import ImageContent, TextContent
 from openhands.sdk.llm.message import content_to_str
 from openhands.sdk.utils.discriminated_union import (
@@ -168,14 +167,6 @@ def from_mcp_schema(
 class ActionBase(Schema, DiscriminatedUnionMixin):
     """Base schema for input action."""
 
-    # NOTE: We make it optional since some weaker
-    # LLMs may not be able to fill it out correctly.
-    # https://github.com/All-Hands-AI/OpenHands/issues/10797
-    security_risk: risk.SecurityRisk = Field(
-        default=risk.SecurityRisk.UNKNOWN,
-        description="The LLM's assessment of the safety risk of this action.",
-    )
-
     @property
     def visualize(self) -> Text:
         """Return Rich Text representation of this action.
@@ -198,23 +189,6 @@ def visualize(self) -> Text:
 
         return content
 
-    @classmethod
-    def to_mcp_schema(cls) -> dict[str, Any]:
-        """Convert to JSON schema format compatible with MCP."""
-        schema = super().to_mcp_schema()
-
-        # We need to move the fields from ActionBase to the END of the properties
-        # We use these properties to generate the llm schema for tool calling
-        # and we want the ActionBase fields to be at the end
-        # e.g. LLM should already outputs the argument for tools
-        # BEFORE it predicts security_risk
-        assert "properties" in schema, "Schema must have properties"
-        for field_name in ActionBase.model_fields.keys():
-            if field_name in schema["properties"]:
-                v = schema["properties"].pop(field_name)
-                schema["properties"][field_name] = v
-        return schema
-
 
 class MCPActionBase(ActionBase):
     """Base schema for MCP input action."""

openhands/sdk/tool/tool.py

@@ -10,6 +10,7 @@
     field_validator,
 )
 
+from openhands.sdk.security import risk
 from openhands.sdk.tool.schema import ActionBase, ObservationBase
 from openhands.sdk.utils.discriminated_union import (
     DiscriminatedUnionMixin,
@@ -105,16 +106,6 @@ def create(cls, *args, **kwargs) -> "Tool | list[Tool]":
         """
         raise NotImplementedError("Tool.create() must be implemented in subclasses")
 
-    @computed_field(return_type=dict[str, Any], alias="input_schema")
-    @property
-    def input_schema(self) -> dict[str, Any]:
-        return self.action_type.to_mcp_schema()
-
-    @computed_field(return_type=dict[str, Any] | None, alias="output_schema")
-    @property
-    def output_schema(self) -> dict[str, Any] | None:
-        return self.observation_type.to_mcp_schema() if self.observation_type else None
-
     @computed_field(return_type=str, alias="title")
     @property
     def title(self) -> str:
@@ -190,24 +181,47 @@ def to_mcp_tool(self) -> dict[str, Any]:
         out = {
             "name": self.name,
             "description": self.description,
-            "inputSchema": self.input_schema,
+            "inputSchema": self.action_type.to_mcp_schema(),
         }
         if self.annotations:
             out["annotations"] = self.annotations
         if self.meta is not None:
             out["_meta"] = self.meta
-        if self.output_schema:
-            out["outputSchema"] = self.output_schema
+        if self.observation_type:
+            out["outputSchema"] = self.observation_type.to_mcp_schema()
         return out
 
-    def to_openai_tool(self) -> ChatCompletionToolParam:
-        """Convert an MCP tool to an OpenAI tool."""
+    def to_openai_tool(
+        self,
+        add_security_risk_prediction: bool = False,
+    ) -> ChatCompletionToolParam:
+        """Convert a Tool to an OpenAI tool.
+
+        Args:
+            add_security_risk_prediction: Whether to add a `security_risk` field
+                to the action schema for LLM to predict. This is useful for
+                tools that may have safety risks, so the LLM can reason about
+                the risk level before calling the tool.
+        """
+
+        class ActionTypeWithRisk(self.action_type):
+            security_risk: risk.SecurityRisk = Field(
+                default=risk.SecurityRisk.UNKNOWN,
+                description="The LLM's assessment of the safety risk of this action.",
+            )
+
+        # We only add security_risk if the tool is not read-only
+        add_security_risk_prediction = add_security_risk_prediction and (
+            self.annotations is None or (not self.annotations.readOnlyHint)
+        )
         return ChatCompletionToolParam(
             type="function",
             function=ChatCompletionToolParamFunctionChunk(
                 name=self.name,
                 description=self.description,
-                parameters=self.input_schema,
+                parameters=ActionTypeWithRisk.to_mcp_schema()
+                if add_security_risk_prediction
+                else self.action_type.to_mcp_schema(),
             ),
         )