gmf_forge_ai_shared_core.observability.tracing

206class TracingProvider(ABC):
207    """
208    Abstract tracing provider.
209
210    Implement this class to add a new backend.  Applications only ever
211    import :class:`TracingProvider` and call :func:`get_tracer` — they
212    have no knowledge of the concrete backend in use.
213    """
214
215    # ── Core operations ───────────────────────────────────────────────────────
216
217    @abstractmethod
218    def trace(self, name: str, **kwargs: Any) -> Span:
219        """
220        Start a new root trace.
221
222        Args:
223            name:   Human-readable name for the operation (e.g. ``"rag_query"``).
224            **kwargs: Optional fields forwarded to the backend:
225                - ``input``   — the operation input (question, prompt, …)
226                - ``user_id`` — end-user identifier for session grouping
227                - ``session_id`` — conversation or session ID
228                - ``metadata`` — dict of arbitrary key/value pairs
229                - ``tags``    — list of string tags
230
231        Returns:
232            A :class:`Span` to use as a context manager.
233        """
234
235    @abstractmethod
236    def score(
237        self,
238        trace_id: str,
239        name: str,
240        value: float,
241        comment: str = "",
242    ) -> None:
243        """
244        Attach a quality score to a completed trace.
245
246        Args:
247            trace_id: ID of the trace to score (``span.trace_id``).
248            name:     Score dimension name, e.g. ``"faithfulness"``, ``"relevance"``.
249            value:    Numeric score — typically 0.0–1.0 but provider-specific.
250            comment:  Optional free-text explanation.
251        """
252
253    @abstractmethod
254    def flush(self) -> None:
255        """Flush any buffered events to the backend. Call before process exit."""
256
257    @abstractmethod
258    def shutdown(self) -> None:
259        """Gracefully shut down the provider and release resources."""
260
261    # ── Convenience helpers (non-abstract, built on trace()) ──────────────────
262
263    def trace_rag_query(
264        self,
265        question: str,
266        user_id: str = "",
267        session_id: str = "",
268        **metadata: Any,
269    ) -> Span:
270        """Shortcut to open a root trace for a RAG question-answering call."""
271        return self.trace(
272            "rag_query",
273            input=question,
274            user_id=user_id or None,
275            session_id=session_id or None,
276            metadata=metadata or None,
277            tags=["rag"],
278        )
279
280    def trace_agent_run(
281        self,
282        agent_name: str,
283        input: Any,
284        user_id: str = "",
285        session_id: str = "",
286        **metadata: Any,
287    ) -> Span:
288        """Shortcut to open a root trace for an agent invocation."""
289        return self.trace(
290            f"agent:{agent_name}",
291            input=input,
292            user_id=user_id or None,
293            session_id=session_id or None,
294            metadata=metadata or None,
295            tags=["agent", agent_name],
296        )

