feat(llmobs): add base datasets and experiments

Author: Kyle-Verhoog

ddtrace/llmobs/__init__.py

@@ -6,8 +6,11 @@
     LLMObs.enable()
 """
 
+from ._experiment import Dataset
+from ._experiment import DatasetRecord
+from ._experiment import Experiment
 from ._llmobs import LLMObs
 from ._llmobs import LLMObsSpan
 
 
-__all__ = ["LLMObs", "LLMObsSpan"]
+__all__ = ["LLMObs", "LLMObsSpan", "Experiment", "Dataset", "DatasetRecord"]

ddtrace/llmobs/_experiment.py

@@ -0,0 +1,42 @@
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import NotRequired
+from typing import Optional
+from typing import TypedDict
+from typing import Union
+
+
+JSONType = Union[str, int, float, bool, None, List["JSONType"], Dict[str, "JSONType"]]
+
+
+class DatasetRecord(TypedDict):
+    input: JSONType
+    expected_output: JSONType
+    record_id: NotRequired[Optional[str]]
+    metadata: NotRequired[Optional[Dict[str, Any]]]
+
+
+class Dataset:
+    _name: str
+    _id: str
+    _data: List[DatasetRecord]
+
+    def __init__(self, name: str, id: str, data: List[DatasetRecord]) -> None:
+        self._name = name
+        self._id = 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,6 +14,7 @@
 from typing import TypedDict
 from typing import Union
 from typing import cast
+from urllib.parse import quote
 
 import ddtrace
 from ddtrace import config
@@ -73,6 +74,8 @@
 from ddtrace.llmobs._constants import TAGS
 from ddtrace.llmobs._context import LLMObsContextProvider
 from ddtrace.llmobs._evaluators.runner import EvaluatorRunner
+from ddtrace.llmobs._experiment import Dataset
+from ddtrace.llmobs._experiment import Experiment
 from ddtrace.llmobs._utils import AnnotationContext
 from ddtrace.llmobs._utils import LinkTracker
 from ddtrace.llmobs._utils import ToolCallTracker
@@ -85,6 +88,7 @@
 from ddtrace.llmobs._utils import validate_prompt
 from ddtrace.llmobs._writer import LLMObsEvalMetricWriter
 from ddtrace.llmobs._writer import LLMObsEvaluationMetricEvent
+from ddtrace.llmobs._writer import LLMObsExperimentsClient
 from ddtrace.llmobs._writer import LLMObsSpanEvent
 from ddtrace.llmobs._writer import LLMObsSpanWriter
 from ddtrace.llmobs._writer import should_use_agentless
@@ -180,6 +184,11 @@ def __init__(
             interval=float(os.getenv("_DD_LLMOBS_EVALUATOR_INTERVAL", 1.0)),
             llmobs_service=self,
         )
+        self._dne_client = LLMObsExperimentsClient(
+            interval=float(os.getenv("_DD_LLMOBS_WRITER_INTERVAL", 1.0)),
+            timeout=float(os.getenv("_DD_LLMOBS_WRITER_TIMEOUT", 5.0)),
+            is_agentless=agentless_enabled,
+        )
 
         forksafe.register(self._child_after_fork)
 
@@ -549,6 +558,13 @@ def enable(
                 config._llmobs_ml_app,
             )
 
+    @classmethod
+    def pull_dataset(cls, name: str) -> Dataset:
+        return cls._instance._dne_client.dataset_pull(name)
+
+    def experiment(self, name: str, dataset: Dataset) -> Experiment:
+        return Experiment(name, dataset)
+
     @classmethod
     def register_processor(cls, processor: Optional[Callable[[LLMObsSpan], LLMObsSpan]] = None) -> None:
         """Register a processor to be called on each LLMObs span.

ddtrace/llmobs/_writer.py

@@ -4,6 +4,7 @@
 from typing import List
 from typing import Optional
 from typing import Union
