Tracing & Observability in Depth

{"label": "ask", "model": "gpt-6-luna", "in": 3120, "out": 140, "usd": 0.000382, "latency_s": 1.9}
{"label": "ask", "model": "gpt-6-luna", "in": 3305, "out": 610, "usd": 0.000636, "latency_s": 8.7}

# script_id: day_042_tracing_observability/atlas/tracing
# atlas/tracing.py (Day 12 + Day 42): one trace per request, nested step spans, numbers from spans.
# Install: uv add opentelemetry-sdk arize-phoenix-otel openinference-instrumentation-openai
# Env: ATLAS_TRACING=0 (no exporter), ATLAS_OTLP_ENDPOINT, ATLAS_OTLP_HEADERS ("k=v,k2=v2")
# Services: `phoenix serve` (optional; without it spans stay in memory only)
from __future__ import annotations

import os
import socket
from collections.abc import Iterator
from contextlib import contextmanager
from urllib.parse import urlparse

from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter

PHOENIX_URL = "http://localhost:6006/v1/traces"
OI_KINDS = {"chain", "retriever", "llm", "tool", "agent", "embedding", "reranker", "guardrail"}


def _reachable(url: str, timeout: float = 0.3) -> bool:
    u = urlparse(url)
    port = u.port or (443 if u.scheme == "https" else 80)
    try:
        with socket.create_connection((u.hostname or "localhost", port), timeout=timeout):
            return True
    except OSError:
        return False


def setup_tracing(project: str = "atlas", *, endpoint: str | None = None,
                  capture_content: bool = False) -> TracerProvider | None:
    """Send spans to Phoenix or any OTLP/HTTP backend. Returns None when no exporter is attached."""
    endpoint = endpoint or os.environ.get("ATLAS_OTLP_ENDPOINT", PHOENIX_URL)
    if os.environ.get("ATLAS_TRACING", "1") == "0" or not _reachable(endpoint):
        print(f"tracing: no exporter ({endpoint} is off or unreachable); in-memory spans still work")
        return None
    if not capture_content:  # read by OpenInference when it patches the SDK
        os.environ.setdefault("OPENINFERENCE_HIDE_INPUTS", "true")
        os.environ.setdefault("OPENINFERENCE_HIDE_OUTPUTS", "true")
    from phoenix.otel import register  # speaks plain OTLP, so the endpoint needn't be Phoenix

    raw = os.environ.get("ATLAS_OTLP_HEADERS", "")  # auth for hosted backends
    headers = dict(h.split("=", 1) for h in raw.split(",") if "=" in h) or None
    provider = register(project_name=project, endpoint=endpoint, headers=headers,
                        protocol="http/protobuf", batch=True, auto_instrument=True, verbose=False)
    print(f"tracing → {endpoint} (project {project!r})")
    return provider

# script_id: day_042_tracing_observability/atlas/tracing
def attach_memory_exporter() -> InMemorySpanExporter:
    """Copy every finished span into memory, beside whatever backend is configured."""
    provider = trace.get_tracer_provider()
    if not isinstance(provider, TracerProvider):  # tracing off: install a bare SDK provider
        provider = TracerProvider()
        trace.set_tracer_provider(provider)
    memory = InMemorySpanExporter()
    provider.add_span_processor(SimpleSpanProcessor(memory))  # in-process, so no network cost
    return memory


@contextmanager
def step_span(name: str, kind: str, **attrs) -> Iterator[trace.Span]:
    """One timed step. `kind` picks Phoenix's view; attrs become searchable span attributes."""
    oi_kind = kind.upper() if kind in OI_KINDS else "UNKNOWN"
    with trace.get_tracer("atlas").start_as_current_span(
            name, attributes={"atlas.kind": kind, "openinference.span.kind": oi_kind, **attrs}) as span:
        yield span  # an exception raised in here is recorded on the span, which is marked ERROR

# script_id: day_042_tracing_observability/atlas/llm
# atlas/llm.py, compact (built on Day 12): one client, prices, and complete() as a traced LLM step.
# Install: uv add "openai>=3.46,<4" opentelemetry-sdk
# Env: OPENAI_API_KEY (unset = offline fake client), ATLAS_MODEL. Services: none.
from __future__ import annotations

import os
import random
import re
import time
from types import SimpleNamespace

from openai import OpenAI

from atlas.tracing import step_span

MODEL = os.environ.get("ATLAS_MODEL", "gpt-6-luna")
# USD per 1M (input, output) tokens. As of 2026-11, verify at https://developers.openai.com/api/docs/pricing
PRICES = {"gpt-6-luna": (0.10, 0.50), "gpt-6.1-terra": (1.00, 5.00)}


