diff --git a/.env.example b/.env.example index e830110..de826f3 100644 --- a/.env.example +++ b/.env.example @@ -36,3 +36,10 @@ LZ_API_TOKEN=user-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx # LogZilla's own documentation uses LOGZILLA_HOST / LOGZILLA_API_KEY in its # examples. This tool accepts those names too, so if you already export them # for `logzilla query`, you can leave the two settings above unset. + +# Optional, for --export-trace: the Phoenix server the LogZilla AI settings send +# traces to (the host of TRACING_SERVER_URL) and the project (TRACING_PROJECT_NAME). +# Both can also be passed as --phoenix-url / --phoenix-project, or read from the +# --deploy-config file. +#LZ_PHOENIX_URL=http://phoenix.example.com +#LZ_PHOENIX_PROJECT=logzilla-ai diff --git a/README.md b/README.md index 7f709b8..a66a469 100644 --- a/README.md +++ b/README.md @@ -191,8 +191,20 @@ output: --raw Include raw WebSocket messages in JSON/saved output. --save-dir DIR Also write a .json and .md file per question. --keep-charts Keep the JSON blocks in terminal output. + --tag TAG Write the saved files under //. -q, --quiet Print only the answer. -v, --verbose Show connection details and extra progress. + +capture: + --deploy-config PATH Copy this ai.yaml to the server as root over SSH, + reload the AI settings, wait, then ask. + --reload-wait SECONDS Seconds to wait after the reload (default: 30). + --export-trace Save each question's Phoenix trace beside its result. + --phoenix-url URL The Phoenix server (TRACING_SERVER_URL's host). + --phoenix-project NAME + The Phoenix project (TRACING_PROJECT_NAME). + --trace-wait SECONDS Seconds to wait for a trace to land in Phoenix + (default: 60). ``` ## JSON output @@ -231,6 +243,42 @@ To pull just the answer out, use `jq`: | `-.json` | Same structure as `--json` output | | `-.md` | Readable transcript: question, answer, tool calls, errors | +## Capturing a run the way LogZilla certifies models + +LogZilla certifies AI models by running this same tool against a test server +and keeping, per question, the result, the AI settings that were active and +the question's trace. Three flags let you capture a run in your own lab the +same way, so a model you test on your hardware can be compared with a +published row like for like. + +```bash +./lz-ai-query.py --file question.txt --save-dir ./results --tag llama-70b-l40s \ + --deploy-config ./ai.yaml --export-trace +``` + +| Flag | What it does | +| --- | --- | +| `--tag TAG` | Writes the files under `results//`, so one directory holds one model or hardware setup. Needs `--save-dir`. | +| `--deploy-config PATH` | Before the first question, copies the file to `/etc/logzilla/settings/ai.yaml` on the server as `root` over SSH (`scp`, then `ssh root@ logzilla settings reload ai`) and waits `--reload-wait` seconds (default 30) for the settings to apply. A failed copy or reload stops the run with the server's own message. With `--save-dir`, a copy of the file is saved as `ai-config.yaml` with every `*_KEY`, `*_TOKEN` and `API_KEY` value replaced by ``. | +| `--export-trace` | After each answer, finds the question's trace on the Phoenix server the AI settings send traces to (`TRACING_ENABLED: true`, `TRACING_SERVER_URL`, `TRACING_PROJECT_NAME`) and saves it as `--trace.json` beside the result: every span with its name, timings, input, output and token counts. The server and project come from `--phoenix-url` and `--phoenix-project`, from `LZ_PHOENIX_URL` and `LZ_PHOENIX_PROJECT` in the environment or `.env`, or from the `--deploy-config` file's `TRACING_*` settings. | + +The trace is matched, never guessed: the tool looks at the root spans that +started inside the question's own time window (five seconds of slack on each +side), exports the candidates and keeps the one whose spans carry the question +text, matched as a whole. Phoenix ingests spans a few seconds after the answer +arrives, so the lookup is repeated every five seconds for up to `--trace-wait` +seconds (default 60) before the trace is reported missing. A question with no +such trace, or with more than one, is reported as a trace error for that +question (`"trace": {"status": "missing"}` or `"ambiguous"` in the JSON, exit +code 2), because a run without its trace cannot be scored; a Phoenix failure +during the lookup is `"status": "error"` with the server's words, the answer +still saved. More than five roots in one window is also ambiguous; keep the +test server to one run at a time. The project named by `TRACING_PROJECT_NAME` +may be new: Phoenix creates it with the first trace that reaches it. + +`--deploy-config` needs root SSH access to the LogZilla server. The other two +flags need only the API token and, for the trace, network access to Phoenix. + ## Exit codes | Code | Meaning | diff --git a/lz-ai-query.py b/lz-ai-query.py index 597f560..d9639ce 100755 --- a/lz-ai-query.py +++ b/lz-ai-query.py @@ -15,10 +15,12 @@ import argparse import asyncio +import functools import json import os import re import ssl +import subprocess import sys import time from contextlib import asynccontextmanager @@ -34,6 +36,28 @@ CONNECT_TIMEOUT = 30 MAX_WS_MESSAGE = 10 * 1024 * 1024 # 10 MB +# --- capture (the flags LogZilla's certification harness runs this tool with) +REMOTE_AI_YAML = "/etc/logzilla/settings/ai.yaml" +RELOAD_WAIT_DEFAULT = 30 # seconds the server takes to apply reloaded AI settings +TRACE_WINDOW_SLACK_S = 5 # a trace root may start this much before/after the question +TRACE_MAX_CANDIDATES = 5 # more roots in one question's window is ambiguous, not a guess +TRACE_MAX_SPANS = 500 +TRACE_WAIT_DEFAULT = 60 # seconds to wait for a question's trace to land in Phoenix +TRACE_POLL_S = 5 # seconds between lookups while waiting +TAG_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,63}$") +# a YAML line whose key (quoted or not) names a secret: its value is replaced in +# the saved copy, and so is any continuation of it on the more-indented lines +# below (a block scalar, a multi-line quoted or plain scalar, a value that starts +# on the next line) +SECRET_SETTING_RE = re.compile( + r"""^(\s*["']?[A-Za-z0-9_]*(?:_KEY|_TOKEN|API_KEY)["']?\s*:)(\s*\S.*)?$""" +) +# a secret key with a value anywhere else on a line (a flow-style mapping, an +# inline value): a form the line-by-line redaction cannot cover, so it is refused +SECRET_ANYWHERE_RE = re.compile( + r"""["']?[A-Za-z0-9_]*(?:_KEY|_TOKEN|API_KEY)["']?\s*:\s*\S""" +) + # -------------------------------------------------------------------------- # Configuration @@ -212,6 +236,7 @@ def __init__(self, query): self.errors = [] self.stop_reason = None self.messages = [] # raw WebSocket messages, for debugging + self.trace = None # {'status', 'trace_ids', 'path'} when --export-trace ran @property def duration_s(self): @@ -246,6 +271,8 @@ def to_dict(self, include_raw=False): "errors": self.errors, "stop_reason": self.stop_reason, } + if self.trace is not None: + data["trace"] = self.trace if include_raw: data["raw_messages"] = self.messages return data @@ -626,21 +653,33 @@ def slugify(text, max_len=48): return (slug[:max_len].rstrip("-") or "query") -def save_result(result, save_dir, include_raw=False): - """Write -.json and .md into save_dir. +def output_dir(save_dir, tag=None): + """save_dir, or save_dir/tag with --tag, created on first use.""" + path = Path(save_dir).expanduser() + if tag: + path = path / tag + path.mkdir(parents=True, exist_ok=True) + return path - UTF-8 is explicit on both writes: answers routinely contain non-ASCII - (em dashes, accented text), which would raise UnicodeEncodeError under a - C/POSIX or cp1252 locale. - """ - save_dir = Path(save_dir).expanduser() - save_dir.mkdir(parents=True, exist_ok=True) + +def result_base(result, save_dir, tag=None): + """The path every file of one question shares: /-.""" stamp = ( datetime.fromtimestamp(result.start_time or time.time(), tz=timezone.utc) .strftime("%Y%m%dT%H%M%SZ") ) - base = save_dir / f"{stamp}-{slugify(result.query)}" + return output_dir(save_dir, tag) / f"{stamp}-{slugify(result.query)}" + + +def save_result(result, save_dir, include_raw=False, tag=None): + """Write -.json and .md into save_dir (or save_dir/tag). + + UTF-8 is explicit on both writes: answers routinely contain non-ASCII + (em dashes, accented text), which would raise UnicodeEncodeError under a + C/POSIX or cp1252 locale. + """ + base = result_base(result, save_dir, tag) json_path = base.with_suffix(".json") with open(json_path, "w", encoding="utf-8") as handle: @@ -654,6 +693,321 @@ def save_result(result, save_dir, include_raw=False): return json_path, md_path +# -------------------------------------------------------------------------- +# Capture: the config deploy, the redacted config copy, the Phoenix trace. +# These are the steps LogZilla's own certification harness wraps around this +# tool, so a customer can capture a run the same way in their own lab. +# -------------------------------------------------------------------------- + +def redact_config_text(text): + """The ai.yaml text with every *_KEY, *_TOKEN and API_KEY value replaced + by , so the saved copy never carries a secret. The key may be + quoted. Whatever follows the key on more-indented lines is part of its + value (a block scalar, a multi-line quoted or plain scalar, a value that + starts on the next line) and is dropped with it; a secret key with no + value at all stays as it is. A secret in a form the line-anchored key + cannot see (a flow-style mapping, an inline value) raises ValueError: + the copy is refused, never saved with the credential in it.""" + lines = text.splitlines() + for number, line in enumerate(lines, 1): + if line.lstrip().startswith("#") or SECRET_SETTING_RE.match(line): + continue + if SECRET_ANYWHERE_RE.search(line): + raise ValueError( + f"line {number}: a secret in a form the redaction does not cover (a " + "flow-style mapping or an inline value); write the setting as " + "KEY: value on its own line" + ) + out = [] + i = 0 + while i < len(lines): + line = lines[i] + match = SECRET_SETTING_RE.match(line) + if not match: + out.append(line) + i += 1 + continue + indent = len(line) - len(line.lstrip()) + end = i + 1 # the first line after the key that is not its continuation + while end < len(lines) and ( + not lines[end].strip() + or len(lines[end]) - len(lines[end].lstrip()) > indent + ): + end += 1 + continued = any(lines[j].strip() for j in range(i + 1, end)) + if match.group(2) is None and not continued: + out.append(line) # `KEY:` with nothing under it: an empty value, kept + i += 1 + continue + out.append(match.group(1) + " ") + i = end if continued else i + 1 + return "\n".join(out) + "\n" + + +def read_config_file(path): + path = Path(path).expanduser() + if not path.is_file(): + die(f"config file not found: {path}") + return path, path.read_text(encoding="utf-8") + + +def yaml_scalar(raw): + """The text of a plain or quoted YAML scalar: a quoted value is what the + quotes enclose, a plain one ends at a trailing ` #` comment (the same + conventions parse_env_line applies to .env values).""" + raw = raw.strip() + if raw[:1] in ("'", '"'): + end = raw.find(raw[0], 1) + return raw[1:end] if end > 0 else raw[1:] + if raw.startswith("#"): + return "" + return raw.split(" #", 1)[0].rstrip() + + +def tracing_from_config(text): + """(Phoenix base URL, project name) from TRACING_SERVER_URL and + TRACING_PROJECT_NAME in the ai.yaml text; None for a missing one. A + trailing comment is not part of the value.""" + url = project = None + for line in text.splitlines(): + match = re.match( + r"^\s*(TRACING_SERVER_URL|TRACING_PROJECT_NAME)\s*:\s*(.+?)\s*$", + line, + ) + if not match: + continue + value = yaml_scalar(match.group(2)) + if not value: + continue + if match.group(1) == "TRACING_SERVER_URL": + url = value + else: + project = value + return (phoenix_base(url) if url else None), (project or None) + + +def phoenix_base(url): + """Scheme and host of the tracing URL (TRACING_SERVER_URL is the OTLP + endpoint, http:///v1/traces; GraphQL lives at /graphql).""" + parsed = urlparse(url if "://" in url else f"http://{url}") + if parsed.scheme not in ("http", "https") or not parsed.hostname: + die(f"Phoenix URL is not a valid URL: {url!r}") + return f"{parsed.scheme}://{parsed.netloc}" + + +def deploy_steps(host, config_path): + """The scp and ssh commands of the deploy. urlparse strips the brackets + of an IPv6 literal; scp needs them back or it reads the first colon as + the host/path separator (ssh takes the bare address).""" + hostname = urlparse(host).hostname + scp_host = f"[{hostname}]" if ":" in hostname else hostname + return [ + ["scp", str(config_path), f"root@{scp_host}:{REMOTE_AI_YAML}"], + ["ssh", f"root@{hostname}", "logzilla", "settings", "reload", "ai"], + ] + + +def deploy_config(host, config_path, wait_s, reporter): + """Copy the config to the server as root over SSH, reload the AI + settings, wait for them to apply. Any failure stops the run with the + remote's own words: a question asked against a half-applied config is + not a result worth keeping.""" + for argv in deploy_steps(host, config_path): + reporter(" ".join(argv), level="detail") + try: + proc = subprocess.run(argv, capture_output=True, text=True) + except OSError as exc: + die(f"cannot run {argv[0]}: {exc}") + if proc.returncode != 0: + die(f"{argv[0]} failed (exit {proc.returncode}) during the config " + f"deploy: {proc.stderr.strip() or proc.stdout.strip() or 'no message'}") + reporter(f"config deployed; waiting {wait_s}s for the settings reload") + time.sleep(wait_s) + + +# Phoenix speaks GraphQL at /graphql with no auth. The declarations +# match its schema (node(id: ID!), Project.spans(first: Int!), +# getProjectByName(name: String!)); a wrong one is a validation error. +PHOENIX_PROJECT_QUERY = ( + "query ProjectByName($name: String!) { getProjectByName(name: $name) " + "{ id name } }" +) +PHOENIX_ROOTS_QUERY = ( + "query Roots($id: ID!, $first: Int!, $start: DateTime, $end: DateTime) { " + "node(id: $id) { ... on Project { " + "spans(first: $first, timeRange: {start: $start, end: $end}, " + 'sort: {col: startTime, dir: desc}, filterCondition: "parent_id is None") ' + "{ edges { node { context { traceId } startTime parentId name } } } } } }" +) +PHOENIX_TRACE_QUERY = ( + "query TraceSpans($traceId: String!, $first: Int!) { " + "getTraceByOtelId(traceId: $traceId) { " + "traceId startTime endTime latencyMs numSpans spans(first: $first) { " + "edges { node { id name spanKind statusCode statusMessage startTime " + "endTime latencyMs parentId context { traceId spanId } " + "input { value } output { value } " + "tokenCountTotal tokenCountPrompt tokenCountCompletion } } } } }" +) + + +class TraceError(Exception): + """A Phoenix step that failed. Raised after the answer is captured, it is + recorded on that question (the answer is still saved, the question fails) + rather than ending the run.""" + + +def phoenix_graphql(base, query, variables, verify_tls=True): + url = f"{base}/graphql" + try: + resp = requests.post(url, json={"query": query, "variables": variables}, + timeout=CONNECT_TIMEOUT, verify=verify_tls) + except requests.exceptions.RequestException as exc: + raise TraceError(f"cannot reach Phoenix at {url}: {exc}") + if resp.status_code != 200: + raise TraceError(f"Phoenix {url} answered HTTP {resp.status_code}: " + f"{resp.text[:200]}") + try: + doc = resp.json() + except ValueError: + raise TraceError(f"Phoenix {url}: the answer is not JSON") + if doc.get("errors"): + raise TraceError(f"Phoenix {url}: {doc['errors']}") + return doc.get("data") or {} + + +def phoenix_project_id(base, name, verify_tls=True): + """The project's id, or None while no trace has reached it: Phoenix + creates a project with its first span, so a fresh TRACING_PROJECT_NAME + exists only after the first question is traced.""" + data = phoenix_graphql(base, PHOENIX_PROJECT_QUERY, {"name": name}, + verify_tls) + node = data.get("getProjectByName") + if isinstance(node, dict) and node.get("id"): + return str(node["id"]) + return None + + +def phoenix_export(base, trace_id, verify_tls=True): + data = phoenix_graphql(base, PHOENIX_TRACE_QUERY, + {"traceId": trace_id, "first": TRACE_MAX_SPANS}, + verify_tls) + trace = data.get("getTraceByOtelId") + if not isinstance(trace, dict): + raise TraceError(f"Phoenix {base}: trace {trace_id} not found") + num = trace.get("numSpans") + if isinstance(num, int) and num > TRACE_MAX_SPANS: + raise TraceError(f"Phoenix {base}: trace {trace_id} has {num} spans, over " + f"the {TRACE_MAX_SPANS} cap") + spans = [edge["node"] for edge in (trace.get("spans") or {}).get("edges") or [] + if edge.get("node")] + head = {key: trace.get(key) + for key in ("traceId", "startTime", "endTime", "latencyMs", "numSpans")} + return {"trace": head, "spans": spans} + + +def _utc_iso(ts): + return datetime.fromtimestamp(ts, tz=timezone.utc).isoformat() + + +def _strings_of(value): + """Every string inside a span value: the value itself, and when it is + JSON, every string nested in it (the question sits inside the rendered + messages as a JSON string, escapes and all).""" + out = [] + if isinstance(value, str): + out.append(value) + try: + parsed = json.loads(value) + except ValueError: + return out + if isinstance(parsed, (dict, list)): + out.extend(_strings_of(parsed)) + elif isinstance(value, dict): + for item in value.values(): + out.extend(_strings_of(item)) + elif isinstance(value, list): + for item in value: + out.extend(_strings_of(item)) + return out + + +def find_trace(base, project, result, verify_tls=True, wait_s=TRACE_WAIT_DEFAULT, + sleep=time.sleep, clock=time.monotonic): + """The one Phoenix trace of this question: a root span that started + inside the question's window and whose spans carry the question text. + Returns (status, trace_ids, trace): 'exported' with the trace, 'missing' + with nothing, or 'ambiguous' with the candidates. It never guesses. + + Phoenix ingests spans after the answer arrives, so a trace that is not + there yet is looked up again every TRACE_POLL_S seconds for up to wait_s + seconds (the question's window stays fixed). The project is resolved + here, not at startup: it exists only once its first trace landed. A + Phoenix failure raises TraceError.""" + deadline = clock() + wait_s + while True: + found = _lookup_trace(base, project, result, verify_tls) + if found[0] != "missing" or clock() >= deadline: + return found + sleep(min(TRACE_POLL_S, deadline - clock())) + + +def _lookup_trace(base, project, result, verify_tls): + project_id = phoenix_project_id(base, project, verify_tls) + if project_id is None: + return "missing", [], None + start = result.start_time - TRACE_WINDOW_SLACK_S + end = (result.end_time or time.time()) + TRACE_WINDOW_SLACK_S + data = phoenix_graphql( + base, PHOENIX_ROOTS_QUERY, + {"id": project_id, "first": 200, "start": _utc_iso(start), + "end": _utc_iso(end)}, + verify_tls, + ) + edges = ((data.get("node") or {}).get("spans") or {}).get("edges") or [] + candidates = [] + for edge in edges: + node = edge.get("node") or {} + when = node.get("startTime") + if node.get("parentId") is not None or not isinstance(when, str): + continue + started = datetime.fromisoformat(when.replace("Z", "+00:00")).timestamp() + if not (start <= started <= end): + continue + trace_id = (node.get("context") or {}).get("traceId") + if trace_id and str(trace_id) not in candidates: + candidates.append(str(trace_id)) + if len(candidates) > TRACE_MAX_CANDIDATES: + return "ambiguous", sorted(candidates), None + needle = " ".join(result.query.split()) + matches = [] + for trace_id in candidates: + trace = phoenix_export(base, trace_id, verify_tls) + texts = [] + for span in trace["spans"]: + for key in ("input", "output"): + texts.extend(_strings_of((span.get(key) or {}).get("value"))) + # the question is a whole message string in the trace (the user turn), + # so the match is on the complete normalized string: a longer question + # that merely contains this one is another question's trace + if any(" ".join(text.split()) == needle for text in texts): + matches.append((trace_id, trace)) + if len(matches) == 1: + return "exported", [matches[0][0]], matches[0][1] + if not matches: + return "missing", [], None + return "ambiguous", sorted(trace_id for trace_id, _trace in matches), None + + +def trace_path_for(base): + return base.with_name(base.name + "-trace.json") + + +def save_trace(trace, path): + with open(path, "w", encoding="utf-8") as handle: + json.dump(trace, handle, indent=2, default=str) + return path + + # -------------------------------------------------------------------------- # Query sources # -------------------------------------------------------------------------- @@ -713,7 +1067,7 @@ async def run_queries(client, queries, args, reporter): reporter.heading(f"Question {i}/{total}: " f"{_preview(query)}") results.append(await chat.ask(query)) - _emit(results[-1], args, reporter, multiple=total > 1) + await _emit(results[-1], args, reporter, multiple=total > 1) return results for i, query in enumerate(queries, 1): @@ -723,7 +1077,7 @@ async def run_queries(client, queries, args, reporter): async with client.session(conv_id) as chat: result = await chat.ask(query) results.append(result) - _emit(result, args, reporter, multiple=total > 1) + await _emit(result, args, reporter, multiple=total > 1) return results @@ -763,7 +1117,7 @@ async def run_interactive(client, args, reporter): print("\n ! aborted", file=sys.stderr) continue results.append(result) - _emit(result, args, reporter, multiple=False) + await _emit(result, args, reporter, multiple=False) if result.errors and not result.answer: # A dead socket cannot be reused; start a new session. break @@ -775,8 +1129,38 @@ def _preview(query, width=70): else single_line -def _emit(result, args, reporter, multiple): - """Print and optionally persist one result.""" +async def _emit(result, args, reporter, multiple): + """Print and optionally persist one result. The trace lookup (blocking + HTTP calls and the ingestion wait) runs in a worker thread: with + --same-conversation the WebSocket stays open across questions, and a + blocked event loop would miss the server's heartbeat and lose it.""" + trace = None + if args.export_trace and result.start_time: + lookup = functools.partial( + find_trace, args.phoenix_base, args.phoenix_project, result, + verify_tls=not args.insecure, wait_s=args.trace_wait, + ) + try: + status, trace_ids, trace = await asyncio.get_running_loop().run_in_executor( + None, lookup) + detail = (", ".join(trace_ids) if trace_ids + else "no root span in the window carries the question") + except TraceError as exc: + # the answer is captured; a Phoenix failure is this question's + # trace error, saved with it, and the next question still runs + status, trace_ids, trace, detail = "error", [], None, str(exc) + base = result_base(result, args.save_dir, args.tag) + result.trace = { + "status": status, + "trace_ids": trace_ids, + "path": str(trace_path_for(base)) if trace is not None else None, + } + if status == "error": + result.trace["error"] = detail + if trace is None: + result.errors.append(f"trace {status}: {detail}") + reporter(f" ! trace {status}: {detail}", level="progress") + if args.json: if not multiple: print(json.dumps(result.to_dict(include_raw=args.raw), indent=2, @@ -786,8 +1170,11 @@ def _emit(result, args, reporter, multiple): if args.save_dir: json_path, md_path = save_result(result, args.save_dir, - include_raw=args.raw) + include_raw=args.raw, tag=args.tag) reporter(f"saved {json_path.name} and {md_path.name}") + if trace is not None: + path = save_trace(trace, Path(result.trace["path"])) + reporter(f"saved {path.name} ({trace['trace']['numSpans']} spans)") # -------------------------------------------------------------------------- @@ -808,6 +1195,8 @@ def build_parser(): %(prog)s "Which devices stopped sending logs?" --timeout 1800 echo "Summarize critical events today" | %(prog)s %(prog)s --interactive + %(prog)s --file q.txt --save-dir ./results --tag llama-70b-l40s \\ + --deploy-config ./ai.yaml --export-trace Configuration (first match wins): --host / --token flags @@ -861,11 +1250,51 @@ def build_parser(): "emits for the web UI. By default they are collapsed " "to a one-line marker in terminal output (JSON and " "saved files always keep the full text).") + out.add_argument("--tag", metavar="TAG", + help="Write the saved files under //, a " + "label for the model or hardware under test " + "(letters, digits, . _ -). Needs --save-dir.") out.add_argument("-q", "--quiet", action="store_true", help="Print only the answer; suppress progress output.") out.add_argument("-v", "--verbose", action="store_true", help="Show connection details and extra progress.") + cap = parser.add_argument_group( + "capture", + "The steps LogZilla's certification harness wraps around this tool; " + "with them a run in your own lab is captured the same way.", + ) + cap.add_argument("--deploy-config", metavar="PATH", + help="Before asking, copy this ai.yaml to " + f"{REMOTE_AI_YAML} on the server as root over SSH, " + "reload the AI settings and wait --reload-wait " + "seconds. With --save-dir, a copy with its secrets " + "redacted is saved as ai-config.yaml.") + cap.add_argument("--reload-wait", type=int, default=RELOAD_WAIT_DEFAULT, + metavar="SECONDS", + help="Seconds to wait after the settings reload " + f"(default: {RELOAD_WAIT_DEFAULT}).") + cap.add_argument("--export-trace", action="store_true", + help="After each answer, find the question's trace in " + "Phoenix and save it beside the result as " + "--trace.json. Needs --save-dir " + "and the Phoenix server and project (flags, " + "LZ_PHOENIX_URL / LZ_PHOENIX_PROJECT, or the " + "--deploy-config file's TRACING_* settings). The " + "lookup waits --trace-wait seconds for Phoenix to " + "ingest the trace; a missing, ambiguous or failed " + "trace is an error for that question.") + cap.add_argument("--phoenix-url", metavar="URL", + help="The Phoenix server (the TRACING_SERVER_URL of the " + "LogZilla AI settings, scheme and host).") + cap.add_argument("--phoenix-project", metavar="NAME", + help="The Phoenix project (TRACING_PROJECT_NAME).") + cap.add_argument("--trace-wait", type=int, default=TRACE_WAIT_DEFAULT, + metavar="SECONDS", + help="Seconds to wait for a question's trace to land in " + f"Phoenix before reporting it missing (default: " + f"{TRACE_WAIT_DEFAULT}).") + parser.add_argument( "--same-conversation", action="store_true", help="Ask all questions in one conversation so later questions can " @@ -887,6 +1316,27 @@ def main(): "--file") if args.timeout <= 0: parser.error("--timeout must be a positive number of seconds") + if args.tag and not TAG_RE.match(args.tag): + parser.error("--tag must be letters, digits, '.', '_' or '-' " + "(up to 64 characters, starting with a letter or digit)") + if args.tag and not args.save_dir: + parser.error("--tag needs --save-dir") + if args.export_trace and not args.save_dir: + parser.error("--export-trace needs --save-dir") + if args.reload_wait < 0: + parser.error("--reload-wait must be zero or more seconds") + if args.trace_wait < 0: + parser.error("--trace-wait must be zero or more seconds") + config_path = config_text = redacted_config = None + if args.deploy_config: + config_path, config_text = read_config_file(args.deploy_config) + if args.save_dir: + # judged before anything is deployed: a config whose secrets the + # redaction cannot cover is refused, not saved with them in it + try: + redacted_config = redact_config_text(config_text) + except ValueError as exc: + die(f"{config_path}: {exc}") mode = "quiet" if args.quiet else ("verbose" if args.verbose else "normal") reporter = Reporter(mode) @@ -902,6 +1352,38 @@ def main(): host, token = resolve_config(args) reporter(f"LogZilla: {host}", level="detail") + if args.export_trace: + phoenix_url = args.phoenix_url or _from_env("LZ_PHOENIX_URL") + phoenix_project = args.phoenix_project or _from_env("LZ_PHOENIX_PROJECT") + if config_text and not (phoenix_url and phoenix_project): + url_from_config, project_from_config = tracing_from_config(config_text) + phoenix_url = phoenix_url or url_from_config + phoenix_project = phoenix_project or project_from_config + if not (phoenix_url and phoenix_project): + die("--export-trace needs the Phoenix server and project: pass " + "--phoenix-url and --phoenix-project, set LZ_PHOENIX_URL and " + "LZ_PHOENIX_PROJECT, or give --deploy-config an ai.yaml that " + "sets TRACING_SERVER_URL and TRACING_PROJECT_NAME") + args.phoenix_base = phoenix_base(phoenix_url) + args.phoenix_project = phoenix_project + # the server must answer before any question is asked (a wrong URL + # is a flag error, not a trace error); the project may not exist + # yet, Phoenix creates it with the first trace that reaches it + try: + phoenix_project_id(args.phoenix_base, phoenix_project, + verify_tls=not args.insecure) + except TraceError as exc: + die(str(exc)) + reporter(f"Phoenix: {args.phoenix_base}, project {phoenix_project}", + level="detail") + + if args.deploy_config: + deploy_config(host, config_path, args.reload_wait, reporter) + if args.save_dir: + copy = output_dir(args.save_dir, args.tag) / "ai-config.yaml" + copy.write_text(redacted_config, encoding="utf-8") + reporter(f"saved {copy.name} (secrets redacted)") + client = LogZillaAI( host, token, verify_tls=not args.insecure, diff --git a/tests/test_capture.py b/tests/test_capture.py new file mode 100644 index 0000000..bcab791 --- /dev/null +++ b/tests/test_capture.py @@ -0,0 +1,372 @@ +"""The capture flags: the redacted config copy, the tracing settings read from +the config, the tag, the trace match. Run: python3 -m unittest discover -s tests -v""" + +import asyncio +import importlib.util +import json +import sys +import tempfile +import threading +import types +import unittest +from pathlib import Path + +TOOL = Path(__file__).resolve().parents[1] / "lz-ai-query.py" +spec = importlib.util.spec_from_file_location("lz_ai_query", TOOL) +tool = importlib.util.module_from_spec(spec) +sys.modules[spec.name] = tool +spec.loader.exec_module(tool) + +CONFIG = """MODEL_TYPE: anthropic +MODEL_NAME: claude-opus-5-5 +API_KEY: sk-live-secret +TRACING_ENABLED: true +TRACING_SERVER_URL: http://phoenix.example.com/v1/traces +TRACING_PROJECT_NAME: lab-1 +TRACING_API_KEY: '' +EXTRA_MCP_URLS: [] +SOME_TOKEN: abc +EMPTY_KEY: +""" + + +class RedactionTests(unittest.TestCase): + def test_secret_values_are_replaced_and_the_rest_kept(self): + out = tool.redact_config_text(CONFIG) + self.assertNotIn("sk-live-secret", out) + self.assertIn("API_KEY: \n", out) + self.assertIn("TRACING_API_KEY: \n", out) + self.assertIn("SOME_TOKEN: \n", out) + self.assertIn("MODEL_NAME: claude-opus-5-5\n", out) + self.assertIn("TRACING_SERVER_URL: http://phoenix.example.com/v1/traces\n", out) + self.assertIn("EMPTY_KEY:\n", out) + + def test_quoted_keys_and_block_scalars_are_redacted_whole(self): + config = ( + "API_KEY: |-\n" + " sk-block-secret\n" + "\n" + " second-line-of-secret\n" + "MODEL_NAME: claude-opus-5-5\n" + "\"SOME_TOKEN\": quoted-key-secret\n" + "'OTHER_KEY': 'single-quoted-secret'\n" + "FOLDED_TOKEN: >\n" + " folded secret\n" + "TRACING_PROJECT_NAME: lab-1\n" + ) + out = tool.redact_config_text(config) + for secret in ("sk-block-secret", "second-line-of-secret", "quoted-key-secret", + "single-quoted-secret", "folded secret"): + self.assertNotIn(secret, out) + self.assertIn("API_KEY: \n", out) + self.assertIn("\"SOME_TOKEN\": \n", out) + self.assertIn("'OTHER_KEY': \n", out) + self.assertIn("FOLDED_TOKEN: \n", out) + self.assertIn("MODEL_NAME: claude-opus-5-5\n", out) + self.assertIn("TRACING_PROJECT_NAME: lab-1\n", out) + + def test_a_value_on_the_lines_below_the_key_is_redacted_whole(self): + config = ( + "API_KEY:\n" + " sk-next-line-secret\n" + "SOME_TOKEN: \"first line\n" + " second-line-of-quoted-secret\"\n" + "OTHER_KEY: plain first\n" + " plain-continuation-secret\n" + "EMPTY_KEY:\n" + "MODEL_NAME: claude-opus-5-5\n" + ) + out = tool.redact_config_text(config) + for secret in ("sk-next-line-secret", "second-line-of-quoted-secret", + "plain-continuation-secret", "first line", "plain first"): + self.assertNotIn(secret, out) + self.assertIn("API_KEY: \n", out) + self.assertIn("SOME_TOKEN: \n", out) + self.assertIn("OTHER_KEY: \n", out) + self.assertIn("EMPTY_KEY:\n", out) + self.assertIn("MODEL_NAME: claude-opus-5-5\n", out) + + def test_a_secret_in_a_form_the_redaction_does_not_cover_is_refused(self): + """A flow-style mapping or an inline value never matches the line-anchored + key, so the copy is refused rather than saved with the credential in it; + a comment that names a key is not a value.""" + for text in ( + '{"MODEL_TYPE": "anthropic", "API_KEY": "sk-live-secret"}\n', + "settings: {MODEL_TYPE: anthropic, API_KEY: sk-live-secret}\n", + ): + with self.assertRaises(ValueError): + tool.redact_config_text(text) + out = tool.redact_config_text("# API_KEY: from the vault\nAPI_KEY: sk-live\n") + self.assertEqual(out, "# API_KEY: from the vault\nAPI_KEY: \n") + + +class TracingSettingsTests(unittest.TestCase): + def test_the_phoenix_base_and_project_come_from_the_config(self): + self.assertEqual(tool.tracing_from_config(CONFIG), + ("http://phoenix.example.com", "lab-1")) + + def test_missing_settings_are_none(self): + self.assertEqual(tool.tracing_from_config("MODEL_TYPE: x\n"), (None, None)) + + def test_a_trailing_comment_is_not_part_of_the_value(self): + config = ( + "TRACING_SERVER_URL: 'http://phoenix.example.com/v1/traces' # the lab\n" + "TRACING_PROJECT_NAME: lab-1 # certification run\n" + ) + self.assertEqual(tool.tracing_from_config(config), + ("http://phoenix.example.com", "lab-1")) + quoted = "TRACING_PROJECT_NAME: 'lab #1' # note\n" + self.assertEqual(tool.tracing_from_config(quoted)[1], "lab #1") + + +class DeployTests(unittest.TestCase): + def test_an_ipv6_host_keeps_its_brackets_for_scp(self): + scp, ssh = tool.deploy_steps("https://[2001:db8::1]", "ai.yaml") + self.assertEqual(scp[-1], "root@[2001:db8::1]:" + tool.REMOTE_AI_YAML) + self.assertEqual(ssh[1], "root@2001:db8::1") + scp, ssh = tool.deploy_steps("http://lz.example.com", "ai.yaml") + self.assertEqual(scp[-1], "root@lz.example.com:" + tool.REMOTE_AI_YAML) + + +class TagTests(unittest.TestCase): + def test_tag_pattern(self): + for ok in ("llama-70b", "rtx5090.q4_k_m", "a"): + self.assertTrue(tool.TAG_RE.match(ok), ok) + for bad in ("", "-x", "a b", "a/b", "x" * 65): + self.assertFalse(tool.TAG_RE.match(bad), bad) + + def test_results_land_under_the_tag(self): + result = tool.QueryResult("Top hosts?") + result.start_time = 1700000000.0 + with tempfile.TemporaryDirectory() as tmp: + base = tool.result_base(result, tmp, "tag-1") + self.assertEqual(base.parent, Path(tmp) / "tag-1") + self.assertTrue(base.parent.is_dir()) + self.assertTrue(base.name.startswith("20231114T221320Z-top-hosts")) + self.assertEqual(tool.trace_path_for(base).name, base.name + "-trace.json") + + +def fake_phoenix(roots, traces, project_absent_calls=0, roots_empty_calls=0): + """Phoenix as a function of the call count: the project is absent for the + first `project_absent_calls` lookups and the roots listing empty for the + first `roots_empty_calls`, the way a trace that has not been ingested looks.""" + calls = {"project": 0, "roots": 0} + + def graphql(base, query, variables, verify_tls=True): + if "getProjectByName" in query: + calls["project"] += 1 + if calls["project"] <= project_absent_calls: + return {"getProjectByName": None} + return {"getProjectByName": {"id": "UHJvamVjdDox", "name": "lab-1"}} + if "Roots" in query: + calls["roots"] += 1 + listed = [] if calls["roots"] <= roots_empty_calls else roots + return {"node": {"spans": {"edges": [{"node": r} for r in listed]}}} + return {"getTraceByOtelId": traces[variables["traceId"]]} + graphql.calls = calls + return graphql + + +class FakeClock: + """A clock that only moves when the code under test sleeps.""" + + def __init__(self): + self.now = 1000.0 + self.sleeps = [] + + def sleep(self, seconds): + self.sleeps.append(seconds) + self.now += seconds + + def clock(self): + return self.now + + +def root(trace_id, start): + return {"context": {"traceId": trace_id}, "startTime": start, + "parentId": None, "name": "FunctionAgent.run"} + + +def trace_doc(trace_id, texts): + spans = [] + for i, text in enumerate(texts): + spans.append({ + "id": f"{trace_id}-{i}", "name": "FunctionAgent.setup_agent", + "parentId": None, "context": {"traceId": trace_id, "spanId": f"s{i}"}, + "input": {"value": "{}"}, + "output": {"value": json.dumps({"messages": [{"content": text}]})}, + }) + return {"traceId": trace_id, "startTime": "2023-11-14T22:13:21+00:00", + "endTime": "2023-11-14T22:13:30+00:00", "latencyMs": 9000, + "numSpans": len(spans), "spans": {"edges": [{"node": s} for s in spans]}} + + +class TraceMatchTests(unittest.TestCase): + def setUp(self): + self.result = tool.QueryResult("Top 10 hosts\nwith errors?") + self.result.start_time = 1700000000.0 # 2023-11-14T22:13:20Z + self.result.end_time = 1700000020.0 + self.original = tool.phoenix_graphql + + def tearDown(self): + tool.phoenix_graphql = self.original + + def test_the_trace_carrying_the_question_is_exported(self): + roots = [root("t1", "2023-11-14T22:13:21+00:00"), + root("t2", "2023-11-14T22:13:25.500000+00:00")] + traces = {"t1": trace_doc("t1", ["something else"]), + "t2": trace_doc("t2", ["Top 10 hosts\nwith errors?"])} + tool.phoenix_graphql = fake_phoenix(roots, traces) + status, ids, trace = tool.find_trace("http://p", "pid", self.result) + self.assertEqual((status, ids), ("exported", ["t2"])) + self.assertEqual(trace["trace"]["traceId"], "t2") + self.assertEqual(len(trace["spans"]), 1) + + def test_a_root_outside_the_window_is_not_a_candidate(self): + roots = [root("t1", "2023-11-14T22:14:30+00:00")] + traces = {"t1": trace_doc("t1", ["Top 10 hosts with errors?"])} + tool.phoenix_graphql = fake_phoenix(roots, traces) + self.assertEqual(tool.find_trace("http://p", "pid", self.result, wait_s=0), + ("missing", [], None)) + + def test_two_traces_with_the_question_are_ambiguous(self): + roots = [root("t1", "2023-11-14T22:13:21+00:00"), + root("t2", "2023-11-14T22:13:22+00:00")] + traces = {t: trace_doc(t, ["Top 10 hosts\nwith errors?"]) for t in ("t1", "t2")} + tool.phoenix_graphql = fake_phoenix(roots, traces) + self.assertEqual(tool.find_trace("http://p", "pid", self.result), + ("ambiguous", ["t1", "t2"], None)) + + def test_more_candidates_than_the_cap_is_ambiguous_without_an_export(self): + roots = [root(f"t{i}", "2023-11-14T22:13:21+00:00") for i in range(6)] + tool.phoenix_graphql = fake_phoenix(roots, {}) + status, ids, trace = tool.find_trace("http://p", "pid", self.result) + self.assertEqual((status, len(ids), trace), ("ambiguous", 6, None)) + + def test_a_longer_question_containing_this_one_is_not_a_match(self): + self.result = tool.QueryResult("Show errors") + self.result.start_time = 1700000000.0 + self.result.end_time = 1700000020.0 + roots = [root("t1", "2023-11-14T22:13:21+00:00")] + traces = {"t1": trace_doc("t1", ["Show errors from yesterday"])} + tool.phoenix_graphql = fake_phoenix(roots, traces) + self.assertEqual(tool.find_trace("http://p", "lab-1", self.result, wait_s=0), + ("missing", [], None)) + + def test_the_lookup_waits_for_ingestion_then_exports(self): + roots = [root("t1", "2023-11-14T22:13:21+00:00")] + traces = {"t1": trace_doc("t1", ["Top 10 hosts\nwith errors?"])} + tool.phoenix_graphql = fake_phoenix(roots, traces, roots_empty_calls=2) + fake = FakeClock() + status, ids, trace = tool.find_trace( + "http://p", "lab-1", self.result, wait_s=60, sleep=fake.sleep, clock=fake.clock) + self.assertEqual((status, ids), ("exported", ["t1"])) + self.assertEqual(fake.sleeps, [tool.TRACE_POLL_S, tool.TRACE_POLL_S]) + + def test_the_wait_is_bounded_and_zero_means_one_lookup(self): + tool.phoenix_graphql = fake_phoenix([], {}, roots_empty_calls=99) + fake = FakeClock() + status, ids, trace = tool.find_trace( + "http://p", "lab-1", self.result, wait_s=12, sleep=fake.sleep, clock=fake.clock) + self.assertEqual((status, ids, trace), ("missing", [], None)) + self.assertEqual(sum(fake.sleeps), 12) + fake = FakeClock() + tool.phoenix_graphql = fake_phoenix([], {}, roots_empty_calls=99) + tool.find_trace("http://p", "lab-1", self.result, wait_s=0, + sleep=fake.sleep, clock=fake.clock) + self.assertEqual(fake.sleeps, []) + + def test_a_project_that_appears_with_its_first_trace_is_resolved_then(self): + roots = [root("t1", "2023-11-14T22:13:21+00:00")] + traces = {"t1": trace_doc("t1", ["Top 10 hosts\nwith errors?"])} + graphql = fake_phoenix(roots, traces, project_absent_calls=1) + tool.phoenix_graphql = graphql + fake = FakeClock() + status, ids, trace = tool.find_trace( + "http://p", "lab-1", self.result, wait_s=60, sleep=fake.sleep, clock=fake.clock) + self.assertEqual((status, ids), ("exported", ["t1"])) + self.assertEqual(graphql.calls["project"], 2) + + def test_a_phoenix_failure_is_a_trace_error_not_an_exit(self): + def failing(base, query, variables, verify_tls=True): + raise tool.TraceError("cannot reach Phoenix at http://p/graphql: refused") + tool.phoenix_graphql = failing + with self.assertRaises(tool.TraceError): + tool.find_trace("http://p", "lab-1", self.result, wait_s=0) + + def test_an_unreachable_phoenix_raises_rather_than_exits(self): + original_post = tool.requests.post + + def refused(*args, **kwargs): + raise tool.requests.exceptions.ConnectionError("refused") + tool.requests.post = refused + try: + with self.assertRaises(tool.TraceError): + tool.phoenix_graphql("http://p", "{ x }", {}) + finally: + tool.requests.post = original_post + + +class EmitTests(unittest.TestCase): + """A Phoenix failure after a captured answer is recorded on the question and + the answer is still saved; the question fails (exit 2), the run goes on.""" + + def setUp(self): + self.original_find = tool.find_trace + self.original_print = tool.print_answer + tool.print_answer = lambda *a, **k: None + + def tearDown(self): + tool.find_trace = self.original_find + tool.print_answer = self.original_print + + def test_a_trace_error_is_recorded_and_the_answer_persisted(self): + def failing(*args, **kwargs): + raise tool.TraceError("Phoenix http://p/graphql answered HTTP 502: bad gateway") + tool.find_trace = failing + result = tool.QueryResult("Top hosts?") + result.start_time = 1700000000.0 + result.end_time = 1700000020.0 + result.answer = "host-1 leads" + with tempfile.TemporaryDirectory() as tmp: + args = types.SimpleNamespace( + export_trace=True, save_dir=tmp, tag=None, json=False, raw=False, + keep_charts=False, insecure=False, phoenix_base="http://p", + phoenix_project="lab-1", trace_wait=0) + asyncio.run(tool._emit(result, args, lambda *a, **k: None, multiple=True)) + saved = list(Path(tmp).glob("*.json")) + self.assertEqual(len(saved), 1) + doc = json.loads(saved[0].read_text(encoding="utf-8")) + self.assertEqual(result.trace["status"], "error") + self.assertIn("HTTP 502", result.trace["error"]) + self.assertFalse(result.ok) + self.assertEqual(doc["trace"]["status"], "error") + self.assertEqual(doc["answer"], "host-1 leads") + + def test_the_lookup_runs_off_the_event_loop_thread(self): + """With --same-conversation the socket stays open across questions; a + lookup that blocks the loop (HTTP calls, the ingestion wait) would let + the server's heartbeat close it, so the lookup runs in a worker thread.""" + seen = {} + + def recording(*args, **kwargs): + seen["thread"] = threading.current_thread() + return "missing", [], None + tool.find_trace = recording + result = tool.QueryResult("Top hosts?") + result.start_time = 1700000000.0 + result.end_time = 1700000020.0 + result.answer = "host-1 leads" + self.assertTrue(asyncio.iscoroutinefunction(tool._emit)) + with tempfile.TemporaryDirectory() as tmp: + args = types.SimpleNamespace( + export_trace=True, save_dir=tmp, tag=None, json=False, raw=False, + keep_charts=False, insecure=False, phoenix_base="http://p", + phoenix_project="lab-1", trace_wait=0) + asyncio.run(tool._emit(result, args, lambda *a, **k: None, multiple=True)) + self.assertIsNot(seen["thread"], threading.main_thread()) + self.assertEqual(result.trace["status"], "missing") + + +if __name__ == "__main__": + unittest.main()