303@dataclass
304class TracingConfig:
305    """
306    Configuration for a tracing provider, loaded from env vars by
307    :func:`get_tracer`.
308
309    All fields are optional — providers only read what they need.
310    """
311    provider: str = "null"
312
313    # Langfuse
314    langfuse_public_key: str = ""
315    langfuse_secret_key: str = ""
316    langfuse_host: str = "https://cloud.langfuse.com"
317
318    # LangSmith
319    langsmith_api_key: str = ""
320    langsmith_project: str = "default"
321    langsmith_endpoint: str = "https://api.smith.langchain.com"
322
323    # Azure AI Foundry
324    azure_ai_connection_string: str = ""   # from AI Foundry project settings
325    azure_ai_project_name: str = ""
326
327    # Splunk (via OpenTelemetry OTLP)
328    splunk_otlp_endpoint: str = ""
329    splunk_access_token: str = ""
330    splunk_service_name: str = "gmf-forge-ai"
331
332    # Common
333    environment: str = "development"
334    release: str = ""
335    extra: Dict[str, Any] = field(default_factory=dict)

 84class Span:
 85    """
 86    A single unit of work inside a trace.
 87
 88    Use as a context manager::
 89
 90        with tracer.trace("rag_query", input=question) as span:
 91            span.set_output(answer)
 92
 93    Instances are returned by :meth:`TracingProvider.trace`,
 94    :meth:`Span.span`, and :meth:`Span.generation`.  Child spans created
 95    via :meth:`Span.span` or :meth:`Span.generation` are automatically
 96    linked to the parent.
 97
 98    Attributes:
 99        id:        Unique identifier for this span/generation.
100        trace_id:  Identifier of the root trace this span belongs to.
101        name:      Human-readable operation name.
102        is_root:   True when this span *is* the root trace.
103    """
104
105    def __init__(
106        self,
107        name: str,
108        trace_id: Optional[str] = None,
109        span_id: Optional[str] = None,
110        is_root: bool = False,
111        kind: str = "span",          # "trace" | "span" | "generation"
112    ):
113        self.name = name
114        self.id: str = span_id or str(uuid.uuid4())
115        self.trace_id: str = trace_id or self.id
116        self.is_root = is_root
117        self.kind = kind
118
119        self._input: Any = None
120        self._output: Any = None
121        self._metadata: Dict[str, Any] = {}
122        self._error: Optional[Exception] = None
123        self._token_usage: Dict[str, int] = {}
124        self._model: Optional[str] = None
125        self._active_token: Any = None  # contextvars token for cleanup
126
127    # ── Mutators ─────────────────────────────────────────────────────────────
128
129    def set_input(self, input: Any) -> "Span":
130        self._input = input
131        return self
132
133    def set_output(self, output: Any) -> "Span":
134        self._output = output
135        return self
136
137    def set_metadata(self, **kwargs: Any) -> "Span":
138        self._metadata.update(kwargs)
139        return self
140
141    def set_error(self, error: Exception) -> "Span":
142        self._error = error
143        return self
144
145    def set_token_usage(
146        self,
147        prompt_tokens: int = 0,
148        completion_tokens: int = 0,
149        total_tokens: Optional[int] = None,
150    ) -> "Span":
151        self._token_usage = {
152            "prompt_tokens": prompt_tokens,
153            "completion_tokens": completion_tokens,
154        }
155        if total_tokens is not None:
156            self._token_usage["total_tokens"] = total_tokens
157        return self
158
159    # ── Child span factories (overridden by backend implementations) ──────────
160
161    def span(self, name: str, **kwargs: Any) -> "Span":
162        """Create a child span under this span. Override in subclasses."""
163        return Span(name=name, trace_id=self.trace_id, kind="span")
164
165    def generation(self, name: str, model: str = "", **kwargs: Any) -> "Span":
166        """Create a child generation span (LLM call). Override in subclasses."""
167        s = Span(name=name, trace_id=self.trace_id, kind="generation")
168        s._model = model
169        return s
170
171    # ── Context manager ───────────────────────────────────────────────────────
172
173    def __enter__(self) -> "Span":
174        # Track this span as active in the current context for decorator nesting
175        self._active_token = _set_active_span(self)
176        return self
177
178    def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
179        if exc_val is not None:
180            self.set_error(exc_val)
181        self._end()
182        # Restore previous active span (only if token was set)
183        if self._active_token is not None:
184            try:
185                _active_span.reset(self._active_token)  # type: ignore[attr-defined]
186            except (AttributeError, LookupError):
187                pass
188        return False  # do not suppress exceptions
189
190    def _end(self) -> None:
191        """Called when the context manager exits. Override to flush to backend."""
192
193    # ── Helpers ───────────────────────────────────────────────────────────────
194
195    def __repr__(self) -> str:
196        return (
197            f"<Span name={self.name!r} id={self.id[:8]} "
198            f"trace_id={self.trace_id[:8]} kind={self.kind}>"
199        )

 68class NullTracer(TracingProvider):
 69    """
 70    Default tracing provider — logs span lifecycle to the console.
 71
 72    Used when ``TRACING_PROVIDER`` is not set (or set to ``"null"`` /
 73    ``"none"``).  Each span emits a *started* and *finished* log line so
 74    the tracing structure is visible during local development and in
 75    example runs, without shipping data to any backend.
 76
 77    Switch to a real backend at any time by setting ``TRACING_PROVIDER``
 78    in ``.env`` — application code stays untouched.
 79
 80    Example::
 81
 82        tracer = NullTracer()
 83
 84        with tracer.trace("rag_query", input=question) as trace:
 85            with trace.span("retrieval", input=question) as span:
 86                docs = retriever.retrieve(question)
 87                span.set_output({"doc_count": len(docs)})
 88    """
 89
 90    def trace(self, name: str, **kwargs: Any) -> _NullSpan:
 91        span = _NullSpan(name=name, is_root=True, kind="trace")
 92        if "input" in kwargs:
 93            span.set_input(kwargs["input"])
 94        return span
 95
 96    def score(self, trace_id: str, name: str, value: float, comment: str = "") -> None:
 97        log_kw: dict[str, Any] = {"trace_id": trace_id[:8], "name": name, "value": value}
 98        if comment:
 99            log_kw["comment"] = comment
100        _logger.info("[trace] score", **log_kw)
101
102    def flush(self) -> None:
103        pass
104
105    def shutdown(self) -> None:
106        pass

