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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of agentx-python might be problematic. Click here for more details.

Files changed (70) hide show
  1. {agentx_python-0.6.22 → agentx_python-0.6.23}/PKG-INFO +3 -1
  2. agentx_python-0.6.23/agentx/integrations/databricks.py +405 -0
  3. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/integrations/moveworks.py +44 -30
  4. agentx_python-0.6.23/agentx/version.py +1 -0
  5. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx_python.egg-info/PKG-INFO +3 -1
  6. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx_python.egg-info/SOURCES.txt +1 -0
  7. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx_python.egg-info/entry_points.txt +1 -0
  8. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx_python.egg-info/requires.txt +3 -0
  9. {agentx_python-0.6.22 → agentx_python-0.6.23}/setup.py +4 -0
  10. agentx_python-0.6.22/agentx/version.py +0 -1
  11. {agentx_python-0.6.22 → agentx_python-0.6.23}/LICENSE +0 -0
  12. {agentx_python-0.6.22 → agentx_python-0.6.23}/README.md +0 -0
  13. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/__init__.py +0 -0
  14. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/agentx.py +0 -0
  15. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/cli.py +0 -0
  16. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/__init__.py +0 -0
  17. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/_term.py +0 -0
  18. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/adapters/__init__.py +0 -0
  19. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/adapters/http_endpoint.py +0 -0
  20. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/adapters/precomputed.py +0 -0
  21. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/adapters/raw.py +0 -0
  22. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/client.py +0 -0
  23. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/datasets.py +0 -0
  24. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/evaluation_settings.py +0 -0
  25. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/models.py +0 -0
  26. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/prompts.py +0 -0
  27. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/redaction.py +0 -0
  28. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/reporting.py +0 -0
  29. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/results.py +0 -0
  30. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/runner.py +0 -0
  31. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/tracing.py +0 -0
  32. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/exceptions.py +0 -0
  33. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/feedback.py +0 -0
  34. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/integrations/__init__.py +0 -0
  35. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/integrations/_traced_call.py +0 -0
  36. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/integrations/anthropic.py +0 -0
  37. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/integrations/autogen.py +0 -0
  38. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/integrations/crewai.py +0 -0
  39. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/integrations/google_adk.py +0 -0
  40. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/integrations/google_genai.py +0 -0
  41. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/integrations/langchain.py +0 -0
  42. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/integrations/litellm.py +0 -0
  43. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/integrations/llamaindex.py +0 -0
  44. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/integrations/openai.py +0 -0
  45. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/integrations/openai_agents.py +0 -0
  46. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/monitor/__init__.py +0 -0
  47. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/monitor/client.py +0 -0
  48. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/monitor/models.py +0 -0
  49. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/monitor/online_evaluators.py +0 -0
  50. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/monitor/patterns.py +0 -0
  51. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/monitor/profile.py +0 -0
  52. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/monitor/signals.py +0 -0
  53. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/outcomes.py +0 -0
  54. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/py.typed +0 -0
  55. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/resources/__init__.py +0 -0
  56. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/resources/agent.py +0 -0
  57. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/resources/conversation.py +0 -0
  58. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/resources/workforce.py +0 -0
  59. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/tracing/__init__.py +0 -0
  60. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/tracing/ci_types.py +0 -0
  61. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/tracing/ingest_client.py +0 -0
  62. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/tracing/tracer.py +0 -0
  63. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/util.py +0 -0
  64. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx_python.egg-info/dependency_links.txt +0 -0
  65. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx_python.egg-info/not-zip-safe +0 -0
  66. {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx_python.egg-info/top_level.txt +0 -0
  67. {agentx_python-0.6.22 → agentx_python-0.6.23}/setup.cfg +0 -0
  68. {agentx_python-0.6.22 → agentx_python-0.6.23}/tests/test_integration.py +0 -0
  69. {agentx_python-0.6.22 → agentx_python-0.6.23}/tests/test_integrations.py +0 -0
  70. {agentx_python-0.6.22 → agentx_python-0.6.23}/tests/test_span_tree.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: agentx-python
3
- Version: 0.6.22
3
+ Version: 0.6.23
4
4
  Summary: Official Python SDK for AgentX (https://www.agentx.so/)
5
5
  Home-page: https://github.com/AgentX-ai/AgentX-python
6
6
  Author: Robin Wang and AgentX Team
@@ -36,6 +36,8 @@ Requires-Dist: llama-index-core>=0.10.0; extra == "llamaindex"
36
36
  Provides-Extra: autogen
37
37
  Requires-Dist: autogen-agentchat>=0.4.0; extra == "autogen"
38
38
  Requires-Dist: autogen-core>=0.4.0; extra == "autogen"
39
+ Provides-Extra: databricks
40
+ Requires-Dist: mlflow>=3.6.0; extra == "databricks"
39
41
  Provides-Extra: all
40
42
  Requires-Dist: langchain-core>=0.1.0; extra == "all"
41
43
  Requires-Dist: crewai>=0.80.0; extra == "all"
@@ -0,0 +1,405 @@
1
+ """
2
+ Databricks / MLflow integration for AgentX.
3
+
4
+ Two complementary paths for agents built on Databricks (Agent Bricks or the Mosaic AI Agent
5
+ Framework - both are auto-instrumented by MLflow 3 Tracing):
6
+
7
+ 1. **Push (live)** - ``enable_mlflow_export()``: point MLflow Tracing's native OTLP exporter at
8
+ the AgentX engine's OTel endpoint. Works anywhere MLflow traces run (notebooks, jobs, Model
9
+ Serving endpoints via environment variables). Dual export keeps Databricks' own MLflow UI and
10
+ inference tables working alongside AgentX::
11
+
12
+ from agentx.integrations.databricks import enable_mlflow_export
13
+
14
+ enable_mlflow_export(
15
+ api_key=os.environ["AGENTX_API_KEY"],
16
+ base_url="http://localhost:4700/api/v1", # your AgentX engine
17
+ service_name="my-databricks-agent",
18
+ )
19
+ # ... then trace as usual (@mlflow.trace, autolog, Agent Framework, ...)
20
+
21
+ On a Model Serving endpoint, set the equivalent environment variables instead (this helper
22
+ prints them with ``dry_run=True``).
23
+
24
+ 2. **Pull (batch)** - ``agentx-databricks sync``: import finished MLflow traces from a Databricks
25
+ (or any MLflow 3) tracking server into AgentX - full span trees, tool calls, sessions -
26
+ deduplicated on deterministic span ids so re-running a window never duplicates, with
27
+ ``--monitor`` / ``--judge-sessions`` mirroring ``agentx-moveworks``::
28
+
29
+ export AGENTX_API_KEY=... # AgentX project key
30
+ export DATABRICKS_HOST=... DATABRICKS_TOKEN=... # or MLFLOW_TRACKING_URI
31
+ agentx-databricks sync --experiment-id 123456 --since 24h
32
+
33
+ Requires (pull path): ``pip install "agentx-python[databricks]"`` (mlflow>=3).
34
+ """
35
+ from __future__ import annotations
36
+
37
+ import argparse
38
+ import json
39
+ import os
40
+ from datetime import datetime, timedelta, timezone
41
+ from pathlib import Path
42
+ from typing import Any, Dict, Iterable, List, Optional
43
+
44
+ from agentx.tracing.ingest_client import IngestClient
45
+ from agentx.version import VERSION
46
+ from agentx.integrations.moveworks import (
47
+ _parse_since,
48
+ _parse_time,
49
+ _read_cursor,
50
+ _write_cursor,
51
+ judge_sessions_via_engine,
52
+ )
53
+
54
+ _DEFAULT_CURSOR_FILE = Path.home() / ".agentx" / "databricks_sync_cursor.json"
55
+
56
+ # MLflow's session grouping metadata key (mlflow.update_current_trace(metadata={...})) - traces
57
+ # sharing it become one AgentX session, judged as a conversation by session-scoped evaluators.
58
+ _SESSION_METADATA_KEYS = ("mlflow.trace.session", "mlflow.trace.session_id", "session_id")
59
+
60
+
61
+ # ----------------------------------------------------------------------------------------------
62
+ # Path 1: push - MLflow OTLP export pointed at the AgentX engine
63
+ # ----------------------------------------------------------------------------------------------
64
+
65
+ def enable_mlflow_export(
66
+ *,
67
+ api_key: str,
68
+ base_url: str,
69
+ service_name: str = "databricks-agent",
70
+ dual: bool = True,
71
+ genai_semconv: bool = False,
72
+ dry_run: bool = False,
73
+ ) -> Dict[str, str]:
74
+ """
75
+ Configure MLflow Tracing's built-in OTLP exporter to send every trace to the AgentX engine.
76
+
77
+ Must run BEFORE the first trace starts (MLflow reads these once, at tracer setup). ``dual``
78
+ keeps MLflow's own tracking export too (Databricks MLflow UI / inference tables keep
79
+ working); ``genai_semconv`` switches the wire format to OTel GenAI semantic conventions -
80
+ AgentX ingests both, and the default (MLflow-native attributes) is the higher-fidelity
81
+ mapping for plain ``@mlflow.trace`` functions. Returns the environment variables set - with
82
+ ``dry_run=True`` nothing is set, so the dict can be copied onto a Databricks Model Serving
83
+ endpoint's environment variables instead.
84
+ """
85
+ env = {
86
+ "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT": f"{base_url.rstrip('/')}/otel/v1/traces",
87
+ "OTEL_EXPORTER_OTLP_TRACES_HEADERS": f"x-api-key={api_key}",
88
+ "OTEL_EXPORTER_OTLP_PROTOCOL": "http/protobuf",
89
+ "OTEL_SERVICE_NAME": service_name,
90
+ }
91
+ if dual:
92
+ env["MLFLOW_TRACE_ENABLE_OTLP_DUAL_EXPORT"] = "true"
93
+ if genai_semconv:
94
+ env["MLFLOW_ENABLE_OTEL_GENAI_SEMCONV"] = "true"
95
+ if not dry_run:
96
+ os.environ.update(env)
97
+ return env
98
+
99
+
100
+ # ----------------------------------------------------------------------------------------------
101
+ # Path 2: pull - MLflow trace search -> AgentX span trees
102
+ # ----------------------------------------------------------------------------------------------
103
+
104
+ class DatabricksSyncReport:
105
+ def __init__(self) -> None:
106
+ self.traces = 0
107
+ self.spans = 0
108
+ self.tool_calls = 0
109
+ self.ingested = 0
110
+ self.failed = 0
111
+ self.skipped_in_progress = 0
112
+ self.session_ids: List[str] = []
113
+ self.sessions_judged = 0
114
+ self.sessions_judge_skipped = 0
115
+ self.sessions_judge_failed = 0
116
+
117
+ def __repr__(self) -> str: # also what the CLI prints
118
+ parts = [
119
+ f"traces={self.traces}",
120
+ f"spans={self.spans}",
121
+ f"tool_calls={self.tool_calls}",
122
+ f"ingested={self.ingested}",
123
+ f"failed={self.failed}",
124
+ f"sessions={len(self.session_ids)}",
125
+ ]
126
+ if self.skipped_in_progress:
127
+ parts.append(f"skipped_in_progress={self.skipped_in_progress}")
128
+ if self.sessions_judged or self.sessions_judge_skipped or self.sessions_judge_failed:
129
+ parts.append(
130
+ f"judged={self.sessions_judged} judge_skipped={self.sessions_judge_skipped} "
131
+ f"judge_failed={self.sessions_judge_failed}"
132
+ )
133
+ return f"DatabricksSyncReport({' '.join(parts)})"
134
+
135
+
136
+ def _span_time_ns(span: Any, attr: str) -> Optional[int]:
137
+ value = getattr(span, attr, None)
138
+ return int(value) if isinstance(value, (int, float)) and value > 0 else None
139
+
140
+
141
+ def _span_error(span: Any) -> Optional[str]:
142
+ status = getattr(span, "status", None)
143
+ code = str(getattr(status, "status_code", "") or "")
144
+ if "ERROR" in code.upper():
145
+ return str(getattr(status, "description", None) or "error")
146
+ return None
147
+
148
+
149
+ def _serialize(value: Any) -> Any:
150
+ if value is None or isinstance(value, (str, int, float, bool)):
151
+ return value
152
+ try:
153
+ return json.loads(json.dumps(value, default=str))
154
+ except (TypeError, ValueError):
155
+ return str(value)
156
+
157
+
158
+ class DatabricksTraceImporter:
159
+ """
160
+ Replays finished MLflow traces (Databricks-hosted or any MLflow 3 tracking server) into
161
+ AgentX as full span trees. ``span_id`` is deterministic (``dbx:<trace_id>[:<span_id>]``), so
162
+ re-syncing a window is idempotent - the engine dedupes on replay and skips re-judging.
163
+ """
164
+
165
+ def __init__(
166
+ self,
167
+ mlflow_client: Any,
168
+ *,
169
+ agentx_api_key: str,
170
+ agentx_base_url: Optional[str] = None,
171
+ agent_name: Optional[str] = None,
172
+ ) -> None:
173
+ self._client = mlflow_client
174
+ self._agent_name = agent_name
175
+ self._ingest = IngestClient(agentx_api_key, sdk_version=VERSION, base_url=agentx_base_url)
176
+
177
+ # -- trace -> wires ------------------------------------------------------------------------
178
+
179
+ def _trace_wires(self, trace: Any) -> "tuple[List[Dict[str, Any]], Optional[str], int]":
180
+ """One MLflow Trace -> [root wire, *child wires], its session id, and its tool-call count."""
181
+ info = trace.info
182
+ trace_id = str(getattr(info, "trace_id", None) or getattr(info, "request_id", ""))
183
+ spans = list(getattr(trace.data, "spans", None) or [])
184
+ if not trace_id or not spans:
185
+ return [], None, 0
186
+
187
+ metadata_bag: Dict[str, Any] = {}
188
+ for source in (getattr(info, "trace_metadata", None), getattr(info, "tags", None)):
189
+ if isinstance(source, dict):
190
+ metadata_bag.update(source)
191
+ session_raw = next((metadata_bag[k] for k in _SESSION_METADATA_KEYS if metadata_bag.get(k)), None)
192
+ session_id = f"dbx_{session_raw}" if session_raw else None
193
+
194
+ root = next((s for s in spans if not getattr(s, "parent_id", None)), spans[0])
195
+ root_span_id = str(getattr(root, "span_id", "") or "root")
196
+ wire_span_id = {root_span_id: f"dbx:{trace_id}"}
197
+ for span in spans:
198
+ sid = str(getattr(span, "span_id", "") or "")
199
+ if sid and sid not in wire_span_id:
200
+ wire_span_id[sid] = f"dbx:{trace_id}:{sid}"
201
+
202
+ tool_calls: List[Dict[str, Any]] = []
203
+ wires: List[Dict[str, Any]] = []
204
+ for span in spans:
205
+ sid = str(getattr(span, "span_id", "") or "")
206
+ is_root = span is root
207
+ start_ns = _span_time_ns(span, "start_time_ns")
208
+ end_ns = _span_time_ns(span, "end_time_ns")
209
+ error = _span_error(span)
210
+ span_type = str(getattr(span, "span_type", "") or "").upper()
211
+ wire: Dict[str, Any] = {
212
+ "name": (self._agent_name if is_root and self._agent_name else str(getattr(span, "name", "span"))),
213
+ "framework": "databricks",
214
+ "span_id": wire_span_id.get(sid, f"dbx:{trace_id}:{sid or 'span'}"),
215
+ }
216
+ if not is_root:
217
+ parent_sid = str(getattr(span, "parent_id", "") or "")
218
+ wire["parent_span_id"] = wire_span_id.get(parent_sid, f"dbx:{trace_id}")
219
+ if session_id:
220
+ wire["session_id"] = session_id
221
+ if start_ns:
222
+ wire["started_at_unix_nano"] = str(start_ns)
223
+ if start_ns and end_ns and end_ns > start_ns:
224
+ wire["latency_ms"] = int((end_ns - start_ns) / 1_000_000)
225
+ inputs = _serialize(getattr(span, "inputs", None))
226
+ outputs = _serialize(getattr(span, "outputs", None))
227
+ if inputs is not None:
228
+ wire["input"] = inputs
229
+ if outputs is not None:
230
+ wire["output"] = outputs
231
+ if error:
232
+ wire["error"] = error
233
+ if is_root:
234
+ wire["metadata"] = {
235
+ "source": "databricks",
236
+ "mlflowTraceId": trace_id,
237
+ **({"experimentId": str(getattr(info, "experiment_id", ""))} if getattr(info, "experiment_id", None) else {}),
238
+ }
239
+ if span_type == "TOOL":
240
+ tool_calls.append(
241
+ {
242
+ "name": str(getattr(span, "name", "tool")),
243
+ "input": inputs,
244
+ "output": outputs,
245
+ "latency_ms": wire.get("latency_ms"),
246
+ "success": not error,
247
+ }
248
+ )
249
+ wires.append(wire)
250
+
251
+ # Root carries the flat tool_calls mirror - what the engine's Tool-failure check and
252
+ # trajectory matching read (same posture as tracer._merge_child_run).
253
+ if tool_calls:
254
+ wires[spans.index(root)]["tool_calls"] = tool_calls
255
+ # Root first so the engine resolves the agent before children arrive.
256
+ wires.sort(key=lambda w: 0 if "parent_span_id" not in w else 1)
257
+ return wires, session_id, len(tool_calls)
258
+
259
+ def _search(self, experiment_ids: List[str], since: datetime, until: datetime) -> Iterable[Any]:
260
+ """Newest-first paginated search, stopping once a page is entirely older than ``since``."""
261
+ page_token: Optional[str] = None
262
+ since_ms = int(since.timestamp() * 1000)
263
+ until_ms = int(until.timestamp() * 1000)
264
+ while True:
265
+ page = self._client.search_traces(
266
+ experiment_ids=experiment_ids,
267
+ max_results=100,
268
+ page_token=page_token,
269
+ order_by=["timestamp_ms DESC"],
270
+ )
271
+ oldest_seen = None
272
+ for trace in page:
273
+ ts = getattr(trace.info, "request_time", None) or getattr(trace.info, "timestamp_ms", None)
274
+ ts_ms = int(ts.timestamp() * 1000) if isinstance(ts, datetime) else (int(ts) if ts else None)
275
+ oldest_seen = ts_ms if ts_ms is not None else oldest_seen
276
+ if ts_ms is not None and (ts_ms < since_ms or ts_ms >= until_ms):
277
+ if ts_ms < since_ms:
278
+ continue
279
+ continue
280
+ yield trace
281
+ page_token = getattr(page, "token", None)
282
+ if not page_token or (oldest_seen is not None and oldest_seen < since_ms):
283
+ return
284
+
285
+ # -- public --------------------------------------------------------------------------------
286
+
287
+ def sync(
288
+ self,
289
+ experiment_ids: List[str],
290
+ since: datetime,
291
+ until: Optional[datetime] = None,
292
+ *,
293
+ monitor: bool = False,
294
+ judge_sessions: bool = False,
295
+ dry_run: bool = False,
296
+ on_payload: Optional[Any] = None,
297
+ ) -> DatabricksSyncReport:
298
+ """
299
+ Import every finished MLflow trace in ``[since, until)`` from the given experiments.
300
+ Safe to re-run over the same window (span_id dedupe). ``monitor=True`` opts every
301
+ imported root trace into the engine's ingest-time checks; ``judge_sessions=True`` judges
302
+ each imported session afterwards (``ifStale`` - never duplicates the engine's own sweep).
303
+ """
304
+ report = DatabricksSyncReport()
305
+ until = until or datetime.now(timezone.utc)
306
+ for trace in self._search(experiment_ids, since, until):
307
+ state = str(getattr(trace.info, "state", "") or "")
308
+ if state and "IN_PROGRESS" in state.upper():
309
+ report.skipped_in_progress += 1
310
+ continue
311
+ wires, session_id, tool_count = self._trace_wires(trace)
312
+ if not wires:
313
+ continue
314
+ report.traces += 1
315
+ report.spans += len(wires)
316
+ report.tool_calls += tool_count
317
+ if session_id and session_id not in report.session_ids:
318
+ report.session_ids.append(session_id)
319
+ for wire in wires:
320
+ if monitor and "parent_span_id" not in wire:
321
+ wire["monitor"] = True
322
+ if on_payload is not None:
323
+ on_payload(wire)
324
+ if dry_run:
325
+ continue
326
+ if self._ingest.send_trace_sync(wire) is not None:
327
+ report.ingested += 1
328
+ else:
329
+ report.failed += 1
330
+ if judge_sessions and not dry_run:
331
+ judged, skipped, failed = judge_sessions_via_engine(self._ingest, report.session_ids)
332
+ report.sessions_judged += judged
333
+ report.sessions_judge_skipped += skipped
334
+ report.sessions_judge_failed += failed
335
+ return report
336
+
337
+
338
+ # ----------------------------------------------------------------------------------------------
339
+ # CLI: agentx-databricks sync
340
+ # ----------------------------------------------------------------------------------------------
341
+
342
+ def cli_main(argv: Optional[List[str]] = None) -> None:
343
+ parser = argparse.ArgumentParser(
344
+ prog="agentx-databricks",
345
+ description="Sync MLflow traces (Databricks agents) into AgentX as span trees/sessions.",
346
+ )
347
+ sub = parser.add_subparsers(dest="command", required=True)
348
+ sync = sub.add_parser("sync", help="Import a time window of MLflow traces")
349
+ sync.add_argument("--experiment-id", action="append", required=True, help="MLflow experiment id (repeatable)")
350
+ sync.add_argument("--since", help="ISO timestamp or relative (7d, 24h). Default: cursor file, else 24h")
351
+ sync.add_argument("--until", help="ISO timestamp (default: now)")
352
+ sync.add_argument("--cursor-file", default=str(_DEFAULT_CURSOR_FILE), help="Incremental cursor path")
353
+ sync.add_argument("--no-cursor", action="store_true", help="Ignore and don't update the cursor file")
354
+ sync.add_argument("--agent-name", help="Attribute every trace to this agent (default: the root span's name)")
355
+ sync.add_argument(
356
+ "--tracking-uri",
357
+ default=os.getenv("MLFLOW_TRACKING_URI", "databricks"),
358
+ help='MLflow tracking URI (default: MLFLOW_TRACKING_URI, else "databricks" - uses DATABRICKS_HOST/TOKEN)',
359
+ )
360
+ sync.add_argument("--monitor", action="store_true", help="Run ingest-time checks on each imported trace")
361
+ sync.add_argument(
362
+ "--judge-sessions",
363
+ action="store_true",
364
+ help="After the sync, judge each imported session with every enabled session-scoped evaluator "
365
+ "(skips sessions that already have an up-to-date verdict)",
366
+ )
367
+ sync.add_argument("--dry-run", action="store_true", help="Map and print payloads without ingesting")
368
+
369
+ args = parser.parse_args(argv)
370
+
371
+ agentx_key = os.getenv("AGENTX_API_KEY", "")
372
+ if not agentx_key and not args.dry_run:
373
+ raise SystemExit("agentx-databricks: set AGENTX_API_KEY (your AgentX project API key)")
374
+
375
+ try:
376
+ from mlflow.client import MlflowClient
377
+ except ImportError as exc: # pragma: no cover
378
+ raise SystemExit('agentx-databricks: pip install "agentx-python[databricks]" (needs mlflow>=3)') from exc
379
+
380
+ cursor_path = Path(args.cursor_file)
381
+ since = (
382
+ _parse_since(args.since)
383
+ if args.since
384
+ else ((None if args.no_cursor else _read_cursor(cursor_path)) or (datetime.now(timezone.utc) - timedelta(hours=24)))
385
+ )
386
+ until = _parse_time(args.until) if args.until else datetime.now(timezone.utc)
387
+
388
+ importer = DatabricksTraceImporter(
389
+ MlflowClient(tracking_uri=args.tracking_uri),
390
+ agentx_api_key=agentx_key,
391
+ agentx_base_url=os.getenv("AGENTX_API_BASE_URL") or None,
392
+ agent_name=args.agent_name,
393
+ )
394
+ report = importer.sync(
395
+ args.experiment_id,
396
+ since,
397
+ until,
398
+ monitor=args.monitor,
399
+ judge_sessions=args.judge_sessions,
400
+ dry_run=args.dry_run,
401
+ on_payload=(lambda wire: print(json.dumps(wire, default=str))) if args.dry_run else None,
402
+ )
403
+ print(report)
404
+ if not args.no_cursor and not args.dry_run and until is not None:
405
+ _write_cursor(cursor_path, until)
@@ -189,6 +189,46 @@ class MoveworksSyncReport:
189
189
  return base + ")"
190
190
 
191
191
 
192
+ def judge_sessions_via_engine(ingest: IngestClient, session_ids: List[str]) -> "tuple[int, int, int]":
193
+ """
194
+ Judge each session with every enabled session-scoped evaluator via the engine's on-demand
195
+ judge route (``ifStale=true``: an up-to-date verdict from the sweep or a previous run is left
196
+ alone). Shared by every pull importer (Moveworks, Databricks). Returns (judged, skipped,
197
+ failed) counts.
198
+ """
199
+ # Same base/key the ingest client already resolved - the judge routes live under
200
+ # /agent-monitoring on the same engine.
201
+ base = ingest._base_url
202
+ http = ingest._session
203
+ try:
204
+ response = http.get(f"{base}/agent-monitoring/online-evaluators", timeout=30)
205
+ response.raise_for_status()
206
+ evaluators = [
207
+ e for e in response.json().get("evaluators", [])
208
+ if e.get("enabled") and e.get("scope") == "session"
209
+ ]
210
+ except requests.RequestException:
211
+ return 0, 0, len(session_ids)
212
+ judged = skipped = failed = 0
213
+ for session_id in session_ids:
214
+ for evaluator in evaluators:
215
+ try:
216
+ result = http.post(
217
+ f"{base}/agent-monitoring/sessions/{session_id}/judge/{evaluator['_id']}",
218
+ params={"ifStale": "true"},
219
+ timeout=120,
220
+ )
221
+ if result.status_code == 200 and result.json().get("skipped"):
222
+ skipped += 1
223
+ elif result.ok:
224
+ judged += 1
225
+ else:
226
+ failed += 1
227
+ except requests.RequestException:
228
+ failed += 1
229
+ return judged, skipped, failed
230
+
231
+
192
232
  class MoveworksImporter:
193
233
  """
194
234
  Pulls Moveworks Data API records and replays them into AgentX as traces.
@@ -386,36 +426,10 @@ class MoveworksImporter:
386
426
  run) is left alone instead of judged again.
387
427
  """
388
428
  report = report or MoveworksSyncReport()
389
- # Same base/key the ingest client already resolved - the judge routes live under
390
- # /agent-monitoring on the same engine.
391
- base = self._ingest._base_url
392
- http = self._ingest._session
393
- try:
394
- response = http.get(f"{base}/agent-monitoring/online-evaluators", timeout=30)
395
- response.raise_for_status()
396
- evaluators = [
397
- e for e in response.json().get("evaluators", [])
398
- if e.get("enabled") and e.get("scope") == "session"
399
- ]
400
- except requests.RequestException:
401
- report.sessions_judge_failed += len(session_ids)
402
- return report
403
- for session_id in session_ids:
404
- for evaluator in evaluators:
405
- try:
406
- result = http.post(
407
- f"{base}/agent-monitoring/sessions/{session_id}/judge/{evaluator['_id']}",
408
- params={"ifStale": "true"},
409
- timeout=120,
410
- )
411
- if result.status_code == 200 and result.json().get("skipped"):
412
- report.sessions_judge_skipped += 1
413
- elif result.ok:
414
- report.sessions_judged += 1
415
- else:
416
- report.sessions_judge_failed += 1
417
- except requests.RequestException:
418
- report.sessions_judge_failed += 1
429
+ judged, skipped, failed = judge_sessions_via_engine(self._ingest, session_ids)
430
+ report.sessions_judged += judged
431
+ report.sessions_judge_skipped += skipped
432
+ report.sessions_judge_failed += failed
419
433
  return report
420
434
 
421
435
 
@@ -0,0 +1 @@
1
+ VERSION = "0.6.23"
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: agentx-python
3
- Version: 0.6.22
3
+ Version: 0.6.23
4
4
  Summary: Official Python SDK for AgentX (https://www.agentx.so/)
5
5
  Home-page: https://github.com/AgentX-ai/AgentX-python
6
6
  Author: Robin Wang and AgentX Team
@@ -36,6 +36,8 @@ Requires-Dist: llama-index-core>=0.10.0; extra == "llamaindex"
36
36
  Provides-Extra: autogen
37
37
  Requires-Dist: autogen-agentchat>=0.4.0; extra == "autogen"
38
38
  Requires-Dist: autogen-core>=0.4.0; extra == "autogen"
39
+ Provides-Extra: databricks
40
+ Requires-Dist: mlflow>=3.6.0; extra == "databricks"
39
41
  Provides-Extra: all
40
42
  Requires-Dist: langchain-core>=0.1.0; extra == "all"
41
43
  Requires-Dist: crewai>=0.80.0; extra == "all"
@@ -31,6 +31,7 @@ agentx/integrations/_traced_call.py
31
31
  agentx/integrations/anthropic.py
32
32
  agentx/integrations/autogen.py
33
33
  agentx/integrations/crewai.py
34
+ agentx/integrations/databricks.py
34
35
  agentx/integrations/google_adk.py
35
36
  agentx/integrations/google_genai.py
36
37
  agentx/integrations/langchain.py
@@ -1,3 +1,4 @@
1
1
  [console_scripts]
2
+ agentx-databricks = agentx.integrations.databricks:cli_main
2
3
  agentx-moveworks = agentx.integrations.moveworks:cli_main
3
4
  agentx-trace-eval = agentx.cli:main
@@ -26,6 +26,9 @@ autogen-core>=0.4.0
26
26
  [crewai]
27
27
  crewai>=0.80.0
28
28
 
29
+ [databricks]
30
+ mlflow>=3.6.0
31
+
29
32
  [google-adk]
30
33
  google-adk>=1.0.0
31
34
 
@@ -35,6 +35,9 @@ setup(
35
35
  # Moveworks Data API -> AgentX traces sync (see agentx/integrations/moveworks.py) -
36
36
  # a pull importer, since Moveworks agents run in their cloud with no in-process hook.
37
37
  "agentx-moveworks=agentx.integrations.moveworks:cli_main",
38
+ # MLflow/Databricks trace sync (see agentx/integrations/databricks.py) - pull importer
39
+ # for agents whose serving endpoints can't set OTLP egress env vars.
40
+ "agentx-databricks=agentx.integrations.databricks:cli_main",
38
41
  ],
39
42
  },
40
43
  install_requires=[
@@ -54,6 +57,7 @@ setup(
54
57
  "litellm": ["litellm>=1.0.0"],
55
58
  "llamaindex": ["llama-index-core>=0.10.0"],
56
59
  "autogen": ["autogen-agentchat>=0.4.0", "autogen-core>=0.4.0"],
60
+ "databricks": ["mlflow>=3.6.0"],
57
61
  "all": [
58
62
  "langchain-core>=0.1.0",
59
63
  "crewai>=0.80.0",
@@ -1 +0,0 @@
1
- VERSION = "0.6.22"
File without changes
File without changes
File without changes