traccia 0.1.2__py3-none-any.whl → 0.1.5__py3-none-any.whl

traccia/instrumentation/decorator.py

@@ -0,0 +1,263 @@
+"""@observe decorator for instrumenting functions."""
+
+from __future__ import annotations
+
+import functools
+import inspect
+import traceback
+from typing import Any, Callable, Dict, Iterable, Optional
+from traccia.tracer.span import SpanStatus
+
+
+def _capture_args(bound_args: inspect.BoundArguments, skip: Iterable[str]) -> Dict[str, Any]:
+    """Capture function arguments, converting complex types to OTel-compatible types."""
+    captured = {}
+    for name, value in bound_args.arguments.items():
+        if name in skip:
+            continue
+        # Skip 'self' - it's an object, not a valid OTel attribute
+        if name == "self":
+            continue
+        # Convert value to OTel-compatible type
+        captured[name] = _convert_to_otel_type(value)
+    return captured
+
+
+def _convert_to_otel_type(value: Any) -> Any:
+    """
+    Convert a value to an OpenTelemetry-compatible type.
+    
+    OTel attributes must be: bool, str, bytes, int, float, or sequences of those.
+    """
+    # Primitive types are fine
+    if isinstance(value, (bool, str, bytes, int, float)) or value is None:
+        return value
+    
+    # For sequences, convert each element
+    if isinstance(value, (list, tuple)):
+        converted = []
+        for item in value:
+            if isinstance(item, (bool, str, bytes, int, float)) or item is None:
+                converted.append(item)
+            else:
+                # Convert complex types to string representation
+                converted.append(str(item)[:1000])  # Truncate long strings
+        return converted[:100]  # Limit sequence length
+    
+    # For dicts and other complex types, convert to JSON string
+    if isinstance(value, dict):
+        try:
+            import json
+            json_str = json.dumps(value, default=str)[:1000]  # Truncate
+            return json_str
+        except Exception:
+            return str(value)[:1000]
+    
+    # For other types, convert to string
+    return str(value)[:1000]  # Truncate long strings
+
+
+def _infer_type_from_attributes(attributes: Dict[str, Any]) -> Optional[str]:
+    """
+    Infer span type from attributes.
+    
+    Returns:
+        - "llm" if LLM-related attributes found
+        - "tool" if tool-related attributes found
+        - None otherwise (will use default "span")
+    """
+    # Check for LLM indicators
+    if any(key in attributes for key in ["llm.model", "llm.vendor", "model"]):
+        return "llm"
+    
+    # Check for tool indicators
+    if any(key in attributes for key in ["tool.name", "tool", "http.url"]):
+        return "tool"
+    
+    return None
+
+
+def _extract_llm_attributes(span_attrs: Dict[str, Any], bound_args: inspect.BoundArguments) -> None:
+    """
+    Extract LLM-related attributes from function arguments.
+    
+    Extracts common LLM parameters like model, temperature, max_tokens, messages.
+    Fails silently if extraction fails or attributes not found.
+    
+    Args:
+        span_attrs: Dictionary to add extracted attributes to
+        bound_args: Bound arguments from the function call
+    """
+    try:
+        args_dict = dict(bound_args.arguments)
+        
+        # Extract model
+        if "model" in args_dict and "llm.model" not in span_attrs:
+            span_attrs["llm.model"] = str(args_dict["model"])
+        
+        # Extract temperature
+        if "temperature" in args_dict and "llm.temperature" not in span_attrs:
+            temp = args_dict["temperature"]
+            if isinstance(temp, (int, float)):
+                span_attrs["llm.temperature"] = temp
+        
+        # Extract max_tokens
+        if "max_tokens" in args_dict and "llm.max_tokens" not in span_attrs:
+            max_tok = args_dict["max_tokens"]
+            if isinstance(max_tok, int):
+                span_attrs["llm.max_tokens"] = max_tok
+        
+        # Extract messages/prompt
+        if "messages" in args_dict and "llm.prompt" not in span_attrs:
+            messages = args_dict["messages"]
+            if isinstance(messages, (list, str)):
+                # Convert messages to string representation
+                prompt_str = _convert_to_otel_type(messages)
+                span_attrs["llm.prompt"] = prompt_str
+        elif "prompt" in args_dict and "llm.prompt" not in span_attrs:
+            prompt = args_dict["prompt"]
+            if isinstance(prompt, str):
+                span_attrs["llm.prompt"] = prompt[:1000]
+    
+    except Exception:
+        # Fail silently - don't interrupt span creation if extraction fails
+        pass
+
+
+def observe(
+    name: Optional[str] = None,
+    *,
+    attributes: Optional[Dict[str, Any]] = None,
+    tags: Optional[Iterable[str]] = None,
+    as_type: str = "span",
+    skip_args: Optional[Iterable[str]] = None,
+    skip_result: bool = False,
+) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+    """
+    Decorate a function to create a span around its execution.
+
+    - Supports sync and async functions.
+    - Captures errors and records exception events.
+    - Optionally captures arguments/results (skip controls).
+    """
+
+    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+        span_name = name or func.__name__
+        arg_names = func.__code__.co_varnames
+        skip_args_set = set(skip_args or [])
+        tags_list = [str(tag) for tag in tags] if tags is not None else []
+
+        is_coro = inspect.iscoroutinefunction(func)
+
+        @functools.wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            tracer = _get_tracer(func.__module__ or "default")
+            bound = inspect.signature(func).bind_partial(*args, **kwargs)
+            bound.apply_defaults()
+
+            span_attrs = dict(attributes or {})
+            if tags_list:
+                span_attrs["span.tags"] = tags_list
+            
+            # Capture function arguments first
+            span_attrs.update(_capture_args(bound, skip_args_set))
+            
+            # Infer type from attributes if not explicitly set (or if set to default "span")
+            inferred_type = as_type
+            if as_type == "span":
+                # Try to infer from attributes
+                detected_type = _infer_type_from_attributes(span_attrs)
+                if detected_type:
+                    inferred_type = detected_type
+            
+            # Set span type
+            span_attrs["span.type"] = inferred_type
+            
+            # Extract LLM attributes if this is an LLM call
+            if inferred_type == "llm":
+                _extract_llm_attributes(span_attrs, bound)
+
+            with tracer.start_as_current_span(span_name, attributes=span_attrs) as span:
+                try:
+                    result = func(*args, **kwargs)
+                    if not skip_result:
+                        # Convert result to OTel-compatible type
+                        otel_result = _convert_to_otel_type(result)
+                        span.set_attribute("result", otel_result)
+                    return result
+                except Exception as exc:
+                    # Record detailed error information
+                    span.record_exception(exc)
+                    span.set_status(SpanStatus.ERROR, str(exc))
+                    
+                    # Add error attributes
+                    span.set_attribute("error.type", type(exc).__name__)
+                    span.set_attribute("error.message", str(exc))
+                    
+                    # Add truncated stack trace
+                    tb = traceback.format_exc()
+                    span.set_attribute("error.stack_trace", tb[:2000])  # Truncate to 2000 chars
+                    
+                    raise
+
+        @functools.wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            tracer = _get_tracer(func.__module__ or "default")
+            bound = inspect.signature(func).bind_partial(*args, **kwargs)
+            bound.apply_defaults()
+
+            span_attrs = dict(attributes or {})
+            if tags_list:
+                span_attrs["span.tags"] = tags_list
+            
+            # Capture function arguments first
+            span_attrs.update(_capture_args(bound, skip_args_set))
+            
+            # Infer type from attributes if not explicitly set (or if set to default "span")
+            inferred_type = as_type
+            if as_type == "span":
+                # Try to infer from attributes
+                detected_type = _infer_type_from_attributes(span_attrs)
+                if detected_type:
+                    inferred_type = detected_type
+            
+            # Set span type
+            span_attrs["span.type"] = inferred_type
+            
+            # Extract LLM attributes if this is an LLM call
+            if inferred_type == "llm":
+                _extract_llm_attributes(span_attrs, bound)
+
+            async with tracer.start_as_current_span(span_name, attributes=span_attrs) as span:
+                try:
+                    result = await func(*args, **kwargs)
+                    if not skip_result:
+                        # Convert result to OTel-compatible type
+                        otel_result = _convert_to_otel_type(result)
+                        span.set_attribute("result", otel_result)
+                    return result
+                except Exception as exc:
+                    # Record detailed error information
+                    span.record_exception(exc)
+                    span.set_status(SpanStatus.ERROR, str(exc))
+                    
+                    # Add error attributes
+                    span.set_attribute("error.type", type(exc).__name__)
+                    span.set_attribute("error.message", str(exc))
+                    
+                    # Add truncated stack trace
+                    tb = traceback.format_exc()
+                    span.set_attribute("error.stack_trace", tb[:2000])  # Truncate to 2000 chars
+                    
+                    raise
+
+        return async_wrapper if is_coro else sync_wrapper
+
+    return decorator
+
+
+def _get_tracer(name: str):
+    import traccia
+
+    return traccia.get_tracer(name)
+