def cost_usd(model: str, input_tokens: int, output_tokens: int) -> float:
    p_in, p_out = PRICES[model]  # a KeyError beats a silent $0; cache discounts ignored here
    return (input_tokens * p_in + output_tokens * p_out) / 1_000_000


class FakeResponses:
    """Offline stand-in (Day 19's fake-client idea): estimated usage, invented latency."""

    def create(self, *, model: str, input: str, instructions: str = "", **_) -> SimpleNamespace:
        n_in = (len(instructions) + len(input)) // 4  # ~4 characters per token
        if instructions.startswith("Classify"):
            text = "question" if "?" in input[:300] else "bug"
        else:  # cite the first source in the context
            text = "See [{}].".format((re.findall(r"^\[([\w-]+)\]", input, re.M) or ["faq"])[0])
        n_out = random.Random(input).randint(40, 160)  # seeded by the input, so reruns match
        time.sleep(0.03 + n_in / 40_000 + n_out / 1_500)
        return SimpleNamespace(model=model, output_text=text,
                               usage=SimpleNamespace(input_tokens=n_in, output_tokens=n_out))


client = OpenAI() if os.environ.get("OPENAI_API_KEY") else SimpleNamespace(responses=FakeResponses())

# script_id: day_042_tracing_observability/atlas/llm
def complete(input: str, *, instructions: str, step: str = "generate", model: str = MODEL) -> str:
    """One model call = one LLM span: standard gen_ai.* attributes plus Atlas's own cost."""
    with step_span(step, "llm", **{"gen_ai.operation.name": "chat",
                                  "gen_ai.provider.name": "openai",
                                  "gen_ai.request.model": model}) as span:
        r = client.responses.create(model=model, instructions=instructions, input=input)
        u = r.usage
        span.set_attributes({
            "gen_ai.response.model": r.model,
            "gen_ai.usage.input_tokens": u.input_tokens,
            "gen_ai.usage.output_tokens": u.output_tokens,
            "atlas.cost_usd": cost_usd(model, u.input_tokens, u.output_tokens),
        })
        return r.output_text

# script_id: day_042_tracing_observability/atlas/rag
# atlas/rag.py: compact mini-R2 (stand-in for Days 26-28 and your R2) with one span per step.
# Install: uv add pydantic opentelemetry-sdk. Env: ORBIT_DATA, ATLAS_TOP_K. Services: none.
from __future__ import annotations

import math
import os
import re
from collections import Counter
from pathlib import Path

from pydantic import BaseModel

from atlas.llm import complete
from atlas.tracing import step_span

DATA_DIR = Path(os.environ.get("ORBIT_DATA", "datasets/orbit"))
TOP_K = int(os.environ.get("ATLAS_TOP_K", "8"))
WORD = re.compile(r"[a-z0-9_]+(?:[.-][a-z0-9_]+)*")  # keeps orbit.toml and orb-1101 whole


class Citation(BaseModel):  # Day 15 / Day 28
    doc_id: str
    quote: str = ""


class Answer(BaseModel):
    text: str
    citations: list[Citation]


class Index:
    """Keyword stand-in for Day 26's Index: one chunk per '## ' section, tf-idf scoring."""

    def __init__(self, docs_dir: Path) -> None:
        self.chunks = [(p.stem, sec) for p in sorted(docs_dir.glob("*.md"))
                       for sec in re.split(r"\n(?=## )", p.read_text(encoding="utf-8"))]
        self.tf = [Counter(WORD.findall(text.lower())) for _, text in self.chunks]
        df = Counter(w for tf in self.tf for w in tf)
        self.idf = {w: math.log(len(self.chunks) / n) for w, n in df.items()}

    def search(self, query: str, k: int = TOP_K) -> list[tuple[str, str, float]]:
        terms = set(WORD.findall(query.lower()))
        scored = sorted(((sum(tf[t] * self.idf.get(t, 0.0) for t in terms), i)
                         for i, tf in enumerate(self.tf)), reverse=True)[:k]
        return [(self.chunks[i][0], self.chunks[i][1], score) for score, i in scored]

# script_id: day_042_tracing_observability/atlas/rag
INSTRUCTIONS = ("Answer from the context only. Cite each source you use as [doc_id]. "
                "If the context does not answer the question, say you don't know.")


