init project

Author: gary-huang

ddtrace/llmobs/_llmobs.py

@@ -183,6 +183,7 @@ class LLMObs(Service):
     enabled = False
     _app_key: str = os.getenv("DD_APP_KEY", "")
     _project_name: str = os.getenv("DD_LLMOBS_PROJECT_NAME", DEFAULT_PROJECT_NAME)
+    _project_id: str = ""
 
     def __init__(
         self,
@@ -212,6 +213,7 @@ def __init__(
             interval=float(os.getenv("_DD_LLMOBS_WRITER_INTERVAL", 1.0)),
             timeout=float(os.getenv("_DD_LLMOBS_WRITER_TIMEOUT", 5.0)),
             _app_key=self._app_key,
+            _default_project_id=self._project_id,
             is_agentless=True,  # agent proxy doesn't seem to work for experiments
         )
 
@@ -613,6 +615,15 @@ def enable(
             cls.enabled = True
             cls._instance.start()
 
+            try:
+                cls._project_id = cls._instance._dne_client.project_create_or_get(cls._project_name)
+            except Exception as e:
+                log.error(
+                    "failed to get project ID with %s, dataset & experiments features may not be functional: %s",
+                    cls._project_name,
+                    e,
+                )
+
             # Register hooks for span events
             core.on("trace.span_start", cls._instance._on_span_start)
             core.on("trace.span_finish", cls._instance._on_span_finish)
@@ -654,15 +665,23 @@ def enable(
             )
 
     @classmethod
-    def pull_dataset(cls, name: str) -> Dataset:
-        ds = cls._instance._dne_client.dataset_get_with_records(name)
+    def pull_dataset(cls, dataset_name: str, project_name: Optional[str] = None) -> Dataset:
+        ds = cls._instance._dne_client.dataset_get_with_records(
+            dataset_name, cls._project_name if project_name is None else project_name
+        )
         return ds
 
     @classmethod
-    def create_dataset(cls, name: str, description: str = "", records: Optional[List[DatasetRecord]] = None) -> Dataset:
+    def create_dataset(
+        cls,
+        dataset_name: str,
+        project_name: Optional[str] = None,
+        description: str = "",
+        records: Optional[List[DatasetRecord]] = None,
+    ) -> Dataset:
         if records is None:
             records = []
-        ds = cls._instance._dne_client.dataset_create(name, description)
+        ds = cls._instance._dne_client.dataset_create(dataset_name, project_name, description)
         for r in records:
             ds.append(r)
         if len(records) > 0:
@@ -678,13 +697,14 @@ def create_dataset_from_csv(
         expected_output_columns: Optional[List[str]] = None,
         metadata_columns: Optional[List[str]] = None,
         csv_delimiter: str = ",",
-        description="",
+        description: str = "",
+        project_name: Optional[str] = None,
     ) -> Dataset:
         if expected_output_columns is None:
             expected_output_columns = []
         if metadata_columns is None:
             metadata_columns = []
-        ds = cls._instance._dne_client.dataset_create(dataset_name, description)
+        ds = cls._instance._dne_client.dataset_create(dataset_name, project_name, description)
 
         # Store the original field size limit to restore it later
         original_field_size_limit = csv.field_size_limit()

ddtrace/llmobs/_writer.py

@@ -140,6 +140,7 @@ def __init__(
         _api_key: str = "",
         _app_key: str = "",
         _override_url: str = "",
+        _default_project_id: str = "",
     ) -> None:
         super(BaseLLMObsWriter, self).__init__(interval=interval)
         self._lock = forksafe.RLock()
@@ -150,6 +151,7 @@ def __init__(
         self._site: str = _site or config._dd_site
         self._app_key: str = _app_key
         self._override_url: str = _override_url or os.environ.get("DD_LLMOBS_OVERRIDE_ORIGIN", "")
+        self._default_project_id: str = _default_project_id
 
         self._agentless: bool = is_agentless
         self._intake: str = self._override_url or (
@@ -362,23 +364,28 @@ def dataset_delete(self, dataset_id: str) -> None:
             raise ValueError(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"
+    def dataset_create(
+        self, dataset_name: str, project_name: Optional[str], description: str,
+    ) -> Dataset:
+        project_id = self.project_create_or_get(project_name)
+        logger.debug("getting records with project ID %s for %s", project_id, project_name)
+
+        path = f"/api/unstable/llm-obs/v1/{project_id}/datasets"
         body: JSONType = {
             "data": {
                 "type": "datasets",
-                "attributes": {"name": name, "description": description},
+                "attributes": {"name": dataset_name, "description": description},
             }
         }
         resp = self.request("POST", path, body)
         if resp.status != 200:
-            raise ValueError(f"Failed to create dataset {name}: {resp.status} {resp.get_json()}")
+            raise ValueError(f"Failed to create dataset {dataset_name}: {resp.status} {resp.get_json()}")
         response_data = resp.get_json()
         dataset_id = response_data["data"]["id"]
         if dataset_id is None or dataset_id == "":
             raise ValueError(f"unexpected dataset state, invalid ID (is None: {dataset_id is None})")
         curr_version = response_data["data"]["attributes"]["current_version"]
-        return Dataset(name, dataset_id, [], description, curr_version, _dne_client=self)
+        return Dataset(dataset_name, dataset_id, [], description, curr_version, _dne_client=self)
 
     @staticmethod
     def _get_record_json(record: Union[UpdatableDatasetRecord, DatasetRecordRaw], is_update: bool) -> JSONType:
@@ -436,16 +443,19 @@ def dataset_batch_update(
         new_record_ids: List[str] = [r["id"] for r in data] if data else []
         return new_version, new_record_ids
 
-    def dataset_get_with_records(self, name: str) -> Dataset:
-        path = f"/api/unstable/llm-obs/v1/datasets?filter[name]={quote(name)}"
+    def dataset_get_with_records(self, dataset_name: str, project_name: Optional[str] = None) -> Dataset:
+        project_id = self.project_create_or_get(project_name)
+        logger.debug("getting records with project ID %s for %s", project_id, project_name)
+
+        path = f"/api/unstable/llm-obs/v1/{project_id}/datasets?filter[name]={quote(dataset_name)}"
         resp = self.request("GET", path)
         if resp.status != 200:
-            raise ValueError(f"Failed to pull dataset {name}: {resp.status}")
+            raise ValueError(f"Failed to pull dataset {dataset_name} from project {project_name}: {resp.status}")
 
         response_data = resp.get_json()
         data = response_data["data"]
         if not data:
-            raise ValueError(f"Dataset '{name}' not found")
+            raise ValueError(f"Dataset '{dataset_name}' not found in project {project_name}")
 
         curr_version = data[0]["attributes"]["current_version"]
         dataset_description = data[0]["attributes"].get("description", "")
@@ -460,7 +470,8 @@ def dataset_get_with_records(self, name: str) -> Dataset:
             resp = self.request("GET", list_path, timeout=self.LIST_RECORDS_TIMEOUT)
             if resp.status != 200:
                 raise ValueError(
-                    f"Failed to pull {page_num}th page of dataset records {name}: {resp.status} {resp.get_json()}"
+                    f"Failed to pull {page_num}th page of dataset records {dataset_name}: "
+                    f"{resp.status} {resp.get_json()}"
                 )
             records_data = resp.get_json()
 
@@ -481,7 +492,7 @@ def dataset_get_with_records(self, name: str) -> Dataset:
                 list_path = f"{list_base_path}?page[cursor]={next_cursor}"
                 logger.debug("next list records request path %s", list_path)
                 page_num += 1
-        return Dataset(name, dataset_id, class_records, dataset_description, curr_version, _dne_client=self)
+        return Dataset(dataset_name, dataset_id, class_records, dataset_description, curr_version, _dne_client=self)
 
     def dataset_bulk_upload(self, dataset_id: str, records: List[DatasetRecord]):
         with tempfile.NamedTemporaryFile(suffix=".csv") as tmp:
@@ -534,7 +545,10 @@ def dataset_bulk_upload(self, dataset_id: str, records: List[DatasetRecord]):
             raise ValueError(f"Failed to upload dataset from file: {resp.status} {resp.get_json()}")
         logger.debug("successfully uploaded with code %d", resp.status)
 
-    def project_create_or_get(self, name: str) -> str:
+    def project_create_or_get(self, name: Optional[str] = None) -> str:
+        if name is None or name == "":
+            return self._default_project_id
+
         path = "/api/unstable/llm-obs/v1/projects"
         resp = self.request(
             "POST",
@@ -544,7 +558,12 @@ def project_create_or_get(self, name: str) -> str:
         if resp.status != 200:
             raise ValueError(f"Failed to create project {name}: {resp.status} {resp.get_json()}")
         response_data = resp.get_json()
-        return response_data["data"]["id"]
+        project_id = response_data["data"]["id"]
+
+        if project_id is None or project_id == "":
+            raise ValueError(f"project ID is required for dataset & experiments features (project name: {name})")
+
+        return project_id
 
     def experiment_create(
         self,