Skip to content
This repository was archived by the owner on Dec 11, 2025. It is now read-only.
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion notdiamond/llms/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@
MissingLLMConfigs,
)
from notdiamond.llms.config import LLMConfig
from notdiamond.llms.request import amodel_select, model_select, report_latency
from notdiamond.llms.request import amodel_select, model_select, report_latency, create_preference_id
from notdiamond.metrics.metric import Metric
from notdiamond.prompts import _curly_escape, inject_system_prompt
from notdiamond.types import NDApiKeyValidator
Expand Down Expand Up @@ -131,6 +131,9 @@ def chat(self):
def completions(self):
return self

def create_preference_id(self, name: Optional[str] = None) -> str:
return create_preference_id(self.api_key, name)

async def amodel_select(
self,
messages: List[Dict[str, str]],
Expand Down
38 changes: 28 additions & 10 deletions notdiamond/llms/request.py
Original file line number Diff line number Diff line change
Expand Up @@ -67,11 +67,7 @@ def model_select_prepare(
if preference_id is not None:
payload["preference_id"] = preference_id

headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {notdiamond_api_key}",
"User-Agent": f"Python-SDK/{settings.VERSION}",
}
headers = _default_headers(notdiamond_api_key)

return url, payload, headers

Expand Down Expand Up @@ -282,11 +278,7 @@ def report_latency(
"feedback": {"tokens_per_second": tokens_per_second},
}

headers = {
"content-type": "application/json",
"Authorization": f"Bearer {notdiamond_api_key}",
"User-Agent": f"Python-SDK/{settings.VERSION}",
}
headers = _default_headers(notdiamond_api_key)

try:
response = requests.post(url, json=payload, headers=headers)
Expand All @@ -297,3 +289,29 @@ def report_latency(
return 500

return response.status_code

def create_preference_id(notdiamond_api_key: str, name: Optional[str] = None) -> str:
"""
Create a preference id with an optional name. The preference name will appear in your
dashboard on Not Diamond.
"""
url = f"{settings.ND_BASE_URL}/v2/preferences/userPreferenceCreate"
headers = _default_headers(notdiamond_api_key)
res = requests.post(
url=url,
headers=headers,
json={"name": name}
)
if res.status_code == 200:
preference_id = res.json()["preference_id"]
else:
raise Exception(f"Error creating preference ID: {res.text}")

return preference_id

def _default_headers(notdiamond_api_key: str) -> Dict[str, str]:
return {
"content-type": "application/json",
"Authorization": f"Bearer {notdiamond_api_key}",
"User-Agent": f"Python-SDK/{settings.VERSION}",
}
15 changes: 15 additions & 0 deletions tests/test_components/test_llms/test_llm.py
Original file line number Diff line number Diff line change
Expand Up @@ -620,6 +620,21 @@ def test_create_unavailable_error(self, ndtarget, NDLLM):
hash_content=True,
)

def test_preference_id(self, ndtarget, NDLLM):
llm_configs = [
"openai/gpt-3.5-turbo",
"openai/gpt-4",
"anthropic/claude-2.1",
"google/gemini-pro",
]
with patch.multiple(
"notdiamond.llms.client", NDApiKeyValidator=Mock(return_value=True)
):
nd_llm = NDLLM(llm_configs=llm_configs)
preference_id = nd_llm.create_preference_id(name="test")
assert preference_id is not None
assert isinstance(preference_id, str)


class Test_OpenAI_style_input:
def test_openai_style_input_invoke(
Expand Down