M19: Tracing & Logging

from dataclasses import dataclass, field, asdict
from typing import Any, Optional
import time, json, traceback

# WHAT: Unified trace event covering all four agent event categories.
# WHY: One dataclass means one JSON schema, one destination, one filter.
# GOTCHA: category must be one of the four defined strings; TraceRecorder
#         enforces this at emit time.

@dataclass
class TraceEvent:
    # ── Common fields (all events) ─────────────────────────────────────
    category: str          # "tool_call" | "llm_turn" | "loop_iter" | "error"
    run_id:   str          # UUID linking all events in one agent run
    ts:       float = field(default_factory=time.time)   # Unix epoch seconds

    # ── Tool call fields ──────────────────────────────────────────────
    tool_name:   Optional[str] = None
    tool_args:   Optional[str] = None   # json.dumps(args, default=str)
    tool_output: Optional[str] = None   # truncated to 2 048 chars
    tool_ok:     Optional[bool] = None
    tool_error:  Optional[str] = None

    # ── LLM turn fields ───────────────────────────────────────────────
    model:             Optional[str] = None
    prompt_tokens:     Optional[int] = None
    completion_tokens: Optional[int] = None
    finish_reason:     Optional[str] = None  # stop | tool_calls | length | error
    turn_index:        Optional[int] = None

    # ── Agent loop iteration fields ───────────────────────────────────
    iteration:     Optional[int] = None
    tools_invoked: Optional[list] = None
    exit_reason:   Optional[str] = None  # max_iterations | goal_achieved | ...

    # ── Error fields ──────────────────────────────────────────────────
    exc_type:    Optional[str] = None
    exc_msg:     Optional[str] = None
    stack_tail:  Optional[str] = None   # last 3 frames only
    retried:     Optional[bool] = None
    retry_count: Optional[int] = None

    # ── Latency (shared across tool calls, LLM turns, loop iters) ────
    latency_ms: Optional[float] = None

    def to_json(self) -> str:
        return json.dumps(
            {k: v for k, v in asdict(self).items() if v is not None},
            default=str
        )


class TraceRecorder:
    """Thin wrapper: build events, write to JSON Lines file."""

    VALID_CATEGORIES = {"tool_call", "llm_turn", "loop_iter", "error"}

    def __init__(self, run_id: str, filepath: str = "agent_trace.jsonl"):
        self.run_id = run_id
        self.filepath = filepath

    def emit(self, event: TraceEvent) -> None:
        if event.category not in self.VALID_CATEGORIES:
            raise ValueError(f"Unknown category: {event.category}")
        with open(self.filepath, "a") as f:
            f.write(event.to_json() + "\n")

    # ── Convenience builders ──────────────────────────────────────────

    def tool_call(self, name: str, args: dict, output: Any,
                  latency_ms: float, ok: bool = True,
                  error: Optional[str] = None) -> None:
        self.emit(TraceEvent(
            category="tool_call", run_id=self.run_id,
            tool_name=name,
            tool_args=json.dumps(args, default=str),
            tool_output=str(output)[:2048],
            tool_ok=ok, tool_error=error,
            latency_ms=latency_ms
        ))

    def llm_turn(self, model: str, prompt_tokens: int,
                 completion_tokens: int, finish_reason: str,
                 latency_ms: float, turn_index: int) -> None:
        self.emit(TraceEvent(
            category="llm_turn", run_id=self.run_id,
            model=model, prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            finish_reason=finish_reason,
            latency_ms=latency_ms, turn_index=turn_index
        ))

    def loop_iter(self, iteration: int, tools: list,
                  exit_reason: str, latency_ms: float) -> None:
        self.emit(TraceEvent(
            category="loop_iter", run_id=self.run_id,
            iteration=iteration, tools_invoked=tools,
            exit_reason=exit_reason, latency_ms=latency_ms
        ))

    def error(self, exc: Exception, retried: bool = False,
              retry_count: int = 0) -> None:
        tb = "".join(traceback.format_tb(exc.__traceback__)[-3:])
        self.emit(TraceEvent(
            category="error", run_id=self.run_id,
            exc_type=type(exc).__name__,
            exc_msg=str(exc),
            stack_tail=tb,
            retried=retried, retry_count=retry_count
        ))

import structlog, logging, time, uuid
from openai import OpenAI  # Ollama's OpenAI-compatible endpoint

