From 28974cc70c907ad89d7365e8782e981b9bc3635a Mon Sep 17 00:00:00 2001 From: Luca Beurer-Kellner Date: Mon, 17 Feb 2025 20:18:21 +0100 Subject: [PATCH] dirty demo hacks --- .env | 2 +- docker-compose.local.yml | 1 + proxy/requirements.txt | 10 ++- proxy/routes/open_ai.py | 24 +++++- proxy/utils/explorer.py | 163 +++++++++++++++++++++++++++++++++++++-- 5 files changed, 189 insertions(+), 11 deletions(-) diff --git a/.env b/.env index 419976f..5684ee0 100644 --- a/.env +++ b/.env @@ -1 +1 @@ -INVARIANT_API_URL=https://explorer.invariantlabs.ai +INVARIANT_API_URL=http://explorer-app-api-1:8000 diff --git a/docker-compose.local.yml b/docker-compose.local.yml index 3096fe1..66f2cd3 100644 --- a/docker-compose.local.yml +++ b/docker-compose.local.yml @@ -4,6 +4,7 @@ services: build: context: ./proxy dockerfile: ../proxy/Dockerfile.proxy + platform: linux/amd64 working_dir: /srv/proxy env_file: - .env diff --git a/proxy/requirements.txt b/proxy/requirements.txt index a564b13..a46823a 100644 --- a/proxy/requirements.txt +++ b/proxy/requirements.txt @@ -2,4 +2,12 @@ fastapi==0.115.7 httpx==0.28.1 invariant-sdk>=0.0.10 starlette-compress==1.4.0 -uvicorn==0.34.0 \ No newline at end of file +uvicorn==0.34.0 +invariant-ai[dev]==0.2 +presidio-analyzer +spacy +transformers +torch +python-dotenv +numpy +invariant_sdk \ No newline at end of file diff --git a/proxy/routes/open_ai.py b/proxy/routes/open_ai.py index ba33d94..eb0d00b 100644 --- a/proxy/routes/open_ai.py +++ b/proxy/routes/open_ai.py @@ -7,7 +7,7 @@ import httpx from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response from starlette.responses import StreamingResponse from utils.constants import CLIENT_TIMEOUT, IGNORED_HEADERS -from utils.explorer import push_trace +from utils.explorer import error_label, push_trace, validate_guardrails ALLOWED_OPEN_AI_ENDPOINTS = {"chat/completions"} @@ -308,21 +308,28 @@ async def push_to_explorer( ) -> None: """Pushes the full trace to the Invariant Explorer""" # Only push the trace to explorer if the message is an end turn message + only_push_if_blocked = False + if ( merged_response.get("choices") and merged_response["choices"][0].get("finish_reason") not in FINISH_REASON_TO_PUSH_TRACE ): - return + only_push_if_blocked = True + # Combine the messages from the request body and the choices from the OpenAI response messages = request_body.get("messages", []) messages += [choice["message"] for choice in merged_response.get("choices", [])] - _ = await push_trace( + + blocked, response = await push_trace( dataset_name=dataset_name, messages=[messages], invariant_authorization=invariant_authorization, + dry=only_push_if_blocked, ) + return blocked[0] + async def handle_non_streaming_response( response: httpx.Response, @@ -332,10 +339,19 @@ async def handle_non_streaming_response( ): """Handles non-streaming OpenAI responses""" json_response = response.json() - await push_to_explorer( + blocked = await push_to_explorer( dataset_name, json_response, request_body_json, invariant_authorization ) + if blocked: + # json_response["c + json_response["choices"][-1]["finish_reason"] = "blocked" + json_response["choices"][-1]["message"]["content"] = ( + f"[Agent execution blocked by guardrail: {error_label(blocked)}]" + ) + # remove tool calls + json_response["choices"][-1]["message"]["tool_calls"] = None + return Response( content=json.dumps(json_response), status_code=response.status_code, diff --git a/proxy/utils/explorer.py b/proxy/utils/explorer.py index 71108d7..852bcf6 100644 --- a/proxy/utils/explorer.py +++ b/proxy/utils/explorer.py @@ -2,9 +2,17 @@ import os from typing import Any, Dict, List +import re from invariant_sdk.async_client import AsyncClient -from invariant_sdk.types.push_traces import PushTracesRequest, PushTracesResponse +from invariant_sdk.types.push_traces import ( + PushTracesRequest, + PushTracesResponse, + AnnotationCreate, +) +import json + +from invariant.analyzer import Policy DEFAULT_API_URL = "https://explorer.invariantlabs.ai" @@ -13,13 +21,14 @@ async def push_trace( messages: List[List[Dict[str, Any]]], dataset_name: str, invariant_authorization: str, + dry: bool = False, ) -> PushTracesResponse: """Pushes traces to the dataset on the Invariant Explorer. If a dataset with the given name does not exist, it will be created. Args: - messages (List[List[Dict[str, Any]]]): List of messages to push. + messages (List[List[Dict[str, Any]]]): List of messages to push dataset_name (str): Name of the dataset. invariant_authorization (str): Value of the invariant-authorization header. @@ -32,13 +41,157 @@ async def push_trace( [{k: v for k, v in msg.items() if v is not None} for msg in msg_list] for msg_list in messages ] - request = PushTracesRequest(messages=update_messages, dataset=dataset_name) + client = AsyncClient( api_url=os.getenv("INVARIANT_API_URL", DEFAULT_API_URL).rstrip("/"), api_key=invariant_authorization.split("Bearer ")[1], ) + + # validate guardrails (and get annotations) + blocked_and_annotations = [ + await validate_guardrails(client, messages, dataset_name) + for messages in update_messages + ] + + annotations = [annotation for _, annotation in blocked_and_annotations] + blocked = [block for block, _ in blocked_and_annotations] + + # for blocked messages histories, apply "system" message with [blocked] content + for i, block in enumerate(blocked): + if block: + update_messages[i].append( + { + "content": f"[Agent execution blocked by guardrail: {error_label(block)}]", + "role": "system", + } + ) + + request = PushTracesRequest( + messages=update_messages, dataset=dataset_name, annotations=annotations + ) + try: - return await client.push_trace(request) + # if dry run, don't push the trace (but still validate guardrails) + if dry and not any(blocked): + result = {"dry_run": True} + else: + result = await client.push_trace(request) except Exception as e: print(f"Failed to push trace: {e}") - return {"error": str(e)} + result = {"error": str(e)} + + return blocked, result + + +async def validate_guardrails( + client: AsyncClient, + messages: List[List[Dict[str, Any]]], + dataset_name: str, +) -> PushTracesResponse: + """Fetches and validates the guardrails for the given dataset. + + Args: + messages (List[List[Dict[str, Any]]]): List of messages to push. + dataset_name (str): Name of the dataset. + invariant_authorization (str): Value of the + invariant-authorization header. + + Returns: + PushTracesResponse: Response containing the trace ID details. + """ + try: + metadata = await client.get_dataset_metadata(dataset_name=dataset_name) + except Exception as e: + print(f"Failed to get dataset metadata: {e}") + return False, [] + guardrails = json.loads(metadata.get("guardrails", "[]")) + + blocked = False + + # preprocess messages + trace = [{**msg} for msg in messages] + + # if content is missing in a msg, set it to empty string + for msg in trace: + if "content" not in msg: + msg["content"] = "" + + annotations = [] + + for guardrail in guardrails: + if not guardrail.get("enabled", False): + print(f"Skipping guardrail {guardrail['name']} as it is disabled") + continue + + policy = Policy.from_string(guardrail["policy"]) + results = policy.analyze(trace) + + for error in results.errors: + label = error_label(error) + # check for action=[block|warn|...] + if "action=" in label: + action = label.split("action=")[1].split()[0] + if action == "block": + blocked = error + + ranges = [range for range in error.ranges] + + if len(ranges) > 1: + prefixfree_ranges = [] + # remove prefixes + for range in ranges: + if not any( + range.json_path.startswith(prefix.json_path) + for prefix in ranges + if range != prefix + ): + prefixfree_ranges.append(range) + ranges = prefixfree_ranges + + for range in ranges: + address = "messages." + range.json_path + old_address = address + try: + # if there is .start and .end, also include them as :start-end + if ( + range.start is not None + and range.end is not None + and "-" not in address + ): + address += f":{range.start}-{range.end}" + # handle special cases + # 1. if the address is messages.[0-9]+, change it to messages.[0-9]+.content:0- + elif re.match(r"messages\.\d+$", address) and "-" not in address: + end = len(trace[int(address.split(".")[1])]["content"]) + address += ".content:0-" + str(end) + # 2. if address is messages.4.tool_calls.0, change it to messages.4.tool_calls.0.function.name:0- + elif ( + re.match(r"messages\.\d+\.tool_calls\.\d+$", address) + and "-" not in address + ): + end = len( + trace[int(address.split(".")[1])]["tool_calls"][ + int(address.split(".")[3]) + ]["function"]["name"] + ) + address += ".function.name:0-" + str(end) + except Exception as e: + print(f"Error while handling special cases: {e}") + pass + + print("fixed", old_address, "->", address) + + annotations.append( + AnnotationCreate( + content=label, + address=address, + extra_metadata={"source": "guardrail"}, + ) + ) + return blocked, annotations + + +def error_label(error): + if str(error).startswith("ErrorInformation("): + return str(error).split("(", 1)[1].rsplit(")", 1)[0] + return str(error)