From ae8e866e4cc6ed9a393de094320d42b5925ab85e Mon Sep 17 00:00:00 2001 From: Mehrin Kiani Date: Tue, 23 Jan 2024 10:25:56 -0500 Subject: [PATCH] Update imports --- python-sdk/pyproject.toml | 1 + python-sdk/rebuff/sdk.py | 9 ++++-- python-sdk/tests/__init__.py | 0 python-sdk/tests/conftest.py | 8 +---- python-sdk/tests/test_sdk.py | 59 ++++++------------------------------ 5 files changed, 18 insertions(+), 59 deletions(-) create mode 100644 python-sdk/tests/__init__.py diff --git a/python-sdk/pyproject.toml b/python-sdk/pyproject.toml index 6febe5a..69a29f3 100644 --- a/python-sdk/pyproject.toml +++ b/python-sdk/pyproject.toml @@ -1,6 +1,7 @@ [tool.poetry] name = "rebuff" version = "0.0.0" +packages = [{ include = "rebuff" }] description = "Rebuff is designed to protect AI applications from prompt injection (PI) attacks through a multi-layered defense." authors = ["ProtectAI "] readme = "README.md" diff --git a/python-sdk/rebuff/sdk.py b/python-sdk/rebuff/sdk.py index 1df2cf8..2095798 100644 --- a/python-sdk/rebuff/sdk.py +++ b/python-sdk/rebuff/sdk.py @@ -4,9 +4,12 @@ from langchain_core.prompts import PromptTemplate from pydantic import BaseModel -from .detect_pi_heuristics import detect_prompt_injection_using_heuristic_on_input -from .detect_pi_openai import call_openai_to_detect_pi, render_prompt_for_pi_detection -from .detect_pi_vectorbase import detect_pi_using_vector_database, init_pinecone +from rebuff.detect_pi_heuristics import detect_prompt_injection_using_heuristic_on_input +from rebuff.detect_pi_openai import ( + call_openai_to_detect_pi, + render_prompt_for_pi_detection, +) +from rebuff.detect_pi_vectorbase import detect_pi_using_vector_database, init_pinecone class RebuffDetectionResponse(BaseModel): diff --git a/python-sdk/tests/__init__.py b/python-sdk/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/python-sdk/tests/conftest.py b/python-sdk/tests/conftest.py index 348fbf2..37b5e28 100644 --- a/python-sdk/tests/conftest.py +++ b/python-sdk/tests/conftest.py @@ -5,13 +5,7 @@ import requests -try: - sys.path.insert( - 0, - os.path.abspath(os.path.join(os.path.dirname(__file__), "../rebuff")), - ) -except NameError: - pass + import subprocess import time diff --git a/python-sdk/tests/test_sdk.py b/python-sdk/tests/test_sdk.py index f4c13c1..6d99724 100644 --- a/python-sdk/tests/test_sdk.py +++ b/python-sdk/tests/test_sdk.py @@ -1,30 +1,15 @@ -import os -import sys -from typing import List, Union, Dict +from typing import List, Dict import pytest - -from sdk import RebuffSdk, RebuffDetectionResponse -from utils import get_environment_variable - -try: - sys.path.insert( - 0, - os.path.abspath(os.path.join(os.path.dirname(__file__), "../rebuff")), - ) -except NameError: - pass +from rebuff.sdk import RebuffSdk, RebuffDetectionResponse +from .utils import get_environment_variable @pytest.fixture() def rebuff() -> RebuffSdk: - openai_apikey = get_environment_variable("OPENAI_API_KEY") - pinecone_apikey = get_environment_variable("PINECONE_API_KEY") - pinecone_index = get_environment_variable("PINECONE_INDEX_NAME") - rb = RebuffSdk( - openai_apikey, - pinecone_apikey, - pinecone_index, + get_environment_variable("OPENAI_API_KEY"), + get_environment_variable("PINECONE_API_KEY"), + get_environment_variable("PINECONE_INDEX_NAME"), ) return rb @@ -67,31 +52,6 @@ def detect_injection_arguments() -> Dict: return detect_injection_arguments -def test_rebuff_detection_response_attributes(): - rebuff_response = RebuffDetectionResponse( - heuristic_score=0.5, - openai_score=0.8, - vector_score=0.9, - run_heuristic_check=True, - run_language_model_check=False, - run_vector_check=True, - max_heuristic_score=0.5, - max_model_score=0.8, - max_vector_score=0.0, - injection_detected=False, - ) - assert hasattr(rebuff_response, "heuristic_score") - assert hasattr(rebuff_response, "openai_score") - assert hasattr(rebuff_response, "vector_score") - assert hasattr(rebuff_response, "run_heuristic_check") - assert hasattr(rebuff_response, "run_language_model_check") - assert hasattr(rebuff_response, "run_vector_check") - assert hasattr(rebuff_response, "max_heuristic_score") - assert hasattr(rebuff_response, "max_model_score") - assert hasattr(rebuff_response, "max_vector_score") - assert hasattr(rebuff_response, "injection_detected") - - def test_add_canary_word(rebuff: RebuffSdk, user_inputs: List[str]): for user_input in user_inputs: prompt = f"Tell me a joke about\n{user_input}" @@ -151,7 +111,7 @@ def test_detect_injection_heuristics( rebuff_response = rebuff.detect_injection(input, **detect_injection_arguments) assert ( rebuff_response.heuristic_score - < detect_injection_arguments["max_heuristic_score"] + <= detect_injection_arguments["max_heuristic_score"] ) assert not rebuff_response.injection_detected @@ -179,7 +139,7 @@ def test_detect_injection_vectorbase( assert ( rebuff_response.vector_score - < detect_injection_arguments["max_vector_score"] + <= detect_injection_arguments["max_vector_score"] ) assert not rebuff_response.injection_detected @@ -205,6 +165,7 @@ def test_detect_injection_llm( rebuff_response = rebuff.detect_injection(input, **detect_injection_arguments) assert ( - rebuff_response.openai_score < detect_injection_arguments["max_model_score"] + rebuff_response.openai_score + <= detect_injection_arguments["max_model_score"] ) assert not rebuff_response.injection_detected