Merge branch 'main' into workspace

invariant/examples/code_agents/swe_agent.ipynb

@@ -105,12 +105,12 @@
    "outputs": [],
    "source": [
     "policy = Policy.from_string(\"\"\"\n",
-    "    from invariant.detectors import pii, code_shield, CodeIssue, secrets\n",
+    "    from invariant.detectors import pii, semgrep, CodeIssue, secrets\n",
     "                                \n",
     "    raise PolicyViolation(\"found unsafe code: \", issue) if:\n",
     "        (call1: ToolCall) -> (call2: ToolCall)\n",
     "        call1.function.name == \"edit\"\n",
-    "        (issue: CodeIssue) in code_shield(call1.function.arguments[\"code\"])\n",
+    "        (issue: CodeIssue) in semgrep(call1.function.arguments[\"code\"])\n",
     "        call2.function.name == \"python\"\n",
     "        \"pickle\" in issue.description\n",
     "    \"\"\")"

invariant/extras.py

@@ -158,11 +158,6 @@ def find_all() -> list["Extra"]:
     "spacy": ExtrasImport("spacy", "spacy", ">=3.7.5")
 })
 
-"""Extra for features that rely on the Llama `codeshield` library."""
-codeshield_extra = Extra("Code Scanning with Llama CodeShield", "Enables the use of CodeShield", {
-    "codeshield.cs": ExtrasImport("codeshield.cs", "codeshield", ">=1.0.1")
-})
-
 """Extra for features that rely on the `semgrep` library."""
 semgrep_extra = Extra("Code Scanning with Semgrep", "Enables the use of Semgrep for code scanning", {
     "semgrep": ExtrasImport("semgrep", "semgrep", ">=1.78.0")

invariant/language/parser.py

@@ -10,6 +10,9 @@
 from invariant.language.typing import typing
 from invariant.runtime.patterns import VALUE_MATCHERS
 
+"""
+Lark EBNF grammar for the Invariant Policy Language.
+"""
 parser = lark.Lark(r"""
     %import common.NUMBER
 
@@ -27,10 +30,11 @@
     def_stmt: "def" func_signature INDENT NEWLINE? statement* DEDENT
     decl_stmt: ( ID | func_signature ) ":=" expr | ( ID | func_signature ) INDENT NEWLINE? expr (NEWLINE expr)* DEDENT
 
-    expr: ID | binary_expr | "(" expr ("," expr)* ")" | block_expr
+    expr: ID | assignment_expr | "(" expr ("," expr)* ")" | block_expr
                    
     block_expr: INDENT expr (NEWLINE expr)* DEDENT
     
+    assignment_expr: ( ID ":=" binary_expr ) | binary_expr
     binary_expr: cmp_expr LOGICAL_OPERATOR cmp_expr | cmp_expr
     cmp_expr: ( term CMP_OPERATOR term ) | term
     term: factor TERM_OPERATOR factor | factor
@@ -76,7 +80,7 @@
     ID.2: /[a-zA-Z_]([a-zA-Z0-9_])*/
     UNARY_OPERATOR.3: /not[\n\t ]/ | "-" | "+"
     LOGICAL_OPERATOR: /and[\n\t ]/ | /or[\n\t ]/
-    CMP_OPERATOR: ":=" | "==" | "!=" | ">" | "<" | ">=" | "<=" | /is[\n\t ]/ | /contains_only[\n\t ]/ | /in[\n\t ]/ | "->"
+    CMP_OPERATOR: "==" | "!=" | ">" | "<" | ">=" | "<=" | /is[\n\t ]/ | /contains_only[\n\t ]/ | /in[\n\t ]/ | "->"
     VALUE_TYPE: /<[a-zA-Z_:]+>/
     
     TERM_OPERATOR: "+" | "-"
@@ -97,8 +101,8 @@ def indent_level(line, unit=1):
     # count the number of leading spaces
     return (len(line) - len(line.lstrip())) // unit
 
-
 def derive_indentation_units(text):
+    # derive the indentation unit from the first non-empty line
     lines = text.split("\n")
     indents = set()
     for line in lines:
@@ -108,8 +112,29 @@ def derive_indentation_units(text):
         return 1
     return min(indents)
 
-
 def parse_indents(text):
+    """
+    This function parses an intended snippet of IPL code and returns a version of the code
+    where indents are replaced by |INDENT| and |DEDENT| tokens.
+
+    This allows our actual language grammar above to be context-free, as it does not need to
+    handle indentation, but can rely on the |INDENT| and |DEDENT| tokens instead.
+
+    |INDENT| and |DEDENT| tokens fulfill the same role as e.g. `{` and `}` in C-like languages.
+
+    Example:
+    ```
+    def foo:
+        bar
+    ```
+
+    is transformed into:
+
+    ```
+    def foo: |INDENT|
+    bar |DEDENT|
+    ```
+    """
     indent_unit = derive_indentation_units(text)
     lines = text.split("\n")
 
@@ -217,6 +242,11 @@ def parameter_decl(self, items):
     def expr(self, items):
         return items[0]
 
+    def assignment_expr(self, items):
+        if len(items) == 1:
+            return items[0]
+        return BinaryExpr(items[0], ":=", items[1]).with_location(self.loc(items))
+
     def block_expr(self, items):
         return self.filter(items)
 
@@ -355,22 +385,10 @@ def NUMBER(self, items):
     def loc(self, items):
         return Location.from_items(items, self.line_mappings, self.source_code)
 
-
-def chain(ind_mappings, ml_mappings):
-    result = {}
-    for k, (line, char) in ind_mappings.items():
-        if line in ml_mappings:
-            result[k] = (ml_mappings[line], char)
-        else:
-            result[k] = (line, char)
-    return result
-
-
 def transform(policy):
     """
-    Basic AST transformations to simplify the AST
+    Basic transformations to simplify the AST
     """
-
     class PostParsingTransformations(Transformation):
         # transforms FunctionCall with a ToolReference target into a SemanticPattern
         def visit_FunctionCall(self, node: FunctionCall):
@@ -388,9 +406,22 @@ def visit_FunctionCall(self, node: FunctionCall):
 
 
 def parse(text, path=None, verbose=True):
+    """
+    Multi-stage parsing process to transform a string of IPL code into an Invariant Policy AST.
+
+    The parsing stages are as follows:
+
+    1. Indentation parsing: The input code is transformed into a version where indents are marked with |INDENT| and |DEDENT| tokens, instead of actual indentation.
+    2. Lark parsing: The indented code is parsed using the Lark parser as defined by the grammar above.
+    3. AST construction: The Lark parse tree is transformed into an AST.
+    4. AST post-processing: The AST is simplified and transformed.
+    5. Type checking: The AST is type-checked.
+    
+    """
+
     # removes common leading indent (e.g. when parsing from an indented multiline string)
     text = textwrap.dedent(text)
-    # creates source code hjandle
+    # creates source code handle
     source_code = SourceCode(text, path=path, verbose=verbose)
 
     # translates an indent-based code into code in which indented 

invariant/runtime/utils/code.py

@@ -6,7 +6,6 @@
 from enum import Enum
 from invariant.runtime.utils.base import BaseDetector, DetectorResult
 from pydantic.dataclasses import dataclass, Field
-from invariant.extras import codeshield_extra
 
 
 class CodeSeverity(str, Enum):
@@ -133,24 +132,6 @@ def detect(self, text: str, ipython_mode=False) -> PythonDetectorResult:
         return ast_visitor.res
 
 
-class CodeShieldDetector(BaseDetector):
-    """Detector which uses Llama CodeShield for safety (currently based on regex and semgrep rules)"""
-
-    async def scan_llm_output(self, llm_output_code):
-        self.CodeShield = codeshield_extra.package("codeshield.cs").import_names("CodeShield")
-        result = await self.CodeShield.scan_code(llm_output_code)
-        return result
-
-    def detect_all(self, text: str) -> list[CodeIssue]:
-        res = asyncio.run(self.scan_llm_output(text))
-        if res.issues_found is None:
-            return []
-        return [
-            CodeIssue(description=issue.description, severity=str(issue.severity).lower())
-            for issue in res.issues_found
-        ]
-
-
 class SemgrepDetector(BaseDetector):
     """Detector which uses Semgrep for safety evaluation."""
 

invariant/stdlib/invariant/detectors/code.py

@@ -2,7 +2,6 @@
 from invariant.runtime.functions import cache
 
 PYTHON_ANALYZER = None
-CODE_SHIELD_DETECTOR = None
 SEMGREP_DETECTOR = None
 
 @cache
@@ -30,25 +29,6 @@ def ipython_code(data: str | list | dict, **config: dict) -> PythonDetectorResul
     """Predicate used to extract entities from IPython cell code."""
     return python_code(data, ipython_mode=True, **config)
 
-@cache
-def code_shield(data: str | list | dict, **config: dict) -> list[CodeIssue]:
-    """Predicate used to run CodeShield on code."""
-
-    global CODE_SHIELD_DETECTOR
-    if CODE_SHIELD_DETECTOR is None:
-        CODE_SHIELD_DETECTOR = CodeShieldDetector()
-
-    if type(data) is str:
-        return CODE_SHIELD_DETECTOR.detect_all(data, **config)
-    if type(data) is not list:
-        data = [data]
-
-    res = []
-    for message in data:
-        if message.content is None:
-            continue
-        res.extend(CODE_SHIELD_DETECTOR.detect_all(message.content, **config))
-    return res
 
 @cache
 def semgrep(data: str | list | dict, **config: dict) -> list[CodeIssue]:

pyproject.toml

@@ -30,7 +30,6 @@ dev-dependencies = [
     "presidio-analyzer>=2.2.354",
     "transformers>=4.41.1",
     "torch>=2.3.0",
-    "codeshield>=1.0.1",
     "python-dotenv>=1.0.1",
 ]
 

requirements-dev.lock

@@ -50,7 +50,6 @@ click-option-group==0.5.6
     # via semgrep
 cloudpathlib==0.18.1
     # via weasel
-codeshield==1.0.1
 colorama==0.4.6
     # via semgrep
 confection==0.1.5
@@ -194,7 +193,6 @@ pygments==2.18.0
 pytest==8.2.2
 python-dotenv==1.0.1
 pyyaml==6.0.1
-    # via codeshield
     # via huggingface-hub
     # via langchain
     # via langchain-core
@@ -234,7 +232,6 @@ ruamel-yaml-clib==0.2.8
 safetensors==0.4.3
     # via transformers
 semgrep==1.78.0
-    # via codeshield
     # via invariant
 setuptools==70.0.0
     # via marisa-trie

tests/test_html_parsing.py

@@ -3,7 +3,7 @@
 from invariant import Policy
 from invariant.policy import analyze_trace
 from invariant.traces import *
-from invariant.extras import extras_available, presidio_extra, transformers_extra, codeshield_extra
+from invariant.extras import extras_available, presidio_extra, transformers_extra
 from invariant.traces import user, assistant, tool, tool_call
 
 class TestHTMLParsing(unittest.TestCase):

tests/test_parser.py

@@ -371,7 +371,6 @@ def test_with_member_access_call_in_addition(self):
             (to: ToolOutput)
             "File " + tc.function.arguments.arg.strip() + " not found" in to.content
         """)
-
         self.assertIsInstance(policy.statements[0].body[1], ast.BinaryExpr)
         self.assertIsInstance(policy.statements[0].body[1].left, ast.BinaryExpr)
         self.assertIsInstance(policy.statements[0].body[1].left.left, ast.BinaryExpr)
@@ -387,5 +386,14 @@ def test_with_member_access_call_in_addition(self):
         self.assertEqual(policy.statements[0].body[1].left.left.right.name.expr.expr.expr.member, "function")
         self.assertIsInstance(policy.statements[0].body[1].left.left.right.name.expr.expr.expr.expr, ast.Identifier)
 
+    def test_assign_in(self):
+        policy = parse("""
+        raise "error" if:
+            (msg: Message)
+            flag := ["a", "b"] in ["a"]
+        """)
+        self.assertEqual(policy.statements[0].body[1].op, ":=")
+
+
 if __name__ == "__main__":
     unittest.main()