fiddler-langgraph 0.1.0rc1__py3-none-any.whl

fiddler_langgraph/tracing/instrumentation.py

@@ -0,0 +1,264 @@
+"""LangGraph instrumentation module for Fiddler."""
+
+from collections.abc import Callable, Collection
+from typing import Any, cast
+
+from langchain_core.callbacks import BaseCallbackManager
+from langchain_core.language_models import BaseLanguageModel
+from langchain_core.retrievers import BaseRetriever
+from langchain_core.runnables import RunnableBinding
+from langchain_core.tools import BaseTool
+from opentelemetry.instrumentation.instrumentor import (  # type: ignore[attr-defined]
+    BaseInstrumentor,
+)
+from pydantic import ConfigDict, validate_call
+from wrapt import wrap_function_wrapper
+
+from fiddler_langgraph.core.attributes import (
+    _CONVERSATION_ID,
+    FIDDLER_METADATA_KEY,
+    FiddlerSpanAttributes,
+)
+from fiddler_langgraph.core.client import FiddlerClient
+from fiddler_langgraph.tracing.callback import _CallbackHandler
+from fiddler_langgraph.tracing.util import _check_langgraph_version, _get_package_version
+
+
+@validate_call(config=ConfigDict(strict=True))
+def set_conversation_id(conversation_id: str) -> None:
+    """Set the conversation ID for the current application invocation.
+    This will remain in use until it is called again with a new conversation ID.
+    Note (Robin 11th Sep 2025): This should be moved to the core.attributes module in the future.
+    """
+    _CONVERSATION_ID.set(conversation_id)
+
+
+@validate_call(config=ConfigDict(strict=True, arbitrary_types_allowed=True))
+def _set_default_metadata(
+    node: BaseLanguageModel | BaseRetriever | BaseTool,
+) -> None:
+    """Ensures a node has the default Fiddler metadata dictionary.
+
+    If `node.metadata` does not exist or is not a dictionary, it will be
+    initialized. This function modifies the node in place.
+
+    Args:
+        node (BaseLanguageModel | BaseRetriever | BaseTool): The node to modify.
+    """
+    if not hasattr(node, 'metadata'):
+        node.metadata = {}
+    if not isinstance(node.metadata, dict):
+        node.metadata = {}
+    metadata = node.metadata
+    if FIDDLER_METADATA_KEY not in metadata:
+        metadata[FIDDLER_METADATA_KEY] = {}
+
+
+@validate_call(config=ConfigDict(strict=True, arbitrary_types_allowed=True))
+def add_span_attributes(
+    node: BaseLanguageModel | BaseRetriever | BaseTool,
+    **kwargs: Any,
+) -> None:
+    """Adds Fiddler-specific attributes to a runnable's metadata.
+
+    This is used for various runnable types like LLM calls, tool
+    calls, and retriever calls.
+
+    Args:
+        node (BaseLanguageModel | BaseRetriever | BaseTool): The runnable node.
+        **kwargs: The attributes to add as key-value pairs.
+    """
+    _set_default_metadata(node)
+    metadata = cast(dict[str, Any], node.metadata)
+    fiddler_attrs = cast(dict[str, Any], metadata.get(FIDDLER_METADATA_KEY, {}))
+    for key, value in kwargs.items():
+        fiddler_attrs[key] = value
+
+
+@validate_call(config=ConfigDict(strict=True))
+def set_llm_context(llm: BaseLanguageModel | RunnableBinding, context: str) -> None:
+    """Sets a context string on a language model instance.
+    If the language model is a RunnableBinding, the context will be set on the bound object.
+
+    https://python.langchain.com/api_reference/core/runnables/langchain_core.runnables.base.RunnableBinding.html
+
+    The bound object of the RunnableBinding must be a BaseLanguageModel.
+    This context can be used to provide additional information about the
+    environment or data that the language model is being used in. This
+    information will be attached to the spans created for this model.
+    In case the user passes a RunnableBinding, the context will be set on the
+    bound object.
+
+    Args:
+        llm (BaseLanguageModel | RunnableBinding): The language model instance. **Required**.
+        context (str): The context string to add. **Required**.
+
+    Examples:
+        >>> from langchain_openai import ChatOpenAI
+        >>> from fiddler_langgraph.tracing.instrumentation import set_llm_context
+        >>>
+        >>> llm = ChatOpenAI()
+        >>> set_llm_context(llm, "This is a test context.")
+        >>>
+        >>> # If you are using a RunnableBinding, you can pass the bound object
+        >>> # directly to set_llm_context.
+        >>> bound_llm = llm.bind(x=1)
+        >>> set_llm_context(bound_llm, "This is a test context.")
+    """
+    if isinstance(llm, RunnableBinding):
+        if not isinstance(llm.bound, BaseLanguageModel):
+            raise TypeError(
+                'llm must be a BaseLanguageModel or a RunnableBinding of a BaseLanguageModel'
+            )
+        # RunnableBinding has config attribute (which can store metadata), however these are not passed
+        # to the callback handlers. So we need to use the bound object directly.
+        _llm = llm.bound
+    else:
+        _llm = llm
+
+    _set_default_metadata(_llm)
+
+    if _llm.metadata is None:
+        _llm.metadata = {}
+    fiddler_attrs = cast(dict[str, Any], _llm.metadata.get(FIDDLER_METADATA_KEY, {}))
+    fiddler_attrs[FiddlerSpanAttributes.LLM_CONTEXT] = context
+
+
+class LangGraphInstrumentor(BaseInstrumentor):
+    """An OpenTelemetry instrumentor for LangGraph applications.
+
+    This class provides automatic instrumentation for applications built with
+    LangGraph. It captures traces from the execution of LangGraph graphs and
+    sends them to the Fiddler platform.
+
+    To use the instrumentor, you first need to create a `FiddlerClient`
+    instance. Then, you can create an instance of `LangGraphInstrumentor` and
+    call the `instrument()` method.
+
+    Examples:
+        >>> from fiddler_langgraph import FiddlerClient
+        >>> from fiddler_langgraph.tracing import LangGraphInstrumentor
+        >>>
+        >>> client = FiddlerClient(api_key="...", application_id="...")
+        >>> instrumentor = LangGraphInstrumentor(client=client)
+        >>> instrumentor.instrument()
+
+    Attributes:
+        _client (FiddlerClient): The FiddlerClient instance used for configuration.
+    """
+
+    def __init__(self, client: FiddlerClient):
+        """Initializes the LangGraphInstrumentor.
+
+        Args:
+            client (FiddlerClient): The `FiddlerClient` instance. **Required**.
+
+        Raises:
+            ImportError: If LangGraph version is incompatible or not installed.
+        """
+        super().__init__()
+        self._client = client
+        self._langgraph_version = _get_package_version('langgraph')
+        self._langchain_version = _get_package_version('langchain_core')
+        self._fiddler_langgraph_version = _get_package_version('fiddler_langgraph')
+
+        self._client.update_resource(
+            {
+                'lib.langgraph.version': self._langgraph_version.public,
+                'lib.langchain_core.version': self._langchain_version.public,
+                'lib.fiddler-langgraph.version': self._fiddler_langgraph_version.public,
+            }
+        )
+        self._tracer: _CallbackHandler | None = None
+        self._original_callback_manager_init: Callable[..., None] | None = None
+
+        # Check LangGraph version compatibility - we don't add this to dependencies
+        # because we leave it to the user to install the correct version of LangGraph
+        # We will check if the user installed version is compatible with the version of fiddler-langgraph
+        _check_langgraph_version(self._langgraph_version)
+
+    def instrumentation_dependencies(self) -> Collection[str]:
+        """Returns the package dependencies required for this instrumentor.
+
+        Returns:
+            Collection[str]: A collection of package dependency strings.
+        """
+        return ('langchain_core >= 0.1.0',)
+
+    def _instrument(self, **kwargs: Any) -> None:
+        """Instruments LangGraph by monkey-patching `BaseCallbackManager`.
+
+        This method injects a custom callback handler into LangGraph's callback
+        system to capture trace data. This is done by wrapping the `__init__`
+        method of `BaseCallbackManager` to inject a `_CallbackHandler`.
+
+        Raises:
+            ValueError: If the tracer is not initialized in the FiddlerClient.
+        """
+        import langchain_core
+
+        tracer = self._client.get_tracer()
+        if tracer is None:
+            raise ValueError('Context tracer is not initialized')
+
+        self._tracer = _CallbackHandler(tracer)
+        self._original_callback_manager_init = langchain_core.callbacks.BaseCallbackManager.__init__
+        wrap_function_wrapper(
+            module='langchain_core.callbacks',
+            name='BaseCallbackManager.__init__',
+            wrapper=_BaseCallbackManagerInit(self._tracer),
+        )
+
+    def _uninstrument(self, **kwargs: Any) -> None:
+        """Removes the instrumentation from LangGraph.
+
+        This is done by restoring the original `__init__` method on the
+        `BaseCallbackManager` class.
+        """
+        import langchain_core
+
+        if self._original_callback_manager_init is not None:
+            setattr(  # noqa: B010
+                langchain_core.callbacks.BaseCallbackManager,
+                '__init__',
+                self._original_callback_manager_init,
+            )
+        self._original_callback_manager_init = None
+        self._tracer = None
+
+
+class _BaseCallbackManagerInit:
+    """A wrapper class for `BaseCallbackManager.__init__` to inject Fiddler's callback handler."""
+
+    __slots__ = ('_callback_handler',)
+
+    def __init__(self, callback_handler: _CallbackHandler):
+        """Initializes the wrapper.
+
+        Args:
+            callback_handler (_CallbackHandler): The Fiddler callback handler instance
+                to be injected into the callback manager.
+        """
+        self._callback_handler = callback_handler
+
+    def __call__(
+        self,
+        wrapped: Callable[..., None],
+        instance: 'BaseCallbackManager',
+        args: Any,
+        kwargs: Any,
+    ) -> None:
+        """Calls the original `__init__` and then adds the Fiddler handler.
+
+        It also ensures that the handler is not added multiple times if it
+        already exists in the list of inheritable handlers.
+        """
+        wrapped(*args, **kwargs)
+        for handler in instance.inheritable_handlers:
+            # Handlers may be copied when new managers are created, so we
+            # don't want to keep adding. E.g. see the following location.
+            # https://github.com/langchain-ai/langchain/blob/5c2538b9f7fb64afed2a918b621d9d8681c7ae32/libs/core/langchain_core/callbacks/manager.py#L1876
+            if isinstance(handler, type(self._callback_handler)):
+                break
+        else:
+            instance.add_handler(self._callback_handler, True)

fiddler_langgraph/tracing/jsonl_capture.py

@@ -0,0 +1,185 @@
+"""JSONL data capture module for simplified span data in structured format."""
+
+import json
+import os
+import threading
+from collections.abc import Sequence
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
+
+from fiddler_langgraph.core.attributes import FiddlerSpanAttributes, SpanType
+
+
+class JSONLSpanCapture:
+    """Captures OpenTelemetry span data and saves it to JSONL format with structured fields."""
+
+    def __init__(self, jsonl_file_path: str | None = None):
+        """Initialize JSONL capture.
+
+        Args:
+            jsonl_file_path: Path to the JSONL file. If None, uses FIDDLER_JSONL_FILE env var (default: 'fiddler_trace_data.jsonl')
+        """
+        if jsonl_file_path is None:
+            jsonl_file_path = os.getenv('FIDDLER_JSONL_FILE', 'fiddler_trace_data.jsonl')
+
+        self.jsonl_file_path = Path(jsonl_file_path)
+        self._lock = threading.Lock()
+        self._ensure_jsonl_file()
+
+    def _ensure_jsonl_file(self) -> None:
+        """Ensure the JSONL file exists and has proper headers."""
+        try:
+            if not self.jsonl_file_path.exists():
+                self.jsonl_file_path.parent.mkdir(parents=True, exist_ok=True)
+                self.jsonl_file_path.touch()
+        except Exception as e:
+            print(f'Warning: Could not create JSONL file {self.jsonl_file_path}: {e}')
+
+    def capture_span(self, span: ReadableSpan) -> None:
+        """Capture a span and write it to JSONL file."""
+        try:
+            span_data = self._convert_span_to_structured_format(span)
+            self._write_span_to_jsonl(span_data)
+        except Exception as e:
+            print(f'Error capturing span to JSONL: {e}')
+
+    def _convert_span_to_structured_format(self, span: ReadableSpan) -> dict[str, Any]:
+        """Convert ReadableSpan to structured format for JSONL export."""
+        # Extract basic span information
+        span_data = {
+            'trace_id': format(span.get_span_context().trace_id, '032x'),
+            'span_id': format(span.get_span_context().span_id, '016x'),
+            'parent_span_id': format(span.parent.span_id, '016x') if span.parent else '',
+            'root_span_id': format(
+                span.get_span_context().trace_id, '032x'
+            ),  # Use trace_id as root_span_id
+            'span_name': span.name,
+            'span_kind': span.kind.name if span.kind else 'CLIENT',
+            'start_time': (
+                datetime.fromtimestamp(span.start_time / 1_000_000_000, tz=timezone.utc).isoformat()
+                if span.start_time is not None
+                else ''
+            ),
+            'end_time': (
+                datetime.fromtimestamp(span.end_time / 1_000_000_000, tz=timezone.utc).isoformat()
+                if span.end_time is not None
+                else ''
+            ),
+            'duration_ms': (
+                int((span.end_time - span.start_time) / 1_000_000)
+                if span.end_time is not None and span.start_time is not None
+                else 0
+            ),
+            'status_code': span.status.status_code.name if span.status else 'OK',
+            'status_message': (
+                span.status.description if span.status and span.status.description else ''
+            ),
+        }
+
+        # Extract attributes and map them to structured fields
+        attributes = dict(span.attributes) if span.attributes else {}
+
+        # Span type and agent info
+        span_data['span_type'] = attributes.get(FiddlerSpanAttributes.TYPE, SpanType.OTHER)
+        span_data['agent_name'] = attributes.get(FiddlerSpanAttributes.AGENT_NAME, '')
+        span_data['agent_id'] = attributes.get(FiddlerSpanAttributes.AGENT_ID, '')
+        span_data['conversation_id'] = attributes.get(FiddlerSpanAttributes.CONVERSATION_ID, '')
+
+        # Model information
+        span_data['model_name'] = attributes.get(FiddlerSpanAttributes.LLM_REQUEST_MODEL, '')
+        span_data['model_provider'] = attributes.get(FiddlerSpanAttributes.LLM_SYSTEM, '')
+
+        # LLM inputs/outputs
+        span_data['llm_input_system'] = attributes.get(FiddlerSpanAttributes.LLM_INPUT_SYSTEM, '')
+        span_data['llm_input_user'] = attributes.get(FiddlerSpanAttributes.LLM_INPUT_USER, '')
+        span_data['llm_output'] = attributes.get(FiddlerSpanAttributes.LLM_OUTPUT, '')
+        span_data['llm_context'] = attributes.get(FiddlerSpanAttributes.LLM_CONTEXT, '')
+
+        # Tool information
+        span_data['tool_name'] = attributes.get(FiddlerSpanAttributes.TOOL_NAME, '')
+        span_data['tool_input'] = attributes.get(FiddlerSpanAttributes.TOOL_INPUT, '')
+        span_data['tool_output'] = attributes.get(FiddlerSpanAttributes.TOOL_OUTPUT, '')
+
+        # Library versions (from resource if available)
+        resource_attributes = (
+            dict(span.resource.attributes) if span.resource and span.resource.attributes else {}
+        )
+        span_data['service_name'] = resource_attributes.get('service.name', '')
+        span_data['service_version'] = resource_attributes.get('service.version', '')
+        span_data['telemetry_sdk_name'] = resource_attributes.get('telemetry.sdk.name', '')
+        span_data['telemetry_sdk_version'] = resource_attributes.get('telemetry.sdk.version', '')
+        span_data['application_id'] = resource_attributes.get('application.id', '')
+
+        # Custom metadata and tags
+        custom_attributes = {}
+        for key, value in attributes.items():
+            if not key.startswith(('gen_ai.', 'fiddler.', 'service.', 'telemetry.')):
+                custom_attributes[key] = value
+
+        span_data['custom_attributes'] = json.dumps(custom_attributes) if custom_attributes else ''
+
+        # Exception information
+        exception_info = []
+        if hasattr(span, 'events') and span.events:
+            for event in span.events:
+                if event.name == 'exception':
+                    event_attrs = dict(event.attributes) if event.attributes else {}
+                    exception_info.append(
+                        {
+                            'type': event_attrs.get('exception.type', ''),
+                            'message': event_attrs.get('exception.message', ''),
+                            'stacktrace': event_attrs.get('exception.stacktrace', ''),
+                        }
+                    )
+
+        span_data['exception_info'] = json.dumps(exception_info) if exception_info else ''
+
+        return span_data
+
+    def _write_span_to_jsonl(self, span_data: dict[str, Any]) -> None:
+        """Write span data to JSONL file."""
+        with self._lock:
+            try:
+                with self.jsonl_file_path.open('a', encoding='utf-8') as f:
+                    json.dump(span_data, f, ensure_ascii=False)
+                    f.write('\n')
+            except Exception as e:
+                print(f'Error writing to JSONL file {self.jsonl_file_path}: {e}')
+
+
+class JSONLSpanExporter(SpanExporter):
+    """SpanExporter that captures spans using JSONLSpanCapture."""
+
+    def __init__(self, jsonl_capture: JSONLSpanCapture):
+        """Initialize the exporter.
+
+        Args:
+            jsonl_capture: The JSONLSpanCapture instance to use for capturing spans
+        """
+        self.jsonl_capture = jsonl_capture
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        """Export spans by capturing them with JSONLSpanCapture."""
+        try:
+            for span in spans:
+                self.jsonl_capture.capture_span(span)
+            return SpanExportResult.SUCCESS
+        except Exception as e:
+            print(f'Error exporting spans to JSONL: {e}')
+            return SpanExportResult.FAILURE
+
+
+def initialize_jsonl_capture(jsonl_file_path: str | None = None) -> JSONLSpanCapture:
+    """Initialize a JSONLSpanCapture instance.
+
+    Args:
+        jsonl_file_path: Path to the JSONL file. If None, uses FIDDLER_JSONL_FILE env var
+
+    Returns:
+        JSONLSpanCapture: The initialized capture instance
+    """
+    return JSONLSpanCapture(jsonl_file_path=jsonl_file_path)

fiddler_langgraph/tracing/util.py

@@ -0,0 +1,83 @@
+import datetime
+import importlib.metadata
+import json
+import logging
+from dataclasses import asdict, is_dataclass
+from typing import Any
+
+from packaging import version as pkg_version
+from pydantic import BaseModel
+
+logger = logging.getLogger(__name__)
+
+
+class _LanggraphJSONEncoder(json.JSONEncoder):
+    """A custom JSON encoder for LangGraph objects.
+
+    This encoder handles the serialization of common LangGraph and Pydantic
+    objects into JSON-serializable formats.
+    """
+
+    # pylint: disable=too-many-return-statements
+    def default(self, o: Any) -> Any:
+        """Serializes an object to a JSON-compatible format.
+
+        This method provides custom serialization for the following types:
+        - Dataclasses
+        - Objects with a `to_json` method
+        - Pydantic models
+        - Datetime objects
+
+        If an object cannot be serialized, it is converted to an empty string.
+
+        Args:
+            o: The object to serialize.
+
+        Returns:
+            A JSON-serializable representation of the object.
+        """
+        if is_dataclass(o):
+            return asdict(o)  # type: ignore[arg-type]
+
+        if hasattr(o, 'to_json'):
+            return o.to_json()
+
+        if isinstance(o, BaseModel) and hasattr(o, 'model_dump_json'):
+            return o.model_dump_json()
+
+        if isinstance(o, datetime.datetime):
+            return o.isoformat()
+
+        try:
+            return str(o)
+        except (TypeError, ValueError) as e:
+            logger.debug('Failed to serialize object of type %s: %s', type(o).__name__, str(e))
+            return ''
+
+
+def _get_package_version(package_name: str) -> pkg_version.Version:
+    """Get the version of a package."""
+    try:
+        version = importlib.metadata.version(package_name)
+        return pkg_version.parse(version)
+    except importlib.metadata.PackageNotFoundError:
+        raise ImportError(f'Package {package_name} is not installed')
+
+
+def _check_langgraph_version(
+    langgraph_version: pkg_version.Version,
+) -> None:
+    """Check if the installed LangGraph version is compatible with the version of fiddler-langgraph."""
+
+    if langgraph_version is None:
+        raise ImportError('Either langgraph or langchain_core should be installed')
+
+    # check compatibility range
+    min_langgraph_version = pkg_version.parse('0.3.28')
+    max_langgraph_version = pkg_version.parse('1.1.0')
+
+    if langgraph_version < min_langgraph_version or langgraph_version > max_langgraph_version:
+        raise ImportError(
+            f'langgraph version {langgraph_version.public} is not compatible. '
+            f'fiddler-langgraph requires langgraph >= 0.3.28 and < 1.1.0. '
+        )