agentx-python 0.6.22__tar.gz → 0.6.24__tar.gz

{agentx_python-0.6.22 → agentx_python-0.6.24}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agentx-python
-Version: 0.6.22
+Version: 0.6.24
 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.22 → agentx_python-0.6.24}/agentx/integrations/_traced_call.py

@@ -13,9 +13,10 @@ from __future__ import annotations
 
 import asyncio
 import inspect
+import json
 from typing import Any, Callable, Dict, Optional
 
-from agentx.tracing.tracer import Tracer
+from agentx.tracing.tracer import Tracer, _safe_serialize
 
 
 def call_and_trace(
@@ -62,6 +63,32 @@ async def _await_and_finish(
     return response
 
 
+# The engine's unregistered-tool surfacing (Tools & MCPs -> Unregistered) reads a trace's
+# metadata "tools" key to show the REAL definition instead of one inferred from observed
+# arguments (see AgentX-trace-eval's toolSchemas.ts draftFromMetadata). Raw-client patches see
+# the request's tools=[...] right in kwargs, so capture it - capped so a huge toolbox never
+# blows up the trace's metadata budget.
+_MAX_TOOL_DEFINITIONS = 20
+_MAX_TOOL_DEFINITIONS_BYTES = 12_000
+
+
+def capture_tool_definitions(tools: Any) -> Optional[list]:
+    """Return a metadata-ready copy of a request's ``tools=[...]`` list, or None."""
+    if not isinstance(tools, list) or not tools:
+        return None
+    # A plain JSON round-trip preserves nested schema objects exactly (default=str catches the
+    # odd SDK object inside); _safe_serialize would repr-stringify nested dicts, turning a
+    # parameters schema into an unusable string.
+    try:
+        text = json.dumps(tools[:_MAX_TOOL_DEFINITIONS], default=str)
+        if len(text) > _MAX_TOOL_DEFINITIONS_BYTES:
+            return None
+        serialized = json.loads(text)
+    except (TypeError, ValueError):
+        return None
+    return serialized if isinstance(serialized, list) else None
+
+
 def finish_llm_call(
     tracer: Tracer,
     *,
@@ -79,6 +106,7 @@ def finish_llm_call(
     error: Optional[str],
     cache_read_tokens: Optional[int] = None,
     cache_write_tokens: Optional[int] = None,
+    tool_definitions: Optional[list] = None,
 ) -> None:
     """
     Close out one raw-client LLM call - shared by the ``on_finish``/exit
@@ -96,8 +124,16 @@ def finish_llm_call(
     """
     latency_ms = int((end_t - start_t) * 1000)
 
+    if tool_definitions:
+        metadata = {**(metadata or {}), "tools": tool_definitions}
+
     active_span = tracer.current_span
     if active_span is not None:
+        # The definitions describe the whole call's toolbox - attach them to the enclosing
+        # span's metadata (first capture wins) so the ROOT trace carries them for the
+        # unregistered-tool listing, same as the standalone-trace path below.
+        if tool_definitions and not (active_span._metadata or {}).get("tools"):
+            active_span._metadata = {**(active_span._metadata or {}), "tools": tool_definitions}
         if error is not None:
             active_span.set_error(error)
         active_span._record_llm_call(

{agentx_python-0.6.22 → agentx_python-0.6.24}/agentx/integrations/anthropic.py

@@ -22,7 +22,7 @@ import time
 from typing import Any, Dict, Optional, Tuple
 
 from agentx.tracing.tracer import Tracer, _safe_serialize
-from agentx.integrations._traced_call import call_and_trace, finish_llm_call
+from agentx.integrations._traced_call import capture_tool_definitions, call_and_trace, finish_llm_call
 
 
 def _extract_output_text(response: Any) -> Optional[str]:
@@ -136,6 +136,7 @@ def _patch_create(
             kwargs.get("system"),
         )
         model = kwargs.get("model")
+        tool_definitions = capture_tool_definitions(kwargs.get("tools"))
 
         input_repr = _safe_serialize(input_messages)
 
@@ -171,6 +172,7 @@ def _patch_create(
                 cache_read_tokens=cache_read_tokens,
                 cache_write_tokens=cache_write_tokens,
                 error=error,
+                tool_definitions=tool_definitions,
             )
 
         return call_and_trace(original, args, kwargs, on_finish)
@@ -199,6 +201,7 @@ def _patch_stream(
         ctx = original_stream(*args, **kwargs)
         input_repr = _safe_serialize(_prepend_system(kwargs.get("messages"), kwargs.get("system")))
         model = kwargs.get("model")
+        tool_definitions = capture_tool_definitions(kwargs.get("tools"))
 
         def build_and_send(end_t: float, error: Optional[str], final_message: Optional[Any]) -> None:
             output = None
@@ -230,6 +233,7 @@ def _patch_stream(
                 cache_read_tokens=cache_read_tokens,
                 cache_write_tokens=cache_write_tokens,
                 error=error,
+                tool_definitions=tool_definitions,
             )
 
         class _TracedStream:

agentx_python-0.6.24/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.22 → agentx_python-0.6.24}/agentx/integrations/google_genai.py

@@ -20,7 +20,7 @@ import time
 from typing import Any, Dict, List, Optional
 
 from agentx.tracing.tracer import Tracer, _safe_serialize
-from agentx.integrations._traced_call import call_and_trace, finish_llm_call
+from agentx.integrations._traced_call import capture_tool_definitions, call_and_trace, finish_llm_call
 
 
 def patch_genai_client(
@@ -90,6 +90,31 @@ def _extract_response_text(response: Any) -> Optional[str]:
     return None
 
 
+
+def _genai_tool_definitions(config: Any) -> Optional[list]:
+    """
+    Flatten a GenAI request's config.tools (each a Tool holding function_declarations) into the
+    flat [{name, description, parameters}, ...] list the engine's unregistered-tool listing
+    reads from trace metadata. Plain dicts and SDK objects both handled.
+    """
+    tools = getattr(config, "tools", None) if config is not None and not isinstance(config, dict) else (config or {}).get("tools") if isinstance(config, dict) else None
+    if not tools:
+        return None
+    flat = []
+    for tool in tools:
+        decls = getattr(tool, "function_declarations", None) if not isinstance(tool, dict) else tool.get("function_declarations")
+        for decl in decls or []:
+            if isinstance(decl, dict):
+                flat.append(decl)
+            else:
+                flat.append({
+                    "name": getattr(decl, "name", None),
+                    "description": getattr(decl, "description", None),
+                    "parameters": getattr(decl, "parameters", None),
+                })
+    return capture_tool_definitions(flat)
+
+
 def _patch_generate_content(
     models: Any,
     tracer: Tracer,
@@ -106,6 +131,7 @@ def _patch_generate_content(
         model = kwargs.get("model") or (args[0] if args else None)
         contents = kwargs.get("contents") or (args[1] if len(args) > 1 else None)
         input_repr = contents if isinstance(contents, str) else _safe_serialize(contents)
+        tool_definitions = _genai_tool_definitions(kwargs.get("config"))
 
         def on_finish(response: Optional[Any], error: Optional[str]) -> None:
             end_t = time.time()
@@ -137,6 +163,7 @@ def _patch_generate_content(
                 output_tokens=output_tokens,
                 cache_read_tokens=cache_read_tokens,
                 error=error,
+                tool_definitions=tool_definitions,
             )
 
         return call_and_trace(original, args, kwargs, on_finish)
@@ -180,6 +207,7 @@ def _patch_sync_generate_content_stream(
     def patched_stream(*args, **kwargs):
         start_t = time.time()
         model, input_repr = _stream_input_and_model(args, kwargs)
+        tool_definitions = _genai_tool_definitions(kwargs.get("config"))
         accumulated_text: List[str] = []
         last_usage_metadata = None
         error: Optional[str] = None
@@ -222,6 +250,7 @@ def _patch_sync_generate_content_stream(
                 output_tokens=output_tokens,
                 cache_read_tokens=cache_read_tokens,
                 error=error,
+                tool_definitions=tool_definitions,
             )
 
     patched_stream._agentx_patched = True
@@ -247,6 +276,7 @@ def _patch_async_generate_content_stream(
         # drives it.
         start_t = time.time()
         model, input_repr = _stream_input_and_model(args, kwargs)
+        tool_definitions = _genai_tool_definitions(kwargs.get("config"))
         inner = await original_stream(*args, **kwargs)
 
         async def traced_agen():
@@ -290,6 +320,7 @@ def _patch_async_generate_content_stream(
                     output_tokens=output_tokens,
                     cache_read_tokens=cache_read_tokens,
                     error=error,
+                    tool_definitions=tool_definitions,
                 )
 
         return traced_agen()

{agentx_python-0.6.22 → agentx_python-0.6.24}/agentx/integrations/langchain.py

@@ -26,6 +26,7 @@ from typing import Any, Dict, List, Optional, Union
 from uuid import UUID
 
 from agentx.tracing.tracer import Tracer, _safe_serialize
+from agentx.integrations._traced_call import capture_tool_definitions
 
 try:
     from langchain_core.callbacks.base import BaseCallbackHandler
@@ -375,6 +376,10 @@ class AgentXCallbackHandler(BaseCallbackHandler):
                 # skipped plumbing runs to the nearest emitted ancestor.
                 "node_runs": {},
                 "chain_parents": {},
+                # The request's tools=[...] as seen on the first LLM call's invocation params -
+                # attached to the root trace's metadata so the engine's unregistered-tool
+                # listing can surface the REAL definition (not one inferred from arguments).
+                "tool_definitions": None,
             }
         else:
             top = self._find_top_ancestor(parent_run_id)
@@ -522,6 +527,8 @@ class AgentXCallbackHandler(BaseCallbackHandler):
 
         active_span = self._tracer.current_span
         if active_span is not None:
+            if state.get("tool_definitions") and not (active_span._metadata or {}).get("tools"):
+                active_span._metadata = {**(active_span._metadata or {}), "tools": state["tool_definitions"]}
             # Part of a `with tracer.trace(...)` block (e.g. an orchestrator
             # spanning several chain/agent/retriever calls) - fold this
             # top-level run into it instead of sending an independent trace.
@@ -540,7 +547,15 @@ class AgentXCallbackHandler(BaseCallbackHandler):
             # Standalone usage (no enclosing `with tracer.trace()`): open a real root span for
             # this chain invocation and let _merge_child_run explode its accumulated
             # execution_steps/tool_calls/retrieval_steps into real child-span rows.
-            with self._tracer.trace(self._name, metadata=self._metadata, session_id=self._session_id) as span:
+            with self._tracer.trace(
+                self._name,
+                metadata=(
+                    {**(self._metadata or {}), "tools": state["tool_definitions"]}
+                    if state.get("tool_definitions")
+                    else self._metadata
+                ),
+                session_id=self._session_id,
+            ) as span:
                 # __enter__ just set _start to "now" - overridden to the chain's real start time,
                 # see llamaindex.py's _send_trace for the identical fix and full rationale.
                 span._start = state["start"]
@@ -577,6 +592,8 @@ class AgentXCallbackHandler(BaseCallbackHandler):
 
         active_span = self._tracer.current_span
         if active_span is not None:
+            if state.get("tool_definitions") and not (active_span._metadata or {}).get("tools"):
+                active_span._metadata = {**(active_span._metadata or {}), "tools": state["tool_definitions"]}
             active_span.set_error(str(error))
             active_span._merge_child_run(
                 tool_calls=state["tool_calls"],
@@ -590,7 +607,15 @@ class AgentXCallbackHandler(BaseCallbackHandler):
             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:
+            with self._tracer.trace(
+                self._name,
+                metadata=(
+                    {**(self._metadata or {}), "tools": state["tool_definitions"]}
+                    if state.get("tool_definitions")
+                    else self._metadata
+                ),
+                session_id=self._session_id,
+            ) as span:
                 span._start = state["start"]
                 span.set_error(str(error))
                 span._merge_child_run(
@@ -629,6 +654,11 @@ class AgentXCallbackHandler(BaseCallbackHandler):
             or serialized.get("name")
         )
         model = str(model) if model and model not in ("None", "none") else None
+        top_for_tools = self._find_top_ancestor(parent_run_id)
+        if top_for_tools and top_for_tools in self._runs and not self._runs[top_for_tools].get("tool_definitions"):
+            captured = capture_tool_definitions(kwargs.get("invocation_params", {}).get("tools"))
+            if captured:
+                self._runs[top_for_tools]["tool_definitions"] = captured
         self._runs[run_id] = {
             "llm_start": time.time(),
             "model": model,

{agentx_python-0.6.22 → agentx_python-0.6.24}/agentx/integrations/litellm.py

@@ -18,7 +18,7 @@ from __future__ import annotations
 from typing import Any, Dict, Optional, Tuple
 
 from agentx.tracing.tracer import Tracer, _safe_serialize
-from agentx.integrations._traced_call import finish_llm_call
+from agentx.integrations._traced_call import capture_tool_definitions, finish_llm_call
 
 try:
     from litellm.integrations.custom_logger import CustomLogger
@@ -98,6 +98,10 @@ class AgentXLiteLLMLogger(CustomLogger):
     def _finish(self, kwargs: Dict[str, Any], response_obj: Any, start_time: Any, end_time: Any, error: Optional[str]) -> None:
         model = kwargs.get("model")
         input_repr = _safe_serialize(kwargs.get("messages"))
+        # LiteLLM surfaces the request's tools under optional_params (and sometimes top-level).
+        tool_definitions = capture_tool_definitions(
+            kwargs.get("tools") or (kwargs.get("optional_params") or {}).get("tools")
+        )
         output = None
         input_tokens = None
         output_tokens = None
@@ -118,6 +122,7 @@ class AgentXLiteLLMLogger(CustomLogger):
             output=output,
             model=model,
             input_tokens=input_tokens,
+            tool_definitions=tool_definitions,
             output_tokens=output_tokens,
             cache_read_tokens=cache_read_tokens,
             error=error,

{agentx_python-0.6.22 → agentx_python-0.6.24}/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.22 → agentx_python-0.6.24}/agentx/integrations/openai.py

@@ -28,7 +28,7 @@ import time
 from typing import Any, Dict, Optional, Tuple
 
 from agentx.tracing.tracer import Tracer, _safe_serialize
-from agentx.integrations._traced_call import call_and_trace, finish_llm_call
+from agentx.integrations._traced_call import capture_tool_definitions, call_and_trace, finish_llm_call
 
 
 def _extract_output_text(response: Any) -> Optional[str]:
@@ -128,6 +128,7 @@ def _patch_chat_completions_create(
         input_messages = kwargs.get("messages") or (args[0] if args else None)
         model = kwargs.get("model")
         input_repr = _safe_serialize(input_messages)
+        tool_definitions = capture_tool_definitions(kwargs.get("tools"))
 
         def on_finish(response: Optional[Any], error: Optional[str]) -> None:
             end_t = time.time()
@@ -159,6 +160,7 @@ def _patch_chat_completions_create(
                 output_tokens=output_tokens,
                 cache_read_tokens=cache_read_tokens,
                 error=error,
+                tool_definitions=tool_definitions,
             )
 
         return call_and_trace(original, args, kwargs, on_finish)

agentx_python-0.6.24/agentx/version.py

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

{agentx_python-0.6.22 → agentx_python-0.6.24}/agentx_python.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agentx-python
-Version: 0.6.22
+Version: 0.6.24
 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.22 → agentx_python-0.6.24}/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.22 → agentx_python-0.6.24}/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.22 → agentx_python-0.6.24}/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.22 → agentx_python-0.6.24}/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.22/agentx/version.py

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