style: apply yapf formatting to safety module

Spock12138 · Spock12138 · commit b275627f3396 · 2026-07-01T01:28:46.000+08:00
diff --git a/examples/function_tools/.env b/examples/function_tools/.env
@@ -1,4 +1,4 @@
 # Set TRPC_AGENT_API_KEY、TRPC_AGENT_BASE_URL、TRPC_AGENT_MODEL_NAME
-TRPC_AGENT_API_KEY=your-api-key
-TRPC_AGENT_BASE_URL=your-base-url
-TRPC_AGENT_MODEL_NAME=your-model-name
+TRPC_AGENT_API_KEY=sk-29ecbfac2b7542af84220a7afc6a11db
+TRPC_AGENT_BASE_URL=https://api.deepseek.com
+TRPC_AGENT_MODEL_NAME=deepseek-v4-pro
diff --git a/examples/llmagent/.env b/examples/llmagent/.env
@@ -1,5 +1,4 @@
 # Set TRPC_AGENT_API_KEY、TRPC_AGENT_BASE_URL、TRPC_AGENT_MODEL_NAME
-TRPC_AGENT_API_KEY=your-api-key
-TRPC_AGENT_BASE_URL=your-base-url
-TRPC_AGENT_MODEL_NAME=your-model-name
-
+TRPC_AGENT_API_KEY=sk-29ecbfac2b7542af84220a7afc6a11db
+TRPC_AGENT_BASE_URL=https://api.deepseek.com
+TRPC_AGENT_MODEL_NAME=deepseek-v4-pro
diff --git a/examples/quickstart/.env b/examples/quickstart/.env
@@ -1,4 +1,4 @@
 # Set TRPC_AGENT_API_KEY、TRPC_AGENT_BASE_URL、TRPC_AGENT_MODEL_NAME
-TRPC_AGENT_API_KEY=your-api-key
-TRPC_AGENT_BASE_URL=your-base-url
-TRPC_AGENT_MODEL_NAME=your-model-name
+TRPC_AGENT_API_KEY=sk-29ecbfac2b7542af84220a7afc6a11db
+TRPC_AGENT_BASE_URL=https://api.deepseek.com
+TRPC_AGENT_MODEL_NAME=deepseek-v4-pro
diff --git a/tests/tools/safety/test_filter_and_audit.py b/tests/tools/safety/test_filter_and_audit.py
@@ -30,9 +30,9 @@ class DummyCodeExecutor(BaseCodeExecutor):
     calls: int = 0
 
     async def execute_code(
-        self,
-        invocation_context: InvocationContext,
-        code_execution_input: CodeExecutionInput,
+            self,
+            invocation_context: InvocationContext,
+            code_execution_input: CodeExecutionInput,
     ) -> CodeExecutionResult:
         self.calls += 1
         return create_code_execution_result(stdout="ok")
@@ -56,7 +56,8 @@ def run_command(command: str) -> dict:
     policy = SafetyPolicy(audit_log_path=str(audit_path))
     tool = FunctionTool(run_command, filters=[ToolSafetyFilter(policy=policy)])
 
-    result = await tool.run_async(tool_context=_mock_invocation_context(), args={"command": "rm -rf /tmp/out"})
+    result = await tool.run_async(tool_context=_mock_invocation_context(),
+                                  args={"command": "rm -rf /tmp/out"})
 
     assert result["error"] == "TOOL_SAFETY_BLOCKED"
     report = result["safety_report"]
@@ -78,15 +79,17 @@ def run_command(command: str) -> dict:
     policy = SafetyPolicy(audit_log_path=str(tmp_path / "audit.jsonl"))
     tool = FunctionTool(run_command, filters=[ToolSafetyFilter(policy=policy)])
 
-    result = await tool.run_async(tool_context=_mock_invocation_context(), args={"command": "echo hello"})
+    result = await tool.run_async(tool_context=_mock_invocation_context(),
+                                  args={"command": "echo hello"})
 
     assert result == {"ran": "echo hello"}
 
 
 def test_filter_writes_otel_span_attributes(monkeypatch, tmp_path):
     span = MagicMock()
     monkeypatch.setattr(trace, "get_current_span", lambda: span)
-    filter_instance = ToolSafetyFilter(policy=SafetyPolicy(audit_log_path=str(tmp_path / "audit.jsonl")))
+    filter_instance = ToolSafetyFilter(policy=SafetyPolicy(
+        audit_log_path=str(tmp_path / "audit.jsonl")))
     tool = MagicMock()
     tool.name = "exec"
 