# ── Step 1: Configure structlog once at startup ──────────────────────
# WHAT: Wire structlog to emit JSON with timestamps and level fields
# WHY: JSON output is queryable; plain text requires regex parsing
# GOTCHA: processors run left-to-right; JSONRenderer must be last

def configure_agent_logging(log_file: str = "agent.log") -> None:
    logging.basicConfig(
        format="%(message)s",
        level=logging.INFO,
        handlers=[
            logging.StreamHandler(),
            logging.FileHandler(log_file, mode="a")
        ]
    )
    structlog.configure(
        processors=[
            structlog.stdlib.add_log_level,
            structlog.stdlib.add_logger_name,
            structlog.processors.TimeStamper(fmt="iso"),
            structlog.processors.StackInfoRenderer(),
            structlog.processors.JSONRenderer()
        ],
        wrapper_class=structlog.stdlib.BoundLogger,
        context_class=dict,
        logger_factory=structlog.stdlib.LoggerFactory()
    )


# ── Step 2: Agent wrapper that emits structured events ────────────────
# WHAT: Wraps every Ollama LLM call and tool call with structured logging
# WHY: One consistent event schema across all events in every run
# GOTCHA: Always log BEFORE raising exceptions so the error event is
#         written even when the caller swallows the exception

class AgentLogger:
    def __init__(self, run_id: str | None = None):
        self.run_id = run_id or str(uuid.uuid4())
        self.log = structlog.get_logger().bind(run_id=self.run_id)
        self.client = OpenAI(
            base_url="http://localhost:11434/v1",
            api_key="ollama"
        )

    def llm_call(self, messages: list, model: str = "mistral") -> str:
        t0 = time.perf_counter()
        try:
            response = self.client.chat.completions.create(
                model=model, messages=messages
            )
            latency_ms = (time.perf_counter() - t0) * 1000
            usage = response.usage
            self.log.info("llm_turn",
                event_type="llm_turn", model=model,
                prompt_tokens=usage.prompt_tokens if usage else None,
                completion_tokens=usage.completion_tokens if usage else None,
                finish_reason=response.choices[0].finish_reason,
                latency_ms=round(latency_ms, 2)
            )
            return response.choices[0].message.content or ""
        except Exception as exc:
            latency_ms = (time.perf_counter() - t0) * 1000
            self.log.error("llm_error",
                event_type="error",
                exc_type=type(exc).__name__,
                exc_msg=str(exc),
                latency_ms=round(latency_ms, 2)
            )
            raise

    def tool_call(self, name: str, args: dict,
                  fn, *fn_args, **fn_kwargs):
        t0 = time.perf_counter()
        try:
            result = fn(*fn_args, **fn_kwargs)
            latency_ms = (time.perf_counter() - t0) * 1000
            self.log.info("tool_call",
                event_type="tool_call", tool_name=name,
                tool_args=str(args)[:512], tool_output=str(result)[:512],
                tool_ok=True, latency_ms=round(latency_ms, 2)
            )
            return result
        except Exception as exc:
            latency_ms = (time.perf_counter() - t0) * 1000
            self.log.warning("tool_error",
                event_type="tool_call", tool_name=name,
                tool_ok=False, tool_error=str(exc),
                latency_ms=round(latency_ms, 2)
            )
            raise

pip install langsmith

export LANGCHAIN_API_KEY="lsv2_pt_..."        # your LangSmith API key
export LANGCHAIN_TRACING_V2="true"            # enable trace shipping
export LANGCHAIN_PROJECT="ollama-c3-agent"    # project name in dashboard

from langsmith import traceable
from openai import OpenAI
import json

# WHAT: OpenAI client pointed at Ollama — @traceable intercepts it
# WHY: LangSmith patches the openai library, so any client call made
#      inside a @traceable function is captured as a child span
# GOTCHA: api_key must be a non-empty string (Ollama ignores the value,
#         but the OpenAI client validates that the field is present)

client = OpenAI(
    base_url="http://localhost:11434/v1",
    api_key="ollama"
)

TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "fuzzy_match_score",
            "description": "Compute fuzzy similarity between two entity names",
            "parameters": {
                "type": "object",
                "properties": {
                    "a": {"type": "string"},
                    "b": {"type": "string"}
                },
                "required": ["a", "b"]
            }
        }
    }
]

