add create/deletion methods, add tests

Author: Kyle-Verhoog

ddtrace/llmobs/_constants.py

@@ -45,8 +45,10 @@
 EVP_SUBDOMAIN_HEADER_NAME = "X-Datadog-EVP-Subdomain"
 SPAN_SUBDOMAIN_NAME = "llmobs-intake"
 EVAL_SUBDOMAIN_NAME = "api"
+EXP_SUBDOMAIN_NAME = "api"
 AGENTLESS_SPAN_BASE_URL = "https://{}".format(SPAN_SUBDOMAIN_NAME)
 AGENTLESS_EVAL_BASE_URL = "https://{}".format(EVAL_SUBDOMAIN_NAME)
+AGENTLESS_EXP_BASE_URL = "https://{}".format(EXP_SUBDOMAIN_NAME)
 
 EVP_PAYLOAD_SIZE_LIMIT = 5 << 20  # 5MB (actual limit is 5.1MB)
 EVP_EVENT_SIZE_LIMIT = (1 << 20) - 1024  # 999KB (actual limit is 1MB)

ddtrace/llmobs/_experiment.py

@@ -18,25 +18,19 @@ class DatasetRecord(TypedDict):
 
 
 class Dataset:
-    _name: str
+    name: str
     _id: str
     _data: List[DatasetRecord]
 
-    def __init__(self, name: str, id: str, data: List[DatasetRecord]) -> None:
-        self._name = name
-        self._id = id
+    def __init__(self, name: str, dataset_id: str, data: List[DatasetRecord]) -> None:
+        self.name = name
+        self._id = dataset_id
         self._data = data
 
-    def __str__(self) -> str:
-        return f"Dataset(name={self._name}, id={self._id}, data={self._data})"
-
 
 class Experiment:
     def __init__(self, name: str, dataset: Dataset, description: str = "") -> None:
         self.name = name
         self._dataset = dataset
         self._experiment_id: Optional[str] = None
         self._project_id: Optional[str] = None
-
-    def run(self):
-        pass

ddtrace/llmobs/_llmobs.py

@@ -14,7 +14,6 @@
 from typing import TypedDict
 from typing import Union
 from typing import cast
-from urllib.parse import quote
 
 import ddtrace
 from ddtrace import config