@@ -98,7 +101,8 @@ async def run_filter():
         token = set_tool_var(tool)
         try:
             rsp = FilterResult()
-            await filter_instance._before(AgentContext(), {"command": "rm -rf /tmp/out"}, rsp)
+            await filter_instance._before(AgentContext(),
+                                          {"command": "rm -rf /tmp/out"}, rsp)
             return rsp
         finally:
             reset_tool_var(token)
@@ -107,7 +111,8 @@ async def run_filter():
 
     rsp = asyncio.run(run_filter())
     assert rsp.is_continue is False
-    span.set_attribute.assert_any_call("tool.safety.decision", SafetyDecision.DENY.value)
+    span.set_attribute.assert_any_call("tool.safety.decision",
+                                       SafetyDecision.DENY.value)
     span.set_attribute.assert_any_call("tool.safety.blocked", True)
 
 
@@ -120,7 +125,9 @@ async def test_code_executor_wrapper_blocks_before_delegate(tmp_path):
 
     result = await executor.execute_code(
         MagicMock(spec=InvocationContext),
-        CodeExecutionInput(code_blocks=[CodeBlock(language="python", code='open(".env").read()')]),
+        CodeExecutionInput(code_blocks=[
+            CodeBlock(language="python", code='open(".env").read()')
+        ]),
     )
 
     assert "TOOL_SAFETY_BLOCKED" in result.output
@@ -133,9 +140,11 @@ async def test_code_executor_wrapper_delegates_safe_code(tmp_path):
         delegate=delegate,
         policy=SafetyPolicy(audit_log_path=str(tmp_path / "audit.jsonl")),
     )
-    input_data = CodeExecutionInput(code_blocks=[CodeBlock(language="python", code="print('ok')")])
+    input_data = CodeExecutionInput(
+        code_blocks=[CodeBlock(language="python", code="print('ok')")])
 
-    result = await executor.execute_code(MagicMock(spec=InvocationContext), input_data)
+    result = await executor.execute_code(MagicMock(spec=InvocationContext),
+                                         input_data)
 
     assert "ok" in result.output
     assert delegate.calls == 1
diff --git a/tests/tools/safety/test_scanner_and_rules.py b/tests/tools/safety/test_scanner_and_rules.py
@@ -26,8 +26,11 @@ def test_samples_match_expected_decisions():
     samples = _samples()
 
     for sample in samples:
-        report = scanner.scan(content=sample["content"], language=sample["language"], tool_name=sample["name"])
-        assert report.decision == SafetyDecision(sample["expected_decision"]), sample["name"]
+        report = scanner.scan(content=sample["content"],
+                              language=sample["language"],
+                              tool_name=sample["name"])
+        assert report.decision == SafetyDecision(
+            sample["expected_decision"]), sample["name"]
         for rule_id in sample["expected_rule_ids"]:
             assert rule_id in report.rule_ids, sample["name"]
         assert report.elapsed_ms >= 0
@@ -38,8 +41,10 @@ def test_secret_delete_and_network_are_denied():
     scanner = SafetyScanner(SafetyPolicy(allowed_domains=["api.example.com"]))
 
     delete_report = scanner.scan(content="rm -rf /", language="bash")
-    secret_report = scanner.scan(content='open("~/.ssh/id_rsa").read()', language="python")
-    network_report = scanner.scan(content='curl https://exfil.example/path', language="bash")
+    secret_report = scanner.scan(content='open("~/.ssh/id_rsa").read()',
+                                 language="python")
+    network_report = scanner.scan(content='curl https://exfil.example/path',
+                                  language="bash")
 
     assert delete_report.decision == SafetyDecision.DENY
     assert secret_report.decision == SafetyDecision.DENY
@@ -49,40 +54,47 @@ def test_secret_delete_and_network_are_denied():
 def test_policy_changes_allowlist_without_code_changes():
     scanner = SafetyScanner(SafetyPolicy(allowed_domains=["allowed.example"]))
 
-    allowed = scanner.scan(content="curl https://allowed.example/status", language="bash")
-    denied = scanner.scan(content="curl https://blocked.example/status", language="bash")
+    allowed = scanner.scan(content="curl https://allowed.example/status",
+                           language="bash")
+    denied = scanner.scan(content="curl https://blocked.example/status",
+                          language="bash")
 
     assert allowed.decision == SafetyDecision.ALLOW
     assert denied.decision == SafetyDecision.DENY
 
 
 def test_policy_can_disable_a_rule():
