From 7fadac8cbc0c74f5979004f7db4351e87374d65e Mon Sep 17 00:00:00 2001
From: Alejandro Companioni <a9@notdiamond.ai>
Date: Tue, 6 Aug 2024 13:17:28 -0400
Subject: [PATCH 1/2] Adding create_preference_id client method.

---
 notdiamond/llms/client.py  |  5 ++++-
 notdiamond/llms/request.py | 38 ++++++++++++++++++++++++++++----------
 2 files changed, 32 insertions(+), 11 deletions(-)

diff --git a/notdiamond/llms/client.py b/notdiamond/llms/client.py
index c6f97e92..a98adafc 100644
--- a/notdiamond/llms/client.py
+++ b/notdiamond/llms/client.py
@@ -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
@@ -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]],
diff --git a/notdiamond/llms/request.py b/notdiamond/llms/request.py
index e0743010..29799396 100644
--- a/notdiamond/llms/request.py
+++ b/notdiamond/llms/request.py
@@ -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
 
@@ -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)
@@ -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}",
+    }
\ No newline at end of file

From d4374a9b767ec0c2fa9e08322db75d5df2577bef Mon Sep 17 00:00:00 2001
From: Alejandro Companioni <a9@notdiamond.ai>
Date: Tue, 6 Aug 2024 14:36:43 -0400
Subject: [PATCH 2/2] Unit test

---
 tests/test_components/test_llms/test_llm.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/tests/test_components/test_llms/test_llm.py b/tests/test_components/test_llms/test_llm.py
index 457bf13c..49de7f92 100644
--- a/tests/test_components/test_llms/test_llm.py
+++ b/tests/test_components/test_llms/test_llm.py
@@ -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(