Add toolchain from agentic system here

This commit is contained in:
Ashwin Bharambe 2024-07-19 12:30:35 -07:00
parent f6b2b2fb39
commit 95781ec85d
71 changed files with 11899 additions and 0 deletions

View file

@ -0,0 +1,28 @@
from codeshield.cs import CodeShield
from termcolor import cprint
from .base import ShieldResponse, TextShield
from toolchain.safety.api.datatypes import * # noqa: F403
class CodeScannerShield(TextShield):
def get_shield_type(self) -> ShieldType:
return BuiltinShield.code_scanner_guard
async def run_impl(self, text: str) -> ShieldResponse:
cprint(f"Running CodeScannerShield on {text[50:]}", color="magenta")
result = await CodeShield.scan_code(text)
if result.is_insecure:
return ShieldResponse(
shield_type=BuiltinShield.code_scanner_guard,
is_violation=True,
violation_type=",".join(
[issue.pattern_id for issue in result.issues_found]
),
violation_return_message="Sorry, I found security concerns in the code.",
)
else:
return ShieldResponse(
shield_type=BuiltinShield.code_scanner_guard, is_violation=False
)