traccia/instrumentation/fastapi.py

@@ -0,0 +1,38 @@
+"""
+FastAPI middleware helpers for tracing HTTP requests with the SDK.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Awaitable, Callable
+
+from traccia.instrumentation import start_server_span
+
+
+def install_http_middleware(app: Any, *, tracer_name: str = "agents-fastapi") -> None:
+    """
+    Attach an HTTP middleware that wraps each FastAPI request in a server span.
+
+    - Propagates incoming context from headers
+    - Records method/path and response status code
+    """
+
+    @app.middleware("http")
+    async def tracing_middleware(request, call_next: Callable[[Any], Awaitable[Any]]):  # type: ignore
+        # Lazy import to avoid circular import when traccia initializes.
+        from traccia import get_tracer
+        tracer = get_tracer(tracer_name)
+        headers = dict(request.headers)
+        attrs = {
+            "http.method": request.method,
+            "http.target": request.url.path,
+        }
+        async with start_server_span(tracer, "http.request", headers, attributes=attrs) as span:
+            response = await call_next(request)
+            try:
+                span.set_attribute("http.status_code", response.status_code)
+            except Exception:
+                pass
+            return response
+
+    return None

traccia/instrumentation/http_client.py

@@ -0,0 +1,21 @@
+"""HTTP client helpers for context propagation."""
+
+from __future__ import annotations
+
+from typing import Dict
+
+from traccia.context import inject_traceparent, inject_tracestate, get_current_span
+
+
+def inject_headers(headers: Dict[str, str]) -> Dict[str, str]:
+    """
+    Inject traceparent/tracestate into the provided headers dict if a current span exists.
+
+    Returns the same headers mapping for convenience.
+    """
+    span = get_current_span()
+    if span:
+        inject_traceparent(headers, span.context)
+        inject_tracestate(headers, span.context)
+    return headers
+

traccia/instrumentation/http_server.py

@@ -0,0 +1,25 @@
+"""HTTP server helpers for extracting context and creating server spans."""
+
+from __future__ import annotations
+
+from typing import Dict, Optional
+
+from traccia.context import extract_trace_context
+from traccia.tracer.tracer import Tracer
+from traccia.tracer.span_context import SpanContext
+
+
+def extract_parent_context(headers: Dict[str, str]) -> Optional[SpanContext]:
+    """Parse traceparent/tracestate from headers and return SpanContext if valid."""
+    return extract_trace_context(headers)
+
+
+def start_server_span(tracer: Tracer, name: str, headers: Dict[str, str], attributes=None):
+    """
+    Convenience helper to start a server span with extracted parent context.
+
+    Returns the span context manager (caller should use 'with' or 'async with').
+    """
+    parent_ctx = extract_parent_context(headers)
+    return tracer.start_as_current_span(name, attributes=attributes, parent_context=parent_ctx)
+

traccia/instrumentation/openai.py

@@ -0,0 +1,178 @@
+"""OpenAI monkey patching for chat completions."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional, Callable
+from traccia.tracer.span import SpanStatus
+
+_patched = False
+
+
+def _safe_get(obj, path: str, default=None):
+    cur = obj
+    for part in path.split("."):
+        if cur is None:
+            return default
+        if isinstance(cur, dict):
+            cur = cur.get(part)
+        else:
+            cur = getattr(cur, part, None)
+    return cur if cur is not None else default
+
+
+def patch_openai() -> bool:
+    """Patch OpenAI chat completions for both legacy and new client APIs."""
+    global _patched
+    if _patched:
+        return True
+    try:
+        import openai
+    except Exception:
+        return False
+
+    def _extract_messages(kwargs, args):
+        messages = kwargs.get("messages")
+        # For new client, first arg after self is messages
+        if messages is None and len(args) >= 2:
+            messages = args[1]
+        if not messages or not isinstance(messages, (list, tuple)):
+            return None
+        # Keep only JSON-friendly, small fields to avoid huge/sensitive payloads.
+        slim = []
+        for m in list(messages)[:50]:
+            if not isinstance(m, dict):
+                continue
+            role = m.get("role")
+            name = m.get("name")
+            content = m.get("content")
+            if isinstance(content, (list, dict)):
+                content = str(content)
+            elif content is not None and not isinstance(content, str):
+                content = str(content)
+            item = {"role": role, "content": content}
+            if name:
+                item["name"] = name
+            slim.append(item)
+        return slim or None
+
+    def _extract_prompt_text(messages_slim) -> Optional[str]:
+        if not messages_slim:
+            return None
+        parts = []
+        for m in messages_slim:
+            role = m.get("role")
+            content = m.get("content")
+            if not content:
+                continue
+            parts.append(f"{role}: {content}" if role else str(content))
+        return "\n".join(parts) if parts else None
+
+    def _extract_prompt(kwargs, args) -> Optional[str]:
+        messages = kwargs.get("messages")
+        if messages is None and len(args) >= 2:
+            messages = args[1]
+        if not messages:
+            return None
+        parts = []
+        for m in messages:
+            content = m.get("content")
+            role = m.get("role")
+            if content:
+                parts.append(f"{role}: {content}" if role else str(content))
+        return "\n".join(parts) if parts else None
+
+    def _wrap(create_fn: Callable):
+        if getattr(create_fn, "_agent_trace_patched", False):
+            return create_fn
+
+        def wrapped_create(*args, **kwargs):
+            tracer = _get_tracer("openai")
+            model = kwargs.get("model") or _safe_get(args, "0.model", None)
+            messages_slim = _extract_messages(kwargs, args)
+            prompt_text = _extract_prompt_text(messages_slim) or _extract_prompt(kwargs, args)
+            attributes: Dict[str, Any] = {"llm.vendor": "openai"}
+            if model:
+                attributes["llm.model"] = model
+            if messages_slim:
+                # Convert messages to JSON string for OTel compatibility
+                import json
+                try:
+                    attributes["llm.openai.messages"] = json.dumps(messages_slim)[:1000]
+                except Exception:
+                    attributes["llm.openai.messages"] = str(messages_slim)[:1000]
+            if prompt_text:
+                attributes["llm.prompt"] = prompt_text
+            with tracer.start_as_current_span("llm.openai.chat.completions", attributes=attributes) as span:
+                try:
+                    resp = create_fn(*args, **kwargs)
+                    # capture model from response if not already set
+                    resp_model = getattr(resp, "model", None) or (_safe_get(resp, "model"))
+                    if resp_model and "llm.model" not in span.attributes:
+                        span.set_attribute("llm.model", resp_model)
+                    usage = getattr(resp, "usage", None) or (resp.get("usage") if isinstance(resp, dict) else None)
+                    if usage:
+                        span.set_attribute("llm.usage.source", "provider_usage")
+                        for k in ("prompt_tokens", "completion_tokens", "total_tokens"):
+                            val = getattr(usage, k, None) if not isinstance(usage, dict) else usage.get(k)
+                            if val is not None:
+                                span.set_attribute(f"llm.usage.{k}", val)
+                        if "llm.usage.prompt_tokens" in span.attributes:
+                            span.set_attribute("llm.usage.prompt_source", "provider_usage")
+                        if "llm.usage.completion_tokens" in span.attributes:
+                            span.set_attribute("llm.usage.completion_source", "provider_usage")
+                    finish_reason = _safe_get(resp, "choices.0.finish_reason")
+                    if finish_reason:
+                        span.set_attribute("llm.finish_reason", finish_reason)
+                    completion = _safe_get(resp, "choices.0.message.content")
+                    if completion:
+                        span.set_attribute("llm.completion", completion)
+                    return resp
+                except Exception as exc:
+                    span.record_exception(exc)
+                    span.set_status(SpanStatus.ERROR, str(exc))
+                    raise
+
+        wrapped_create._agent_trace_patched = True
+        return wrapped_create
+
+    patched_any = False
+
+    # Legacy: openai.ChatCompletion.create
+    target_legacy = getattr(openai, "ChatCompletion", None) or getattr(openai, "chat", None)
+    if target_legacy:
+        create_fn = getattr(target_legacy, "create", None)
+        if create_fn:
+            setattr(target_legacy, "create", _wrap(create_fn))
+            patched_any = True
+
+    # New client: OpenAI.chat.completions.create
+    new_client_cls = getattr(openai, "OpenAI", None)
+    if new_client_cls and hasattr(new_client_cls, "chat"):
+        chat = getattr(new_client_cls, "chat", None)
+        if chat and hasattr(chat, "completions"):
+            completions = getattr(chat, "completions")
+            if hasattr(completions, "create"):
+                patched = _wrap(completions.create)
+                setattr(completions, "create", patched)
+                patched_any = True
+
+    # New client resource class: openai.resources.chat.completions.Completions
+    try:
+        from openai.resources.chat.completions import Completions  # type: ignore
+
+        if hasattr(Completions, "create"):
+            Completions.create = _wrap(Completions.create)
+            patched_any = True
+    except Exception:
+        pass
+
+    if patched_any:
+        _patched = True
+    return patched_any
+
+
+def _get_tracer(name: str):
+    import traccia
+
+    return traccia.get_tracer(name)
+