def ask(question: str, *, index: Index, k: int = TOP_K) -> Answer:
    """One request = one trace: ask → search → build_context → generate."""
    with step_span("ask", "chain", **{"atlas.top_k": k}) as root:
        with step_span("search", "retriever") as span:
            hits = index.search(question, k)
            for i, (doc_id, _, score) in enumerate(hits):  # OpenInference's flattened document list
                span.set_attribute(f"retrieval.documents.{i}.document.id", doc_id)
                span.set_attribute(f"retrieval.documents.{i}.document.score", score)
        with step_span("build_context", "chain") as span:
            context = "\n\n".join(f"[{doc_id}]\n{text}" for doc_id, text, _ in hits)
            span.set_attribute("atlas.context_chars", len(context))
        text = complete(f"Context:\n{context}\n\nQuestion: {question}", instructions=INSTRUCTIONS)
        cited = list(dict.fromkeys(re.findall(r"\[([\w-]+)\]", text)))  # ordered, no repeats
        root.set_attribute("atlas.cited", cited)
        return Answer(text=text, citations=[Citation(doc_id=d) for d in cited])

# script_id: day_042_tracing_observability/atlas/workflows
# atlas/workflows.py: compact Day 41 route_issue() with spans (subset of Day 41's RouteResult).
# Install: uv add pydantic opentelemetry-sdk. Env: as atlas/llm. Services: none.
from __future__ import annotations

import asyncio
import re
import time
from dataclasses import dataclass

from atlas.llm import complete
from atlas.rag import Index
from atlas.tracing import step_span

RISK = re.compile(r"ignore (all|previous) instructions|api[_-]?key|password", re.I)


@dataclass
class RouteResult:
    issue_id: int
    route: str
    needs_human: bool
    latency_s: float


async def find_duplicates(issue: dict, others: list[dict]) -> list[int]:
    with step_span("find_duplicates", "tool"):
        words = set(issue["title"].lower().split())
        return [o["id"] for o in others if o["id"] != issue["id"]
                and len(words & set(o["title"].lower().split())) >= 6]


async def match_docs(issue: dict, index: Index) -> list[str]:
    with step_span("match_docs", "retriever"):
        return [doc_id for doc_id, _, _ in index.search(issue["title"], 3)]


async def screen_risk(issue: dict) -> bool:
    with step_span("screen_risk", "guardrail"):
        return bool(RISK.search(issue["body"]))

# script_id: day_042_tracing_observability/atlas/workflows
async def route_issue(issue: dict, *, index: Index, others: list[dict]) -> RouteResult:
    start = time.perf_counter()
    with step_span("route_issue", "chain", **{"atlas.issue_id": issue["id"]}) as root:
        kind = await asyncio.to_thread(  # the sync client runs in a thread; to_thread copies context
            complete, f"{issue['title']}\n\n{issue['body'][:2000]}", step="classify",
            instructions="Classify this GitHub issue. Reply with one word: bug, question, feature or docs.")
        dups, docs, risky = await asyncio.gather(
            find_duplicates(issue, others), match_docs(issue, index), screen_risk(issue))
        route = "needs_human" if risky else ("duplicate" if dups else kind.strip(" .\n").lower())
        root.set_attributes({"atlas.route": route, "atlas.doc_ids": docs})
        return RouteResult(issue["id"], route, risky, time.perf_counter() - start)

# script_id: day_042_tracing_observability/atlas/tracing
from collections import defaultdict
from collections.abc import Iterable
from dataclasses import dataclass
from statistics import quantiles

from opentelemetry.sdk.trace import ReadableSpan
from opentelemetry.trace import StatusCode


@dataclass
class StepStats:
    name: str
    n: int
    p50_ms: float
    p95_ms: float
    input_tokens: int
    output_tokens: int
    cost_usd: float
    error_rate: float


def _pct(values: list[float], p: int) -> float:
    return values[0] if len(values) == 1 else quantiles(values, n=100, method="inclusive")[p - 1]


def summarize_spans(spans: Iterable[ReadableSpan]) -> list[StepStats]:
    """Group Atlas spans by name: p50/p95 latency, tokens, $ and error rate. Slowest p95 first."""
    groups: dict[str, list[ReadableSpan]] = defaultdict(list)
    for s in spans:
        if "atlas.kind" in (s.attributes or {}):
            groups[s.name].append(s)
    rows = []
    for name, group in groups.items():
        ms = [(s.end_time - s.start_time) / 1e6 for s in group]  # OTel timestamps are nanoseconds
        attrs = [s.attributes for s in group]
        rows.append(StepStats(
            name, len(group), _pct(ms, 50), _pct(ms, 95),
            sum(int(a.get("gen_ai.usage.input_tokens", 0)) for a in attrs),
            sum(int(a.get("gen_ai.usage.output_tokens", 0)) for a in attrs),
            sum(float(a.get("atlas.cost_usd", 0.0)) for a in attrs),
            sum(s.status.status_code is StatusCode.ERROR for s in group) / len(group)))
    return sorted(rows, key=lambda r: r.p95_ms, reverse=True)