@@ -562,6 +561,14 @@ def enable(
     def pull_dataset(cls, name: str) -> Dataset:
         return cls._instance._dne_client.dataset_pull(name)
 
+    @classmethod
+    def create_dataset(cls, name: str, description: str) -> Dataset:
+        return cls._instance._dne_client.dataset_create(name, description)
+
+    @classmethod
+    def _delete_dataset(cls, dataset_id: str) -> None:
+        return cls._instance._dne_client.dataset_delete(dataset_id)
+
     def experiment(self, name: str, dataset: Dataset) -> Experiment:
         return Experiment(name, dataset)
 

ddtrace/llmobs/_writer.py

@@ -1,4 +1,5 @@
 import atexit
+import json
 from typing import Any
 from typing import Dict
 from typing import List
@@ -25,6 +26,7 @@
 from ddtrace.internal.utils.retry import fibonacci_backoff_with_jitter
 from ddtrace.llmobs import _telemetry as telemetry
 from ddtrace.llmobs._constants import AGENTLESS_EVAL_BASE_URL
+from ddtrace.llmobs._constants import AGENTLESS_EXP_BASE_URL
 from ddtrace.llmobs._constants import AGENTLESS_SPAN_BASE_URL
 from ddtrace.llmobs._constants import DROPPED_IO_COLLECTION_ERROR
 from ddtrace.llmobs._constants import DROPPED_VALUE_TEXT
@@ -38,6 +40,7 @@
 from ddtrace.llmobs._constants import SPAN_SUBDOMAIN_NAME
 from ddtrace.llmobs._experiment import Dataset
 from ddtrace.llmobs._experiment import DatasetRecord
+from ddtrace.llmobs._experiment import JSONType
 from ddtrace.llmobs._utils import safe_json
 from ddtrace.settings._agent import config as agent_config
 
@@ -270,57 +273,73 @@ def _data(self, events: List[LLMObsEvaluationMetricEvent]) -> Dict[str, Any]:
 
 class LLMObsExperimentsClient(BaseLLMObsWriter):
 
-    def request(self, method: str, path: str, body: bytes = b"") -> Response:
+    AGENTLESS_BASE_URL = AGENTLESS_EXP_BASE_URL
+
+    def request(self, method: str, path: str, body: JSONType = None) -> Response:
         headers = {
             "Content-Type": "application/json",
             "DD-API-KEY": self._api_key,
             "DD-APPLICATION-KEY": self._app_key,
         }
-        site = self._site
-        if site == "datad0g.com":
-            base = "https://dd.datad0g.com"
-        else:
-            base = f"https://api.{site}"
-
-        conn = get_connection(base)
+        body = json.dumps(body).encode("utf-8") if body else b""
+        conn = get_connection(self._intake)
         try:
-            url = base + path
+            url = self._intake + path
             logger.debug("requesting %s", url)
             conn.request(method, url, body, headers)
             resp = conn.getresponse()
-            if resp.status >= 300:
-                raise ValueError(f"Failed to {method} {path}: {resp.status}")
             return Response.from_http_response(resp)
         finally:
             conn.close()
 
-    def dataset_pull(self, name: str) -> Dataset:
+    def dataset_delete(self, dataset_id: str) -> None:
+        path = "/api/unstable/llm-obs/v1/datasets/delete"
+        resp = self.request("POST", path, body={
+            "data": {
+                "type": "datasets",
+                "attributes": {
+                    "type": "soft",
+                    "dataset_ids": [dataset_id],
+                },
+            },
+        })
+        assert resp.status == 200, f"Failed to delete dataset {id}: {resp.get_json()}"
+        return None
+
+    def dataset_create(self, name: str, description: str) -> Dataset:
+        path = "/api/unstable/llm-obs/v1/datasets"
+        body = {
+            "data": {
+                "type": "datasets",
+                "attributes": {"name": name, "description": description},
+            }
+        }
+        resp = self.request("POST", path, body)
+        response_data = resp.get_json()
+        dataset_id = response_data["data"]["id"]
+        return Dataset(name, dataset_id, [])
 
+    def dataset_pull(self, name: str) -> Dataset:
         path = f"/api/unstable/llm-obs/v1/datasets?filter[name]={quote(name)}"
         resp = self.request("GET", path)
 
         response_data = resp.get_json()
-        datasets = response_data.get("data", [])
-
-        if not datasets:
+        data = response_data["data"]
+        if not data:
             raise ValueError(f"Dataset '{name}' not found")
 
-        dataset_id = datasets[0]["id"]
+        dataset_id = data[0]["id"]
         url = f"/api/unstable/llm-obs/v1/datasets/{dataset_id}/records"
-        try:
-            resp = self.request("GET", url)
-            records_data = resp.get_json()
-        except ValueError as e:
-            if "404" in str(e):
-                raise ValueError(f"Dataset '{name}' not found") from e
-            raise
+        resp = self.request("GET", url)
+        if resp.status == 404:
+            raise ValueError(f"Dataset '{name}' not found")
+        records_data = resp.get_json()
 
         class_records: List[DatasetRecord] = []
         for record in records_data.get("data", []):
             attrs = record.get("attributes", {})
             input_data = attrs.get("input")
             expected_output = attrs.get("expected_output")
-
             class_records.append(
                 {
                     "record_id": record.get("id"),
@@ -329,7 +348,6 @@ def dataset_pull(self, name: str) -> Dataset:
                     **attrs.get("metadata", {}),
                 }
             )
-
         return Dataset(name, dataset_id, class_records)
 
 

tests/llmobs/_utils.py

@@ -24,7 +24,7 @@
         cassette_library_dir=os.path.join(os.path.dirname(__file__), "llmobs_cassettes/"),
         record_mode="once",
         match_on=["path"],
-        filter_headers=["authorization", "OpenAI-Organization", "api-key", "x-api-key", ("DD-API-KEY", "XXXXXX")],
+        filter_headers=["authorization", "OpenAI-Organization", "api-key", "x-api-key", ("DD-API-KEY", "XXXXXX"), ("DD-APPLICATION-KEY", "XXXXXX")],
         # Ignore requests to the agent
         ignore_localhost=True,
     )

tests/llmobs/conftest.py

@@ -273,7 +273,9 @@ def llmobs(
         llmobs_service.enable(_tracer=tracer, **llmobs_enable_opts)
         llmobs_service._instance._llmobs_span_writer = llmobs_span_writer
         llmobs_service._instance._llmobs_span_writer.start()
+        llmobs_service._instance._dne_client._intake = "http://localhost:9126/vcr/datadog"
         yield llmobs_service
+    tracer.shutdown()
     llmobs_service.disable()
 
 

tests/llmobs/test_experiments.py

@@ -1,5 +1,33 @@
+import os
 
+import pytest
 
-def test_dataset_pull(llmobs):
-    dataset = llmobs.pull_dataset(name="kyle-test")
-    assert dataset._id == "929531d1-3cd2-473d-ab4e-2423b40c5db5"
+
+@pytest.fixture
+def test_dataset(llmobs):
+    ds = llmobs.create_dataset(name="test-dataset", description="A test dataset")
+
+    # When recording the requests, we need to wait for the dataset to be queryable.
+    if os.environ.get("RECORD_REQUESTS"):
+        import time
+        time.sleep(1)
+
+    yield ds
+
+    llmobs._delete_dataset(dataset_id=ds._id)
+
+
+def test_dataset_create_delete(llmobs):
+    dataset = llmobs.create_dataset(name="test-dataset-2", description="A second test dataset")
+    assert dataset._id is not None
+    llmobs._delete_dataset(dataset_id=dataset._id)
+
+
+def test_dataset_pull_non_existent(llmobs):
+    with pytest.raises(ValueError):
+        llmobs.pull_dataset(name="test-dataset-non-existent")
+
+
+def test_dataset_pull(llmobs, test_dataset):
+    dataset = llmobs.pull_dataset(name=test_dataset.name)
+    assert dataset._id is not None