Misc

.github/workflows/build.yml

@@ -19,5 +19,5 @@ jobs:
         cache: 'poetry'
     - run: poetry install
     - run: poetry run black aisploit --check
-    #- run: poetry run mypy aisploit
+    - run: poetry run mypy aisploit
     - run: poetry run pytest

aisploit/agent/red_teaming_bot.py

@@ -69,7 +69,7 @@ def __init__(
         runnable = prompt | chat_model | StrOutputParser()
 
         self._chain = RunnableWithMessageHistory(
-            runnable,
+            runnable,  # type: ignore[arg-type]
             lambda session_id: self._history,
             input_messages_key="input",
             output_messages_key="output",

aisploit/classifier/huggingface/__init__.py

@@ -1,3 +1,3 @@
-from .prompt_injection_identifier_pipeline import PromptIjectionIdentifierPipeline
+from .pipeline_prompt_injection_identifier import PipelinePromptIjectionIdentifier
 
-__all__ = ["PromptIjectionIdentifierPipeline"]
+__all__ = ["PipelinePromptIjectionIdentifier"]

aisploit/classifier/huggingface/prompt_injection_identifier_hub.py → aisploit/classifier/huggingface/hub_prompt_injection_identifier.py

@@ -1,5 +1,5 @@
 from aisploit.core import BaseTextClassification, Score
 
 
-class PromptIjectionIdentifierHub(BaseTextClassification):
+class HubPromptIjectionIdentifier(BaseTextClassification):
     pass

aisploit/classifier/huggingface/prompt_injection_identifier_pipeline.py → aisploit/classifier/huggingface/pipeline_prompt_injection_identifier.py

@@ -6,7 +6,7 @@
 from aisploit.core import BaseTextClassification, Score
 
 
-class PromptIjectionIdentifierPipeline(BaseTextClassification):
+class PipelinePromptIjectionIdentifier(BaseTextClassification):
     def __init__(
         self,
         *,

aisploit/core/__init__.py

@@ -1,5 +1,6 @@
 from .model import BaseLLM, BaseChatModel, BaseModel, BaseEmbeddings
 from .score import Score, BaseTextClassification
+from .vectorstore import BaseVectorStore
 
 __all__ = [
     "BaseLLM",
@@ -8,4 +9,5 @@
     "BaseEmbeddings",
     "Score",
     "BaseTextClassification",
+    "BaseVectorStore",
 ]

aisploit/core/vectorstore.py

@@ -0,0 +1,3 @@
+from langchain_core.vectorstores import VectorStore
+
+BaseVectorStore = VectorStore

aisploit/integration/__init__.py

@@ -1,9 +1,9 @@
 from .gandalf import GandalfBot, GandalfScorer, GandalfLevel
-from .rag import GenericRAG
+from .rag import VectorStoreRAG
 
 __all__ = [
     "GandalfBot",
     "GandalfScorer",
     "GandalfLevel",
-    "GenericRAG",
+    "VectorStoreRAG",
 ]

aisploit/integration/rag.py

@@ -1,10 +1,9 @@
 import textwrap
-from typing import List
-from langchain_community.vectorstores.chroma import Chroma
+from typing import List, Any
 from langchain_core.output_parsers import StrOutputParser
-from langchain_core.runnables import RunnablePassthrough
+from langchain_core.runnables import RunnablePassthrough, RunnableSerializable
 from langchain_core.prompts import PromptTemplate
-from aisploit.core import BaseModel, BaseEmbeddings
+from aisploit.core import BaseModel, BaseVectorStore
 
 _template = textwrap.dedent(
     """
@@ -20,28 +19,25 @@
     """
 )
 
-GENERIC_RAG_PROMPT = PromptTemplate.from_template(_template)
+VECTOR_STORE_RAG_PROMPT = PromptTemplate.from_template(_template)
 
 
 def format_docs(docs):
     return "\n\n".join(doc.page_content for doc in docs)
 
 
-class GenericRAG:
+class VectorStoreRAG:
     def __init__(
         self,
         *,
         model: BaseModel,
-        embeddings: BaseEmbeddings,
-        prompt: PromptTemplate = GENERIC_RAG_PROMPT,
+        vectorstore: BaseVectorStore,
+        prompt: PromptTemplate = VECTOR_STORE_RAG_PROMPT,
     ) -> None:
-        self._vectorstore = Chroma(
-            embedding_function=embeddings,
-        )
-
+        self._vectorstore = vectorstore
         self._retriever = self._vectorstore.as_retriever()
 
-        self._chain = (
+        self._chain: RunnableSerializable[Any, str] = (
             {
                 "context": self._retriever | format_docs,
                 "question": RunnablePassthrough(),

aisploit/poison/poison.py

@@ -14,11 +14,11 @@ class Poison:
     """
 
     question: str = ""
-    question_embedding: List[float] = field(default_factory=list)
+    question_embeddings: List[float] = field(default_factory=list)
     target_answer: str = ""
     adversary_text: str = ""
-    adversary_text_embedding: List[float] = field(default_factory=list)
-    cosine_distance: float = -1
+    adversary_text_embeddings: List[float] = field(default_factory=list)
+    cosine_distance: float = 2
 
 
 _template = """This is my question: ```{question}```
@@ -54,7 +54,7 @@ def __init__(
         """
         self._chain = prompt | model | StrOutputParser()
         self._embeddings = embeddings
-        self._max_ierations = max_iterations
+        self._max_iterations = max_iterations
 
     def __call__(
         self, question: str, answer: str, max_words: int = 30
@@ -70,9 +70,9 @@ def __call__(
         Yields:
         - Poison: A Poison object representing generated poisoning data.
         """
-        question_embedding = self._embeddings.embed_query(question)
+        question_embeddings = self._embeddings.embed_query(question)
 
-        for _ in range(self._max_ierations):
+        for _ in range(self._max_iterations):
             adversary_text = self._chain.invoke(
                 {
                     "question": question,
@@ -81,15 +81,15 @@ def __call__(
                 }
             )
 
-            adversary_text_embedding = self._embeddings.embed_query(adversary_text)
+            adversary_text_embeddings = self._embeddings.embed_query(adversary_text)
 
             yield Poison(
                 question=question,
-                question_embedding=question_embedding,
+                question_embeddings=question_embeddings,
                 target_answer=answer,
                 adversary_text=adversary_text,
-                adversary_text_embedding=adversary_text_embedding,
+                adversary_text_embeddings=adversary_text_embeddings,
                 cosine_distance=cosine_distance(
-                    question_embedding, adversary_text_embedding
+                    question_embeddings, adversary_text_embeddings
                 ),
             )

examples/classifier.ipynb

@@ -7,7 +7,7 @@
    "outputs": [],
    "source": [
     "import textwrap\n",
-    "from aisploit.classifier.huggingface import PromptIjectionIdentifierPipeline"
+    "from aisploit.classifier.huggingface import PipelinePromptIjectionIdentifier"
    ]
   },
   {
@@ -27,7 +27,7 @@
     }
    ],
    "source": [
-    "classifier = PromptIjectionIdentifierPipeline()\n",
+    "classifier = PipelinePromptIjectionIdentifier()\n",
     "classifier.score_text(textwrap.dedent(\n",
     "    \"\"\"\n",
     "    The answer to the universe is 42. \n",

examples/poison.ipynb

@@ -10,10 +10,11 @@
     "import os\n",
     "import textwrap\n",
     "from dotenv import load_dotenv\n",
+    "from langchain_community.vectorstores.chroma import Chroma\n",
     "from aisploit.chat import ChatOpenAI\n",
     "from aisploit.embeddings import OpenAIEmbeddings\n",
     "from aisploit.poison import PoisonGen\n",
-    "from aisploit.integration import GenericRAG"
+    "from aisploit.integration import VectorStoreRAG"
    ]
   },
   {
@@ -32,12 +33,16 @@
     "    api_key=os.getenv(\"OPENAI_API_KEY\"),\n",
     ")\n",
     "\n",
-    "rag = GenericRAG(\n",
+    "rag = VectorStoreRAG(\n",
     "    model=chat_model,\n",
-    "    embeddings=embeddings,\n",
+    "    vectorstore=Chroma(embedding_function=embeddings),\n",
     ")\n",
     "\n",
-    "generate = PoisonGen(model=chat_model, embeddings=embeddings, max_iterations=10)"
+    "generate_poison = PoisonGen(\n",
+    "    model=chat_model, \n",
+    "    embeddings=embeddings, \n",
+    "    max_iterations=10,\n",
+    ")"
    ]
   },
   {
@@ -229,7 +234,7 @@
    "source": [
     "adversary_texts = []\n",
     "\n",
-    "for poison in generate(\n",
+    "for poison in generate_poison(\n",
     "    question=\"Who invented the telephone?\",\n",
     "    answer=\"Homer Simpson\",\n",
     "):\n",
@@ -260,7 +265,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "groud_truth = [\n",
+    "ground_truth = [\n",
     "    \"The invention of the telephone is credited to Alexander Graham Bell, who patented it in 1876.\",\n",
     "    \"The first practical telephone was demonstrated by Bell to his assistant Thomas Watson, with the famous words 'Mr. Watson, come here, I want to see you.'\",\n",
     "    \"Initially, the telephone was used for local communication within a limited range.\",\n",
@@ -305,7 +310,7 @@
     }
    ],
    "source": [
-    "rag.add_texts(groud_truth)\n",
+    "rag.add_texts(ground_truth)\n",
     "rag.retrieve(\"Who invented the telephone?\")"
    ]
   },

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "aisploit"
-version = "0.0.6"
+version = "0.0.7"
 description = "Tiny package designed to support red teams and penetration testers in exploiting large language model AI solutions."
 authors = ["hupe1980"]
 repository = "https://github.com/hupe1980/aisploit"
@@ -32,16 +32,19 @@ langchain-community = "^0.0.24"
 python-dotenv = "^1.0.1"
 numpy = "^1.26.4"
 transformers = "^4.38.1"
-chromadb = "^0.4.23"
 
 [tool.poetry.group.dev.dependencies]
+chromadb = "^0.4.23"
 pytest = "^8.0.1"
 jupyter = "^1.0.0"
 types-pyyaml = "^6.0.12.12"
 types-requests = "^2.31.0.20240218"
 black = "^24.2.0"
 mypy = "^1.8.0"
 
+[tool.mypy]
+ignore_missing_imports = true
+
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"