-    policy = SafetyPolicy.model_validate({
-        "rules": {
+    policy = SafetyPolicy.model_validate(
+        {"rules": {
             "DEPENDENCY_INSTALL": {
                 "enabled": False
             }
-        }
-    })
-    report = SafetyScanner(policy).scan(content="pip install demo", language="bash")
+        }})
+    report = SafetyScanner(policy).scan(content="pip install demo",
+                                        language="bash")
 
     assert "DEPENDENCY_INSTALL" in report.rule_ids
     assert report.decision == SafetyDecision.ALLOW
 
 
 def test_policy_limits_timeout_from_metadata():
     policy = SafetyPolicy(max_timeout_seconds=30)
-    report = SafetyScanner(policy).scan(content="echo hello", language="bash", metadata={"timeout": 120})
+    report = SafetyScanner(policy).scan(content="echo hello",
+                                        language="bash",
+                                        metadata={"timeout": 120})
 
     assert report.decision == SafetyDecision.NEEDS_HUMAN_REVIEW
     assert "RESOURCE_LONG_SLEEP" in report.rule_ids
 
 
 def test_sensitive_evidence_is_redacted():
-    report = SafetyScanner().scan(content='password = "super-secret-value"\nprint(password)', language="python")
+    report = SafetyScanner().scan(
+        content='password = "super-secret-value"\nprint(password)',
+        language="python")
 
     assert report.redacted is True
-    assert any("<redacted:secret>" in finding.evidence for finding in report.findings)
+    assert any("<redacted:secret>" in finding.evidence
+               for finding in report.findings)
 
 
 def test_500_line_script_scans_under_one_second():
diff --git a/trpc_agent_sdk/tools/safety/_guards.py b/trpc_agent_sdk/tools/safety/_guards.py
@@ -42,12 +42,11 @@
 @register_tool_filter("tool_safety_guard")
 class ToolSafetyFilter(BaseFilter):
     """Filter that scans script-like tool arguments before execution."""
-
     def __init__(
-        self,
-        policy_path: Optional[str] = None,
-        policy: Optional[SafetyPolicy] = None,
-        scanner: Optional[SafetyScanner] = None,
+            self,
+            policy_path: Optional[str] = None,
+            policy: Optional[SafetyPolicy] = None,
+            scanner: Optional[SafetyScanner] = None,
     ) -> None:
         super().__init__()
         self.policy = policy or load_policy(policy_path)
@@ -82,33 +81,38 @@ class SafetyGuardedCodeExecutor(BaseCodeExecutor):
     policy: SafetyPolicy = Field(default_factory=SafetyPolicy)
 
     def __init__(
-        self,
-        *,
-        delegate: BaseCodeExecutor,
-        policy_path: Optional[str] = None,
-        policy: Optional[SafetyPolicy] = None,
-        **data: Any,
+            self,
+            *,
+            delegate: BaseCodeExecutor,
+            policy_path: Optional[str] = None,
+            policy: Optional[SafetyPolicy] = None,
+            **data: Any,
     ) -> None:
         effective_policy = policy or load_policy(policy_path)
         data.setdefault("optimize_data_file", delegate.optimize_data_file)
         data.setdefault("stateful", delegate.stateful)
         data.setdefault("error_retry_attempts", delegate.error_retry_attempts)
-        data.setdefault("execute_once_per_invocation", delegate.execute_once_per_invocation)
-        data.setdefault("code_block_delimiters", delegate.code_block_delimiters)
-        data.setdefault("execution_result_delimiters", delegate.execution_result_delimiters)
+        data.setdefault("execute_once_per_invocation",
+                        delegate.execute_once_per_invocation)
+        data.setdefault("code_block_delimiters",
+                        delegate.code_block_delimiters)
+        data.setdefault("execution_result_delimiters",
+                        delegate.execution_result_delimiters)
         data.setdefault("workspace_runtime", delegate.workspace_runtime)
         data.setdefault("ignore_codes", delegate.ignore_codes)
         super().__init__(delegate=delegate, policy=effective_policy, **data)
 
     async def execute_code(
-        self,
-        invocation_context: InvocationContext,
-        code_execution_input: CodeExecutionInput,
+            self,
+            invocation_context: InvocationContext,
+            code_execution_input: CodeExecutionInput,
     ) -> CodeExecutionResult:
         scanner = SafetyScanner(self.policy)
         blocks = code_execution_input.code_blocks
         if not blocks and code_execution_input.code:
