notdiamond-python/tests/test_documentation/test_getting_started.py at 4698582e1ec982b5ed7cbba04159356d91a2c869 · 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
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
import pytest

from notdiamond.llms.client import NotDiamond
from notdiamond.llms.config import LLMConfig


@pytest.fixture
def start_prompt():
    return [
        {
            "role": "system",
            "content": "You are a world class software developer.",
        },
        {"role": "user", "content": "Write a merge sort in Python"},
    ]


@pytest.mark.vcr
def test_main_example():
    model_list = [
        "openai/gpt-3.5-turbo",
        "openai/gpt-4",
        "openai/gpt-4-1106-preview",
        "openai/gpt-4-turbo-preview",
        "anthropic/claude-3-haiku-20240307",
        "anthropic/claude-sonnet-4-5-20250929",
        "anthropic/claude-3-opus-20240229",
    ]

    client = NotDiamond()

    # After fuzzy hashing the inputs, the best LLM is determined by the ND API and the LLM is called client-side
    result, session_id, provider = client.chat.completions.create(
        messages=[
            {"role": "system", "content": "You are a world class programmer."},
            {"role": "user", "content": "Write a merge sort in Python"},
        ],
        model=model_list,
        tradeoff="cost",
    )

    print(
        "ND session ID: ", session_id
    )  # A unique ID of the model call. Important for personalizing ND to your use-case
    print("LLM called: ", provider.model)  # The LLM routed to
    print("LLM output: ", result.content)  # The LLM response


@pytest.mark.vcr
def test_pass_array_of_messages(start_prompt):
    # Define the available LLMs you'd like to route between
    llm_configs = [
        "openai/gpt-3.5-turbo",
        "openai/gpt-4",
        "openai/gpt-4-1106-preview",
        "openai/gpt-4-turbo-preview",
        "anthropic/claude-3-haiku-20240307",
        "anthropic/claude-sonnet-4-5-20250929",
        "anthropic/claude-3-opus-20240229",
    ]

    # Create the NDLLM object -> like a 'meta-LLM' combining all of the specified models
    nd_llm = NotDiamond(
        llm_configs=llm_configs,
        tradeoff="cost",
    )  # Define preferences

    # After fuzzy hashing the inputs, the best LLM is determined by the ND API and the LLM is called client-side
    result, session_id, provider = nd_llm.invoke(messages=start_prompt)

    print(
        "ND session ID: ", session_id
    )  # A unique ID of the invoke. Important for personalizing ND to your use-case
    print("LLM called: ", provider.model)  # The LLM routed to
    print("LLM output: ", result.content)  # The LLM response


@pytest.mark.vcr
def test_programatic_define_llm_configs(start_prompt):
    # Define the available LLMs you'd like to route between
    llm_configs = [
        LLMConfig(
            provider="openai",
            model="gpt-3.5-turbo",
            temperature=0.8,
            max_tokens=256,
        ),
        LLMConfig(
            provider="anthropic",
            model="claude-sonnet-4-5-20250929",
            temperature=0.8,
            max_tokens=256,
        ),
    ]

    # Create the NDLLM object -> like a 'meta-LLM' combining all of the specified models
    nd_llm = NotDiamond(
        llm_configs=llm_configs,
        tradeoff="cost",
    )  # Define preferences

    # After fuzzy hashing the inputs, the best LLM is determined by the ND API and the LLM is called client-side
    result, session_id, provider = nd_llm.invoke(start_prompt)

    print(
        "ND session ID: ", session_id
    )  # A unique ID of the invoke. Important for personalizing ND to your use-case
    print("LLM called: ", provider.model)  # The LLM routed to
    print("LLM output: ", result.content)  # The LLM response


@pytest.mark.vcr
def test_model_select(start_prompt):
    # Define the available LLMs you'd like to route between
    llm_configs = [
        "openai/gpt-3.5-turbo",
        "openai/gpt-4",
        "openai/gpt-4-1106-preview",
        "openai/gpt-4-turbo-preview",
        "anthropic/claude-3-haiku-20240307",
        "anthropic/claude-sonnet-4-5-20250929",
        "anthropic/claude-3-opus-20240229",
    ]

    # Create the NDLLM object -> like a 'meta-LLM' combining all of the specified models
    nd_llm = NotDiamond()

    session_id, provider = nd_llm.model_select(
        messages=start_prompt,
        model=llm_configs,
        tradeoff="cost",
    )

    print(
        "ND session ID: ", session_id
    )  # A unique ID of the model_select. Important for personalizing ND to your use-case
    print("LLM called: ", provider.model)  # The LLM routed to