# ── Tool handlers as child spans ──────────────────────────────────────
@traceable(name="fuzzy-match-score")
def fuzzy_match_score(a: str, b: str) -> float:
    from rapidfuzz import fuzz
    return fuzz.token_sort_ratio(a.lower(), b.lower()) / 100.0

@traceable(name="search-filings")
def search_filings(query: str, limit: int = 3) -> list[dict]:
    # Replace with your actual vector store / SQLite query
    return [{"filing_id": "UCC-001", "debtor": query, "score": 0.91}]

# ── Root span ─────────────────────────────────────────────────────────
@traceable(name="entity-resolution-agent")
def run_agent(query: str) -> str:
    """Full ReAct agent loop, traced end-to-end via LangSmith."""
    messages = [
        {"role": "system", "content": "You are an entity resolution specialist."},
        {"role": "user",   "content": query}
    ]
    tool_map = {
        "fuzzy_match_score": fuzzy_match_score,
        "search_filings": search_filings
    }
    for _ in range(10):
        response = client.chat.completions.create(
            model="mistral", messages=messages,
            tools=TOOLS, tool_choice="auto"
        )
        choice = response.choices[0]
        messages.append(choice.message)
        if choice.finish_reason == "stop":
            return choice.message.content or "No answer"
        if choice.finish_reason == "tool_calls":
            for tc in (choice.message.tool_calls or []):
                fn_name = tc.function.name
                fn_args = json.loads(tc.function.arguments)
                result = tool_map[fn_name](**fn_args)
                messages.append({
                    "role": "tool",
                    "tool_call_id": tc.id,
                    "content": json.dumps(result)
                })
    return "Max iterations reached"


if __name__ == "__main__":
    answer = run_agent(
        "Are 'Acme Logistics LLC' and 'ACME LOGISTICS, L.L.C.' the same entity?"
    )
    print(answer)
    # View trace at: https://smith.langchain.com
    # Projects -> ollama-c3-agent -> Filter: model = mistral

pip install \
  opentelemetry-api \
  opentelemetry-sdk \
  opentelemetry-exporter-otlp

# Start Jaeger all-in-one (traces stored in memory)
docker run -d --name jaeger \
  -p 4317:4317 \
  -p 16686:16686 \
  jaegertracing/all-in-one:latest

# View traces at: http://localhost:16686

from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.trace import StatusCode
import time, json
from openai import OpenAI

# ── Bootstrap OTEL — run once at app start ───────────────────────────
# WHAT: Wire TracerProvider to Jaeger via OTLP gRPC
# WHY: BatchSpanProcessor buffers + flushes every 5s, adding <1ms overhead
# GOTCHA: OTLPSpanExporter silently drops spans when Jaeger is not running

def setup_otel(service_name: str = "entity-resolution-agent") -> trace.Tracer:
    exporter = OTLPSpanExporter(endpoint="http://localhost:4317")
    provider = TracerProvider()
    provider.add_span_processor(BatchSpanProcessor(exporter))
    trace.set_tracer_provider(provider)
    return trace.get_tracer(service_name)


tracer = setup_otel()
client = OpenAI(base_url="http://localhost:11434/v1", api_key="ollama")

def run_agent_otel(query: str) -> str:
    with tracer.start_as_current_span("entity-resolution-run") as root:
        root.set_attribute("query", query[:256])
        messages = [
            {"role": "system", "content": "You are an entity resolution specialist."},
            {"role": "user",   "content": query}
        ]
        try:
            for iteration in range(10):
                root.set_attribute("iterations", iteration + 1)
                t0 = time.perf_counter()
                with tracer.start_as_current_span(f"llm-turn-{iteration+1}") as llm_span:
                    response = client.chat.completions.create(
                        model="mistral", messages=messages,
                        tools=TOOLS, tool_choice="auto"
                    )
                    latency_ms = (time.perf_counter() - t0) * 1000
                    usage = response.usage
                    llm_span.set_attribute("model", "mistral")
                    llm_span.set_attribute("latency_ms", round(latency_ms, 1))
                    if usage:
                        llm_span.set_attribute("prompt_tokens", usage.prompt_tokens)
                        llm_span.set_attribute("completion_tokens", usage.completion_tokens)
                    llm_span.set_attribute("finish_reason", response.choices[0].finish_reason)

                choice = response.choices[0]
                messages.append(choice.message)

                if choice.finish_reason == "stop":
                    root.set_attribute("exit_reason", "goal_achieved")
                    return choice.message.content or "No answer"

                if choice.finish_reason == "tool_calls":
                    for tc in (choice.message.tool_calls or []):
                        fn_name = tc.function.name
                        fn_args = json.loads(tc.function.arguments)
                        t1 = time.perf_counter()
                        with tracer.start_as_current_span(f"tool:{fn_name}") as ts:
                            ts.set_attribute("tool.name", fn_name)
                            ts.set_attribute("tool.args", str(fn_args)[:256])
                            try:
                                result = TOOL_MAP[fn_name](**fn_args)
                                ts.set_attribute("tool.ok", True)
                                ts.set_attribute("tool.output", str(result)[:256])
                            except Exception as e:
                                ts.set_attribute("tool.ok", False)
                                ts.set_attribute("tool.error", str(e))
                                ts.set_status(StatusCode.ERROR, str(e))
                                raise
                            finally:
                                ts.set_attribute("latency_ms",
                                    round((time.perf_counter() - t1) * 1000, 1))
                        messages.append({
                            "role": "tool",
                            "tool_call_id": tc.id,
                            "content": json.dumps(result)
                        })
            root.set_attribute("exit_reason", "max_iterations")
            return "Max iterations reached"
        except Exception as e:
            root.set_status(StatusCode.ERROR, str(e))
            raise

