agent-observability-trace-cli 0.1.2__py3-none-any.whl

agent_trace/integrations/haystack.py

@@ -0,0 +1,254 @@
+"""
+Haystack integration — implements Haystack's native ``tracing.Tracer`` /
+``tracing.Span`` interface so ``Pipeline.run()`` and ``Component.run()``
+execution is captured as agent-trace spans.
+
+Haystack 2.x ships its own instrumentation surface
+(``haystack.tracing.Tracer`` / ``haystack.tracing.Span``, see
+``haystack/tracing/tracer.py``) that ``Pipeline.run()``
+(``haystack/core/pipeline/pipeline.py``) and ``PipelineBase._run_component``
+(``haystack/core/pipeline/base.py``) call directly — there is no callback
+list to pass in, unlike LangGraph.  A tracer implementation is registered
+globally via ``haystack.tracing.enable_tracing(...)`` and every pipeline/
+component run after that point is traced through it until
+``haystack.tracing.disable_tracing()`` is called.
+
+Two spans are produced per pipeline run:
+
+- ``haystack.pipeline.run`` — one span for the whole ``Pipeline.run()`` call,
+  tagged with the pipeline's input/output data and metadata.
+- ``haystack.component.run`` — one span per component invocation (nested
+  under the pipeline span), tagged with the component's name, type, and I/O
+  socket spec.  When Haystack's own content tracing is enabled (see below)
+  the component's *actual received arguments* and *actual returned output*
+  are attached too — this is the exact capability gap issue #4574 exposes:
+  a caller-to-callee argument-propagation bug that occurs entirely
+  in-process, before any HTTP request, so agent-trace's httpx/requests
+  interceptor is structurally the wrong layer to catch it.
+
+Haystack gates raw input/output content behind its own
+``HAYSTACK_CONTENT_TRACING_ENABLED`` environment variable (default
+``false``) because pipeline inputs/outputs can contain arbitrary user
+content — set it to ``"true"`` (or call
+``haystack.tracing.tracer.tracer.is_content_tracing_enabled = True``) to
+capture that content in agent-trace spans as well.
+
+Usage::
+
+    import haystack.tracing
+    from agent_trace import Tracer
+    from agent_trace.integrations.haystack import HaystackTracer
+
+    t = Tracer()
+    with t.start_trace("my_pipeline", record=True) as trace:
+        haystack.tracing.enable_tracing(HaystackTracer(tracer=t, trace=trace))
+        try:
+            result = pipeline.run({"component_name": {"value": "hello"}})
+        finally:
+            haystack.tracing.disable_tracing()
+"""
+
+from __future__ import annotations
+
+import contextlib
+import logging
+import threading
+from typing import TYPE_CHECKING, Any
+
+from agent_trace.core.span import Span, SpanStatus
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+    from agent_trace import Trace, Tracer
+
+__all__ = ["HaystackTracer"]
+
+logger = logging.getLogger(__name__)
+
+_INSTALL_HINT = (
+    "HaystackTracer requires haystack-ai.\n"
+    "Install it with:\n\n"
+    "    pip install haystack-ai\n"
+)
+
+# agent_trace.core.span.Span attributes only accept str | int | float | bool;
+# Haystack tags/content-tags are frequently dicts, lists, dataclasses, or
+# Haystack domain objects (Document, ChatMessage, ...).  Truncate the
+# stringified form so one component's raw output can't blow up trace.json.
+_MAX_TAG_LEN = 4000
+
+
+def _require_haystack_tracing() -> Any:
+    """Lazy import guard — raises a clear error if haystack-ai is absent."""
+    try:
+        import haystack.tracing
+
+        return haystack.tracing
+    except ImportError as exc:
+        raise ImportError(_INSTALL_HINT) from exc
+
+
+def _coerce_tag_value(value: Any) -> str | int | float | bool:
+    """Coerce an arbitrary Haystack tag value to a Span-attribute-safe primitive."""
+    if isinstance(value, str | int | float | bool):
+        return value
+    try:
+        text = repr(value)
+    except Exception:
+        text = f"<unrepr-able {type(value).__name__}>"
+    if len(text) > _MAX_TAG_LEN:
+        text = text[:_MAX_TAG_LEN] + "...<truncated>"
+    return text
+
+
+# Module-level singletons: the concrete Tracer/Span implementations are built
+# once (the first time _get_classes() is called) so that Haystack's
+# tracing.Tracer / tracing.Span are real bases at class-definition time
+# rather than spliced in later.
+_HaystackTracerImplClass: type | None = None
+_HaystackSpanImplClass: type | None = None
+_classes_lock: threading.Lock = threading.Lock()
+
+
+def _get_classes() -> tuple[type, type]:
+    """Return (and lazily build) the concrete Tracer/Span implementations."""
+    global _HaystackTracerImplClass, _HaystackSpanImplClass  # noqa: PLW0603
+    if _HaystackTracerImplClass is not None and _HaystackSpanImplClass is not None:
+        return _HaystackTracerImplClass, _HaystackSpanImplClass
+
+    with _classes_lock:
+        if _HaystackTracerImplClass is not None and _HaystackSpanImplClass is not None:
+            return _HaystackTracerImplClass, _HaystackSpanImplClass
+
+        htracing = _require_haystack_tracing()
+        base_span_cls: type = htracing.Span
+        base_tracer_cls: type = htracing.Tracer
+
+        class _AgentTraceHaystackSpan(base_span_cls):  # type: ignore[misc]
+            """Wraps an agent_trace Span behind Haystack's Span interface."""
+
+            def __init__(self, agent_span: Span) -> None:
+                self._agent_span = agent_span
+
+            def set_tag(self, key: str, value: Any) -> None:
+                self._agent_span.set_attribute(key, _coerce_tag_value(value))
+
+            def raw_span(self) -> Any:
+                """Expose the underlying agent_trace Span for direct access."""
+                return self._agent_span
+
+        class _AgentTraceHaystackTracerImpl(base_tracer_cls):  # type: ignore[misc]
+            """Concrete implementation — see HaystackTracer for public docs."""
+
+            def __init__(self, tracer: Tracer, trace: Trace) -> None:
+                self._tracer: Tracer = tracer
+                self._trace: Trace = trace
+                # LIFO stack of open agent_trace spans, used to infer the
+                # parent span when Haystack doesn't pass one explicitly.
+                self._span_stack: list[Span] = []
+                self._lock: threading.Lock = threading.Lock()
+
+            @contextlib.contextmanager
+            def trace(
+                self,
+                operation_name: str,
+                tags: dict[str, Any] | None = None,
+                parent_span: Any | None = None,
+            ) -> Iterator[Any]:
+                """Open a span for *operation_name*, yield it, then close it.
+
+                Honors an explicit *parent_span* (Haystack passes the
+                pipeline-level span in when tracing a component run — see
+                ``PipelineBase._create_component_span``); falls back to the
+                innermost currently-open span on this tracer otherwise.
+                """
+                parent_id: str | None = None
+                if parent_span is not None:
+                    raw = getattr(parent_span, "raw_span", lambda: None)()
+                    if isinstance(raw, Span):
+                        parent_id = raw.span_id
+                else:
+                    with self._lock:
+                        if self._span_stack:
+                            parent_id = self._span_stack[-1].span_id
+
+                agent_span = self._tracer.start_span(
+                    operation_name, parent_id=parent_id
+                )
+                wrapped = _AgentTraceHaystackSpan(agent_span)
+                if tags:
+                    wrapped.set_tags(tags)
+
+                with self._lock:
+                    self._span_stack.append(agent_span)
+                try:
+                    yield wrapped
+                except Exception as exc:
+                    agent_span.record_exception(exc)
+                    if agent_span.end_time is None:
+                        agent_span.end(SpanStatus.ERROR)
+                    raise
+                else:
+                    if agent_span.end_time is None:
+                        agent_span.end(SpanStatus.OK)
+                finally:
+                    with self._lock:
+                        if self._span_stack and self._span_stack[-1] is agent_span:
+                            self._span_stack.pop()
+                        else:
+                            # Defensive: out-of-order close (shouldn't happen
+                            # with well-behaved context managers) — drop it
+                            # wherever it is rather than corrupting the stack.
+                            with contextlib.suppress(ValueError):
+                                self._span_stack.remove(agent_span)
+
+            def current_span(self) -> Any | None:
+                """Return the innermost currently-open span, or None."""
+                with self._lock:
+                    if not self._span_stack:
+                        return None
+                    return _AgentTraceHaystackSpan(self._span_stack[-1])
+
+        _HaystackTracerImplClass = _AgentTraceHaystackTracerImpl
+        _HaystackSpanImplClass = _AgentTraceHaystackSpan
+        return _HaystackTracerImplClass, _HaystackSpanImplClass
+
+
+class HaystackTracer:
+    """Haystack ``tracing.Tracer`` implementation that emits agent-trace spans.
+
+    Register an instance globally with ``haystack.tracing.enable_tracing(...)``
+    before calling ``pipeline.run()`` / ``pipeline.run_async()``; every
+    ``haystack.pipeline.run`` and ``haystack.component.run`` span Haystack
+    creates internally is forwarded to agent-trace for the lifetime of the
+    registration.  Call ``haystack.tracing.disable_tracing()`` (or register a
+    fresh ``NullTracer``) when done so later pipeline runs outside the
+    ``with tracer.start_trace(...)`` block aren't captured into a closed
+    trace.
+
+    ``haystack-ai`` is imported lazily — importing this module succeeds even
+    when it is not installed.
+
+    Parameters
+    ----------
+    tracer:
+        The active :class:`~agent_trace.Tracer` instance.
+    trace:
+        The :class:`~agent_trace.Trace` that spans will be registered on.
+    """
+
+    def __new__(cls, tracer: Tracer, trace: Trace) -> HaystackTracer:
+        # Construct the concrete impl directly so Python's normal
+        # type.__call__ runs _AgentTraceHaystackTracerImpl.__init__
+        # automatically.  See LangGraphTracer.__new__ for why this two-step
+        # __new__/__init__ split is required rather than mutating __bases__.
+        impl_cls, _ = _get_classes()
+        return impl_cls(tracer, trace)  # type: ignore[no-any-return]
+
+    def __init__(self, tracer: Tracer, trace: Trace) -> None:
+        # __init__ is called on the instance whose __class__ is already the
+        # concrete impl class (set by __new__).  This path is only reached if
+        # someone subclasses HaystackTracer directly; normal construction
+        # goes through the impl class __init__.
+        pass  # pragma: no cover

