Setting the file. One moment. scripts/publish_dataset.py
scripts/publish_dataset.py
Python·122 lines·4 KB
15 --dataset-name my_seed_20260529 \\
16 --project-name my-llm-app \\
17 [--env-file /abs/path/to/extra.env]
18
19Prints `OK dataset_name=... record_count=... url=...` on success.
20"""
21from __future__ import annotations
22
23import argparse
24import json
25import os
26import pathlib
27import sys
28
29# Make sibling load_env.py importable regardless of how the script is invoked.
30sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent))
31from load_env import load_env_files # noqa: E402
32
33
34def normalize_tags(raw):
35 """Return (normalized_tags, fix_count).
36
37 Wraps bare strings as `tag:<value>`, drops empty/None, preserves
38 malformed leading/trailing colons by wrapping the original as
39 `tag:<original>` so the SDK's `validate_tags_list` cannot reject
40 the record. See agent-observability/agent-observability-eval-bootstrap/SKILL.md Phase 3D
41 "Tag normalization" for the rationale.
42 """
43 fixed = []
44 fix_count = 0
45 for t in raw or []:
46 if not isinstance(t, str):
47 fix_count += 1
48 continue
49 t = t.strip()
50 if not t:
51 fix_count += 1
52 continue
53 if ":" in t:
54 k, _, v = t.partition(":")
55 if k and v:
56 fixed.append(t)
57 continue
58 fixed.append(f"tag:{t}")
59 fix_count += 1
60 continue
61 fixed.append(f"tag:{t}")
62 fix_count += 1
63 return fixed, fix_count
64
65
66def main() -> int:
67 parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
68 parser.add_argument("--records", required=True, help="Absolute path to the DatasetRecordRaw[] JSON file.")
69 parser.add_argument("--dataset-name", required=True, help="Name to publish the dataset under in Datadog.")
70 parser.add_argument("--project-name", required=True, help="Datadog project to publish the dataset under. Created lazily.")
71 parser.add_argument("--env-file", action="append", default=[], help="Extra .env path(s) to load FIRST (repeatable).")
72 args = parser.parse_args()
73
74 loaded = load_env_files(args.env_file)
75 if loaded:
76 print(f"Loaded credentials from: {', '.join(loaded)}")
77
78 api_key = os.getenv("DD_API_KEY")
79 app_key = os.getenv("DD_APPLICATION_KEY") or os.getenv("DD_APP_KEY")
80 if not api_key:
81 print("ERROR: DD_API_KEY is not set. Export it in your shell or add it to a discovered .env file.", file=sys.stderr)
82 return 2
83 if not app_key:
84 print("ERROR: DD_APPLICATION_KEY (or DD_APP_KEY) is not set. Same fallback paths as above.", file=sys.stderr)
85 return 2
86
87 from ddtrace.llmobs import LLMObs # imported AFTER env load so credentials are picked up
88
89 LLMObs.enable(
90 api_key=api_key,
91 app_key=app_key,
92 site=os.getenv("DD_SITE", "datadoghq.com"),
93 project_name=args.project_name, # project is created lazily here if it does not exist
94 agentless_enabled=True,
95 )
96
97 with open(args.records) as f:
98 records = json.load(f)
99
100 total_fixes = 0
101 for r in records:
102 if "tags" in r:
103 r["tags"], n = normalize_tags(r.get("tags"))
104 total_fixes += n
105 if total_fixes:
106 print(
107 f"WARNING: normalized {total_fixes} malformed tag(s) before publish "
108 "(bare strings wrapped as 'tag:<value>'; empties dropped)."
109 )
110
111 dataset = LLMObs.create_dataset(
112 dataset_name=args.dataset_name,
113 description=f"Seed dataset for {args.dataset_name} (eval-pipeline flow, sampled from production traces).",
114 records=records,
115 )
116 url = dataset.url if hasattr(dataset, "url") else "<inspect in UI>"
117 print(f"OK dataset_name={args.dataset_name} record_count={len(records)} url={url}")
118 return 0
119
120
121if __name__ == "__main__":
122 sys.exit(main())