traccia/instrumentation/requests.py

@@ -0,0 +1,68 @@
+"""Requests monkey patching for HTTP client tracing."""
+
+from __future__ import annotations
+
+from typing import Any, Dict
+from traccia.tracer.span import SpanStatus
+from traccia.context import get_current_span, inject_traceparent, inject_tracestate
+
+_patched = False
+
+
+def patch_requests() -> bool:
+    """Patch requests.Session.request; returns True if patched, False otherwise."""
+    global _patched
+    if _patched:
+        return True
+    try:
+        import requests
+    except Exception:
+        return False
+
+    original_request = requests.sessions.Session.request
+    if getattr(original_request, "_agent_trace_patched", False):
+        _patched = True
+        return True
+
+    def wrapped_request(self, method, url, *args, **kwargs):
+        # Skip instrumentation for OTLP exporter endpoints to prevent feedback loop
+        url_str = str(url) if url else ""
+        if "/v1/traces" in url_str or "/api/v1/traces" in url_str:
+            # This is likely an exporter endpoint - don't instrument it
+            import requests
+            return original_request(self, method, url, *args, **kwargs)
+        
+        tracer = _get_tracer("requests")
+        attributes: Dict[str, Any] = {
+            "http.method": method,
+            "http.url": url,
+        }
+        headers = kwargs.get("headers")
+        if headers is None:
+            headers = {}
+            kwargs["headers"] = headers
+        current = get_current_span()
+        if current:
+            inject_traceparent(headers, current.context)
+            inject_tracestate(headers, current.context)
+        with tracer.start_as_current_span("http.client", attributes=attributes) as span:
+            try:
+                resp = original_request(self, method, url, *args, **kwargs)
+                span.set_attribute("http.status_code", getattr(resp, "status_code", None))
+                return resp
+            except Exception as exc:
+                span.record_exception(exc)
+                span.set_status(SpanStatus.ERROR, str(exc))
+                raise
+
+    wrapped_request._agent_trace_patched = True
+    requests.sessions.Session.request = wrapped_request
+    _patched = True
+    return True
+
+
+def _get_tracer(name: str):
+    import traccia
+
+    return traccia.get_tracer(name)
+