138class SplunkTracer(TracingProvider):
139    """
140    Tracing provider that exports to Splunk Observability Cloud via
141    OpenTelemetry OTLP/HTTP.
142
143    Traces appear in **Splunk APM** under the configured service name.
144    Token usage, model info, and scores are stored as span attributes
145    and can be surfaced via Splunk's MetricSets or custom dashboards.
146
147    Example::
148
149        import os
150        os.environ["TRACING_PROVIDER"]       = "splunk"
151        os.environ["SPLUNK_OTLP_ENDPOINT"]   = "https://ingest.us0.signalfx.com/v2/trace/otlp"
152        os.environ["SPLUNK_ACCESS_TOKEN"]    = "your-token"
153        os.environ["SPLUNK_SERVICE_NAME"]    = "policy-rag"
154
155        from gmf_forge_ai_shared_core.observability.tracing import get_tracer
156
157        tracer = get_tracer()
158
159        with tracer.trace("rag_query", input=question) as trace:
160            with trace.span("retrieval", input=question) as span:
161                docs = retriever.retrieve(question)
162                span.set_output({"doc_count": len(docs)})
163            with trace.generation("llm", model="gpt-4o", input=prompt) as gen:
164                reply = llm.complete(prompt)
165                gen.set_output(reply.text)
166                gen.set_token_usage(prompt_tokens=120, completion_tokens=80)
167            trace.set_output(reply.text)
168
169        tracer.flush()
170    """
171
172    def __init__(self, config: TracingConfig):
173        self._config = config
174        self._otel_tracer: Optional[Any] = None
175        self._init_client()
176
177    def _init_client(self) -> None:
178        try:
179            from opentelemetry import trace as otel_trace                           # type: ignore
180            from opentelemetry.sdk.trace import TracerProvider                      # type: ignore
181            from opentelemetry.sdk.trace.export import BatchSpanProcessor           # type: ignore
182            from opentelemetry.sdk.resources import Resource                        # type: ignore
183            from opentelemetry.exporter.otlp.proto.http.trace_exporter import (    # type: ignore
184                OTLPSpanExporter,
185            )
186        except ImportError:
187            _logger.warning(
188                "opentelemetry-sdk or opentelemetry-exporter-otlp-proto-http not found — falling back to no-op tracing",
189                install_cmd="pip install opentelemetry-sdk opentelemetry-exporter-otlp-proto-http",
190            )
191            return
192
193        endpoint    = self._config.splunk_otlp_endpoint
194        token       = self._config.splunk_access_token
195        service     = self._config.splunk_service_name or "gmf-forge-ai"
196        environment = self._config.environment or "development"
197
198        if not endpoint or not token:
199            _logger.warning(
200                "SPLUNK_OTLP_ENDPOINT / SPLUNK_ACCESS_TOKEN not set — falling back to no-op tracing",
201                missing=[v for v, val in [("SPLUNK_OTLP_ENDPOINT", endpoint), ("SPLUNK_ACCESS_TOKEN", token)] if not val],
202            )
203            return
204
205        try:
206            resource = Resource.create({
207                "service.name":        service,
208                "deployment.environment": environment,
209            })
210            exporter = OTLPSpanExporter(
211                endpoint=endpoint,
212                headers={"X-SF-Token": token},
213            )
214            provider = TracerProvider(resource=resource)
215            provider.add_span_processor(BatchSpanProcessor(exporter))
216            otel_trace.set_tracer_provider(provider)
217            self._otel_tracer = provider.get_tracer(service)
218        except Exception as exc:
219            _logger.error("Splunk tracer init failed", error=str(exc))
220
221    # ── TracingProvider interface ─────────────────────────────────────────────
222
223    def trace(self, name: str, **kwargs: Any) -> Span:
224        from .null_tracer import _NullSpan
225
226        if self._otel_tracer is None:
227            return _NullSpan(name=name, is_root=True, kind="trace")
228
229        input_data = kwargs.get("input")
230        user_id    = kwargs.get("user_id")
231        session_id = kwargs.get("session_id")
232        tags       = kwargs.get("tags") or []
233
234        try:
235            otel_span = self._otel_tracer.start_span(name)
236            if input_data is not None:
237                otel_span.set_attribute("input", str(input_data))
238            if user_id:
239                otel_span.set_attribute("user_id", str(user_id))
240            if session_id:
241                otel_span.set_attribute("session_id", str(session_id))
242            if tags:
243                otel_span.set_attribute("tags", ",".join(tags))
244
245            trace_id = format(otel_span.get_span_context().trace_id, "032x")
246            span = _SplunkSpan(
247                otel_span=otel_span,
248                otel_tracer=self._otel_tracer,
249                trace_id=trace_id,
250                name=name,
251                kind="trace",
252            )
253            if input_data is not None:
254                span._input = input_data
255            return span
256        except Exception as exc:
257            _logger.error("Splunk trace() failed", error=str(exc))
258            return _NullSpan(name=name, is_root=True, kind="trace")
259
260    def score(self, trace_id: str, name: str, value: float, comment: str = "") -> None:
261        # Splunk APM doesn't have a native scoring API; scores can be
262        # emitted as custom metrics via the Splunk metrics ingest endpoint.
263        # For now we record it as a log-level attribute on the active span.
264        pass
265
266    def flush(self) -> None:
267        try:
268            from opentelemetry import trace as otel_trace  # type: ignore
269            provider = otel_trace.get_tracer_provider()
270            if hasattr(provider, "force_flush"):
271                provider.force_flush()
272        except Exception:
273            pass
274
275    def shutdown(self) -> None:
276        try:
277            from opentelemetry import trace as otel_trace  # type: ignore
278            provider = otel_trace.get_tracer_provider()
279            if hasattr(provider, "shutdown"):
280                provider.shutdown()
281        except Exception:
282            pass

