upstream guardrails updates (#50)

Author: knielsen404

invariant/analyzer/examples/agent_bugs/traceset.py

@@ -115,9 +115,9 @@ def analyze(self, policy):
     def filter(
         self,
         invariant_condition: str,
-        max_items: int = None,
-        python: str = None,
-        prefix: str = None,
+        max_items: int | None = None,
+        python: str | None = None,
+        prefix: str | None = None,
     ) -> "TraceSet":
         max_items = self.get_max_items(max_items)
 
@@ -155,7 +155,7 @@ def get_max_items(self, max_items):
 
         return max_items
 
-    def prepare_policy(self, invariant_condition: str, prefix: str = None):
+    def prepare_policy(self, invariant_condition: str, prefix: str | None = None):
         from invariant.analyzer import Policy
 
         # construct makeshift policy

invariant/analyzer/examples/agent_flan/run.py

@@ -52,7 +52,7 @@ def has_bash(conv):
             else:
                 tool_called = True
                 match = re.search(r"```bash\n(.+?)\n```", msg["content"], re.DOTALL)
-                bash_cmd = match.group(1)
+                bash_cmd = match.group(1) if match else None
                 # print("-> ", bash_cmd)
                 ipl_messages += [
                     {

invariant/analyzer/language/scope.py

@@ -33,6 +33,8 @@
     "tool_call",
     "text",
     "image",
+    "Tool",
+    "ToolParameter",
 ]
 
 

invariant/analyzer/language/typing.py

@@ -158,15 +158,6 @@ def visit_MemberAccess(self, node: MemberAccess):
             node.type = UnknownType()
             return node
 
-        def has_member(obj_type, member):
-            if isinstance(obj_type, type) and hasattr(obj_type, "__dataclass_fields__"):
-                # check on type reference if attribute exists
-                return member in obj_type.__dataclass_fields__
-            return False
-
-        if not has_member(node.expr.type, node.member):
-            raise PolicyError(f"Type {node.expr}: {node.expr.type} has no member {node.member}")
-
         node.type = UnknownType()
 
         return node
@@ -245,7 +236,7 @@ def visit_ListComprehension(self, node: ListComprehension):
         node.scope.parent = self.context.scope
 
         # Add the iteration variable to the scope
-        var_name = node.var_name.name if hasattr(node.var_name, 'name') else node.var_name
+        var_name = node.var_name.name if hasattr(node.var_name, "name") else node.var_name
         var_decl = VariableDeclaration(var_name, UnknownType())
         node.scope.declarations = {var_name: var_decl}
 

invariant/analyzer/policy.py

@@ -70,7 +70,7 @@ def errors(self):
         return self.policy_root.errors
 
     @classmethod
-    def from_file(cls, path: str) -> "Policy":
+    def from_file(cls, path: str) -> "LocalPolicy":
         return cls(parse_file(path))
 
     @classmethod
@@ -80,14 +80,14 @@ def from_string(
         path: str | None = None,
         optimize: bool = False,
         symbol_table: Optional[SymbolTable] = None,
-    ) -> "Policy":
+    ) -> "LocalPolicy":
         return cls(parse(string, path, optimize_rules=optimize), symbol_table=symbol_table)
 
     def add_error_to_result(self, error, analysis_result: AnalysisResult):
         """Implements how errors are added to an analysis result (e.g. as handled or non-handled errors)."""
         analysis_result.errors.append(error)
 
-    def analyze(self, input: dict, raise_unhandled=False, **policy_parameters):
+    def analyze(self, input: list[dict], raise_unhandled=False, **policy_parameters):
         return asyncio.run(self.a_analyze(input, raise_unhandled, **policy_parameters))
 
     async def a_analyze(

invariant/analyzer/remote_policy.py

@@ -9,6 +9,7 @@
 from invariant.analyzer.runtime.runtime_errors import (
     ExcessivePolicyError,
     InvariantAttributeError,
+    InvariantInputValidationError,
     MissingPolicyParameter,
     PolicyExecutionError,
 )
@@ -111,7 +112,7 @@ async def a_preload(self):
                 if "errors" in result:
                     raise ValueError("Invalid response from policy service: " + str(result))
 
-    def analyze(self, input: dict, raise_unhandled=False, **policy_parameters):
+    def analyze(self, input: list[dict], raise_unhandled=False, **policy_parameters):
         return asyncio.run(self.a_analyze(input, raise_unhandled, **policy_parameters))
 
     def get_json_policy_parameters(self, **policy_parameters):
@@ -261,6 +262,16 @@ async def handle_error_response(response: aiohttp.ClientResponse):
             except Exception as e:
                 text = str(e)
             raise InvariantAttributeError(text)
+        elif InvariantInputValidationError.catchphrase in text:
+            # get 'detail' from json
+            try:
+                error_json = await response.json()
+                text = error_json.get("detail", str(error_json))
+                # remove 'Invariant attribute error' from text
+                text = text.replace(InvariantInputValidationError.catchphrase, "")
+            except Exception as e:
+                text = str(e)
+            raise InvariantInputValidationError(text)
         else:
             raise ValueError("Invalid response from policy service: " + text)
 

invariant/analyzer/runtime/evaluation.py

@@ -5,10 +5,12 @@
 import sys
 from dataclasses import dataclass  # noqa: I001
 from itertools import product
-from typing import AsyncGenerator
+from typing import AsyncGenerator, Awaitable, Callable, ParamSpec, TypeVar
 
 import termcolor
 
+P = ParamSpec("P")
+R = TypeVar("R")
 from invariant.analyzer.language.ast import (
     ArrayLiteral,
     BinaryExpr,
@@ -786,10 +788,6 @@ def _is_unknown(self, value):
     async def visit_FunctionCall(self, node: FunctionCall):
         function = await self.visit(node.name)
         args = await asyncio.gather(*[self.visit(arg) for arg in node.args])
-
-        # only call functions, once all parameters are known
-        if function is Unknown or any(self._is_unknown(arg) for arg in args):
-            return Unknown
         kwarg_items = await asyncio.gather(
             *[
                 asyncio.gather(self.visit(entry.key), self.visit(entry.value))
@@ -798,6 +796,14 @@ async def visit_FunctionCall(self, node: FunctionCall):
         )
         kwargs = {k: v for k, v in kwarg_items}
 
+        # only call functions, once all parameters are known
+        if (
+            function is Unknown
+            or any(self._is_unknown(arg) for arg in args)
+            or any(self._is_unknown(v) for v in kwargs.values())
+        ):
+            return Unknown
+
         if isinstance(function, Declaration):
             return await self.visit_PredicateCall(function, args, **kwargs)
         else:
@@ -932,6 +938,9 @@ async def visit_ListComprehension(self, node: ListComprehension):
         if iterable is None:
             return []
 
+        if iterable is Unknown:
+            return Unknown
+
         var_name = node.var_name.id if hasattr(node.var_name, "id") else node.var_name
         results = []
         original_vars = self.variable_store.copy()
@@ -953,14 +962,21 @@ async def visit_ListComprehension(self, node: ListComprehension):
         self.variable_store = original_vars
         return results
 
-    async def acall_function(self, function, *args, **kwargs):
+    async def acall_function(
+        self,
+        function: Callable[P, Awaitable[R]] | Callable[P, R] | CachedFunctionWrapper,
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> R:
         ctx: EvaluationContext = self.evaluation_context
         # if function is a cached function wrapper, unwrap it (at this point we
         # are already caching it)
-        if isinstance(function, CachedFunctionWrapper):
-            function = function.func
-        linked_function = ctx.link(function, None)
+        func = function.func if isinstance(function, CachedFunctionWrapper) else function
+        linked_function = ctx.link(func, None)
         # also unwrap linked function, if a cached function wrapper
-        if isinstance(linked_function, CachedFunctionWrapper):
-            linked_function = linked_function.func
-        return await ctx.acall_function(linked_function, args, **kwargs)
+        linked_function = (
+            linked_function.func
+            if isinstance(linked_function, CachedFunctionWrapper)
+            else linked_function
+        )
+        return await ctx.acall_function(linked_function, *args, **kwargs)

invariant/analyzer/runtime/evaluation_context.py

@@ -4,11 +4,13 @@
 In a separate file, for better separation of dependencies.
 """
 
-from typing import Optional
+from typing import Optional, TypeVar
 
 from invariant.analyzer.runtime.input import Input
 from invariant.analyzer.runtime.symbol_table import SymbolTable
+from invariant.analyzer.language.ast import Node
 
+R = TypeVar("R")
 
 class EvaluationContext:
     """
@@ -28,7 +30,7 @@ def call_function(self, function, args, **kwargs):
     async def acall_function(self, function, args, **kwargs):
         raise NotImplementedError("EvaluationContext must implement acall_function()")
 
-    def link(self, function, node):
+    def link(self, function: R, node: Node | None) -> R:
         if self.symbol_table:
             return self.symbol_table.link(function, node)
         else:

invariant/analyzer/runtime/function_cache.py

@@ -1,4 +1,8 @@
 import inspect
+from typing import Awaitable, Callable, ParamSpec, TypeVar
+
+P = ParamSpec("P")
+R = TypeVar("R")
 
 
 class FunctionCache:
@@ -43,7 +47,12 @@ async def get(self, key):
     async def set(self, key, value):
         self.cache[key] = value
 
-    async def acall(self, function, args, **kwargs):
+    async def acall(
+        self,
+        function: Callable[P, Awaitable[R]] | Callable[P, R],
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ) -> R:
         # if function is not marked with @cached we just call it directly (see ./functions.py module)
         if not hasattr(function, "__invariant_cache__"):
             return await call_either_way(function, *args, **kwargs)
@@ -59,8 +68,11 @@ async def acall(self, function, args, **kwargs):
             return value
 
 
-async def call_either_way(fct, *args, **kwargs):
-    if inspect.iscoroutinefunction(fct):
-        return await fct(*args, **kwargs)
+async def call_either_way(
+    func: Callable[P, Awaitable[R]] | Callable[P, R], *args: P.args, **kwargs: P.kwargs
+) -> R:
+    if inspect.iscoroutinefunction(func):
+        return await func(*args, **kwargs)
     else:
-        return fct(*args, **kwargs)
+        print([func, args, kwargs], flush=True)
+        return func(*args, **kwargs)  # type: ignore

invariant/analyzer/runtime/functions.py

@@ -3,9 +3,12 @@
 with special runtime attributes, relevant in the context of the
 invariant agent analyzer.
 """
+from typing import Callable, TypeVar, ParamSpec, Awaitable, Generic
 
+P = ParamSpec("P")
+R = TypeVar("R")
 
-def cached(func):
+def cached(func: Callable[P, Awaitable[R]] | Callable[P, R]) -> Callable[P, Awaitable[R]]:
     """
     Decorator to mark a guardrailing function or built-in predicate as interpreter-cached.
 
@@ -38,19 +41,18 @@ def guardrailing_function(...):
 
     return CachedFunctionWrapper(func)
 
-
-class CachedFunctionWrapper:
+class CachedFunctionWrapper(Generic[P, R]):
     """
     Wraps a function such that is is always called via the current Interpreter instance.
 
     This enables caching and other runtime features like function re-linking in a server context.
     """
 
-    def __init__(self, func):
+    def __init__(self, func: Callable[P, Awaitable[R]] | Callable[P, R]):
         self.func = func
-        self.func.__invariant_cache__ = True
+        self.func.__invariant_cache__ = True # type: ignore
 
-    def __call__(self, *args, **kwargs):
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Awaitable[R]:
         from invariant.analyzer.runtime.evaluation import Interpreter
 
         return Interpreter.current().acall_function(self.func, *args, **kwargs)