def print_table(rows: list[StepStats]) -> None:
    print(f"{'step':<16}{'n':>4}{'p50 ms':>9}{'p95 ms':>9}{'in tok':>9}{'out tok':>9}{'$':>10}{'err':>6}")
    for r in rows:
        print(f"{r.name:<16}{r.n:>4}{r.p50_ms:>9.0f}{r.p95_ms:>9.0f}{r.input_tokens:>9,}"
              f"{r.output_tokens:>9,}{r.cost_usd:>10.5f}{r.error_rate:>6.0%}")


if __name__ == "__main__":  # offline smoke test: three fake requests, no backend needed
    import time

    setup_tracing()
    memory = attach_memory_exporter()
    for delay_ms in (40, 60, 400):
        with step_span("ask", "chain"), step_span("search", "retriever"):
            time.sleep(delay_ms / 1000)
    print_table(summarize_spans(memory.get_finished_spans()))

# script_id: day_042_tracing_observability/replay
# Replay 20 qa_gold questions + 10 issues through traced Atlas; print the per-step dashboard.
# Install: uv add "openai>=3.46,<4" pydantic opentelemetry-sdk arize-phoenix-otel openinference-instrumentation-openai
# Env: OPENAI_API_KEY (optional: offline fake client without it), ORBIT_DATA. Services: `phoenix serve` (optional)
from __future__ import annotations

import asyncio
import json

from atlas.tracing import attach_memory_exporter, print_table, setup_tracing, summarize_spans

setup_tracing("atlas")  # first, so auto-instrumentation is in place before the client exists
memory = attach_memory_exporter()

from atlas.rag import DATA_DIR, Index, ask  # noqa: E402
from atlas.workflows import route_issue  # noqa: E402


def load(rel: str) -> list[dict]:
    return [json.loads(line) for line in (DATA_DIR / rel).read_text(encoding="utf-8").splitlines()]


QA, ISSUES = load("eval/qa_gold.jsonl")[:20], load("issues/issues.jsonl")
index = Index(DATA_DIR / "docs")


def run_ask(k: int) -> tuple[list, int]:
    """Replay QA at top_k=k. Returns its spans and how many answers cite a gold doc (Day 30-style)."""
    memory.clear()
    cites_gold = 0
    for case in QA:
        answer = ask(case["question"], index=index, k=k)
        cites_gold += any(c.doc_id in case["citations"] for c in answer.citations)
    return list(memory.get_finished_spans()), cites_gold


async def run_routes() -> list:
    memory.clear()
    for issue in ISSUES[:10]:
        await route_issue(issue, index=index, others=ISSUES)
    return list(memory.get_finished_spans())

# script_id: day_042_tracing_observability/replay
def request_row(label: str, spans: list, cites_gold: int) -> str:
    rows = {r.name: r for r in summarize_spans(spans)}
    req, gen = rows["ask"], rows["generate"]
    ctx_tokens = sum(s.attributes.get("atlas.context_chars", 0) for s in spans) / 4  # ≈ tokens
    return (f"{label:<9}{req.p95_ms:>8.0f}{gen.input_tokens / req.n:>12,.0f}"
            f"{ctx_tokens / gen.input_tokens:>9.0%}{gen.cost_usd / req.n * 1000:>11.4f}"
            f"{cites_gold:>8}/{req.n}")


if __name__ == "__main__":
    before, gold_before = run_ask(8)
    print_table(summarize_spans(before + asyncio.run(run_routes())))
    after, gold_after = run_ask(4)
    print(f"\n{'run':<9}{'p95 ms':>8}{'in tok/req':>12}{'context':>9}{'$/1k req':>11}{'cites gold':>11}")
    print(request_row("top_k=8", before, gold_before))
    print(request_row("top_k=4", after, gold_after))
    print("guard:", "OK" if gold_after >= gold_before else "FAILED, the cheaper run cites gold docs less")

tracing: no exporter (http://localhost:6006/v1/traces is off or unreachable); in-memory spans still work
step               n   p50 ms   p95 ms   in tok  out tok         $   err
ask               20      199      236        0        0   0.00000    0%
generate          20      198      235   67,917    2,154   0.00787    0%
route_issue       10      112      134        0        0   0.00000    0%
classify          10      109      130    2,013      991   0.00070    0%
find_duplicates   10        0        2        0        0   0.00000    0%
search            20        1        1        0        0   0.00000    0%
match_docs        10        1        1        0        0   0.00000    0%
screen_risk       10        0        0        0        0   0.00000    0%
build_context     20        0        0        0        0   0.00000    0%

run        p95 ms  in tok/req  context   $/1k req cites gold
top_k=8       236       3,396      98%     0.3934      13/20
top_k=4       192       2,179      98%     0.2666      13/20
guard: OK