169class LangfuseTracer(TracingProvider):
170    """Tracing provider that exports traces to Langfuse (SDK v3+)."""
171
172    def __init__(self, config: TracingConfig):
173        self._config = config
174        self._client: Optional[Any] = None
175        self._init_client()
176
177    def _init_client(self) -> None:
178        try:
179            from langfuse import Langfuse  # type: ignore
180        except ImportError:
181            _logger.warning(
182                "langfuse package not found — falling back to no-op tracing",
183                install_cmd="pip install langfuse",
184            )
185            return
186
187        public_key = self._config.langfuse_public_key
188        secret_key = self._config.langfuse_secret_key
189        host = self._config.langfuse_host or "https://cloud.langfuse.com"
190
191        missing = [
192            name
193            for name, value in [
194                ("LANGFUSE_PUBLIC_KEY", public_key),
195                ("LANGFUSE_SECRET_KEY", secret_key),
196            ]
197            if not value
198        ]
199        if missing:
200            _logger.warning(
201                "LANGFUSE_PUBLIC_KEY / LANGFUSE_SECRET_KEY not set — falling back to no-op tracing",
202                missing=missing,
203            )
204            return
205
206        try:
207            self._client = Langfuse(
208                public_key=public_key,
209                secret_key=secret_key,
210                host=host,
211            )
212        except Exception as exc:
213            _logger.error("Langfuse tracer init failed", error=str(exc))
214
215    def trace(self, name: str, **kwargs: Any) -> Span:
216        from .null_tracer import _NullSpan
217
218        if self._client is None:
219            return _NullSpan(name=name, is_root=True, kind="trace")
220
221        input_data = kwargs.get("input")
222        user_id = kwargs.get("user_id")
223        session_id = kwargs.get("session_id")
224        tags = kwargs.get("tags")
225
226        # Merge caller metadata with trace-level fields not supported as direct
227        # parameters on start_as_current_observation in SDK v3+.
228        metadata: dict[str, Any] = dict(kwargs.get("metadata") or {})
229        if user_id:
230            metadata["user_id"] = user_id
231        if session_id:
232            metadata["session_id"] = session_id
233        if tags:
234            metadata["tags"] = tags
235        if self._config.environment:
236            metadata["environment"] = self._config.environment
237
238        try:
239            from langfuse.types import TraceContext  # type: ignore
240
241            # Pre-allocate a deterministic trace ID so callers can reference it
242            # before the observation is flushed.
243            trace_id = self._client.create_trace_id()
244
245            lf_ctx = self._client.start_as_current_observation(
246                trace_context=TraceContext(trace_id=trace_id),
247                name=name,
248                as_type="span",
249                input=input_data,
250                metadata=metadata or None,
251                version=self._config.release or None,
252            )
253            span = _LangfuseSpan(
254                lf_client=self._client,
255                lf_ctx=lf_ctx,
256                trace_id=trace_id,
257                name=name,
258                kind="trace",
259            )
260            span.id = trace_id
261            span.is_root = True
262            if input_data is not None:
263                span.set_input(input_data)
264            return span
265        except Exception as exc:
266            _logger.error("Langfuse trace() failed", error=str(exc))
267            return _NullSpan(name=name, is_root=True, kind="trace")
268
269    def score(self, trace_id: str, name: str, value: float, comment: str = "") -> None:
270        if self._client is None:
271            return
272        try:
273            self._client.create_score(
274                trace_id=trace_id,
275                name=name,
276                value=value,
277                comment=comment or None,
278            )
279        except Exception:
280            pass
281
282    def flush(self) -> None:
283        if self._client is None:
284            return
285        try:
286            if hasattr(self._client, "flush"):
287                self._client.flush()
288        except Exception:
289            pass
290
291    def shutdown(self) -> None:
292        if self._client is None:
293            return
294        try:
295            if hasattr(self._client, "shutdown"):
296                self._client.shutdown()
297            elif hasattr(self._client, "flush"):
298                self._client.flush()
299        except Exception:
300            pass