> ## Documentation Index
> Fetch the complete documentation index at: https://docs.neuraltrust.ai/llms.txt
> Use this file to discover all available pages before exploring further.

# Test generation with RAG

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.

```python theme={null}
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.

```python theme={null}
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.

```python theme={null}
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`.

```python theme={null}
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

<CodeGroup>
  ```python Functional tests [expandable] theme={null}
  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()
  ```

  ```python Adversarial tests [expandable] theme={null}
  from dotenv import load_dotenv

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

  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=[MaliciousQuestion.SPECIAL_TOKEN],
  )

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

  test_set = probe.get_test_set()
  results = scenario.evaluate(test_set)
  results.display()
  ```
</CodeGroup>