import functools, inspect, json, time, uuid
from pathlib import Path
from typing import Callable

# WHAT: @trace_agent wraps any agent function, writing run_start,
#       sub-events, and run_end to a JSONL file
# WHY: Zero deps — works anywhere Python 3.9+ runs
# GOTCHA: Non-serializable arguments are converted to str() automatically

def trace_agent(
    output_dir: str = "traces",
    log_llm_calls: bool = True,
    log_tool_calls: bool = True
) -> Callable:
    Path(output_dir).mkdir(parents=True, exist_ok=True)

    def decorator(fn: Callable) -> Callable:
        @functools.wraps(fn)
        def wrapper(*args, **kwargs):
            run_id = str(uuid.uuid4())[:8]
            trace_file = Path(output_dir) / f"trace_{run_id}.jsonl"

            def emit(event: dict) -> None:
                event["run_id"] = run_id
                event["ts"] = time.time()
                with open(trace_file, "a") as f:
                    f.write(json.dumps(event, default=str) + "\n")

            emit({"category": "run_start", "fn_name": fn.__name__,
                  "fn_args": str(args)[:512], "fn_kwargs": str(kwargs)[:512]})
            t0 = time.perf_counter()
            sig = inspect.signature(fn)
            if "_emit" in sig.parameters:
                kwargs["_emit"] = emit
            try:
                result = fn(*args, **kwargs)
                emit({"category": "run_end", "exit_reason": "success",
                      "latency_ms": round((time.perf_counter() - t0) * 1000, 2),
                      "output_preview": str(result)[:256]})
                return result
            except Exception as exc:
                emit({"category": "run_end", "exit_reason": "error",
                      "latency_ms": round((time.perf_counter() - t0) * 1000, 2),
                      "exc_type": type(exc).__name__, "exc_msg": str(exc)})
                raise
        return wrapper
    return decorator


@trace_agent(output_dir="traces")
def run_agent(query: str, _emit=None) -> str:
    from openai import OpenAI
    client = OpenAI(base_url="http://localhost:11434/v1", api_key="ollama")
    t_llm = time.perf_counter()
    response = client.chat.completions.create(
        model="mistral",
        messages=[{"role": "user", "content": query}]
    )
    if _emit:
        _emit({
            "category": "llm_turn", "model": "mistral",
            "latency_ms": round((time.perf_counter() - t_llm) * 1000, 2),
            "finish_reason": response.choices[0].finish_reason
        })
    return response.choices[0].message.content or ""


def replay_trace(trace_file: str) -> dict:
    """Load JSONL trace file into a summary dict."""
    events = []
    with open(trace_file) as f:
        for line in f:
            line = line.strip()
            if line:
                events.append(json.loads(line))
    start = next((e for e in events if e["category"] == "run_start"), None)
    end   = next((e for e in events if e["category"] == "run_end"),   None)
    return {
        "run_id":      events[0].get("run_id") if events else None,
        "fn_name":     start.get("fn_name") if start else None,
        "total_ms":    end.get("latency_ms") if end else None,
        "exit_reason": end.get("exit_reason") if end else "incomplete",
        "events":      events
    }

