diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4250f0f..0731e53 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -6,40 +6,66 @@ on: pull_request: branches: [main] +concurrency: + group: ci-${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + jobs: test: runs-on: ubuntu-latest + timeout-minutes: 15 strategy: matrix: - python-version: ["3.9", "3.10", "3.11", "3.12"] + python-version: ["3.9", "3.10", "3.11", "3.12", "3.13", "3.14"] steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v7 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 + uses: actions/setup-python@v7 with: python-version: ${{ matrix.python-version }} + cache: pip - name: Install dependencies run: | python -m pip install --upgrade pip - pip install -e ".[all]" + pip install -e "." pip install pytest pytest-cov - name: Run tests run: | pytest tests/ -v --cov=agent_trace --cov-report=xml - name: Upload coverage - if: matrix.python-version == '3.12' - uses: codecov/codecov-action@v4 + if: matrix.python-version == '3.14' + uses: codecov/codecov-action@v7 + with: + files: ./coverage.xml + fail_ci_if_error: false + + optional-dependencies: + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - uses: actions/checkout@v7 + - uses: actions/setup-python@v7 with: - file: ./coverage.xml + python-version: "3.13" + cache: pip + - name: Install all optional dependencies + run: | + python -m pip install --upgrade pip + pip install -e ".[all]" + pip install pytest pytest-cov + - name: Run tests with all optional dependencies + run: pytest tests/ -q lint: runs-on: ubuntu-latest + timeout-minutes: 10 steps: - - uses: actions/checkout@v4 - - uses: actions/setup-python@v5 + - uses: actions/checkout@v7 + - uses: actions/setup-python@v7 with: - python-version: "3.12" + python-version: "3.14" + cache: pip - name: Install dependencies run: pip install ruff - name: Lint diff --git a/README.md b/README.md index 772e741..071e07b 100644 --- a/README.md +++ b/README.md @@ -322,7 +322,7 @@ traceweave/ │ ├── exporters/ # Export to JSON, Chrome Trace, etc. │ └── cli.py # CLI: traceweave tui|dashboard|demo|export ├── examples/ # Demo scripts (no API keys needed) -└── tests/ # 39 tests, 100% passing +└── tests/ # Automated test suite ``` | Technology | Purpose | @@ -378,7 +378,7 @@ MIT License — see [LICENSE](LICENSE) for details. [license-shield]: https://img.shields.io/badge/license-MIT-369eff?labelColor=black&style=flat-square [license-link]: https://opensource.org/licenses/MIT [downloads-shield]: https://img.shields.io/pypi/dm/traceweave?color=369eff&labelColor=black&style=flat-square -[test-shield]: https://img.shields.io/badge/tests-39%20passed-369eff?labelColor=black&logo=pytest&logoColor=white&style=flat-square +[test-shield]: https://github.com/weivwang/trace-wave/actions/workflows/ci.yml/badge.svg [test-link]: https://github.com/weivwang/trace-wave/actions [docs-link]: https://github.com/weivwang/trace-wave#-quick-start [issues-link]: https://github.com/weivwang/trace-wave/issues diff --git a/agent_trace/__init__.py b/agent_trace/__init__.py index 5361d24..9605ab1 100644 --- a/agent_trace/__init__.py +++ b/agent_trace/__init__.py @@ -1,19 +1,19 @@ """traceweave: Distributed tracing and observability for AI agents.""" -__version__ = "0.1.0" +__version__ = "0.1.2" +from agent_trace.core.context import get_current_span, get_current_trace +from agent_trace.core.decorators import trace_agent, trace_llm, trace_tool from agent_trace.core.models import ( + SpanData, + SpanEvent, SpanKind, SpanStatus, TokenUsage, - SpanEvent, - SpanData, TraceData, ) from agent_trace.core.span import Span -from agent_trace.core.context import get_current_span, get_current_trace from agent_trace.core.tracer import AgentTracer, tracer -from agent_trace.core.decorators import trace_agent, trace_tool, trace_llm __all__ = [ "__version__", diff --git a/agent_trace/cli.py b/agent_trace/cli.py index 8c7ca75..a371d12 100644 --- a/agent_trace/cli.py +++ b/agent_trace/cli.py @@ -11,13 +11,14 @@ traceweave export trace.json --format chrome -o trace.chrome.json traceweave demo """ + import click -import json -import sys + +from agent_trace import __version__ @click.group() -@click.version_option(version="0.1.0", prog_name="traceweave") +@click.version_option(version=__version__, prog_name="traceweave") def main(): """🔍 traceweave: Distributed tracing and observability for AI agents.""" pass diff --git a/agent_trace/core/__init__.py b/agent_trace/core/__init__.py index 1367c5a..97f3a08 100644 --- a/agent_trace/core/__init__.py +++ b/agent_trace/core/__init__.py @@ -1,20 +1,20 @@ """Core tracing components for traceweave.""" +from agent_trace.core.context import ( + get_current_span, + get_current_trace, + set_current_span, + set_current_trace, +) from agent_trace.core.models import ( + SpanData, + SpanEvent, SpanKind, SpanStatus, TokenUsage, - SpanEvent, - SpanData, TraceData, ) from agent_trace.core.span import Span -from agent_trace.core.context import ( - get_current_span, - set_current_span, - get_current_trace, - set_current_trace, -) from agent_trace.core.tracer import AgentTracer, tracer __all__ = [ diff --git a/agent_trace/core/context.py b/agent_trace/core/context.py index 478b6a2..95977cd 100644 --- a/agent_trace/core/context.py +++ b/agent_trace/core/context.py @@ -13,6 +13,7 @@ # ... nested code sees my_span via get_current_span() ... reset_current_span(token) """ + from __future__ import annotations import contextvars @@ -32,6 +33,7 @@ # ── Public helpers ─────────────────────────────────────────────────────── + def get_current_span() -> Optional[Span]: """Return the currently active span, or ``None`` if outside a trace.""" return _current_span.get() diff --git a/agent_trace/core/decorators.py b/agent_trace/core/decorators.py index dcb8da0..420d9ca 100644 --- a/agent_trace/core/decorators.py +++ b/agent_trace/core/decorators.py @@ -15,13 +15,14 @@ def search(query: str) -> list[str]: def generate(prompt: str) -> str: return call_llm(prompt) """ + import functools import inspect from contextlib import contextmanager from typing import Any, Callable, Optional, TypeVar -from agent_trace.core.models import SpanKind from agent_trace.core.context import get_current_trace +from agent_trace.core.models import SpanKind from agent_trace.core.tracer import tracer as default_tracer F = TypeVar("F", bound=Callable[..., Any]) @@ -296,9 +297,7 @@ def _serialize_output(value: Any) -> Any: if hasattr(value, "model_dump"): return value.model_dump() if hasattr(value, "__dict__"): - return { - k: str(v) for k, v in value.__dict__.items() if not k.startswith("_") - } + return {k: str(v) for k, v in value.__dict__.items() if not k.startswith("_")} except Exception: pass return str(value)[:1000] @@ -327,6 +326,4 @@ def _try_extract_usage(span: Any, result: Any) -> None: usage.get("completion_tokens", 0) if isinstance(usage, dict) else 0 ) if prompt_tokens or completion_tokens: - span.set_token_usage( - prompt_tokens=prompt_tokens, completion_tokens=completion_tokens - ) + span.set_token_usage(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens) diff --git a/agent_trace/core/models.py b/agent_trace/core/models.py index c280e79..c0faca3 100644 --- a/agent_trace/core/models.py +++ b/agent_trace/core/models.py @@ -4,6 +4,7 @@ tracing system: spans, traces, token usage, and events. These models form the data layer that all other components build upon. """ + from __future__ import annotations import uuid @@ -174,4 +175,4 @@ def _count_spans(self, span: SpanData) -> int: model_config = { "arbitrary_types_allowed": True, "protected_namespaces": (), - } \ No newline at end of file + } diff --git a/agent_trace/core/span.py b/agent_trace/core/span.py index 37d91b1..a82f596 100644 --- a/agent_trace/core/span.py +++ b/agent_trace/core/span.py @@ -11,13 +11,14 @@ The span automatically records timing, captures exceptions, and notifies the parent :class:`~agent_trace.core.tracer.AgentTracer` on completion. """ + from __future__ import annotations import traceback from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, Optional -from agent_trace.core.models import SpanData, SpanKind, SpanStatus, SpanEvent, TokenUsage +from agent_trace.core.models import SpanData, SpanEvent, SpanKind, SpanStatus, TokenUsage if TYPE_CHECKING: from agent_trace.core.tracer import AgentTracer diff --git a/agent_trace/core/tracer.py b/agent_trace/core/tracer.py index 2c3514b..f8eb4ce 100644 --- a/agent_trace/core/tracer.py +++ b/agent_trace/core/tracer.py @@ -15,24 +15,25 @@ result = plan() span.set_output(result) """ + from __future__ import annotations import threading import uuid +from contextlib import contextmanager from datetime import datetime, timezone from typing import Any, Callable, Optional -from contextlib import contextmanager -from agent_trace.core.models import SpanData, SpanKind, SpanStatus, TraceData -from agent_trace.core.span import Span from agent_trace.core.context import ( get_current_span, - set_current_span, - reset_current_span, get_current_trace, - set_current_trace, + reset_current_span, reset_current_trace, + set_current_span, + set_current_trace, ) +from agent_trace.core.models import SpanData, SpanKind, TraceData +from agent_trace.core.span import Span # Type alias for event listener callbacks EventListener = Callable[[str, dict[str, Any]], None] @@ -177,7 +178,8 @@ def start_span( The new :class:`Span`. """ parent_span = get_current_span() - trace_id = get_current_trace() or uuid.uuid4().hex + current_trace_id = get_current_trace() + trace_id = current_trace_id or uuid.uuid4().hex span_data = SpanData( trace_id=trace_id, @@ -193,6 +195,24 @@ def start_span( if parent_span: parent_span._data.children.append(span_data) + # Auto-instrumentations call start_span() directly. When there is no + # surrounding trace, retain that operation as a one-span trace instead + # of silently discarding it when the span ends. + standalone_trace = None + trace_token = None + if parent_span is None and current_trace_id is None: + standalone_trace = TraceData( + trace_id=trace_id, + name=name, + start_time=span_data.start_time, + root_span=span_data, + metadata={}, + ) + with self._lock: + self._traces[trace_id] = standalone_trace + trace_token = set_current_trace(trace_id) + self._emit("trace_start", {"trace_id": trace_id, "name": name}) + with self._lock: self._active_spans[span_data.span_id] = span @@ -218,6 +238,20 @@ def start_span( span.end() # Restore previous span context using proper reset reset_current_span(old_span_token) + if standalone_trace is not None: + standalone_trace.end_time = span_data.end_time + if trace_token is not None: + reset_current_trace(trace_token) + self._emit( + "trace_end", + { + "trace_id": trace_id, + "duration_ms": standalone_trace.total_duration_ms, + "total_tokens": standalone_trace.total_tokens, + "total_cost": standalone_trace.total_cost, + "span_count": standalone_trace.span_count, + }, + ) # ── Internal callbacks ─────────────────────────────────────────── @@ -235,9 +269,7 @@ def _on_span_end(self, span: Span) -> None: "status": span._data.status.value, "duration_ms": span._data.duration_ms, "token_usage": ( - span._data.token_usage.model_dump() - if span._data.token_usage - else None + span._data.token_usage.model_dump() if span._data.token_usage else None ), "error": span._data.error, }, diff --git a/agent_trace/dashboard/__init__.py b/agent_trace/dashboard/__init__.py index 7c29ee5..0c2c868 100644 --- a/agent_trace/dashboard/__init__.py +++ b/agent_trace/dashboard/__init__.py @@ -1,7 +1,8 @@ """Dashboard components for traceweave.""" + from agent_trace.dashboard.server import run_server -from agent_trace.dashboard.tui import TraceDashboard, print_trace, run_tui from agent_trace.dashboard.trace_viewer import view_trace_file +from agent_trace.dashboard.tui import TraceDashboard, print_trace, run_tui __all__ = [ "run_server", diff --git a/agent_trace/dashboard/server.py b/agent_trace/dashboard/server.py index e53843b..f24182f 100644 --- a/agent_trace/dashboard/server.py +++ b/agent_trace/dashboard/server.py @@ -1,12 +1,12 @@ """Lightweight web dashboard server for traceweave.""" + import json -import threading from http.server import HTTPServer, SimpleHTTPRequestHandler -from typing import Optional -from agent_trace.core.tracer import tracer as default_tracer +from agent_trace import __version__ +from agent_trace.core.tracer import tracer as default_tracer -DASHBOARD_HTML = ''' +DASHBOARD_HTML = """ @@ -75,7 +75,7 @@

🔍 traceweave

- v0.1.0 + v__TRACEWEAVE_VERSION__ 0 traces ● Live
@@ -207,7 +207,7 @@ def my_agent(query): poll(); -''' +""".replace("__TRACEWEAVE_VERSION__", __version__) class DashboardHandler(SimpleHTTPRequestHandler): @@ -216,15 +216,15 @@ class DashboardHandler(SimpleHTTPRequestHandler): tracer = None def do_GET(self): - if self.path == '/' or self.path == '/index.html': + if self.path == "/" or self.path == "/index.html": self.send_response(200) - self.send_header('Content-Type', 'text/html') + self.send_header("Content-Type", "text/html") self.end_headers() self.wfile.write(DASHBOARD_HTML.encode()) - elif self.path == '/api/traces': + elif self.path == "/api/traces": self.send_response(200) - self.send_header('Content-Type', 'application/json') - self.send_header('Access-Control-Allow-Origin', '*') + self.send_header("Content-Type", "application/json") + self.send_header("Access-Control-Allow-Origin", "*") self.end_headers() traces = self.tracer.get_all_traces() if self.tracer else [] data = [t.model_dump(mode="json") for t in traces] diff --git a/agent_trace/dashboard/trace_viewer.py b/agent_trace/dashboard/trace_viewer.py index 2bbd7dd..e4acae9 100644 --- a/agent_trace/dashboard/trace_viewer.py +++ b/agent_trace/dashboard/trace_viewer.py @@ -1,6 +1,9 @@ """View saved trace files in the terminal.""" + import json + from rich.console import Console + from agent_trace.core.models import TraceData from agent_trace.dashboard.tui import print_trace diff --git a/agent_trace/dashboard/tui.py b/agent_trace/dashboard/tui.py index c7c0aa2..298ba2b 100644 --- a/agent_trace/dashboard/tui.py +++ b/agent_trace/dashboard/tui.py @@ -1,20 +1,17 @@ """Beautiful terminal UI for traceweave visualization using Rich.""" + import time from typing import Optional + +from rich import box from rich.console import Console +from rich.live import Live +from rich.panel import Panel from rich.table import Table from rich.tree import Tree -from rich.panel import Panel -from rich.layout import Layout -from rich.live import Live -from rich.text import Text -from rich.progress_bar import ProgressBar -from rich.columns import Columns -from rich import box from agent_trace.core.models import SpanData, SpanKind, SpanStatus, TraceData - # Color scheme COLORS = { SpanKind.AGENT: "bold cyan", @@ -47,9 +44,9 @@ def format_duration(ms: Optional[float]) -> str: if ms < 1000: return f"{ms:.0f}ms" elif ms < 60000: - return f"{ms/1000:.1f}s" + return f"{ms / 1000:.1f}s" else: - return f"{ms/60000:.1f}m" + return f"{ms / 60000:.1f}m" def format_tokens(tokens: int) -> str: @@ -57,7 +54,7 @@ def format_tokens(tokens: int) -> str: return "-" if tokens < 1000: return str(tokens) - return f"{tokens/1000:.1f}k" + return f"{tokens / 1000:.1f}k" def format_cost(cost: float) -> str: @@ -96,7 +93,10 @@ def _build_span_tree(span: SpanData, tree: Tree, max_duration: float) -> None: else: bar_color = COLORS.get(span.kind, "white").replace("bold ", "") - waterfall = f"[{bar_color}]{bar_char * bar_fill}[/{bar_color}][dim]{bar_empty * (bar_width - bar_fill)}[/dim]" + waterfall = ( + f"[{bar_color}]{bar_char * bar_fill}[/{bar_color}]" + f"[dim]{bar_empty * (bar_width - bar_fill)}[/dim]" + ) # Model badge model_badge = f" [dim]({span.model_name})[/dim]" if span.model_name else "" @@ -123,8 +123,7 @@ def render_trace(trace: TraceData) -> Panel: max_duration = trace.total_duration_ms or 1.0 tree = Tree( - f"[bold white]📍 Trace: {trace.name}[/bold white]" - f" [dim]id={trace.trace_id[:8]}...[/dim]" + f"[bold white]📍 Trace: {trace.name}[/bold white] [dim]id={trace.trace_id[:8]}...[/dim]" ) _build_span_tree(trace.root_span, tree, max_duration) @@ -141,21 +140,21 @@ def render_trace(trace: TraceData) -> Panel: layout = Table.grid(padding=(1, 0)) layout.add_row(tree) - layout.add_row(Panel(summary, title="[bold]Summary[/bold]", border_style="dim", box=box.ROUNDED)) + layout.add_row( + Panel(summary, title="[bold]Summary[/bold]", border_style="dim", box=box.ROUNDED) + ) return Panel( layout, - title=f"[bold cyan]🔍 traceweave[/bold cyan]", + title="[bold cyan]🔍 traceweave[/bold cyan]", subtitle=f"[dim]{trace.start_time.strftime('%H:%M:%S')}[/dim]", border_style="cyan", box=box.DOUBLE, ) -def render_multi_trace(traces: list[TraceData]) -> Layout: +def render_multi_trace(traces: list[TraceData]) -> Table: """Render multiple traces with an overview panel.""" - layout = Layout() - # Overview table overview = Table( title="[bold]📋 Traces Overview[/bold]", @@ -190,6 +189,7 @@ class TraceDashboard: def __init__(self, tracer=None): from agent_trace.core.tracer import tracer as default_tracer + self.tracer = tracer or default_tracer self.console = Console() self._live: Optional[Live] = None @@ -218,7 +218,12 @@ def _render(self) -> Panel: grid = Table.grid(padding=(1, 0)) grid.add_row(overview) grid.add_row(trace_panel) - return Panel(grid, title="[bold cyan]🔍 traceweave dashboard[/bold cyan]", border_style="cyan", box=box.DOUBLE) + return Panel( + grid, + title="[bold cyan]🔍 traceweave dashboard[/bold cyan]", + border_style="cyan", + box=box.DOUBLE, + ) return trace_panel diff --git a/agent_trace/examples/demo_runner.py b/agent_trace/examples/demo_runner.py index 48789d2..0449024 100644 --- a/agent_trace/examples/demo_runner.py +++ b/agent_trace/examples/demo_runner.py @@ -5,16 +5,15 @@ traceweave demo """ -import sys + import os +import sys def run_demo(): """Run the multi-agent research team demo.""" # Ensure the project root is importable so we can reach examples/ - project_root = os.path.dirname( - os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - ) + project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) sys.path.insert(0, project_root) from examples.multi_agent_demo import main diff --git a/agent_trace/exporters/__init__.py b/agent_trace/exporters/__init__.py index 6f5ab21..9d995f8 100644 --- a/agent_trace/exporters/__init__.py +++ b/agent_trace/exporters/__init__.py @@ -1,6 +1,7 @@ """Exporters for traceweave data.""" -from agent_trace.exporters.json_exporter import export_json, export_json_multi + from agent_trace.exporters.chrome_exporter import export_chrome from agent_trace.exporters.export import export_trace +from agent_trace.exporters.json_exporter import export_json, export_json_multi __all__ = ["export_json", "export_json_multi", "export_chrome", "export_trace"] diff --git a/agent_trace/exporters/chrome_exporter.py b/agent_trace/exporters/chrome_exporter.py index b559b76..16687b4 100644 --- a/agent_trace/exporters/chrome_exporter.py +++ b/agent_trace/exporters/chrome_exporter.py @@ -1,7 +1,9 @@ """Export traces to Chrome Trace Event format for chrome://tracing visualization.""" + import json from pathlib import Path from typing import Union + from agent_trace.core.models import SpanData, TraceData @@ -13,36 +15,40 @@ def _span_to_chrome_events(span: SpanData, pid: int = 1) -> list[dict]: duration_us = int(span.duration_ms * 1000) if span.duration_ms else 0 # Duration event - events.append({ - "name": span.name, - "cat": span.kind.value, - "ph": "X", # Complete event - "ts": start_us, - "dur": duration_us, - "pid": pid, - "tid": span.trace_id[:8], - "args": { - "kind": span.kind.value, - "status": span.status.value, - "span_id": span.span_id, - **({"model": span.model_name} if span.model_name else {}), - **({"tokens": span.token_usage.total_tokens} if span.token_usage else {}), - **({"error": span.error} if span.error else {}), + events.append( + { + "name": span.name, + "cat": span.kind.value, + "ph": "X", # Complete event + "ts": start_us, + "dur": duration_us, + "pid": pid, + "tid": span.trace_id[:8], + "args": { + "kind": span.kind.value, + "status": span.status.value, + "span_id": span.span_id, + **({"model": span.model_name} if span.model_name else {}), + **({"tokens": span.token_usage.total_tokens} if span.token_usage else {}), + **({"error": span.error} if span.error else {}), + }, } - }) + ) # Add events as instant events for event in span.events: - events.append({ - "name": event.name, - "cat": "event", - "ph": "i", - "ts": int(event.timestamp.timestamp() * 1_000_000), - "pid": pid, - "tid": span.trace_id[:8], - "s": "t", - "args": event.attributes, - }) + events.append( + { + "name": event.name, + "cat": "event", + "ph": "i", + "ts": int(event.timestamp.timestamp() * 1_000_000), + "pid": pid, + "tid": span.trace_id[:8], + "s": "t", + "args": event.attributes, + } + ) # Recurse into children for child in span.children: diff --git a/agent_trace/exporters/export.py b/agent_trace/exporters/export.py index 57d6092..ce2dda2 100644 --- a/agent_trace/exporters/export.py +++ b/agent_trace/exporters/export.py @@ -1,6 +1,8 @@ """Main export dispatcher.""" + import json from pathlib import Path + from agent_trace.core.models import TraceData @@ -17,9 +19,11 @@ def export_trace(trace_file: str, fmt: str = "json", output: str = None) -> str: if fmt == "json": from agent_trace.exporters.json_exporter import export_json + return export_json(trace, output) elif fmt == "chrome": from agent_trace.exporters.chrome_exporter import export_chrome + return export_chrome(trace, output) elif fmt == "otlp": # Future: OpenTelemetry export diff --git a/agent_trace/exporters/json_exporter.py b/agent_trace/exporters/json_exporter.py index 7ceb7a9..23164d2 100644 --- a/agent_trace/exporters/json_exporter.py +++ b/agent_trace/exporters/json_exporter.py @@ -1,8 +1,10 @@ """Export traces to JSON format.""" -import json -import orjson + from pathlib import Path from typing import Union + +import orjson + from agent_trace.core.models import TraceData diff --git a/agent_trace/integrations/__init__.py b/agent_trace/integrations/__init__.py index 27c22ef..3afcee2 100644 --- a/agent_trace/integrations/__init__.py +++ b/agent_trace/integrations/__init__.py @@ -28,6 +28,7 @@ def instrument_all(): try: from agent_trace.integrations.openai_integration import instrument_openai + instrument_openai() instrumented.append("openai") except ImportError: @@ -35,6 +36,7 @@ def instrument_all(): try: from agent_trace.integrations.anthropic_integration import instrument_anthropic + instrument_anthropic() instrumented.append("anthropic") except ImportError: diff --git a/agent_trace/integrations/anthropic_integration.py b/agent_trace/integrations/anthropic_integration.py index 6bb96c7..464934e 100644 --- a/agent_trace/integrations/anthropic_integration.py +++ b/agent_trace/integrations/anthropic_integration.py @@ -7,6 +7,7 @@ client = anthropic.Anthropic() response = client.messages.create(...) # auto-traced """ + import functools from typing import Any @@ -68,12 +69,14 @@ def traced_create(self: Any, *args: Any, **kwargs: Any) -> Any: with tracer.start_span(f"anthropic.{model}", SpanKind.LLM) as span: span._data.model_name = model - span.set_input({ - "model": model, - "message_count": len(messages), - "max_tokens": kwargs.get("max_tokens"), - "system": str(kwargs.get("system", ""))[:200] if kwargs.get("system") else None, - }) + span.set_input( + { + "model": model, + "message_count": len(messages), + "max_tokens": kwargs.get("max_tokens"), + "system": str(kwargs.get("system", ""))[:200] if kwargs.get("system") else None, + } + ) try: response = _original_create(self, *args, **kwargs) diff --git a/agent_trace/integrations/langchain_integration.py b/agent_trace/integrations/langchain_integration.py index 84b721d..350812d 100644 --- a/agent_trace/integrations/langchain_integration.py +++ b/agent_trace/integrations/langchain_integration.py @@ -9,7 +9,8 @@ # Or attach to any chain / agent: chain.invoke({"input": "..."}, config={"callbacks": [handler]}) """ -from typing import Any, Optional, Union + +from typing import Any, Optional from uuid import UUID from agent_trace.core.models import SpanKind diff --git a/agent_trace/integrations/openai_integration.py b/agent_trace/integrations/openai_integration.py index de9aaee..869a6fb 100644 --- a/agent_trace/integrations/openai_integration.py +++ b/agent_trace/integrations/openai_integration.py @@ -8,6 +8,7 @@ client = openai.OpenAI() response = client.chat.completions.create(...) # auto-traced """ + import functools from typing import Any @@ -39,7 +40,9 @@ def _get_cost(model: str) -> tuple[float, float]: return (0.01, 0.03) # conservative default -def _truncated_messages(messages: list[dict], max_messages: int = 5, max_chars: int = 200) -> list[dict]: +def _truncated_messages( + messages: list[dict], max_messages: int = 5, max_chars: int = 200 +) -> list[dict]: """Return a compact representation of the first *max_messages* messages.""" return [ { @@ -81,12 +84,14 @@ def traced_create(self: Any, *args: Any, **kwargs: Any) -> Any: with tracer.start_span(f"openai.{model}", SpanKind.LLM) as span: span._data.model_name = model - span.set_input({ - "model": model, - "messages": _truncated_messages(messages), - "temperature": kwargs.get("temperature"), - "max_tokens": kwargs.get("max_tokens"), - }) + span.set_input( + { + "model": model, + "messages": _truncated_messages(messages), + "temperature": kwargs.get("temperature"), + "max_tokens": kwargs.get("max_tokens"), + } + ) try: response = _original_create(self, *args, **kwargs) @@ -141,9 +146,11 @@ async def traced_async_create(self: Any, *args: Any, **kwargs: Any) -> Any: ) if hasattr(response, "choices") and response.choices: - span.set_output({ - "content": str(response.choices[0].message.content)[:500], - }) + span.set_output( + { + "content": str(response.choices[0].message.content)[:500], + } + ) return response except Exception as e: diff --git a/agent_trace/py.typed b/agent_trace/py.typed new file mode 100644 index 0000000..1242d43 --- /dev/null +++ b/agent_trace/py.typed @@ -0,0 +1 @@ +# Marker file for PEP 561. diff --git a/pyproject.toml b/pyproject.toml index b275a8d..e3e2af7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "traceweave" -version = "0.1.2" +dynamic = ["version"] description = "Distributed tracing and observability for AI agents. Datadog for AI Agents." readme = "README.md" license = {text = "MIT"} @@ -31,6 +31,7 @@ classifiers = [ "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", "Topic :: Scientific/Engineering :: Artificial Intelligence", "Topic :: System :: Monitoring", "Topic :: System :: Distributed Computing", @@ -75,11 +76,13 @@ all = [ traceweave = "agent_trace.cli:main" [project.urls] -Homepage = "https://github.com/traceweave/traceweave" -Documentation = "https://github.com/traceweave/traceweave#readme" -Repository = "https://github.com/traceweave/traceweave" -Issues = "https://github.com/traceweave/traceweave/issues" -Changelog = "https://github.com/traceweave/traceweave/blob/main/CHANGELOG.md" +Homepage = "https://github.com/weivwang/trace-wave" +Documentation = "https://github.com/weivwang/trace-wave#readme" +Repository = "https://github.com/weivwang/trace-wave" +Issues = "https://github.com/weivwang/trace-wave/issues" + +[tool.hatch.version] +path = "agent_trace/__init__.py" [tool.hatch.build.targets.wheel] packages = ["agent_trace"] @@ -90,6 +93,10 @@ line-length = 100 [tool.ruff.lint] select = ["E", "F", "I", "N", "W", "UP"] +ignore = ["UP045"] + +[tool.ruff.lint.per-file-ignores] +"agent_trace/dashboard/server.py" = ["E501"] [tool.pytest.ini_options] testpaths = ["tests"] diff --git a/tests/test_package_metadata.py b/tests/test_package_metadata.py new file mode 100644 index 0000000..8f5d75f --- /dev/null +++ b/tests/test_package_metadata.py @@ -0,0 +1,30 @@ +"""Tests for installed package metadata.""" + +from importlib import metadata, resources + +import agent_trace +from click.testing import CliRunner + +from agent_trace.cli import main +from agent_trace.dashboard.server import DASHBOARD_HTML + + +def test_runtime_version_matches_distribution_metadata(): + assert agent_trace.__version__ == metadata.version("traceweave") + + +def test_user_facing_versions_match_runtime_version(): + result = CliRunner().invoke(main, ["--version"]) + + assert result.exit_code == 0 + assert agent_trace.__version__ in result.output + assert f"v{agent_trace.__version__}" in DASHBOARD_HTML + assert "__TRACEWEAVE_VERSION__" not in DASHBOARD_HTML + + +def test_public_project_metadata_and_typing_marker(): + project_urls = metadata.metadata("traceweave").get_all("Project-URL") or [] + + assert project_urls + assert all("github.com/weivwang/trace-wave" in url for url in project_urls) + assert resources.files("agent_trace").joinpath("py.typed").is_file() diff --git a/tests/test_tracer.py b/tests/test_tracer.py index 516377c..198aae9 100644 --- a/tests/test_tracer.py +++ b/tests/test_tracer.py @@ -27,6 +27,30 @@ def test_nested_spans(self): trace = self.tracer.get_all_traces()[-1] assert trace.span_count >= 3 + def test_standalone_span_creates_queryable_trace(self): + with self.tracer.start_span("openai.gpt-test", SpanKind.LLM) as span: + span.set_output({"content": "ok"}) + with self.tracer.start_span("nested-tool", SpanKind.TOOL): + pass + + traces = self.tracer.get_all_traces() + assert len(traces) == 1 + assert traces[0].name == "openai.gpt-test" + assert traces[0].root_span is span.data + assert traces[0].end_time is not None + assert traces[0].span_count == 2 + + def test_standalone_span_restores_context_and_records_errors(self): + with pytest.raises(RuntimeError, match="provider failed"): + with self.tracer.start_span("anthropic.test", SpanKind.LLM): + raise RuntimeError("provider failed") + + trace = self.tracer.get_all_traces()[-1] + assert trace.root_span.status == SpanStatus.ERROR + assert trace.end_time is not None + assert get_current_span() is None + assert get_current_trace() is None + def test_span_error_handling(self): with pytest.raises(ValueError): with self.tracer.start_trace("error-test") as root: