gatewaysdk 0.3.2__py3-none-any.whl

gatewaysdk/tracing/span_builder.py

@@ -0,0 +1,356 @@
+# Copyright (c) Gateway. All rights reserved.
+
+"""OTel span builder for JSON push-to-trace.
+
+Creates OpenTelemetry spans from resolved mapping data. Isolated from
+the mapping engine -- takes already-validated, already-resolved field
+dicts and produces real spans with proper parent-child linking.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import time
+from typing import Any, Dict, List, Optional
+
+from opentelemetry import context as otel_context
+from opentelemetry import trace as otel_trace
+
+from . import semconv
+from .mapping import ChildMapping, coerce_timestamp, resolve_path
+from .push_models import SpanPreview, ValidationIssue
+
+logger = logging.getLogger(__name__)
+
+# Default timestamp: "now" in nanoseconds
+_NANOS_PER_SECOND = 1_000_000_000
+
+
+def _now_ns() -> int:
+    return int(time.time() * _NANOS_PER_SECOND)
+
+
+def _safe_json(value: Any) -> str:
+    """Serialize a value to JSON string for span attributes."""
+    if isinstance(value, str):
+        return value
+    try:
+        return json.dumps(value)
+    except (TypeError, ValueError):
+        logger.warning(
+            "Value of type %s is not JSON-serializable, using str() fallback: %s",
+            type(value).__name__,
+            repr(value)[:100],
+        )
+        return str(value)
+
+
+# =============================================================================
+# Span Building
+# =============================================================================
+
+
+def build_trace_spans(
+    data: dict,
+    fields: Dict[str, Any],
+    child_mapping: Optional[ChildMapping],
+    tracer: Optional[otel_trace.Tracer] = None,
+    extra_attributes: Optional[Dict[str, Any]] = None,
+    timestamp_unit: str = "auto",
+) -> tuple[str, int, List[ValidationIssue]]:
+    """Create OTel spans for a trace and its children.
+
+    Args:
+        data: The original JSON data dict (for resolving child fields).
+        fields: Merged trace-level fields (name, input, output, etc.).
+        child_mapping: Optional mapping for children observations.
+        tracer: OTel tracer to use. Defaults to the global gatewaysdk tracer.
+        extra_attributes: Additional OTel span attributes to set on the root span.
+
+    Returns:
+        (trace_id, observation_count, warnings) -- the hex trace ID, number of
+        child observations created, and any runtime warnings.
+    """
+    if tracer is None:
+        tracer = otel_trace.get_tracer("gatewaysdk.tracing.push")
+
+    start_ns = coerce_timestamp(fields.get("timestamp"), unit=timestamp_unit) or _now_ns()
+
+    # Build root span attributes
+    root_attrs = _build_trace_attributes(fields)
+    if extra_attributes:
+        root_attrs.update(extra_attributes)
+
+    root_span = tracer.start_span(
+        name=fields.get("name") or "push",
+        start_time=start_ns,
+        attributes=root_attrs,
+    )
+
+    observation_count = 0
+    runtime_warnings: List[ValidationIssue] = []
+
+    # Build children
+    if child_mapping is not None:
+        found, children_data = resolve_path(data, child_mapping.path)
+        if found and isinstance(children_data, list):
+            root_ctx = otel_trace.set_span_in_context(root_span)
+            observation_count, child_warnings = _build_children(
+                children_data,
+                child_mapping,
+                tracer,
+                root_ctx,
+                start_ns,
+                timestamp_unit=timestamp_unit,
+            )
+            runtime_warnings.extend(child_warnings)
+
+    # End root span
+    end_ns = _now_ns()
+    root_span.end(end_time=end_ns)
+
+    # Extract trace ID
+    span_context = root_span.get_span_context()
+    trace_id = format(span_context.trace_id, "032x")
+
+    return trace_id, observation_count, runtime_warnings
+
+
+def _build_trace_attributes(fields: Dict[str, Any]) -> Dict[str, Any]:
+    """Build OTel span attributes for the root trace span."""
+    attrs: Dict[str, Any] = {}
+
+    if fields.get("name"):
+        attrs[semconv.TRACE_NAME] = fields["name"]
+    if fields.get("input") is not None:
+        attrs[semconv.TRACE_INPUT] = _safe_json(fields["input"])
+    if fields.get("output") is not None:
+        attrs[semconv.TRACE_OUTPUT] = _safe_json(fields["output"])
+    if fields.get("metadata") is not None:
+        attrs[semconv.TRACE_METADATA] = _safe_json(fields["metadata"])
+    if fields.get("tags"):
+        attrs[semconv.TRACE_TAGS] = _safe_json(fields["tags"])
+    if fields.get("user_id"):
+        attrs[semconv.USER_ID] = fields["user_id"]
+    if fields.get("session_id"):
+        attrs[semconv.SESSION_ID] = fields["session_id"]
+    if fields.get("run_id"):
+        attrs[semconv.RUN_ID] = fields["run_id"]
+    if fields.get("environment"):
+        attrs[semconv.ENVIRONMENT] = fields["environment"]
+
+    return attrs
+
+
+def _build_children(
+    children_data: List[dict],
+    child_mapping: ChildMapping,
+    tracer: otel_trace.Tracer,
+    parent_ctx: otel_context.Context,
+    parent_start_ns: int,
+    timestamp_unit: str = "auto",
+) -> tuple[int, List[ValidationIssue]]:
+    """Create child observation spans under the parent context.
+
+    Returns (observation_count, warnings).
+    """
+    count = 0
+    warnings: List[ValidationIssue] = []
+
+    for i, child_dict in enumerate(children_data):
+        if not isinstance(child_dict, dict):
+            warnings.append(
+                ValidationIssue(
+                    field="children",
+                    path=f"{child_mapping.path}[{i}]",
+                    issue="invalid_child",
+                    message=f"Children item at index {i} is {type(child_dict).__name__}, expected dict. Skipped.",
+                )
+            )
+            continue
+
+        child_fields, field_warnings = _extract_child_fields(child_dict, child_mapping, timestamp_unit=timestamp_unit)
+        warnings.extend(field_warnings)
+
+        child_start = child_fields.get("start_time_ns") or (parent_start_ns + i)
+        child_end = child_fields.get("end_time_ns")
+
+        child_attrs = _build_observation_attributes(child_fields)
+
+        child_span = tracer.start_span(
+            name=child_fields.get("name") or f"observation-{i}",
+            context=parent_ctx,
+            start_time=child_start,
+            attributes=child_attrs,
+        )
+
+        # Recurse into nested children if the mapping defines them
+        nested_count = 0
+        if child_mapping.children is not None:
+            found, nested_data = resolve_path(child_dict, child_mapping.children.path)
+            if found and isinstance(nested_data, list):
+                child_ctx = otel_trace.set_span_in_context(child_span)
+                nested_count, nested_warnings = _build_children(
+                    nested_data,
+                    child_mapping.children,
+                    tracer,
+                    child_ctx,
+                    child_start,
+                    timestamp_unit=timestamp_unit,
+                )
+                warnings.extend(nested_warnings)
+
+        child_span.end(end_time=child_end or _now_ns())
+        count += 1 + nested_count
+
+    return count, warnings
+
+
+def _extract_child_fields(
+    child_dict: dict,
+    mapping: ChildMapping,
+    timestamp_unit: str = "auto",
+) -> tuple[Dict[str, Any], List[ValidationIssue]]:
+    """Extract mapped fields from a single child dict.
+
+    Returns (fields, warnings).
+    """
+    fields: Dict[str, Any] = {}
+    warnings: List[ValidationIssue] = []
+
+    simple_fields = [
+        "name",
+        "input",
+        "output",
+        "type",
+        "model",
+        "metadata",
+        "level",
+        "usage",
+        "cost",
+        "model_parameters",
+        "status_message",
+        "prompt_name",
+        "prompt_version",
+        "attributes",
+    ]
+    for field in simple_fields:
+        dot_path = getattr(mapping, field, None)
+        if dot_path is not None:
+            found, value = resolve_path(child_dict, dot_path)
+            if found:
+                fields[field] = value
+
+    # Timestamps need coercion
+    for ts_field in ("start_time", "end_time"):
+        dot_path = getattr(mapping, ts_field, None)
+        if dot_path is not None:
+            found, value = resolve_path(child_dict, dot_path)
+            if found:
+                coerced = coerce_timestamp(value, unit=timestamp_unit)
+                if coerced is None and value is not None:
+                    warnings.append(
+                        ValidationIssue(
+                            field=f"children.{ts_field}",
+                            path=dot_path,
+                            issue="timestamp_coercion_failed",
+                            message=f"Could not parse timestamp value: {value!r}. Using fallback.",
+                        )
+                    )
+                fields[f"{ts_field}_ns"] = coerced
+
+    return fields, warnings
+
+
+def _build_observation_attributes(fields: Dict[str, Any]) -> Dict[str, Any]:
+    """Build OTel span attributes for an observation span."""
+    attrs: Dict[str, Any] = {}
+
+    if fields.get("input") is not None:
+        attrs[semconv.OBSERVATION_INPUT] = _safe_json(fields["input"])
+    if fields.get("output") is not None:
+        attrs[semconv.OBSERVATION_OUTPUT] = _safe_json(fields["output"])
+    if fields.get("type"):
+        attrs[semconv.OBSERVATION_TYPE] = fields["type"].upper()
+    if fields.get("model"):
+        attrs[semconv.MODEL_NAME] = fields["model"]
+        # Presence of model implies GENERATION type (unless explicitly set)
+        if not fields.get("type"):
+            attrs[semconv.OBSERVATION_TYPE] = "GENERATION"
+    if fields.get("metadata") is not None:
+        attrs[semconv.OBSERVATION_METADATA] = _safe_json(fields["metadata"])
+    if fields.get("level"):
+        attrs[semconv.OBSERVATION_LEVEL] = fields["level"].upper()
+    if fields.get("status_message"):
+        attrs[semconv.OBSERVATION_STATUS_MESSAGE] = fields["status_message"]
+    if fields.get("model_parameters") is not None:
+        attrs[semconv.MODEL_PARAMETERS] = _safe_json(fields["model_parameters"])
+
+    # Usage and cost use the langfuse-compatible attributes
+    if fields.get("usage") is not None:
+        attrs["gateway.observation.usage_details"] = _safe_json(fields["usage"])
+    if fields.get("cost") is not None:
+        attrs["gateway.observation.cost_details"] = _safe_json(fields["cost"])
+
+    # Prompt reference
+    if fields.get("prompt_name"):
+        attrs["gateway.observation.prompt.name"] = fields["prompt_name"]
+    if fields.get("prompt_version") is not None:
+        attrs["gateway.observation.prompt.version"] = fields["prompt_version"]
+
+    # Extra attributes escape hatch
+    extra = fields.get("attributes")
+    if isinstance(extra, dict):
+        attrs.update(extra)
+
+    return attrs
+
+
+# =============================================================================
+# Preview Building (no OTel dependency)
+# =============================================================================
+
+
+def build_trace_preview(
+    data: dict,
+    fields: Dict[str, Any],
+    child_mapping: Optional[ChildMapping],
+) -> List[SpanPreview]:
+    """Build SpanPreview objects for children (no OTel spans created).
+
+    Returns list of SpanPreview for the children.
+    """
+    if child_mapping is None:
+        return []
+
+    found, children_data = resolve_path(data, child_mapping.path)
+    if not found or not isinstance(children_data, list):
+        return []
+
+    previews = []
+    for i, child_dict in enumerate(children_data):
+        if not isinstance(child_dict, dict):
+            continue
+
+        child_fields, _ = _extract_child_fields(child_dict, child_mapping)
+
+        nested = []
+        if child_mapping.children is not None:
+            nested = build_trace_preview(child_dict, child_fields, child_mapping.children)
+
+        previews.append(
+            SpanPreview(
+                name=child_fields.get("name") or f"observation-{i}",
+                type=child_fields.get("type"),
+                input=child_fields.get("input"),
+                output=child_fields.get("output"),
+                model=child_fields.get("model"),
+                metadata=child_fields.get("metadata"),
+                start_time=str(child_fields.get("start_time_ns")) if child_fields.get("start_time_ns") else None,
+                end_time=str(child_fields.get("end_time_ns")) if child_fields.get("end_time_ns") else None,
+                children=nested,
+            )
+        )
+
+    return previews

gatewaysdk/trainer/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from .init_utils import build_component
+from .trainer import Trainer
+
+__all__ = ["Trainer", "build_component"]

gatewaysdk/trainer/init_utils.py

@@ -0,0 +1,263 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Utility helpers for dynamic component initialization within the trainer."""
+
+from __future__ import annotations
+
+import importlib
+import inspect
+from typing import Any, Callable, Dict, Optional, TypeVar, Union, cast, overload
+
+OptionalDefaults = Dict[str, Callable[[], Any] | Any]
+T = TypeVar("T")
+
+
+def load_class(path: str) -> type[Any]:
+    """Load a class from its fully qualified import path."""
+    module_name, class_name = path.rsplit(".", 1)
+    module = importlib.import_module(module_name)
+    return getattr(module, class_name)
+
+
+def instantiate_component(
+    cls: type[Any],
+    provided_kwargs: Optional[Dict[str, Any]] = None,
+    optional_defaults: Optional[OptionalDefaults] = None,
+) -> Any:
+    """Instantiate `cls`, filling optional kwargs when the constructor accepts them."""
+    kwargs = dict(provided_kwargs or {})
+    if optional_defaults:
+        signature = inspect.signature(cls.__init__)
+        for name, value in optional_defaults.items():
+            if name in kwargs or name not in signature.parameters:
+                continue
+            kwargs[name] = value() if callable(value) else value
+    return cls(**kwargs)
+
+
+def instantiate_from_spec(
+    spec: Union[str, Dict[str, Any]],
+    *,
+    spec_name: str,
+    optional_defaults: Optional[OptionalDefaults] = None,
+    dict_requires_type: bool = True,
+    dict_default_cls: type[Any] | None = None,
+    registry: Optional[Dict[str, str]] = None,
+) -> Any:
+    """Instantiate a component from a string or dict spec."""
+    if isinstance(spec, str):
+        type_path = registry.get(spec, spec) if registry else spec
+        cls = load_class(type_path)
+        return instantiate_component(cls, optional_defaults=optional_defaults)
+
+    if isinstance(spec, dict):  # pyright: ignore[reportUnnecessaryIsInstance]
+        spec_conf = dict(spec)
+        type_path = spec_conf.pop("type", None)
+        if type_path is None and registry and "name" in spec_conf:
+            type_path = registry.get(spec_conf.pop("name"))
+        elif registry and type_path is not None:
+            type_path = registry.get(type_path, type_path)
+        if type_path is None:
+            if dict_requires_type:
+                raise ValueError(f"{spec_name} dict must have a 'type' key with the class full name")
+            if dict_default_cls is None:
+                raise ValueError(f"{spec_name} dict missing 'type' and no default class provided")
+            cls = dict_default_cls
+        else:
+            cls = load_class(type_path)
+        return instantiate_component(cls, spec_conf, optional_defaults)
+
+    raise TypeError(f"{spec_name} spec must be a string or dict (got {type(spec)}).")
+
+
+def _ensure_expected_type(
+    instance: Any,
+    expected_type: type[T],
+    spec_name: str,
+    type_error_fmt: str | None,
+) -> T:
+    if not isinstance(instance, expected_type):
+        type_name = str(type(instance))  # type: ignore
+        if type_error_fmt:
+            raise TypeError(type_error_fmt.format(type_name=type_name, expected_type=expected_type.__name__))
+        raise TypeError(f"{spec_name} factory returned {type_name}, which is not a {expected_type.__name__} subclass.")
+    return instance
+
+
+@overload
+def build_component(
+    spec: Union[T, str, Dict[str, Any], type[T], Callable[[], T], None],
+    *,
+    expected_type: type[T],
+    spec_name: str,
+    default_factory: Callable[[], T],
+    allow_none: bool = ...,
+    optional_defaults: Optional[OptionalDefaults] = ...,
+    dict_requires_type: bool = ...,
+    dict_default_cls: type[T] | None = ...,
+    type_error_fmt: str | None = ...,
+    invalid_spec_error_fmt: str | None = ...,
+    registry: Optional[Dict[str, str]] = ...,
+) -> T: ...
+
+
+@overload
+def build_component(
+    spec: Union[T, str, Dict[str, Any], type[T], Callable[[], T], None],
+    *,
+    expected_type: type[T],
+    spec_name: str,
+    default_factory: None = ...,
+    allow_none: bool,
+    optional_defaults: Optional[OptionalDefaults] = ...,
+    dict_requires_type: bool = ...,
+    dict_default_cls: type[T] | None = ...,
+    type_error_fmt: str | None = ...,
+    invalid_spec_error_fmt: str | None = ...,
+    registry: Optional[Dict[str, str]] = ...,
+) -> T | None: ...
+
+
+@overload
+def build_component(
+    spec: Union[T, str, Dict[str, Any], type[T], Callable[[], T], None],
+    *,
+    expected_type: type[T],
+    spec_name: str,
+    default_factory: None = ...,
+    allow_none: bool = ...,
+    optional_defaults: Optional[OptionalDefaults] = ...,
+    dict_requires_type: bool = ...,
+    dict_default_cls: type[T] | None = ...,
+    type_error_fmt: str | None = ...,
+    invalid_spec_error_fmt: str | None = ...,
+    registry: Optional[Dict[str, str]] = ...,
+) -> T | None: ...
+
+
+def build_component(
+    spec: Union[T, str, Dict[str, Any], type[T], Callable[[], T], None],
+    *,
+    expected_type: type[T],
+    spec_name: str,
+    default_factory: Callable[[], T] | None = None,
+    allow_none: bool = False,
+    optional_defaults: Optional[OptionalDefaults] = None,
+    dict_requires_type: bool = True,
+    dict_default_cls: type[T] | None = None,
+    type_error_fmt: str | None = None,
+    invalid_spec_error_fmt: str | None = None,
+    registry: Optional[Dict[str, str]] = None,
+) -> T | None:
+    """Build and return a component instance from a flexible specification.
+
+    This function provides a flexible way to create component instances from various
+    input formats including direct instances, class types, factory functions, import
+    paths, or configuration dictionaries.
+
+    Args:
+        spec: The component specification. Can be:
+            - An instance of expected_type (returned as-is)
+            - A string import path (e.g., 'module.Class') or registry key
+            - A dict with 'type' key (import path or registry key) and constructor kwargs
+            - A class type (will be instantiated)
+            - A factory function (will be called)
+            - None (uses default_factory or returns None if allow_none=True)
+        expected_type: The type that the resulting instance must be or inherit from.
+        spec_name: Descriptive name for the spec, used in error messages.
+        default_factory: Optional factory function called when spec is None.
+        allow_none: If True, allows None to be returned when spec is None and
+            no default_factory is provided.
+        optional_defaults: Dict mapping parameter names to default values or factory
+            functions that will be injected if the constructor accepts them.
+        dict_requires_type: If True, dict specs must include a 'type' key.
+        dict_default_cls: Default class to use for dict specs without a 'type' key
+            (only used when dict_requires_type=False).
+        type_error_fmt: Custom format string for type validation errors. Should include
+            {type_name} and {expected_type} placeholders.
+        invalid_spec_error_fmt: Custom format string for invalid spec type errors.
+            Should include {actual_type} and {expected_type} placeholders.
+        registry: Optional mapping of short names to fully qualified import paths.
+            When provided, string specs or dict 'type'/'name' entries are first
+            resolved through this registry before attempting to import.
+
+    Returns:
+        An instance of expected_type, or None if allow_none=True and spec is None
+        without a default_factory.
+
+    Raises:
+        TypeError: If the instantiated object is not an instance of expected_type.
+        ValueError: If spec is None and neither default_factory nor allow_none is set,
+            or if spec type is invalid, or if dict spec is invalid.
+
+    Examples:
+        >>> # Direct instance
+        >>> optimizer = build_component(AdamW(), expected_type=Optimizer, spec_name='optimizer')
+        >>>
+        >>> # String import path
+        >>> optimizer = build_component('torch.optim.AdamW', expected_type=Optimizer, spec_name='optimizer')
+        >>>
+        >>> # Dict with type and kwargs
+        >>> spec = {'type': 'torch.optim.AdamW', 'lr': 0.001}
+        >>> optimizer = build_component(spec, expected_type=Optimizer, spec_name='optimizer')
+        >>>
+        >>> # Class type
+        >>> optimizer = build_component(AdamW, expected_type=Optimizer, spec_name='optimizer')
+        >>>
+        >>> # Factory function
+        >>> optimizer = build_component(lambda: AdamW(lr=0.001), expected_type=Optimizer,
+        ...                            spec_name='optimizer')
+    """
+    if isinstance(spec, expected_type):
+        return cast(T, spec)
+
+    if spec is None:
+        if default_factory is not None:
+            instance = default_factory()
+            return _ensure_expected_type(instance, expected_type, spec_name, type_error_fmt)
+        if allow_none:
+            return None
+        raise ValueError(
+            invalid_spec_error_fmt.format(actual_type=type(spec), expected_type=expected_type.__name__)
+            if invalid_spec_error_fmt
+            else f"{spec_name} cannot be None."
+        )
+
+    if isinstance(spec, type) and issubclass(spec, expected_type):
+        instance = instantiate_component(spec, optional_defaults=optional_defaults)
+        return _ensure_expected_type(instance, expected_type, spec_name, type_error_fmt)
+
+    if callable(spec) and not isinstance(spec, type):  # type: ignore
+        instance = spec()
+        return _ensure_expected_type(instance, expected_type, spec_name, type_error_fmt)
+
+    if isinstance(spec, str):
+        instance = instantiate_from_spec(
+            spec,
+            spec_name=spec_name,
+            optional_defaults=optional_defaults,
+            dict_requires_type=dict_requires_type,
+            dict_default_cls=dict_default_cls,
+            registry=registry,
+        )
+        return _ensure_expected_type(instance, expected_type, spec_name, type_error_fmt)
+
+    if isinstance(spec, dict):
+        instance = instantiate_from_spec(
+            spec,  # type: ignore
+            spec_name=spec_name,
+            optional_defaults=optional_defaults,
+            dict_requires_type=dict_requires_type,
+            dict_default_cls=dict_default_cls,
+            registry=registry,
+        )
+        return _ensure_expected_type(instance, expected_type, spec_name, type_error_fmt)
+
+    if invalid_spec_error_fmt:
+        raise ValueError(invalid_spec_error_fmt.format(actual_type=type(spec), expected_type=expected_type.__name__))  # type: ignore
+
+    type_name = str(type(spec))  # type: ignore
+    raise ValueError(f"Invalid {spec_name} type: {type_name}. Expected {expected_type.__name__}, str, dict, or None.")
+
+
+__all__ = ["OptionalDefaults", "build_component", "instantiate_component", "instantiate_from_spec", "load_class"]