#!/usr/bin/env python3
"""trace_viewer.py — CLI pretty-printer for agent JSONL trace files.

Usage:
    python trace_viewer.py traces/trace_a3f7b2.jsonl
    python trace_viewer.py traces/trace_a3f7b2.jsonl --slow-threshold 1000
"""
import json, sys, argparse
from pathlib import Path

RED = "\033[91m"; YELLOW = "\033[93m"; GREEN = "\033[92m"
CYAN = "\033[96m"; DIM = "\033[2m"; RESET = "\033[0m"; BOLD = "\033[1m"

CATEGORY_COLOR = {
    "run_start": CYAN, "llm_turn": YELLOW,
    "tool_call": GREEN, "error": RED, "run_end": CYAN,
}

def load_events(filepath: str) -> list[dict]:
    events = []
    with open(filepath) as f:
        for line in f:
            line = line.strip()
            if line:
                try:
                    events.append(json.loads(line))
                except json.JSONDecodeError:
                    pass
    return events

def print_call_tree(events: list[dict], slow_ms: float = 2000) -> None:
    print(f"\n{BOLD}{'─'*60}{RESET}")
    print(f"{BOLD}TRACE: {events[0].get('run_id','unknown')}{RESET}")
    print('─'*60)
    indent = 0
    for ev in events:
        cat     = ev.get("category", "unknown")
        color   = CATEGORY_COLOR.get(cat, RESET)
        latency = ev.get("latency_ms", 0) or 0
        slow    = f" {RED}⚠ SLOW{RESET}" if latency > slow_ms else ""

        if cat == "run_start":
            print(f"\n{color}▶ {ev.get('fn_name','?')}(){RESET}")
            indent = 2
        elif cat == "llm_turn":
            print(f"{' '*indent}{color}LLM [{ev.get('model','?')}]{RESET}  "
                  f"{DIM}{ev.get('prompt_tokens','?')}→{ev.get('completion_tokens','?')} tok  "
                  f"{ev.get('finish_reason','?')}{RESET}  "
                  f"{YELLOW}{latency:.0f}ms{RESET}{slow}")
        elif cat == "tool_call":
            ok     = ev.get("tool_ok", True)
            status = f"{GREEN}✓{RESET}" if ok else f"{RED}✗ {ev.get('tool_error','')}{RESET}"
            print(f"{' '*(indent+2)}{color}TOOL {ev.get('tool_name','?')}{RESET}  "
                  f"{status}  {YELLOW}{latency:.0f}ms{RESET}{slow}")
        elif cat == "error":
            print(f"{' '*indent}{RED}ERROR {ev.get('exc_type','?')}: "
                  f"{str(ev.get('exc_msg',''))[:80]}{RESET}")
        elif cat == "run_end":
            reason = ev.get("exit_reason", "?")
            col    = GREEN if reason == "success" else RED
            print(f"\n{col}{reason.upper()}  {YELLOW}{latency:.0f}ms total{RESET}\n")
    print('─'*60 + '\n')

def load_trace(filepath: str):
    events = load_events(filepath)
    try:
        import pandas as pd
        df = pd.DataFrame(events)
        df["latency_ms"] = pd.to_numeric(df.get("latency_ms"), errors="coerce")
        return df
    except ImportError:
        print("pandas not installed — returning list of dicts")
        return events

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("tracefile")
    parser.add_argument("--slow-threshold", type=float, default=2000)
    args = parser.parse_args()

    if not Path(args.tracefile).exists():
        print(f"File not found: {args.tracefile}", file=sys.stderr)
        sys.exit(1)

    events = load_events(args.tracefile)
    if not events:
        print("Trace file is empty.", file=sys.stderr)
        sys.exit(1)

    print_call_tree(events, slow_ms=args.slow_threshold)

    df = load_trace(args.tracefile)
    if hasattr(df, "groupby"):
        print(f"{BOLD}Category breakdown:{RESET}")
        summary = df.groupby("category").agg(
            count=("category","count"),
            avg_latency_ms=("latency_ms","mean")
        ).round(1)
        print(summary.to_string())
        slow = df[df["latency_ms"] > args.slow_threshold]
        if not slow.empty:
            print(f"\n{RED}Slow steps (>{args.slow_threshold:.0f}ms):{RESET}")
            cols = [c for c in ["category","tool_name","model","latency_ms"] if c in slow.columns]
            print(slow[cols].to_string(index=False))

if __name__ == "__main__":
    main()