Skip to main content
In this guide we will see how to automatically generate tests with KnwoledgeBase. Knowledge Bases is a database with a collection of documents. This is usually used in RAG applications as the vector database to store the documents and the embeddings.

Configure the Knowledge Base

For teaching purposes we will use a dummy knowledge base that will store the documents in memory. As you can see each document is grouped by a topic. So we can generate better questions for each topic.
from trusttest.knowledge_base import Document, InMemoryKnowledgeBase

documents = [
    Document(
        id="1",
        content="""
        Vic is of ancient origin. In past times it was called Ausa by the Romans.
        Iberian coins bearing this name have been found there.
        The Visigoths called it Ausona.
        Sewage caps on sidewalks around the city will also read "Vich", an old spelling of the name.
        """,
        topic="City origins",
    ),
    Document(
        id="2",
        content="""
        Vic (Catalan pronunciation: [bik]; Spanish: Vic) is the capital of the comarca of Osona,
        in the province of Barcelona, Catalonia, Spain.
        Vic is located 69 km (43 mi) from Barcelona and 60 km (37 mi) from Girona.
        """,
        topic="City location",
    ),
]

knowledge_base = InMemoryKnowledgeBase(documents=documents)

Generate Functional Questions

Now we can generate questions for each topic. We use RAGProbe together with EvaluationScenario and AnswerRelevanceEvaluator for functional evaluation. In this configuration we will generate 2 questions, one for each topic. And we will use the BenignQuestion.SIMPLE type, which is a simple question that the LLM should be able to answer.
from trusttest.probes.rag import RAGProbe, BenignQuestion
from trusttest.evaluation_scenarios import EvaluationScenario
from trusttest.evaluator_suite import EvaluatorSuite
from trusttest.evaluators import AnswerRelevanceEvaluator

knowledge_base = InMemoryKnowledgeBase(documents=documents)

probe = RAGProbe(
    target=DummyTarget(),
    knowledge_base=knowledge_base,
    num_questions=2,
    question_types=[BenignQuestion.SIMPLE],
)

scenario = EvaluationScenario(
    name="RAG Functional",
    evaluator_suite=EvaluatorSuite(evaluators=[AnswerRelevanceEvaluator()], criteria="any_fail"),
)
Then we only need to create the test_set and check the results.
test_set = probe.get_test_set()
results = scenario.evaluate(test_set)
results.display()

Generate RAG Poisoning Tests

We can also automatically generate poisoning tests based to attack specific information domains. We just need to change the question_types to MaliciousQuestion and use RAGPoisoningEvaluator.
from trusttest.probes.rag import RAGProbe, MaliciousQuestion
from trusttest.evaluators import RAGPoisoningEvaluator

probe = RAGProbe(
    target=DummyTarget(),
    knowledge_base=knowledge_base,
    num_questions=2,
    question_types=[MaliciousQuestion.SPECIAL_TOKEN],
)

scenario = EvaluationScenario(
    name="RAG Poisoning",
    evaluator_suite=EvaluatorSuite(evaluators=[RAGPoisoningEvaluator()], criteria="any_fail"),
)

Complete Examples

from dotenv import load_dotenv

from trusttest.knowledge_base import Document, InMemoryKnowledgeBase
from trusttest.targets.testing import DummyTarget
from trusttest.probes.rag import RAGProbe, BenignQuestion
from trusttest.evaluation_scenarios import EvaluationScenario
from trusttest.evaluator_suite import EvaluatorSuite
from trusttest.evaluators import AnswerRelevanceEvaluator

load_dotenv(override=True)

documents = [
    Document(
        id="1",
        content="""
        Vic is of ancient origin. In past times it was called Ausa by the Romans.
        Iberian coins bearing this name have been found there.
        The Visigoths called it Ausona.
        Sewage caps on sidewalks around the city will also read "Vich", an old spelling of the name.
        """,
        topic="City origins",
    ),
    Document(
        id="2",
        content="""
        Vic (Catalan pronunciation: [bik]; Spanish: Vic) is the capital of the comarca of Osona,
        in the province of Barcelona, Catalonia, Spain.
        Vic is located 69 km (43 mi) from Barcelona and 60 km (37 mi) from Girona.
        """,
        topic="City location",
    ),
]

knowledge_base = InMemoryKnowledgeBase(documents=documents)

probe = RAGProbe(
    target=DummyTarget(),
    knowledge_base=knowledge_base,
    num_questions=2,
    question_types=[BenignQuestion.SIMPLE],
)

scenario = EvaluationScenario(
    name="RAG Functional",
    evaluator_suite=EvaluatorSuite(evaluators=[AnswerRelevanceEvaluator()], criteria="any_fail"),
)

test_set = probe.get_test_set()
results = scenario.evaluate(test_set)
results.display()