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.
- {agentx_python-0.6.22 → agentx_python-0.6.23}/PKG-INFO +3 -1
- agentx_python-0.6.23/agentx/integrations/databricks.py +405 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/integrations/moveworks.py +44 -30
- agentx_python-0.6.23/agentx/version.py +1 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx_python.egg-info/PKG-INFO +3 -1
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx_python.egg-info/SOURCES.txt +1 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx_python.egg-info/entry_points.txt +1 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx_python.egg-info/requires.txt +3 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/setup.py +4 -0
- agentx_python-0.6.22/agentx/version.py +0 -1
- {agentx_python-0.6.22 → agentx_python-0.6.23}/LICENSE +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/README.md +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/__init__.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/agentx.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/cli.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/__init__.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/_term.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/adapters/__init__.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/adapters/http_endpoint.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/adapters/precomputed.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/adapters/raw.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/client.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/datasets.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/evaluation_settings.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/models.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/prompts.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/redaction.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/reporting.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/results.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/runner.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/evaluations/tracing.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/exceptions.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/feedback.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/integrations/__init__.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/integrations/_traced_call.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/integrations/anthropic.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/integrations/autogen.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/integrations/crewai.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/integrations/google_adk.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/integrations/google_genai.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/integrations/langchain.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/integrations/litellm.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/integrations/llamaindex.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/integrations/openai.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/integrations/openai_agents.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/monitor/__init__.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/monitor/client.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/monitor/models.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/monitor/online_evaluators.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/monitor/patterns.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/monitor/profile.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/monitor/signals.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/outcomes.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/py.typed +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/resources/__init__.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/resources/agent.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/resources/conversation.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/resources/workforce.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/tracing/__init__.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/tracing/ci_types.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/tracing/ingest_client.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/tracing/tracer.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx/util.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx_python.egg-info/dependency_links.txt +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx_python.egg-info/not-zip-safe +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/agentx_python.egg-info/top_level.txt +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/setup.cfg +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/tests/test_integration.py +0 -0
- {agentx_python-0.6.22 → agentx_python-0.6.23}/tests/test_integrations.py +0 -0
- {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.
|
|
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
|
-
|
|
390
|
-
|
|
391
|
-
|
|
392
|
-
|
|
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.
|
|
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
|
|
@@ -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
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|