traccia/integrations/__init__.py

@@ -0,0 +1,22 @@
+"""Traccia integrations for frameworks like LangChain, CrewAI, LlamaIndex."""
+
+__all__ = []
+
+# Lazy imports for optional dependencies
+def _import_langchain():
+    try:
+        from traccia.integrations.langchain import TracciaCallbackHandler
+        return TracciaCallbackHandler
+    except ImportError as e:
+        raise ModuleNotFoundError(
+            "LangChain integration requires langchain-core. "
+            "Install with: pip install traccia[langchain]"
+        ) from e
+
+
+# Make available if imported
+try:
+    from traccia.integrations.langchain import TracciaCallbackHandler
+    __all__.append("TracciaCallbackHandler")
+except ImportError:
+    pass

traccia/integrations/langchain/__init__.py

@@ -0,0 +1,14 @@
+"""Traccia LangChain integration via callback handler."""
+
+try:
+    from traccia.integrations.langchain.callback import TracciaCallbackHandler
+    
+    # Convenience alias: from traccia.integrations.langchain import CallbackHandler
+    CallbackHandler = TracciaCallbackHandler
+    
+    __all__ = ["TracciaCallbackHandler", "CallbackHandler"]
+except ImportError as e:
+    raise ModuleNotFoundError(
+        "LangChain integration requires langchain-core. "
+        "Install with: pip install traccia[langchain]"
+    ) from e