notdiamond-python/tests/test_toolkit/rag/conftest.py at 1709fb4ded52f9ca00ac9f1e7008b907c6d248ff · Not-Diamond/notdiamond-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
import pandas as pd
import pytest
from llama_index.core import SimpleDirectoryReader

from notdiamond.llms.config import EmbeddingConfig, LLMConfig
from notdiamond.toolkit.rag.evaluation_dataset import (
    RAGEvaluationDataset,
    RAGSample,
)


def format_prompt(user_input: str, retrieved_contexts: list[str]) -> str:
    context = "\n".join(retrieved_contexts)
    prompt = f"""
    Use the following context to answer the question.

    Context: {context}

    Question: {user_input}
    """
    return prompt


@pytest.fixture
def user_input():
    return "What's the capital of France?"


@pytest.fixture
def retrieved_contexts():
    return ["Paris is the capital and most populous city of France."]


@pytest.fixture
def response():
    return "The capital of France is Paris."


@pytest.fixture
def reference():
    return "Paris"


@pytest.fixture
def gpt_4o():
    return LLMConfig.from_string("openai/gpt-4o")


@pytest.fixture
def sonnet_3_5():
    # Deprecated; replacing with supported claude-3-5-haiku-20241022
    return LLMConfig.from_string("anthropic/claude-3-5-haiku-20241022")


@pytest.fixture
def openai_embedding():
    return EmbeddingConfig.from_string("openai/text-embedding-3-large")


@pytest.fixture
def dataset(gpt_4o, user_input, retrieved_contexts, response, reference):
    samples = [
        RAGSample(
            user_input=user_input,
            retrieved_contexts=retrieved_contexts,
            response=response,
            reference=reference,
            generation_prompt=format_prompt(user_input, retrieved_contexts),
            generator_llm=str(gpt_4o),
        ),
        RAGSample(
            user_input=user_input,
            retrieved_contexts=retrieved_contexts,
            response=response,
            reference=reference,
            generation_prompt=format_prompt(user_input, retrieved_contexts),
            generator_llm=str(gpt_4o),
        ),
    ]
    return RAGEvaluationDataset(samples)


@pytest.fixture
def pandas_dataset(
    gpt_4o, user_input, retrieved_contexts, response, reference
):
    data = {
        "user_input": [user_input, user_input],
        "retrieved_contexts": [retrieved_contexts, retrieved_contexts],
        "response": [response, response],
        "reference": [reference, reference],
        "generation_prompt": [
            format_prompt(user_input, retrieved_contexts),
            format_prompt(user_input, retrieved_contexts),
        ],
        "generator_llm": [str(gpt_4o), str(gpt_4o)],
    }
    df = pd.DataFrame(data=data)
    return df


@pytest.fixture
def llamaindex_documents():
    loader = SimpleDirectoryReader(
        input_files=[
            "tests/static/airbnb_tos.md",
        ]
    )
    docs = loader.load_data()
    return docs


@pytest.fixture
def test_queries():
    return ["summarize airbnb's ToS", "What are the cancellation policies?"]