From RAG

from trusttest.knowledge_base import InMemoryKnowledgeBase
from trusttest.probes.rag import RAGProbe
from trusttest.targets.http import HttpTarget, PayloadConfig
from trusttest.evaluators import CorrectnessEvaluator
from trusttest.evaluator_suite import EvaluatorSuite
from trusttest.evaluation_scenarios import EvaluationScenario

# Your document chunks
documents = [
    "TrustTest is a framework for testing AI models for safety and reliability.",
    "TrustTest supports multiple knowledge base connectors including Azure, Neo4j, and PostgreSQL.",
    "Probes in TrustTest generate test cases to evaluate model behavior.",
]

# Create knowledge base
kb = InMemoryKnowledgeBase(documents=documents)

# Configure target
target = HttpTarget(
    url="https://your-rag-endpoint.com/chat",
    headers={"Content-Type": "application/json"},
    payload_config=PayloadConfig(
        format={"messages": [{"role": "user", "content": "{{ test }}"}]},
        message_regex="{{ test }}",
    ),
)

# Create RAG probe
probe = RAGProbe(
    target=target,
    knowledge_base=kb,
    num_questions=20,
)

# Evaluate with correctness judge
evaluator = CorrectnessEvaluator()
suite = EvaluatorSuite(evaluators=[evaluator])
scenario = EvaluationScenario(evaluator_suite=suite)

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

from trusttest.knowledge_base import AzureSearchKnowledgeBase

kb = AzureSearchKnowledgeBase(
    endpoint="https://your-search.search.windows.net",
    index_name="your-index",
    api_key="your-api-key",
)

probe = RAGProbe(
    target=target,
    knowledge_base=kb,
    num_questions=50,
)

from trusttest.knowledge_base import PgVectorKnowledgeBase

kb = PgVectorKnowledgeBase(
    connection_string="postgresql://user:pass@localhost/db",
    table_name="documents",
    embedding_column="embedding",
    content_column="content",
)

probe = RAGProbe(
    target=target,
    knowledge_base=kb,
    num_questions=50,
)

from trusttest.evaluators import (
    CorrectnessEvaluator,
    CompletenessEvaluator,
    RAGPoisoningEvaluator,
)

evaluators = [
    CorrectnessEvaluator(),      # Is the answer factually correct?
    CompletenessEvaluator(),     # Does it cover all relevant points?
    RAGPoisoningEvaluator(),     # Is the response grounded in context?
]