Add optional Monocle observability (opt-in, NAT telemetry plugin)

Signed-off-by: Mohammed Ansari <mohammed.ansari@okahu.ai>

Author: imohammedansari

deploy/.env.example

@@ -118,6 +118,12 @@ DASK_DISTRIBUTED__LOGGING__DISTRIBUTED=warning
 # LANGCHAIN_PROJECT=aiq-research
 # WANDB_API_KEY=
 
+# Monocle (optional): OpenTelemetry tracing for agents. Requires the `monocle`
+# extra (`uv sync --extra monocle`). See docs/source/deployment/observability.md.
+# MONOCLE_TRACING=true
+# MONOCLE_EXPORTERS=file          # file, console, okahu, s3, blob, gcs (default: file)
+# OKAHU_API_KEY=okh_xxxxxxxx      # required only for the `okahu` exporter
+
 # -----------------------------------------------------------------------------
 # Evaluation (optional)
 # -----------------------------------------------------------------------------

docs/source/deployment/observability.md

@@ -255,6 +255,67 @@ general:
         shutdown_timeout: 30000
 ```
 
+## Monocle
+
+[Monocle](https://github.com/monocle2ai/monocle) is an open-source, OpenTelemetry-based tracer for agentic applications. It instruments the frameworks already in use (LangChain, LangGraph, and others) and records each run end-to-end -- LLM calls, agent steps, and tool and MCP invocations, with their inputs, outputs, timings, and token counts. It is packaged as an **optional** tracing backend and is disabled unless you both install it and enable it -- the default install and default behavior are unchanged.
+
+Each run writes one trace file to `.monocle/` in the working directory; open it in the [Monocle VS Code extension](https://marketplace.visualstudio.com/items?itemName=OkahuAI.monocle-apptrace) to inspect the span timeline and token counts. Connect to [Okahu](https://www.okahu.ai), an agent-observability platform, to analyze traces across runs and run trace-based and agentic evaluations (via the `okahu` exporter).
+
+### Install
+
+Monocle ships as an optional extra, so a default install does not pull the OpenTelemetry stack:
+
+```bash
+uv sync --extra monocle
+# or, in an existing environment:
+pip install 'aiq-agent[monocle]'
+```
+
+Enabling Monocle without the extra installed fails fast with an actionable error naming the install command; when Monocle is not enabled, `monocle_apptrace` is never imported.
+
+### Enable
+
+There are two ways to opt in. Both forward the exporter list to `setup_monocle_telemetry(workflow_name="nvidia-aiq", monocle_exporters_list=...)` and Monocle initializes at most once per process.
+
+**Environment gate (matches the other demos).** Add the following to your `.env` file (`deploy/.env` for the CLI):
+
+```bash
+MONOCLE_TRACING=true
+MONOCLE_EXPORTERS=file          # file, console, okahu, s3, blob, gcs (default: file)
+OKAHU_API_KEY=okh_xxxxxxxx      # required only for the `okahu` exporter
+```
+
+The CLI reads these at startup and initializes Monocle before the workflow runs.
+
+**NAT telemetry exporter (canonical).** Add the `monocle` exporter to your workflow YAML; it initializes when the workflow builds:
+
+```yaml
+general:
+  telemetry:
+    tracing:
+      monocle:
+        _type: monocle
+        workflow_name: nvidia-aiq   # optional, stamped on spans
+        exporters: file             # optional; comma-separated, e.g. "file" or "file,okahu"
+```
+
+`OKAHU_API_KEY` is always read from the environment. `MONOCLE_EXPORTERS` is the default for the YAML `exporters` field, so the two paths agree unless you override it in YAML.
+
+**Precedence.** The environment gate runs at CLI startup, before the workflow (and any YAML `monocle` exporter) is built. When both are set, the env gate wins and the YAML exporter finds Monocle already initialized. To drive the exporter list purely from YAML, leave `MONOCLE_TRACING` unset.
+
+### Data handling
+
+Traces capture span inputs and outputs verbatim -- prompts, tool arguments, and model responses -- plus token usage and timings. The `file` exporter keeps them on local disk and never rotates or cleans them up, so prune `.monocle/` periodically; the remote exporters (`okahu`, `s3`, `blob`, `gcs`) send that same data off-box, so enable only destinations you trust. An unknown exporter name or a missing `OKAHU_API_KEY` fails fast with a clear error before any instrumentation.
+
+### Configuration Reference
+
+| Setting | Env / field | Default | Description |
+| :-- | :-- | :-- | :-- |
+| Enable | `MONOCLE_TRACING` (env) | off | Truthy gate for the environment opt-in path. |
+| Exporters | `MONOCLE_EXPORTERS` (env) / `exporters` (YAML) | `file` | Comma-separated Monocle exporters, validated against `file, console, okahu, s3, blob, gcs`. |
+| Okahu key | `OKAHU_API_KEY` (env) | -- | Required only when the `okahu` exporter is selected. |
+| Workflow name | `workflow_name` (YAML) | `nvidia-aiq` | Workflow name Monocle stamps onto emitted spans. |
+
 ## Verbose Logging
 
 For quick debugging without any external services, enable the built-in verbose callback logger. This prints detailed agent execution information directly to the console.

frontends/cli/cli.py

@@ -418,6 +418,12 @@ def main() -> None:
         except Exception as e:
             print(f"Warning: Failed to load .env file: {e}")
 
+    # Optional Monocle observability via the deer-flow-style MONOCLE_TRACING gate.
+    # No-op (and never imports monocle_apptrace) unless MONOCLE_TRACING is truthy.
+    from aiq_agent.observability.monocle_exporter import setup_monocle_tracing_if_enabled
+
+    setup_monocle_tracing_if_enabled()
+
     # Validate LLM API keys based on config
     try:
         config_path = Path(args.config_file)
@@ -455,6 +461,11 @@ async def _run():
         loop = asyncio.get_event_loop()
         loop.run_until_complete(_run())
     finally:
+        # os._exit skips normal interpreter shutdown, so flush pending Monocle/OTel
+        # spans first — otherwise the last spans of a run are dropped on exit.
+        from aiq_agent.observability.monocle_exporter import flush_monocle_if_enabled
+
+        flush_monocle_if_enabled()
         os._exit(0)
 
 

pyproject.toml

@@ -57,6 +57,13 @@ dependencies = [
 s3 = [
     "boto3>=1.35.0,<2",
 ]
+# Optional Monocle observability backend, exposed as the `monocle` NAT tracing
+# exporter (see src/aiq_agent/observability/monocle_exporter.py). Not required
+# for normal operation; install only when you want Monocle tracing:
+#   uv sync --extra monocle   (or)   pip install 'aiq-agent[monocle]'
+monocle = [
+    "monocle_apptrace>=0.8.8,<0.9",
+]
 dev = [
     "pytest>=7.4.0",
     "pytest-asyncio>=0.21.0",

src/aiq_agent/agents/chat_researcher/register.py

@@ -31,6 +31,7 @@
 from aiq_agent.common.citation_verification import get_or_create_session_registry
 from aiq_agent.common.citation_verification import reset_session_registry
 from aiq_agent.common.citation_verification import set_session_registry
+from aiq_agent.observability.monocle_exporter import ensure_registered as _ensure_monocle_registered
 from aiq_agent.observability.otel_header_redaction_exporter import (
     ensure_registered as _ensure_otel_redaction_registered,
 )
@@ -55,6 +56,7 @@
 _REPORT_ASK_TIMEOUT_S = 120
 
 _ensure_otel_redaction_registered()
+_ensure_monocle_registered()
 
 
 def _build_report_ask_prompt(

src/aiq_agent/observability/monocle_exporter.py

@@ -0,0 +1,245 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Optional Monocle observability for AI-Q.
+
+Monocle (``monocle_apptrace``) instruments supported frameworks (LangChain,
+LangGraph, etc.) in place and manages its own OpenTelemetry pipeline. It is not
+a hard dependency of AI-Q: it activates only when a user opts in, and only if
+``monocle_apptrace`` is installed (``uv sync --extra monocle``).
+
+There are two opt-in paths, both funnelling through :func:`_setup_monocle_once`
+so Monocle initializes at most once per process:
+
+* **NAT telemetry exporter (canonical).** Add a ``monocle`` exporter under
+  ``general.telemetry.tracing`` in the workflow YAML. This is the NAT-idiomatic
+  mechanism; the exporter is built (and Monocle initialized) only when the
+  workflow that references it is loaded.
+* **Environment gate (deer-flow-style).** Set ``MONOCLE_TRACING=true`` and the
+  CLI initializes Monocle at startup via :func:`setup_monocle_tracing_if_enabled`,
+  reading ``MONOCLE_EXPORTERS`` / ``OKAHU_API_KEY``. Handy for enabling tracing
+  without editing the config.
+
+Precedence: the env gate runs at CLI startup, before the workflow (and thus any
+YAML ``monocle`` exporter) is built, so when both are set the env gate wins and
+the YAML exporter finds Monocle already initialized and does nothing further. To
+drive the exporter list purely from YAML, leave ``MONOCLE_TRACING`` unset.
+
+When neither path is taken, this module registers a config type with the NAT
+registry but never imports ``monocle_apptrace`` -- default behavior is unchanged.
+"""
+
+import logging
+import os
+
+from pydantic import Field
+
+from nat.builder.builder import Builder
+from nat.cli.register_workflow import register_telemetry_exporter
+from nat.data_models.telemetry_exporter import TelemetryExporterBaseConfig
+from nat.observability.exporter.base_exporter import BaseExporter
+
+logger = logging.getLogger(__name__)
+
+# Default workflow name stamped onto Monocle spans.
+_DEFAULT_WORKFLOW_NAME = "nvidia-aiq"
+
+# Manual mirror of monocle_apptrace's supported exporters, kept local so a typo
+# fails fast with a clear message instead of an opaque upstream error. Update
+# this tuple when a monocle_apptrace bump adds or renames an exporter.
+_MONOCLE_EXPORTERS = ("file", "console", "okahu", "s3", "blob", "gcs")
+
+_TRUTHY_VALUES = {"1", "true", "yes", "on"}
+
+# Guard so global Monocle instrumentation is initialized at most once per process,
+# regardless of which opt-in path fires first.
+_MONOCLE_INITIALIZED = False
+
+
+def _env_flag(name: str) -> bool:
+    """Whether env var ``name`` is set to a truthy value."""
+    value = os.environ.get(name)
+    return bool(value) and value.strip().lower() in _TRUTHY_VALUES
+
+
+def _parse_exporters(exporters: str) -> list[str]:
+    """Split a comma-separated exporter string, dropping blanks."""
+    return [e.strip() for e in exporters.split(",") if e.strip()]
+
+
+def _resolve_exporters(exporter_list: list[str], okahu_api_key: str | None) -> list[str]:
+    """Validate and return the effective exporter list.
+
+    An unknown exporter is a config typo, so it still fails fast with an actionable
+    ``ValueError`` before any instrumentation runs. A missing secret, however, must
+    degrade gracefully: when ``okahu`` is selected without ``OKAHU_API_KEY`` it is
+    dropped (with a warning) and the remaining exporters continue. The returned list
+    may be empty, in which case the caller skips Monocle cleanly rather than crashing.
+    """
+    unknown = [e for e in exporter_list if e not in _MONOCLE_EXPORTERS]
+    if unknown:
+        raise ValueError(
+            f"MONOCLE_EXPORTERS has unknown exporter(s): {', '.join(unknown)}. "
+            f"Allowed: {', '.join(_MONOCLE_EXPORTERS)}."
+        )
+    if "okahu" in exporter_list and not okahu_api_key:
+        logger.warning(
+            "Monocle 'okahu' exporter is selected but OKAHU_API_KEY is not set; "
+            "skipping the okahu exporter and continuing with the remaining exporters."
+        )
+        exporter_list = [e for e in exporter_list if e != "okahu"]
+    return exporter_list
+
+
+def flush_monocle_if_enabled() -> None:
+    """Force-flush pending Monocle/OpenTelemetry spans; a no-op if Monocle was never set up.
+
+    Call this before a hard process exit (e.g. ``os._exit``) so pending spans that the
+    batch span processor has not yet exported are not silently dropped.
+    """
+    if not _MONOCLE_INITIALIZED:
+        return
+    try:
+        from opentelemetry import trace
+
+        force_flush = getattr(trace.get_tracer_provider(), "force_flush", None)
+        if callable(force_flush):
+            force_flush()
+    except Exception:  # pragma: no cover - best-effort flush during shutdown
+        logger.debug("Monocle span flush on exit failed.", exc_info=True)
+
+
+def _setup_monocle_once(workflow_name: str, exporters: str) -> None:
+    """Initialize Monocle telemetry once per process.
+
+    Imports ``monocle_apptrace`` lazily and raises a clear ``RuntimeError`` with
+    the install hint when the optional dependency is missing. ``exporters`` is a
+    comma-separated string passed to ``monocle_exporters_list`` as-is.
+    """
+    global _MONOCLE_INITIALIZED
+    if _MONOCLE_INITIALIZED:
+        return
+
+    try:
+        from monocle_apptrace import setup_monocle_telemetry
+    except ImportError as exc:
+        raise RuntimeError(
+            "Monocle observability is enabled but the optional 'monocle_apptrace' package is not installed. "
+            "Install the 'monocle' extra: `uv sync --extra monocle` (or `pip install 'aiq-agent[monocle]'`)."
+        ) from exc
+
+    setup_monocle_telemetry(workflow_name=workflow_name, monocle_exporters_list=exporters or None)
+    _MONOCLE_INITIALIZED = True
+    logger.info(
+        "Monocle telemetry initialized (workflow_name=%s, exporters=%s).",
+        workflow_name,
+        exporters or "<from environment>",
+    )
+
+
+def setup_monocle_tracing_if_enabled(workflow_name: str = _DEFAULT_WORKFLOW_NAME) -> bool:
+    """Initialize Monocle from the environment when ``MONOCLE_TRACING`` is truthy.
+
+    Mirrors the env surface documented across the sibling demos:
+
+    * ``MONOCLE_TRACING`` -- truthy gate, off by default.
+    * ``MONOCLE_EXPORTERS`` -- comma-separated exporter list, default ``file``.
+    * ``OKAHU_API_KEY`` -- required only when the ``okahu`` exporter is selected.
+
+    A no-op returning ``False`` when the gate is off. Called from the CLI startup
+    so embedded/other entry points can call it themselves. Validates before
+    instrumenting; a bad value raises ``ValueError``.
+    """
+    if not _env_flag("MONOCLE_TRACING"):
+        return False
+
+    exporters = (os.environ.get("MONOCLE_EXPORTERS") or "file").strip() or "file"
+    exporter_list = _resolve_exporters(_parse_exporters(exporters), os.environ.get("OKAHU_API_KEY"))
+    if not exporter_list:
+        logger.warning("MONOCLE_TRACING is enabled but no usable exporters remain; skipping Monocle.")
+        return False
+    # Off-box exporters move prompts, tool I/O, and completions beyond local disk.
+    off_box = [e for e in exporter_list if e not in ("file", "console")]
+    if off_box:
+        logger.warning(
+            "Monocle is exporting trace data (prompts, tool inputs/outputs, completions) beyond the local "
+            ".monocle/ directory via: %s. Make sure that destination is trusted.",
+            ", ".join(off_box),
+        )
+
+    _setup_monocle_once(workflow_name=workflow_name, exporters=",".join(exporter_list))
+    return True
+
+
+def ensure_registered() -> None:
+    """Import side effect: registers the ``monocle`` telemetry exporter config type.
+
+    Importing this module runs the ``@register_telemetry_exporter`` decorator
+    below, which is all that is needed for ``_type: monocle`` to be resolvable in
+    workflow config. Kept as an explicit no-op call to mirror the sibling
+    ``otel_header_redaction_exporter`` registration idiom.
+    """
+    return None
+
+
+class MonocleTelemetryExporter(TelemetryExporterBaseConfig, name="monocle"):
+    """Optional Monocle observability backend.
+
+    Enable by adding this exporter under ``general.telemetry.tracing`` in a
+    workflow config. Requires the ``monocle`` optional dependency to be
+    installed; if it is missing, building the exporter raises a clear error
+    instead of failing at import time.
+    """
+
+    workflow_name: str = Field(
+        default=_DEFAULT_WORKFLOW_NAME,
+        description="Workflow name Monocle stamps onto emitted spans.",
+    )
+    exporters: str = Field(
+        # Defaults to the same MONOCLE_EXPORTERS env var the CLI env gate reads,
+        # so the two opt-in paths agree; falls back to 'file'.
+        default_factory=lambda: (os.environ.get("MONOCLE_EXPORTERS") or "file").strip() or "file",
+        description="Comma-separated Monocle exporters (passed as monocle_exporters_list), "
+        "e.g. 'file', 'okahu', or 'file,okahu'. Validated against: " + ", ".join(_MONOCLE_EXPORTERS) + ".",
+    )
+
+
+class _MonocleInstrumentationExporter(BaseExporter):
+    """No-op NAT exporter that represents the Monocle instrumentation lifecycle.
+
+    Monocle exports spans through its own pipeline, so this exporter does not
+    consume NAT intermediate steps. It exists so Monocle initialization is tied
+    to NAT's telemetry-exporter lifecycle (built only when configured).
+    """
+
+    def export(self, event) -> None:  # noqa: D102 - inherited semantics; intentionally a no-op
+        return None
+
+
+@register_telemetry_exporter(config_type=MonocleTelemetryExporter)
+async def monocle_telemetry_exporter(config: MonocleTelemetryExporter, _builder: Builder):
+    """Initialize Monocle telemetry when the ``monocle`` exporter is configured."""
+    # If the env gate (or a prior workflow load) already initialized Monocle, the YAML
+    # settings are documented as a no-op — return the exporter without re-validating.
+    if _MONOCLE_INITIALIZED:
+        yield _MonocleInstrumentationExporter()
+        return
+    # OKAHU_API_KEY is always sourced from the environment, matching the env gate.
+    exporter_list = _resolve_exporters(_parse_exporters(config.exporters), os.environ.get("OKAHU_API_KEY"))
+    if exporter_list:
+        _setup_monocle_once(workflow_name=config.workflow_name, exporters=",".join(exporter_list))
+    else:
+        logger.warning("Monocle exporter configured but no usable exporters remain; skipping Monocle.")
+    yield _MonocleInstrumentationExporter()