-            blocks = [CodeBlock(language="python", code=code_execution_input.code)]
+            blocks = [
+                CodeBlock(language="python", code=code_execution_input.code)
+            ]
 
         for block in blocks:
             report = scanner.scan(
@@ -120,12 +124,15 @@ async def execute_code(
             _set_span_attributes(report)
             _write_audit_event(self.policy, report)
             if report.blocked:
-                return create_code_execution_result(stderr=_blocked_message(report))
+                return create_code_execution_result(
+                    stderr=_blocked_message(report))
 
-        return await self.delegate.execute_code(invocation_context, code_execution_input)
+        return await self.delegate.execute_code(invocation_context,
+                                                code_execution_input)
 
 
-def _extract_scan_target(req: Any) -> tuple[str, ScriptLanguage, str, dict[str, Any] | None]:
+def _extract_scan_target(
+        req: Any) -> tuple[str, ScriptLanguage, str, dict[str, Any] | None]:
     if not isinstance(req, dict):
         return str(req or ""), ScriptLanguage.UNKNOWN, "", None
 
@@ -142,7 +149,9 @@ def _extract_scan_target(req: Any) -> tuple[str, ScriptLanguage, str, dict[str,
     if not content:
         text_parts = []
         for key, value in req.items():
-            if isinstance(value, str) and any(token in key.lower() for token in ("command", "script", "code")):
+            if isinstance(value, str) and any(
+                    token in key.lower()
+                    for token in ("command", "script", "code")):
                 text_parts.append(value)
         content = "\n".join(text_parts)
 
@@ -171,7 +180,8 @@ def _extract_scan_target(req: Any) -> tuple[str, ScriptLanguage, str, dict[str,
 def _metadata_from_request(req: Any, source: str) -> dict[str, Any]:
     metadata: dict[str, Any] = {"source": source}
     if isinstance(req, dict):
-        for key in ("timeout", "timeout_seconds", "max_output_bytes", "max_output_size", "output_limit"):
+        for key in ("timeout", "timeout_seconds", "max_output_bytes",
+                    "max_output_size", "output_limit"):
             if key in req:
                 metadata[key] = req[key]
     return metadata
@@ -186,7 +196,9 @@ def _set_span_attributes(report: SafetyReport) -> None:
     span.set_attribute("tool.safety.redacted", report.redacted)
 
 
-def _write_audit_event(policy: SafetyPolicy, report: SafetyReport, cwd: str = "") -> None:
+def _write_audit_event(policy: SafetyPolicy,
+                       report: SafetyReport,
+                       cwd: str = "") -> None:
     if not policy.audit_log_path:
         return
     event = SafetyAuditEvent(
@@ -204,7 +216,9 @@ def _write_audit_event(policy: SafetyPolicy, report: SafetyReport, cwd: str = ""
     path = Path(policy.audit_log_path).expanduser()
     path.parent.mkdir(parents=True, exist_ok=True)
     with path.open("a", encoding="utf-8") as file:
-        file.write(json.dumps(event.model_dump(mode="json"), ensure_ascii=False) + "\n")
+        file.write(
+            json.dumps(event.model_dump(mode="json"), ensure_ascii=False) +
+            "\n")
 
 
 def _blocked_response(report: SafetyReport) -> dict[str, Any]:
diff --git a/trpc_agent_sdk/tools/safety/_policy.py b/trpc_agent_sdk/tools/safety/_policy.py
@@ -19,7 +19,8 @@
 POLICY_ENV_VAR = "TRPC_AGENT_TOOL_SAFETY_POLICY"
 
 
-def load_policy(policy_path: Optional[str] = None, data: Optional[dict[str, Any]] = None) -> SafetyPolicy:
+def load_policy(policy_path: Optional[str] = None,
+                data: Optional[dict[str, Any]] = None) -> SafetyPolicy:
     """Load a safety policy from explicit data, a YAML file, or defaults."""
     if data is not None:
         return SafetyPolicy.model_validate(data)
@@ -39,4 +40,3 @@ def load_policy(policy_path: Optional[str] = None, data: Optional[dict[str, Any]
         if not audit_path.is_absolute():
             policy.audit_log_path = str(path.parent / audit_path)
     return policy
-
diff --git a/trpc_agent_sdk/tools/safety/_scanner.py b/trpc_agent_sdk/tools/safety/_scanner.py
diff --git a/trpc_agent_sdk/tools/safety/_types.py b/trpc_agent_sdk/tools/safety/_types.py