agent_trace/integrations/langchain_core.py

@@ -0,0 +1,361 @@
+"""
+Generic, framework-agnostic LangChain callback integration.
+
+``LangGraphTracer`` (``agent_trace.integrations.langgraph``) only attaches
+to a compiled LangGraph ``StateGraph`` — it hard-codes LangGraph-specific
+span-naming conventions (``node:<name>``, ``branch:dispatch``, checkpointer
+tracking, ...) and is documented as a LangGraph integration. The much
+larger population of plain-LangChain users — anyone calling
+``Runnable.invoke()``/``.batch()``/``.ainvoke()`` directly (a document
+compressor pipeline, a bare `Chain`, a `Retriever`, ...) with no LangGraph
+graph anywhere in sight — has had *no* chain-level integration at all: only
+the raw HTTP interceptor sees anything, and an application-level exception
+raised inside a Runnable's own Python code (not during an HTTP call) is
+invisible everywhere, confirmed via reproduction against issue #31192 (a
+`DocumentCompressorPipeline`'s `IndexError` in `_parse_ranking` attaches to
+nothing — ``Tracer.start_trace``'s exception handler only records onto
+spans already present in ``trace.spans``, and with no chain-level
+integration wired in, none exist).
+
+``LangChainTracer`` closes that gap: a plain ``BaseCallbackHandler``
+implementation that works with *any* LangChain `Runnable`, not just a
+LangGraph graph — pass it into ``config={"callbacks": [...]}`` on any
+``.invoke()``/``.batch()``/``.ainvoke()`` call. On error, it calls
+``Span.record_exception()`` on a span it opens itself for that Runnable
+invocation, the same mechanism LangGraphTracer uses, but without any
+LangGraph-specific assumptions about span naming or graph structure.
+
+Usage::
+
+    from agent_trace import Tracer
+    from agent_trace.integrations.langchain_core import LangChainTracer
+
+    t = Tracer(trace_dir=...)
+    with t.start_trace("my-runnable", record=True) as trace:
+        cb = LangChainTracer(tracer=t, trace=trace)
+        result = my_runnable.invoke(input, config={"callbacks": [cb]})
+
+Requires ``langchain-core`` (``pip install agent-trace[langchain]``) —
+notably *not* ``langgraph``, unlike ``LangGraphTracer``.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import threading
+import uuid
+from typing import TYPE_CHECKING, Any
+
+from agent_trace.core.span import Span, SpanStatus
+
+if TYPE_CHECKING:
+    from agent_trace import Trace, Tracer
+
+__all__ = ["LangChainTracer"]
+
+logger = logging.getLogger(__name__)
+
+_INSTALL_HINT = (
+    "LangChainTracer requires langchain-core.\n"
+    "Install it with:\n\n"
+    "    pip install agent-trace[langchain]\n"
+)
+
+_MAX_ATTR_LEN = 8_000
+
+
+def _to_attr_string(value: Any, *, max_len: int = _MAX_ATTR_LEN) -> str:
+    """Best-effort, bounded string form of an arbitrary LangChain value
+    (inputs/outputs dicts, message lists, ...) — mirrors
+    agent_trace.integrations.langgraph._to_attr_string, duplicated (not
+    imported) so this module has zero dependency on langgraph.py and,
+    transitively, on the ``langgraph`` package itself."""
+    try:
+        text = json.dumps(value, default=str, ensure_ascii=False)
+    except Exception:
+        text = str(value)
+    if len(text) > max_len:
+        text = text[:max_len] + "...<truncated>"
+    return text
+
+
+def _require_langchain_core() -> Any:
+    try:
+        import langchain_core
+
+        return langchain_core
+    except ImportError as exc:
+        raise ImportError(_INSTALL_HINT) from exc
+
+
+def _base_callback_handler() -> Any:
+    _require_langchain_core()
+    from langchain_core.callbacks import BaseCallbackHandler
+
+    return BaseCallbackHandler
+
+
+_LangChainTracerClass: type | None = None
+_tracer_class_lock: threading.Lock = threading.Lock()
+
+
+def _get_tracer_class() -> type:
+    """Return (and lazily build) the concrete LangChainTracer
+    implementation, with BaseCallbackHandler as a genuine base class at
+    definition time — same pattern as
+    agent_trace.integrations.langgraph._get_tracer_class."""
+    global _LangChainTracerClass  # noqa: PLW0603
+    if _LangChainTracerClass is not None:
+        return _LangChainTracerClass
+
+    with _tracer_class_lock:
+        if _LangChainTracerClass is not None:
+            return _LangChainTracerClass
+
+        BaseCallbackHandler = _base_callback_handler()
+
+        class _LangChainTracerImpl(BaseCallbackHandler):  # type: ignore[misc,valid-type]
+            """Generic BaseCallbackHandler — see module docstring."""
+
+            def __init__(self, tracer: Tracer, trace: Trace) -> None:
+                self._tracer = tracer
+                self._trace = trace
+                self._spans: dict[str, Span] = {}
+                self._lock: threading.Lock = threading.Lock()
+
+            # -- span lifecycle -------------------------------------------
+
+            def _open_span(
+                self,
+                run_id: uuid.UUID | str,
+                name: str,
+                parent_run_id: uuid.UUID | str | None = None,
+            ) -> Span:
+                run_key = str(run_id)
+                parent_span_id: str | None = None
+                if parent_run_id is not None:
+                    with self._lock:
+                        parent_span = self._spans.get(str(parent_run_id))
+                    parent_span_id = (
+                        parent_span.span_id if parent_span is not None else None
+                    )
+                span = self._tracer.start_span(name, parent_id=parent_span_id)
+                with self._lock:
+                    self._spans[run_key] = span
+                return span
+
+            def _close_span(
+                self, run_id: uuid.UUID | str, status: SpanStatus = SpanStatus.OK
+            ) -> Span | None:
+                run_key = str(run_id)
+                with self._lock:
+                    span = self._spans.pop(run_key, None)
+                if span is not None and span.end_time is None:
+                    span.end(status)
+                return span
+
+            def _close_span_with_exception(
+                self, run_id: uuid.UUID | str, error: BaseException
+            ) -> None:
+                run_key = str(run_id)
+                with self._lock:
+                    span = self._spans.pop(run_key, None)
+                if span is None:
+                    return
+                span.record_exception(error)
+                if span.end_time is None:
+                    span.end(SpanStatus.ERROR)
+
+            # -- chain callbacks -------------------------------------------
+
+            def on_chain_start(
+                self,
+                serialized: dict[str, Any] | None,
+                inputs: dict[str, Any],
+                *,
+                run_id: uuid.UUID,
+                parent_run_id: uuid.UUID | None = None,
+                tags: list[str] | None = None,
+                metadata: dict[str, Any] | None = None,
+                **kwargs: Any,
+            ) -> None:
+                ser = serialized or {}
+                name = (
+                    kwargs.get("name")
+                    or ser.get("name")
+                    or (ser.get("id") or [None])[-1]
+                    or "chain"
+                )
+                span = self._open_span(run_id, f"chain:{name}", parent_run_id)
+                if tags:
+                    span.set_attribute("chain.tags", ",".join(tags))
+                if metadata:
+                    span.set_attribute("chain.metadata", _to_attr_string(metadata))
+                if inputs:
+                    span.set_attribute("chain.inputs", _to_attr_string(inputs))
+
+            def on_chain_end(
+                self,
+                outputs: Any,
+                *,
+                run_id: uuid.UUID,
+                parent_run_id: uuid.UUID | None = None,
+                **kwargs: Any,
+            ) -> None:
+                run_key = str(run_id)
+                with self._lock:
+                    span = self._spans.get(run_key)
+                if span is not None and outputs:
+                    try:
+                        span.set_attribute("chain.outputs", _to_attr_string(outputs))
+                    except Exception:
+                        logger.debug(
+                            "agent-trace: failed to record chain outputs for "
+                            "run %r",
+                            run_key,
+                            exc_info=True,
+                        )
+                self._close_span(run_id, SpanStatus.OK)
+
+            def on_chain_error(
+                self,
+                error: BaseException,
+                *,
+                run_id: uuid.UUID,
+                parent_run_id: uuid.UUID | None = None,
+                **kwargs: Any,
+            ) -> None:
+                self._close_span_with_exception(run_id, error)
+
+            # -- LLM callbacks -----------------------------------------------
+
+            def on_llm_start(
+                self,
+                serialized: dict[str, Any] | None,
+                prompts: list[str],
+                *,
+                run_id: uuid.UUID,
+                parent_run_id: uuid.UUID | None = None,
+                **kwargs: Any,
+            ) -> None:
+                ser = serialized or {}
+                model_name = (
+                    (ser.get("kwargs") or {}).get("model_name")
+                    or (ser.get("kwargs") or {}).get("model")
+                    or ser.get("name")
+                    or "llm"
+                )
+                span = self._open_span(run_id, f"llm:{model_name}", parent_run_id)
+                span.set_attribute("llm.model", str(model_name))
+
+            def on_chat_model_start(
+                self,
+                serialized: dict[str, Any] | None,
+                messages: list[list[Any]],
+                *,
+                run_id: uuid.UUID,
+                parent_run_id: uuid.UUID | None = None,
+                **kwargs: Any,
+            ) -> None:
+                ser = serialized or {}
+                model_name = (
+                    (ser.get("kwargs") or {}).get("model_name")
+                    or (ser.get("kwargs") or {}).get("model")
+                    or ser.get("name")
+                    or "llm"
+                )
+                span = self._open_span(run_id, f"llm:{model_name}", parent_run_id)
+                span.set_attribute("llm.model", str(model_name))
+                try:
+                    span.set_attribute("llm.messages", _to_attr_string(messages))
+                except Exception:
+                    logger.debug(
+                        "agent-trace: failed to record chat model messages "
+                        "for run %r",
+                        str(run_id),
+                        exc_info=True,
+                    )
+
+            def on_llm_end(
+                self,
+                response: Any,
+                *,
+                run_id: uuid.UUID,
+                parent_run_id: uuid.UUID | None = None,
+                **kwargs: Any,
+            ) -> None:
+                self._close_span(run_id, SpanStatus.OK)
+
+            def on_llm_error(
+                self,
+                error: BaseException,
+                *,
+                run_id: uuid.UUID,
+                parent_run_id: uuid.UUID | None = None,
+                **kwargs: Any,
+            ) -> None:
+                self._close_span_with_exception(run_id, error)
+
+            # -- tool callbacks -------------------------------------------
+
+            def on_tool_start(
+                self,
+                serialized: dict[str, Any] | None,
+                input_str: str,
+                *,
+                run_id: uuid.UUID,
+                parent_run_id: uuid.UUID | None = None,
+                **kwargs: Any,
+            ) -> None:
+                ser = serialized or {}
+                tool_name = kwargs.get("name") or ser.get("name") or "tool"
+                span = self._open_span(run_id, f"tool:{tool_name}", parent_run_id)
+                span.set_attribute("tool.name", str(tool_name))
+                if input_str:
+                    span.set_attribute("tool.input", _to_attr_string(input_str))
+
+            def on_tool_end(
+                self,
+                output: Any,
+                *,
+                run_id: uuid.UUID,
+                parent_run_id: uuid.UUID | None = None,
+                **kwargs: Any,
+            ) -> None:
+                run_key = str(run_id)
+                with self._lock:
+                    span = self._spans.get(run_key)
+                if span is not None and output is not None:
+                    try:
+                        span.set_attribute("tool.output", _to_attr_string(output))
+                    except Exception:
+                        logger.debug(
+                            "agent-trace: failed to record tool output for "
+                            "run %r",
+                            run_key,
+                            exc_info=True,
+                        )
+                self._close_span(run_id, SpanStatus.OK)
+
+            def on_tool_error(
+                self,
+                error: BaseException,
+                *,
+                run_id: uuid.UUID,
+                parent_run_id: uuid.UUID | None = None,
+                **kwargs: Any,
+            ) -> None:
+                self._close_span_with_exception(run_id, error)
+
+        _LangChainTracerClass = _LangChainTracerImpl
+        return _LangChainTracerClass
+
+
+def LangChainTracer(tracer: Tracer, trace: Trace) -> Any:  # noqa: N802
+    """Construct a generic LangChain ``BaseCallbackHandler`` — see module
+    docstring for usage. Works with any ``Runnable``, not just a compiled
+    LangGraph graph (that's ``LangGraphTracer``, a separate, LangGraph-
+    specific integration)."""
+    cls = _get_tracer_class()
+    return cls(tracer=tracer, trace=trace)