agentx-python 0.6.21__tar.gz → 0.6.23__tar.gz

{agentx_python-0.6.21 → agentx_python-0.6.23}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agentx-python
-Version: 0.6.21
+Version: 0.6.23
 Summary: Official Python SDK for AgentX (https://www.agentx.so/)
 Home-page: https://github.com/AgentX-ai/AgentX-python
 Author: Robin Wang and AgentX Team
@@ -36,6 +36,8 @@ Requires-Dist: llama-index-core>=0.10.0; extra == "llamaindex"
 Provides-Extra: autogen
 Requires-Dist: autogen-agentchat>=0.4.0; extra == "autogen"
 Requires-Dist: autogen-core>=0.4.0; extra == "autogen"
+Provides-Extra: databricks
+Requires-Dist: mlflow>=3.6.0; extra == "databricks"
 Provides-Extra: all
 Requires-Dist: langchain-core>=0.1.0; extra == "all"
 Requires-Dist: crewai>=0.80.0; extra == "all"

{agentx_python-0.6.21 → agentx_python-0.6.23}/agentx/evaluations/datasets.py

@@ -86,9 +86,19 @@ class DatasetBuilder:
         judge_guideline: Optional[str] = None,
         smoke_test_count: Optional[int] = None,
         smoke_test_guidance: Optional[str] = None,
+        expected_tools: Optional[List[str]] = None,
+        trajectory_match_mode: str = "strict",
     ) -> "DatasetBuilder":
         """Add a case. `judge_guideline` is optional grading guidance specific to this question.
 
+        `expected_tools` declares the tool calls a correct run of this case should make. When a
+        result links its trace (return `{"output": ..., "trace_id": span.trace_id}` from the
+        agent function), the engine matches the trace's actual tool-call sequence against it and
+        reports a pass/fail "Trajectory match" scorer row on the result. `trajectory_match_mode`
+        follows agentevals semantics: "strict" (same calls, same order), "unordered" (same calls,
+        any order), "superset" (all expected present, extras allowed), or "subset" (no unexpected
+        calls, missing allowed).
+
         `smoke_test_count`, when set (1-10), asks this question that many extra ways each
         evaluation run, LLM-paraphrased server-side, to catch agents that are brittle to phrasing
         rather than genuinely wrong. `smoke_test_guidance` optionally steers what kind of variants
@@ -111,6 +121,8 @@ class DatasetBuilder:
             main["smokeTest"] = {"enabled": True, "count": smoke_test_count}
             if smoke_test_guidance:
                 main["smokeTest"]["guidance"] = smoke_test_guidance
+        if expected_tools:
+            main["expectedTrajectory"] = {"tools": expected_tools, "mode": trajectory_match_mode}
         self._payload["questions"].append(
             {
                 "main_question": main,

{agentx_python-0.6.21 → agentx_python-0.6.23}/agentx/evaluations/runner.py

@@ -271,6 +271,11 @@ class EvaluationRunContext:
     # UI reads, computed once in the API.
     # ------------------------------------------------------------------
 
+    @property
+    def run_id(self) -> str:
+        """The server-side run id - handy for fetching the run's full results afterwards."""
+        return self._run.run_id
+
     @property
     def rated_count(self) -> int:
         """Number of submitted results that have received a rating so far."""

{agentx_python-0.6.21 → agentx_python-0.6.23}/agentx/integrations/autogen.py

@@ -178,8 +178,13 @@ class AgentXAutoGenObserver:
             if input_tokens is not None or out_tokens is not None:
                 total_input_tokens += input_tokens or 0
                 total_output_tokens += out_tokens or 0
