projects/repobot/v4/tracing.py

"""一个极简的追踪器：把一次任务里的每次模型调用、每次工具调用都记成一行 JSON。

同一个任务的所有记录共享一个 trace_id；每条记录有自己的 span_id，以及它的上级 parent_id。
这样事后能还原出"这个问题先调了模型，模型调了哪些工具，每一步花了多久、多少钱"。
"""
import json
import threading
import time
import uuid
from contextlib import contextmanager
from pathlib import Path

_lock = threading.Lock()
_local = threading.local()  # 记录当前线程正在进行的 span，用来自动找到上级


class Tracer:
    def __init__(self, path):
        self.path = Path(path)

    def _write(self, record):
        with _lock, self.path.open("a") as f:
            f.write(json.dumps(record, ensure_ascii=False) + "\n")

    @contextmanager
    def span(self, kind, name, **attrs):
        """用法：with tracer.span("llm", "chat") as s: ...; s["tokens"] = 100"""
        parent = getattr(_local, "current", None)
        record = {
            "trace_id": parent["trace_id"] if parent else uuid.uuid4().hex[:12],
            "span_id": uuid.uuid4().hex[:8],
            "parent_id": parent["span_id"] if parent else None,
            "kind": kind,
            "name": name,
            "start": time.strftime("%Y-%m-%d %H:%M:%S"),
            **attrs,
        }
        _local.current = record
        start = time.time()
        try:
            yield record
            record.setdefault("status", "ok")
        except Exception as e:
            record["status"] = "error"
            record["error"] = f"{type(e).__name__}: {e}"
            raise
        finally:
            record["ms"] = round((time.time() - start) * 1000)
            _local.current = parent
            self._write(record)


def load(path):
    return [json.loads(line) for line in Path(path).read_text().splitlines() if line.strip()]