+from urllib.parse import quote
 
 
 # TypedDict was added to typing in python 3.8
@@ -35,6 +36,8 @@
 from ddtrace.llmobs._constants import EVP_SUBDOMAIN_HEADER_NAME
 from ddtrace.llmobs._constants import SPAN_ENDPOINT
 from ddtrace.llmobs._constants import SPAN_SUBDOMAIN_NAME
+from ddtrace.llmobs._experiment import Dataset
+from ddtrace.llmobs._experiment import DatasetRecord
 from ddtrace.llmobs._utils import safe_json
 from ddtrace.settings._agent import config as agent_config
 
@@ -109,6 +112,7 @@ def __init__(
         is_agentless: bool,
         _site: str = "",
         _api_key: str = "",
+        _app_key: str = "",
         _override_url: str = "",
     ) -> None:
         super(BaseLLMObsWriter, self).__init__(interval=interval)
@@ -118,6 +122,7 @@ def __init__(
         self._timeout: float = timeout
         self._api_key: str = _api_key or config._dd_api_key
         self._site: str = _site or config._dd_site
+        self._app_key: str = _app_key or config._dd_app_key
         self._override_url: str = _override_url
 
         self._agentless: bool = is_agentless
@@ -263,6 +268,71 @@ def _data(self, events: List[LLMObsEvaluationMetricEvent]) -> Dict[str, Any]:
         return {"data": {"type": "evaluation_metric", "attributes": {"metrics": events}}}
 
 
+class LLMObsExperimentsClient(BaseLLMObsWriter):
+
+    def request(self, method: str, path: str, body: bytes = b"") -> 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)
+        try:
+            url = base + 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:
+
+        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:
+            raise ValueError(f"Dataset '{name}' not found")
+
+        dataset_id = datasets[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
+
+        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"),
+                    "input": input_data,
+                    "expected_output": expected_output,
+                    **attrs.get("metadata", {}),
+                }
+            )
+
+        return Dataset(name, dataset_id, class_records)
+
+
 class LLMObsSpanWriter(BaseLLMObsWriter):
     """Writes span events to the LLMObs Span Endpoint."""
 

ddtrace/settings/_config.py

@@ -645,6 +645,7 @@ def __init__(self):
         self._telemetry_install_time = _get_config("DD_INSTRUMENTATION_INSTALL_TYPE")
 
         self._dd_api_key = _get_config("DD_API_KEY")
+        self._dd_app_key = _get_config("DD_APP_KEY")
         self._dd_site = _get_config("DD_SITE", "datadoghq.com")
 
         self._llmobs_enabled = _get_config("DD_LLMOBS_ENABLED", False, asbool)

tests/llmobs/conftest.py

@@ -99,7 +99,11 @@ def ddtrace_global_config():
 
 
 def default_global_config():
-    return {"_dd_api_key": "<not-a-real-api_key>", "_llmobs_ml_app": "unnamed-ml-app", "service": "tests.llmobs"}
+    return {
+        "_dd_api_key": os.environ.get("DD_API_KEY", "<not-a-real-api_key>"),
+        "_llmobs_ml_app": "unnamed-ml-app",
+        "service": "tests.llmobs",
+    }
 
 
 @pytest.fixture
@@ -172,7 +176,7 @@ def tracer():
 @pytest.fixture
 def llmobs_env():
     return {
-        "DD_API_KEY": "<default-not-a-real-key>",
+        "DD_API_KEY": os.environ.get("DD_API_KEY", "<default-not-a-real-key>"),
         "DD_LLMOBS_ML_APP": "unnamed-ml-app",
     }
 

tests/llmobs/test_experiments.py

@@ -0,0 +1,5 @@
+
+
+def test_dataset_pull(llmobs):
+    dataset = llmobs.pull_dataset(name="kyle-test")
+    assert dataset._id == "929531d1-3cd2-473d-ab4e-2423b40c5db5"