+                # Prefer the speaking agent's name (message.source) over a generic counter -
+                # in a multi-agent team the step sequence then reads as the actual
+                # agent-turn trajectory ("planner" -> "coder" -> "reviewer"), which is what
+                # trajectory evaluation cares about.
+                source = getattr(message, "source", None)
                 execution_steps.append({
-                    "name": f"LLM Call {len(execution_steps) + 1}",
+                    "name": str(source) if source else f"LLM Call {len(execution_steps) + 1}",
                     "duration_ms": (end_t - start_t) * 1000,
                     "start_time": start_t,
                     "end_time": end_t,

agentx_python-0.6.23/agentx/integrations/databricks.py

@@ -0,0 +1,405 @@
+"""
+Databricks / MLflow integration for AgentX.
+
+Two complementary paths for agents built on Databricks (Agent Bricks or the Mosaic AI Agent
+Framework - both are auto-instrumented by MLflow 3 Tracing):
+
+1. **Push (live)** - ``enable_mlflow_export()``: point MLflow Tracing's native OTLP exporter at
+   the AgentX engine's OTel endpoint. Works anywhere MLflow traces run (notebooks, jobs, Model
+   Serving endpoints via environment variables). Dual export keeps Databricks' own MLflow UI and
+   inference tables working alongside AgentX::
+
+       from agentx.integrations.databricks import enable_mlflow_export
+
+       enable_mlflow_export(
+           api_key=os.environ["AGENTX_API_KEY"],
+           base_url="http://localhost:4700/api/v1",   # your AgentX engine
+           service_name="my-databricks-agent",
+       )
+       # ... then trace as usual (@mlflow.trace, autolog, Agent Framework, ...)
+
+   On a Model Serving endpoint, set the equivalent environment variables instead (this helper
+   prints them with ``dry_run=True``).
+
+2. **Pull (batch)** - ``agentx-databricks sync``: import finished MLflow traces from a Databricks
+   (or any MLflow 3) tracking server into AgentX - full span trees, tool calls, sessions -
+   deduplicated on deterministic span ids so re-running a window never duplicates, with
+   ``--monitor`` / ``--judge-sessions`` mirroring ``agentx-moveworks``::
+
+       export AGENTX_API_KEY=...            # AgentX project key
+       export DATABRICKS_HOST=... DATABRICKS_TOKEN=...   # or MLFLOW_TRACKING_URI
+       agentx-databricks sync --experiment-id 123456 --since 24h
+
+Requires (pull path): ``pip install "agentx-python[databricks]"`` (mlflow>=3).
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from datetime import datetime, timedelta, timezone
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Optional
+
+from agentx.tracing.ingest_client import IngestClient
+from agentx.version import VERSION
+from agentx.integrations.moveworks import (
+    _parse_since,
+    _parse_time,
+    _read_cursor,
+    _write_cursor,
+    judge_sessions_via_engine,
+)
+
+_DEFAULT_CURSOR_FILE = Path.home() / ".agentx" / "databricks_sync_cursor.json"
+
+# MLflow's session grouping metadata key (mlflow.update_current_trace(metadata={...})) - traces
+# sharing it become one AgentX session, judged as a conversation by session-scoped evaluators.
+_SESSION_METADATA_KEYS = ("mlflow.trace.session", "mlflow.trace.session_id", "session_id")
+
+
+# ----------------------------------------------------------------------------------------------
+# Path 1: push - MLflow OTLP export pointed at the AgentX engine
+# ----------------------------------------------------------------------------------------------
+
+def enable_mlflow_export(
+    *,
+    api_key: str,
+    base_url: str,
+    service_name: str = "databricks-agent",
+    dual: bool = True,
+    genai_semconv: bool = False,
+    dry_run: bool = False,
+) -> Dict[str, str]:
+    """
+    Configure MLflow Tracing's built-in OTLP exporter to send every trace to the AgentX engine.
+
+    Must run BEFORE the first trace starts (MLflow reads these once, at tracer setup). ``dual``
+    keeps MLflow's own tracking export too (Databricks MLflow UI / inference tables keep
+    working); ``genai_semconv`` switches the wire format to OTel GenAI semantic conventions -
+    AgentX ingests both, and the default (MLflow-native attributes) is the higher-fidelity
+    mapping for plain ``@mlflow.trace`` functions. Returns the environment variables set - with
+    ``dry_run=True`` nothing is set, so the dict can be copied onto a Databricks Model Serving
+    endpoint's environment variables instead.
+    """
+    env = {
+        "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT": f"{base_url.rstrip('/')}/otel/v1/traces",
+        "OTEL_EXPORTER_OTLP_TRACES_HEADERS": f"x-api-key={api_key}",
+        "OTEL_EXPORTER_OTLP_PROTOCOL": "http/protobuf",
+        "OTEL_SERVICE_NAME": service_name,
+    }
+    if dual:
+        env["MLFLOW_TRACE_ENABLE_OTLP_DUAL_EXPORT"] = "true"
+    if genai_semconv:
+        env["MLFLOW_ENABLE_OTEL_GENAI_SEMCONV"] = "true"
+    if not dry_run:
+        os.environ.update(env)
+    return env
+
+
+# ----------------------------------------------------------------------------------------------
+# Path 2: pull - MLflow trace search -> AgentX span trees
+# ----------------------------------------------------------------------------------------------
+
+class DatabricksSyncReport:
+    def __init__(self) -> None:
+        self.traces = 0
+        self.spans = 0
+        self.tool_calls = 0
+        self.ingested = 0
+        self.failed = 0
+        self.skipped_in_progress = 0
+        self.session_ids: List[str] = []
+        self.sessions_judged = 0
+        self.sessions_judge_skipped = 0
+        self.sessions_judge_failed = 0
+
+    def __repr__(self) -> str:  # also what the CLI prints
+        parts = [
+            f"traces={self.traces}",
+            f"spans={self.spans}",
+            f"tool_calls={self.tool_calls}",
+            f"ingested={self.ingested}",
+            f"failed={self.failed}",
+            f"sessions={len(self.session_ids)}",
+        ]
+        if self.skipped_in_progress:
+            parts.append(f"skipped_in_progress={self.skipped_in_progress}")
+        if self.sessions_judged or self.sessions_judge_skipped or self.sessions_judge_failed:
+            parts.append(
+                f"judged={self.sessions_judged} judge_skipped={self.sessions_judge_skipped} "
+                f"judge_failed={self.sessions_judge_failed}"
+            )
+        return f"DatabricksSyncReport({' '.join(parts)})"
+
+
+def _span_time_ns(span: Any, attr: str) -> Optional[int]:
+    value = getattr(span, attr, None)
+    return int(value) if isinstance(value, (int, float)) and value > 0 else None
+
+
+def _span_error(span: Any) -> Optional[str]:
+    status = getattr(span, "status", None)
+    code = str(getattr(status, "status_code", "") or "")
+    if "ERROR" in code.upper():
+        return str(getattr(status, "description", None) or "error")
+    return None
+
+
+def _serialize(value: Any) -> Any:
+    if value is None or isinstance(value, (str, int, float, bool)):
+        return value
+    try:
+        return json.loads(json.dumps(value, default=str))
+    except (TypeError, ValueError):
+        return str(value)
+
+
+class DatabricksTraceImporter:
+    """
+    Replays finished MLflow traces (Databricks-hosted or any MLflow 3 tracking server) into
+    AgentX as full span trees. ``span_id`` is deterministic (``dbx:<trace_id>[:<span_id>]``), so
+    re-syncing a window is idempotent - the engine dedupes on replay and skips re-judging.
+    """
+
+    def __init__(
+        self,
+        mlflow_client: Any,
+        *,
+        agentx_api_key: str,
+        agentx_base_url: Optional[str] = None,
+        agent_name: Optional[str] = None,
+    ) -> None:
+        self._client = mlflow_client
+        self._agent_name = agent_name
+        self._ingest = IngestClient(agentx_api_key, sdk_version=VERSION, base_url=agentx_base_url)
+
+    # -- trace -> wires ------------------------------------------------------------------------
+
+    def _trace_wires(self, trace: Any) -> "tuple[List[Dict[str, Any]], Optional[str], int]":
+        """One MLflow Trace -> [root wire, *child wires], its session id, and its tool-call count."""
+        info = trace.info
+        trace_id = str(getattr(info, "trace_id", None) or getattr(info, "request_id", ""))
+        spans = list(getattr(trace.data, "spans", None) or [])
+        if not trace_id or not spans:
+            return [], None, 0
+
+        metadata_bag: Dict[str, Any] = {}
+        for source in (getattr(info, "trace_metadata", None), getattr(info, "tags", None)):
+            if isinstance(source, dict):
+                metadata_bag.update(source)
+        session_raw = next((metadata_bag[k] for k in _SESSION_METADATA_KEYS if metadata_bag.get(k)), None)
+        session_id = f"dbx_{session_raw}" if session_raw else None
+
+        root = next((s for s in spans if not getattr(s, "parent_id", None)), spans[0])
+        root_span_id = str(getattr(root, "span_id", "") or "root")
+        wire_span_id = {root_span_id: f"dbx:{trace_id}"}
+        for span in spans:
+            sid = str(getattr(span, "span_id", "") or "")
+            if sid and sid not in wire_span_id:
+                wire_span_id[sid] = f"dbx:{trace_id}:{sid}"
+
+        tool_calls: List[Dict[str, Any]] = []
+        wires: List[Dict[str, Any]] = []
+        for span in spans:
+            sid = str(getattr(span, "span_id", "") or "")
+            is_root = span is root
+            start_ns = _span_time_ns(span, "start_time_ns")
+            end_ns = _span_time_ns(span, "end_time_ns")
+            error = _span_error(span)
+            span_type = str(getattr(span, "span_type", "") or "").upper()
+            wire: Dict[str, Any] = {
+                "name": (self._agent_name if is_root and self._agent_name else str(getattr(span, "name", "span"))),
+                "framework": "databricks",
+                "span_id": wire_span_id.get(sid, f"dbx:{trace_id}:{sid or 'span'}"),
+            }
+            if not is_root:
+                parent_sid = str(getattr(span, "parent_id", "") or "")
+                wire["parent_span_id"] = wire_span_id.get(parent_sid, f"dbx:{trace_id}")
+            if session_id:
+                wire["session_id"] = session_id
+            if start_ns:
+                wire["started_at_unix_nano"] = str(start_ns)
+            if start_ns and end_ns and end_ns > start_ns:
+                wire["latency_ms"] = int((end_ns - start_ns) / 1_000_000)
+            inputs = _serialize(getattr(span, "inputs", None))
+            outputs = _serialize(getattr(span, "outputs", None))
+            if inputs is not None:
+                wire["input"] = inputs
+            if outputs is not None:
+                wire["output"] = outputs
+            if error:
+                wire["error"] = error
+            if is_root:
+                wire["metadata"] = {
+                    "source": "databricks",
+                    "mlflowTraceId": trace_id,
+                    **({"experimentId": str(getattr(info, "experiment_id", ""))} if getattr(info, "experiment_id", None) else {}),
+                }
+            if span_type == "TOOL":
+                tool_calls.append(
+                    {
+                        "name": str(getattr(span, "name", "tool")),
+                        "input": inputs,
+                        "output": outputs,
+                        "latency_ms": wire.get("latency_ms"),
+                        "success": not error,
+                    }
+                )
+            wires.append(wire)
+
+        # Root carries the flat tool_calls mirror - what the engine's Tool-failure check and
+        # trajectory matching read (same posture as tracer._merge_child_run).
+        if tool_calls:
+            wires[spans.index(root)]["tool_calls"] = tool_calls
+        # Root first so the engine resolves the agent before children arrive.
+        wires.sort(key=lambda w: 0 if "parent_span_id" not in w else 1)
+        return wires, session_id, len(tool_calls)
+
+    def _search(self, experiment_ids: List[str], since: datetime, until: datetime) -> Iterable[Any]:
+        """Newest-first paginated search, stopping once a page is entirely older than ``since``."""
+        page_token: Optional[str] = None
+        since_ms = int(since.timestamp() * 1000)
+        until_ms = int(until.timestamp() * 1000)
+        while True:
+            page = self._client.search_traces(
+                experiment_ids=experiment_ids,
+                max_results=100,
+                page_token=page_token,
+                order_by=["timestamp_ms DESC"],
+            )
+            oldest_seen = None
+            for trace in page:
+                ts = getattr(trace.info, "request_time", None) or getattr(trace.info, "timestamp_ms", None)
+                ts_ms = int(ts.timestamp() * 1000) if isinstance(ts, datetime) else (int(ts) if ts else None)
+                oldest_seen = ts_ms if ts_ms is not None else oldest_seen
+                if ts_ms is not None and (ts_ms < since_ms or ts_ms >= until_ms):
+                    if ts_ms < since_ms:
+                        continue
+                    continue
+                yield trace
+            page_token = getattr(page, "token", None)
+            if not page_token or (oldest_seen is not None and oldest_seen < since_ms):
+                return
+
+    # -- public --------------------------------------------------------------------------------
+
+    def sync(
+        self,
+        experiment_ids: List[str],
+        since: datetime,
+        until: Optional[datetime] = None,
+        *,
+        monitor: bool = False,
+        judge_sessions: bool = False,
+        dry_run: bool = False,
+        on_payload: Optional[Any] = None,
+    ) -> DatabricksSyncReport:
+        """
+        Import every finished MLflow trace in ``[since, until)`` from the given experiments.
+        Safe to re-run over the same window (span_id dedupe). ``monitor=True`` opts every
+        imported root trace into the engine's ingest-time checks; ``judge_sessions=True`` judges
+        each imported session afterwards (``ifStale`` - never duplicates the engine's own sweep).
+        """
+        report = DatabricksSyncReport()
+        until = until or datetime.now(timezone.utc)
+        for trace in self._search(experiment_ids, since, until):
+            state = str(getattr(trace.info, "state", "") or "")
+            if state and "IN_PROGRESS" in state.upper():
+                report.skipped_in_progress += 1
+                continue
+            wires, session_id, tool_count = self._trace_wires(trace)
+            if not wires:
+                continue
+            report.traces += 1
+            report.spans += len(wires)
+            report.tool_calls += tool_count
+            if session_id and session_id not in report.session_ids:
+                report.session_ids.append(session_id)
+            for wire in wires:
+                if monitor and "parent_span_id" not in wire:
+                    wire["monitor"] = True
+                if on_payload is not None:
+                    on_payload(wire)
+                if dry_run:
+                    continue
+                if self._ingest.send_trace_sync(wire) is not None:
+                    report.ingested += 1
+                else:
+                    report.failed += 1
+        if judge_sessions and not dry_run:
+            judged, skipped, failed = judge_sessions_via_engine(self._ingest, report.session_ids)
+            report.sessions_judged += judged
+            report.sessions_judge_skipped += skipped
+            report.sessions_judge_failed += failed
+        return report
+
+
+# ----------------------------------------------------------------------------------------------
+# CLI: agentx-databricks sync
+# ----------------------------------------------------------------------------------------------
+
+def cli_main(argv: Optional[List[str]] = None) -> None:
+    parser = argparse.ArgumentParser(
+        prog="agentx-databricks",
+        description="Sync MLflow traces (Databricks agents) into AgentX as span trees/sessions.",
+    )
+    sub = parser.add_subparsers(dest="command", required=True)
+    sync = sub.add_parser("sync", help="Import a time window of MLflow traces")
+    sync.add_argument("--experiment-id", action="append", required=True, help="MLflow experiment id (repeatable)")
+    sync.add_argument("--since", help="ISO timestamp or relative (7d, 24h). Default: cursor file, else 24h")
+    sync.add_argument("--until", help="ISO timestamp (default: now)")
+    sync.add_argument("--cursor-file", default=str(_DEFAULT_CURSOR_FILE), help="Incremental cursor path")
+    sync.add_argument("--no-cursor", action="store_true", help="Ignore and don't update the cursor file")
+    sync.add_argument("--agent-name", help="Attribute every trace to this agent (default: the root span's name)")
+    sync.add_argument(
+        "--tracking-uri",
+        default=os.getenv("MLFLOW_TRACKING_URI", "databricks"),
+        help='MLflow tracking URI (default: MLFLOW_TRACKING_URI, else "databricks" - uses DATABRICKS_HOST/TOKEN)',
+    )
+    sync.add_argument("--monitor", action="store_true", help="Run ingest-time checks on each imported trace")
+    sync.add_argument(
+        "--judge-sessions",
+        action="store_true",
+        help="After the sync, judge each imported session with every enabled session-scoped evaluator "
+        "(skips sessions that already have an up-to-date verdict)",
+    )
+    sync.add_argument("--dry-run", action="store_true", help="Map and print payloads without ingesting")
+
+    args = parser.parse_args(argv)
+
+    agentx_key = os.getenv("AGENTX_API_KEY", "")
+    if not agentx_key and not args.dry_run:
+        raise SystemExit("agentx-databricks: set AGENTX_API_KEY (your AgentX project API key)")
+
+    try:
+        from mlflow.client import MlflowClient
+    except ImportError as exc:  # pragma: no cover
+        raise SystemExit('agentx-databricks: pip install "agentx-python[databricks]" (needs mlflow>=3)') from exc
+
+    cursor_path = Path(args.cursor_file)
+    since = (
+        _parse_since(args.since)
+        if args.since
+        else ((None if args.no_cursor else _read_cursor(cursor_path)) or (datetime.now(timezone.utc) - timedelta(hours=24)))
+    )
+    until = _parse_time(args.until) if args.until else datetime.now(timezone.utc)
+
+    importer = DatabricksTraceImporter(
+        MlflowClient(tracking_uri=args.tracking_uri),
+        agentx_api_key=agentx_key,
+        agentx_base_url=os.getenv("AGENTX_API_BASE_URL") or None,
+        agent_name=args.agent_name,
+    )
+    report = importer.sync(
+        args.experiment_id,
+        since,
+        until,
+        monitor=args.monitor,
+        judge_sessions=args.judge_sessions,
+        dry_run=args.dry_run,
+        on_payload=(lambda wire: print(json.dumps(wire, default=str))) if args.dry_run else None,
+    )
+    print(report)
+    if not args.no_cursor and not args.dry_run and until is not None:
+        _write_cursor(cursor_path, until)

{agentx_python-0.6.21 → agentx_python-0.6.23}/agentx/integrations/langchain.py

@@ -231,6 +231,39 @@ def _extract_llm_output(response: "LLMResult") -> Optional[str]:
     return None
 
 
+# Nested chain runs that are plumbing, not agent structure - LCEL composition wrappers,
+# prompt/parse steps, and LangGraph's internal channel machinery. Skipped when deciding which
+# chain runs become "node" child spans; their LLM/tool descendants re-parent to the nearest
+# non-noise ancestor (see _emit_span_tree's resolve_parent walk).
+_NOISE_CHAIN_NAMES = frozenset({
+    "RunnableSequence",
+    "RunnableParallel",
+    "RunnableLambda",
+    "RunnableAssign",
+    "RunnablePick",
+    "RunnableBinding",
+    "RunnableWithFallbacks",
+    "RunnableWithMessageHistory",
+    "RunnableBranch",
+    "ChatPromptTemplate",
+    "PromptTemplate",
+    "StrOutputParser",
+    "JsonOutputParser",
+    "ToolsAgentOutputParser",
+    "OpenAIFunctionsAgentOutputParser",
+    "LangGraph",
+    "CompiledStateGraph",
+    "Prompt",
+    "_Exception",
+})
+
+_NOISE_CHAIN_PREFIXES = ("ChannelWrite", "ChannelRead", "Branch<", "RunnableParallel<", "_")
+
+
+def _is_noise_chain(name: str) -> bool:
+    return name in _NOISE_CHAIN_NAMES or name.startswith(_NOISE_CHAIN_PREFIXES)
+
+
 class AgentXCallbackHandler(BaseCallbackHandler):
     """
     LangChain callback handler that captures the top-level chain run and all
@@ -238,6 +271,11 @@ class AgentXCallbackHandler(BaseCallbackHandler):
 
     Compatible with AgentExecutor, LCEL chains, and LangGraph agents
     (``create_agent``, ``create_react_agent``).
+
+    The trace is a real span tree: LangGraph graph nodes (and named non-plumbing sub-chains)
+    become child spans, and each LLM call / tool call / retrieval becomes a span parented under
+    the node that ran it - so the engine's Execution Timeline shows the actual graph trajectory
+    (which nodes ran, in what order, and what each did), not a flat step list.
     """
 
     def __init__(
@@ -331,7 +369,122 @@ class AgentXCallbackHandler(BaseCallbackHandler):
                 "retrieval_steps": pending,
                 "input_tokens": 0,
                 "output_tokens": 0,
+                # Graph structure captured under this top-level run: named nested chain runs
+                # (LangGraph nodes, sub-agents) keyed by run_id, and the FULL nested-chain
+                # parent map (noise chains included) so _emit_span_tree can walk through
+                # skipped plumbing runs to the nearest emitted ancestor.
+                "node_runs": {},
+                "chain_parents": {},
             }
+        else:
+            top = self._find_top_ancestor(parent_run_id)
+            if top is None:
+                return
+            state = self._runs.get(top)
+            if state is None:
+                return
+            name = kwargs.get("name") or (serialized or {}).get("name")
+            # LangGraph stamps its node runs with metadata.langgraph_node - trust that over the
+            # noise heuristic when present (a user's node could legitimately be named
+            # "RunnableLambda"-style by a wrapper).
+            run_meta = kwargs.get("metadata") or {}
+            is_node = bool(name) and (run_meta.get("langgraph_node") == name or not _is_noise_chain(str(name)))
+            with self._state_lock:
+                state["chain_parents"][run_id] = parent_run_id
+                if is_node:
+                    state["node_runs"][run_id] = {
+                        "name": str(name),
+                        "start": time.time(),
+                        "end": None,
+                        "input": _extract_input(inputs),
+                        "output": None,
+                        "error": None,
+                        "parent": parent_run_id,
+                    }
+
+    def _finalize_node(self, run_id: UUID, *, output: Any = None, error: Optional[str] = None) -> None:
+        """Close a nested chain run's node record (if it became one) with end time + result."""
+        top = self._find_top_ancestor(run_id)
+        if top is None:
+            return
+        state = self._runs.get(top)
+        if state is None:
+            return
+        with self._state_lock:
+            node = state.get("node_runs", {}).get(run_id)
+            if node is None:
+                return
+            node["end"] = time.time()
+            if output is not None:
+                node["output"] = output
+            if error is not None:
+                node["error"] = error
+
+    def _emit_span_tree(self, span, state: Dict[str, Any], tool_calls: List[Dict[str, Any]]) -> None:
+        """
+        Emit the recorded run as a hierarchical span tree under ``span``: node runs first (in
+        start order, parented to their nearest emitted ancestor), then every LLM call, tool
+        call, and retrieval parented under the node that ran it. Records whose parent chain is
+        entirely noise (or missing, e.g. pre-run retrievals) land directly under the root span.
+        """
+        emitted: Dict[Any, Any] = {}
+        chain_parents: Dict[Any, Any] = state.get("chain_parents", {})
+        node_runs: Dict[Any, Dict[str, Any]] = state.get("node_runs", {})
+
+        def resolve_parent(parent_id: Any):
+            seen: set = set()
+            current = parent_id
+            while current is not None and current not in seen:
+                seen.add(current)
+                if current in emitted:
+                    return emitted[current]
+                current = chain_parents.get(current)
+            return span
+
+        for run_id, node in sorted(node_runs.items(), key=lambda kv: kv[1]["start"]):
+            parent_span = resolve_parent(node.get("parent"))
+            emitted[run_id] = parent_span.child_span(
+                node["name"],
+                start_time=node["start"],
+                end_time=node.get("end") or node["start"],
+                input=node.get("input"),
+                output=node.get("output"),
+                error=node.get("error"),
+            )
+
+        llm_count = 0
+        for step in state.get("execution_steps", []):
+            llm_count += 1
+            resolve_parent(step.get("parent_run_id")).child_span(
+                step.get("name") or f"LLM Call {llm_count}",
+                start_time=step.get("start_time"),
+                end_time=step.get("end_time"),
+                duration_ms=step.get("duration_ms"),
+                input=step.get("input"),
+                output=step.get("output"),
+                model=step.get("model"),
+                input_tokens=step.get("inputTokenSize"),
+                output_tokens=step.get("outputTokenSize"),
+            )
+        for tc in tool_calls:
+            resolve_parent(tc.get("parent_run_id")).child_span(
+                tc.get("name") or "Tool call",
+                start_time=tc.get("start_time"),
+                end_time=tc.get("end_time"),
+                duration_ms=tc.get("latency_ms"),
+                input=tc.get("input"),
+                output=tc.get("output"),
+                error=None if tc.get("success", True) else str(tc.get("output") or "Tool call failed"),
+            )
+        for step in state.get("retrieval_steps", []):
+            resolve_parent(step.get("parent_run_id")).child_span(
+                step.get("name") or "Retrieval",
+                start_time=step.get("start_time"),
+                end_time=step.get("end_time"),
+                duration_ms=step.get("duration_ms"),
+                input=step.get("query"),
+                output=step.get("output"),
+            )
 
     def on_chain_end(
         self,
@@ -347,9 +500,13 @@ class AgentXCallbackHandler(BaseCallbackHandler):
         # entries behind here would leak forever in a long-lived singleton
         # handler, since nothing else ever cleans up a non-top-level run_id.
         is_top = self._top_level.pop(run_id, None)
-        self._parents.pop(run_id, None)
         if not is_top:
+            # Close the node record before dropping this run's _parents entry - the top-ancestor
+            # walk inside _finalize_node still needs it.
+            self._finalize_node(run_id, output=_extract_output(outputs))
+            self._parents.pop(run_id, None)
             return
+        self._parents.pop(run_id, None)
         state = self._runs.pop(run_id, None)
         if state is None:
             return
@@ -369,16 +526,16 @@ class AgentXCallbackHandler(BaseCallbackHandler):
             # spanning several chain/agent/retriever calls) - fold this
             # top-level run into it instead of sending an independent trace.
             active_span._merge_child_run(
-                execution_steps=state["execution_steps"],
                 tool_calls=tool_calls,
-                retrieval_steps=state["retrieval_steps"],
                 input=state["input"],
                 output=output,
                 model=state.get("model"),
                 framework="langchain",
                 input_tokens=state["input_tokens"] or None,
                 output_tokens=state["output_tokens"] or None,
+                emit_steps=False,
             )
+            self._emit_span_tree(active_span, state, tool_calls)
         else:
             # Standalone usage (no enclosing `with tracer.trace()`): open a real root span for
             # this chain invocation and let _merge_child_run explode its accumulated
@@ -388,16 +545,16 @@ class AgentXCallbackHandler(BaseCallbackHandler):
                 # see llamaindex.py's _send_trace for the identical fix and full rationale.
                 span._start = state["start"]
                 span._merge_child_run(
-                    execution_steps=state["execution_steps"],
                     tool_calls=tool_calls,
-                    retrieval_steps=state["retrieval_steps"],
                     input=state["input"],
                     output=output,
                     model=state.get("model"),
                     framework="langchain",
                     input_tokens=state["input_tokens"] or None,
                     output_tokens=state["output_tokens"] or None,
+                    emit_steps=False,
                 )
+                self._emit_span_tree(span, state, tool_calls)
 
     def on_chain_error(
         self,
@@ -409,9 +566,11 @@ class AgentXCallbackHandler(BaseCallbackHandler):
     ) -> None:
         # See on_chain_end's comment - pop for every chain run, not just top-level.
         is_top = self._top_level.pop(run_id, None)
-        self._parents.pop(run_id, None)
         if not is_top:
+            self._finalize_node(run_id, error=str(error))
+            self._parents.pop(run_id, None)
             return
+        self._parents.pop(run_id, None)
         state = self._runs.pop(run_id, None)
         if state is None:
             return
@@ -420,30 +579,30 @@ class AgentXCallbackHandler(BaseCallbackHandler):
         if active_span is not None:
             active_span.set_error(str(error))
             active_span._merge_child_run(
-                execution_steps=state["execution_steps"],
                 tool_calls=state["tool_calls"],
-                retrieval_steps=state["retrieval_steps"],
                 input=state["input"],
                 model=state.get("model"),
                 framework="langchain",
                 input_tokens=state["input_tokens"] or None,
                 output_tokens=state["output_tokens"] or None,
+                emit_steps=False,
             )
+            self._emit_span_tree(active_span, state, state["tool_calls"])
         else:
             # See on_chain_end's matching branch - same standalone-usage handling.
             with self._tracer.trace(self._name, metadata=self._metadata, session_id=self._session_id) as span:
                 span._start = state["start"]
                 span.set_error(str(error))
                 span._merge_child_run(
-                    execution_steps=state["execution_steps"],
                     tool_calls=state["tool_calls"],
-                    retrieval_steps=state["retrieval_steps"],
                     input=state["input"],
                     model=state.get("model"),
                     framework="langchain",
                     input_tokens=state["input_tokens"] or None,
                     output_tokens=state["output_tokens"] or None,
+                    emit_steps=False,
                 )
+                self._emit_span_tree(span, state, state["tool_calls"])
 
     # ------------------------------------------------------------------
     # LLM lifecycle
@@ -545,6 +704,7 @@ class AgentXCallbackHandler(BaseCallbackHandler):
                 if start_t is not None and top and top in self._runs:
                     steps = self._runs[top]["execution_steps"]
                     steps.append({
+                        "parent_run_id": parent_run_id,
                         "name": f"LLM Call {len(steps) + 1}",
                         "duration_ms": (end_t - start_t) * 1000,
                         "start_time": start_t,
@@ -592,6 +752,7 @@ class AgentXCallbackHandler(BaseCallbackHandler):
         start_t = state["start"]
         latency_ms = int((end_t - start_t) * 1000)
         tool_call = {
+            "parent_run_id": parent_run_id,
             "name": state["tool_name"],
             "input": state["tool_input"],
             "output": str(output),
@@ -623,6 +784,7 @@ class AgentXCallbackHandler(BaseCallbackHandler):
         end_t = time.time()
         start_t = state.get("start", end_t)
         tool_call = {
+            "parent_run_id": parent_run_id,
             "name": state.get("tool_name", "unknown"),
             "input": state.get("tool_input"),
             "output": f"ERROR: {error}",
@@ -669,6 +831,7 @@ class AgentXCallbackHandler(BaseCallbackHandler):
         query: Optional[str] = state["query"] or None
         doc_count = len(documents) if hasattr(documents, "__len__") else None
         step: Dict[str, Any] = {
+            "parent_run_id": parent_run_id,
             "name": "Retrieval 1",  # renumbered below
             "duration_ms": (end_t - start_t) * 1000,
             "start_time": start_t,

{agentx_python-0.6.21 → agentx_python-0.6.23}/agentx/integrations/moveworks.py

@@ -189,6 +189,46 @@ class MoveworksSyncReport:
         return base + ")"
 
 
+def judge_sessions_via_engine(ingest: IngestClient, session_ids: List[str]) -> "tuple[int, int, int]":
+    """
+    Judge each session with every enabled session-scoped evaluator via the engine's on-demand
+    judge route (``ifStale=true``: an up-to-date verdict from the sweep or a previous run is left
+    alone). Shared by every pull importer (Moveworks, Databricks). Returns (judged, skipped,
+    failed) counts.
+    """
+    # Same base/key the ingest client already resolved - the judge routes live under
+    # /agent-monitoring on the same engine.
+    base = ingest._base_url
+    http = ingest._session
+    try:
+        response = http.get(f"{base}/agent-monitoring/online-evaluators", timeout=30)
+        response.raise_for_status()
+        evaluators = [
+            e for e in response.json().get("evaluators", [])
+            if e.get("enabled") and e.get("scope") == "session"
+        ]
+    except requests.RequestException:
+        return 0, 0, len(session_ids)
+    judged = skipped = failed = 0
+    for session_id in session_ids:
+        for evaluator in evaluators:
+            try:
+                result = http.post(
+                    f"{base}/agent-monitoring/sessions/{session_id}/judge/{evaluator['_id']}",
+                    params={"ifStale": "true"},
+                    timeout=120,
+                )
+                if result.status_code == 200 and result.json().get("skipped"):
+                    skipped += 1
+                elif result.ok:
+                    judged += 1
+                else:
+                    failed += 1
+            except requests.RequestException:
+                failed += 1
+    return judged, skipped, failed
+
+
 class MoveworksImporter:
     """
     Pulls Moveworks Data API records and replays them into AgentX as traces.
@@ -386,36 +426,10 @@ class MoveworksImporter:
         run) is left alone instead of judged again.
         """
         report = report or MoveworksSyncReport()
-        # Same base/key the ingest client already resolved - the judge routes live under
-        # /agent-monitoring on the same engine.
-        base = self._ingest._base_url
-        http = self._ingest._session
-        try:
-            response = http.get(f"{base}/agent-monitoring/online-evaluators", timeout=30)
-            response.raise_for_status()
-            evaluators = [
-                e for e in response.json().get("evaluators", [])
-                if e.get("enabled") and e.get("scope") == "session"
-            ]
-        except requests.RequestException:
-            report.sessions_judge_failed += len(session_ids)
-            return report
-        for session_id in session_ids:
-            for evaluator in evaluators:
-                try:
-                    result = http.post(
-                        f"{base}/agent-monitoring/sessions/{session_id}/judge/{evaluator['_id']}",
-                        params={"ifStale": "true"},
-                        timeout=120,
-                    )
-                    if result.status_code == 200 and result.json().get("skipped"):
-                        report.sessions_judge_skipped += 1
-                    elif result.ok:
-                        report.sessions_judged += 1
-                    else:
-                        report.sessions_judge_failed += 1
-                except requests.RequestException:
-                    report.sessions_judge_failed += 1
+        judged, skipped, failed = judge_sessions_via_engine(self._ingest, session_ids)
+        report.sessions_judged += judged
+        report.sessions_judge_skipped += skipped
+        report.sessions_judge_failed += failed
         return report
 
 

{agentx_python-0.6.21 → agentx_python-0.6.23}/agentx/tracing/tracer.py

@@ -338,6 +338,7 @@ class _TraceSpan:
         output_tokens: Optional[int] = None,
         cache_read_tokens: Optional[int] = None,
         cache_write_tokens: Optional[int] = None,
+        emit_steps: bool = True,
     ) -> None:
         """
         Explode a whole auto-instrumented sub-run (e.g. one top-level LangChain
@@ -348,9 +349,15 @@ class _TraceSpan:
         tool_calls entries here carry no start/end timing (only latency_ms, see callers e.g.
         langchain.py) - their child span falls back to offset-0 positioning in the tree panel;
         execution_steps (LLM calls) always carry real timing and position correctly.
+
+        ``emit_steps=False`` skips the flat child-span emission but keeps everything else (the
+        flat tool_calls mirror for the Monitor tool-failure check, and the summary adoption) -
+        for callers that emit their own HIERARCHICAL span tree instead (langchain.py's
+        _emit_span_tree parents steps under the graph node that ran them, rather than flat
+        under this span).
         """
         with self._merge_lock:
-            for step in execution_steps or []:
+            for step in [] if not emit_steps else (execution_steps or []):
                 self._child_span_count += 1
                 self.child_span(
                     step.get("name") or f"LLM Call {self._child_span_count}",
@@ -371,14 +378,16 @@ class _TraceSpan:
                 # start_time/end_time; prefer those over latency_ms alone when both know duration
                 # since they also let this child span position correctly in the tree panel
                 # instead of defaulting to offset 0.
-                self.child_span(
-                    tc.get("name") or "Tool call",
-                    start_time=tc.get("start_time"),
-                    end_time=tc.get("end_time"),
-                    duration_ms=tc.get("latency_ms"),
-                    input=tc.get("input"),
-                    output=tc.get("output"),
-                )
+                if emit_steps:
+                    self.child_span(
+                        tc.get("name") or "Tool call",
+                        start_time=tc.get("start_time"),
+                        end_time=tc.get("end_time"),
+                        duration_ms=tc.get("latency_ms"),
+                        input=tc.get("input"),
+                        output=tc.get("output"),
+                        error=None if tc.get("success", True) else str(tc.get("output") or "Tool call failed"),
+                    )
                 # Also mirror onto this span's own flat tool_calls list, sent in this span's own
                 # wire payload on __exit__ (see tool_calls=self.tool_calls or None below). The
                 # child span above is only for the trace detail view's span tree; the engine's
@@ -393,7 +402,7 @@ class _TraceSpan:
                     "latency_ms": tc.get("latency_ms"),
                     "success": tc.get("success", True),
                 })
-            for step in retrieval_steps or []:
+            for step in [] if not emit_steps else (retrieval_steps or []):
                 self._child_span_count += 1
                 self.child_span(
                     step.get("name") or f"Retrieval {self._child_span_count}",

agentx_python-0.6.23/agentx/version.py

@@ -0,0 +1 @@
+VERSION = "0.6.23"

{agentx_python-0.6.21 → agentx_python-0.6.23}/agentx_python.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agentx-python
-Version: 0.6.21
+Version: 0.6.23
 Summary: Official Python SDK for AgentX (https://www.agentx.so/)
 Home-page: https://github.com/AgentX-ai/AgentX-python
 Author: Robin Wang and AgentX Team
@@ -36,6 +36,8 @@ Requires-Dist: llama-index-core>=0.10.0; extra == "llamaindex"
 Provides-Extra: autogen
 Requires-Dist: autogen-agentchat>=0.4.0; extra == "autogen"
 Requires-Dist: autogen-core>=0.4.0; extra == "autogen"
+Provides-Extra: databricks
+Requires-Dist: mlflow>=3.6.0; extra == "databricks"
 Provides-Extra: all
 Requires-Dist: langchain-core>=0.1.0; extra == "all"
 Requires-Dist: crewai>=0.80.0; extra == "all"

{agentx_python-0.6.21 → agentx_python-0.6.23}/agentx_python.egg-info/SOURCES.txt

@@ -31,6 +31,7 @@ agentx/integrations/_traced_call.py
 agentx/integrations/anthropic.py
 agentx/integrations/autogen.py
 agentx/integrations/crewai.py
+agentx/integrations/databricks.py
 agentx/integrations/google_adk.py
 agentx/integrations/google_genai.py
 agentx/integrations/langchain.py

{agentx_python-0.6.21 → agentx_python-0.6.23}/agentx_python.egg-info/entry_points.txt

@@ -1,3 +1,4 @@
 [console_scripts]
+agentx-databricks = agentx.integrations.databricks:cli_main
 agentx-moveworks = agentx.integrations.moveworks:cli_main
 agentx-trace-eval = agentx.cli:main

{agentx_python-0.6.21 → agentx_python-0.6.23}/agentx_python.egg-info/requires.txt

@@ -26,6 +26,9 @@ autogen-core>=0.4.0
 [crewai]
 crewai>=0.80.0
 
+[databricks]
+mlflow>=3.6.0
+
 [google-adk]
 google-adk>=1.0.0
 

{agentx_python-0.6.21 → agentx_python-0.6.23}/setup.py

@@ -35,6 +35,9 @@ setup(
             # Moveworks Data API -> AgentX traces sync (see agentx/integrations/moveworks.py) -
             # a pull importer, since Moveworks agents run in their cloud with no in-process hook.
             "agentx-moveworks=agentx.integrations.moveworks:cli_main",
+            # MLflow/Databricks trace sync (see agentx/integrations/databricks.py) - pull importer
+            # for agents whose serving endpoints can't set OTLP egress env vars.
+            "agentx-databricks=agentx.integrations.databricks:cli_main",
         ],
     },
     install_requires=[
@@ -54,6 +57,7 @@ setup(
         "litellm": ["litellm>=1.0.0"],
         "llamaindex": ["llama-index-core>=0.10.0"],
         "autogen": ["autogen-agentchat>=0.4.0", "autogen-core>=0.4.0"],
+        "databricks": ["mlflow>=3.6.0"],
         "all": [
             "langchain-core>=0.1.0",
             "crewai>=0.80.0",

agentx_python-0.6.21/agentx/version.py

@@ -1 +0,0 @@
-VERSION = "0.6.21"