snowglobe 0.4.9__py3-none-any.whl → 0.4.10__py3-none-any.whl

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.
Files changed (7) hide show
  1. snowglobe/client/src/telemetry.py +31 -17
  2. {snowglobe-0.4.9.dist-info → snowglobe-0.4.10.dist-info}/METADATA +1 -1
  3. {snowglobe-0.4.9.dist-info → snowglobe-0.4.10.dist-info}/RECORD +7 -7
  4. {snowglobe-0.4.9.dist-info → snowglobe-0.4.10.dist-info}/WHEEL +0 -0
  5. {snowglobe-0.4.9.dist-info → snowglobe-0.4.10.dist-info}/entry_points.txt +0 -0
  6. {snowglobe-0.4.9.dist-info → snowglobe-0.4.10.dist-info}/licenses/LICENSE +0 -0
  7. {snowglobe-0.4.9.dist-info → snowglobe-0.4.10.dist-info}/top_level.txt +0 -0
snowglobe/client/src/telemetry.py CHANGED
@@ -8,6 +8,7 @@ from snowglobe.client.src.models import CompletionRequest, RiskEvaluationRequest
8
8
  try:
9
9
  import mlflow
10
10
  import mlflow.tracing
11
+ from databricks.sdk import WorkspaceClient
11
12
 
12
13
  mlflow.tracing.enable()
13
14
  except ImportError:
@@ -28,11 +29,17 @@ def trace_completion_fn(
28
29
  def trace_decorator(completion_fn: Callable):
29
30
  disable_mlflow = os.getenv("SNOWGLOBE_DISABLE_MLFLOW_TRACING") or ""
30
31
  if mlflow and disable_mlflow.lower() != "true":
32
+ w = WorkspaceClient()
33
+ current_user = w.current_user.me()
34
+
35
+ formatted_sim_name = simulation_name.lower().replace(" ", "_")
36
+ default_experiment_name = f"/Users/{current_user.user_name}/{formatted_sim_name}"
37
+
31
38
  mlflow_experiment_name = (
32
- os.getenv("MLFLOW_EXPERIMENT_NAME") or simulation_name
39
+ os.getenv("MLFLOW_EXPERIMENT_NAME") or default_experiment_name
33
40
  )
34
41
  mlflow.set_experiment(mlflow_experiment_name)
35
-
42
+
36
43
  mlflow_active_model_id = os.getenv("MLFLOW_ACTIVE_MODEL_ID")
37
44
  if mlflow_active_model_id:
38
45
  mlflow.set_active_model(model_id=mlflow_active_model_id)
@@ -42,9 +49,9 @@ def trace_completion_fn(
42
49
  span_attributes = {
43
50
  "snowglobe.version": SNOWGLOBE_VERSION,
44
51
  "type": span_type,
45
- "session_id": session_id,
46
- "conversation_id": conversation_id,
47
- "message_id": message_id,
52
+ "session_id": str(session_id),
53
+ "conversation_id": str(conversation_id),
54
+ "message_id": str(message_id),
48
55
  "simulation_name": simulation_name,
49
56
  "agent_name": agent_name,
50
57
  }
@@ -58,11 +65,11 @@ def trace_completion_fn(
58
65
  async def completion_fn_wrapper(test_request: CompletionRequest):
59
66
  try:
60
67
  mlflow.update_current_trace(
61
- metadata={"mlflow.trace.session": session_id},
68
+ metadata={"mlflow.trace.session": str(session_id)},
62
69
  tags={
63
- "session_id": session_id,
64
- "conversation_id": conversation_id,
65
- "message_id": message_id,
70
+ "session_id": str(session_id),
71
+ "conversation_id": str(conversation_id),
72
+ "message_id": str(message_id),
66
73
  "simulation_name": simulation_name,
67
74
  "agent_name": agent_name,
68
75
  },
@@ -92,8 +99,14 @@ def trace_risk_evaluation_fn(
92
99
  def trace_decorator(risk_evaluation_fn: Callable):
93
100
  disable_mlflow = os.getenv("SNOWGLOBE_DISABLE_MLFLOW_TRACING") or ""
94
101
  if mlflow and disable_mlflow.lower() != "true":
102
+ w = WorkspaceClient()
103
+ current_user = w.current_user.me()
104
+
105
+ formatted_sim_name = simulation_name.lower().replace(" ", "_")
106
+ default_experiment_name = f"/Users/{current_user.user_name}/{formatted_sim_name}"
107
+
95
108
  mlflow_experiment_name = (
96
- os.getenv("MLFLOW_EXPERIMENT_NAME") or simulation_name
109
+ os.getenv("MLFLOW_EXPERIMENT_NAME") or default_experiment_name
97
110
  )
98
111
  mlflow.set_experiment(mlflow_experiment_name)
99
112
 
@@ -102,12 +115,13 @@ def trace_risk_evaluation_fn(
102
115
  mlflow.set_active_model(model_id=mlflow_active_model_id)
103
116
  else:
104
117
  mlflow.set_active_model(name=agent_name)
118
+
105
119
  span_attributes = {
106
120
  "snowglobe.version": SNOWGLOBE_VERSION,
107
121
  "type": span_type,
108
- "session_id": session_id,
109
- "conversation_id": conversation_id,
110
- "message_id": message_id,
122
+ "session_id": str(session_id),
123
+ "conversation_id": str(conversation_id),
124
+ "message_id": str(message_id),
111
125
  "simulation_name": simulation_name,
112
126
  "agent_name": agent_name,
113
127
  "risk_name": risk_name,
@@ -124,11 +138,11 @@ def trace_risk_evaluation_fn(
124
138
  ):
125
139
  try:
126
140
  mlflow.update_current_trace(
127
- metadata={"mlflow.trace.session": session_id},
141
+ metadata={"mlflow.trace.session": str(session_id)},
128
142
  tags={
129
- "session_id": session_id,
130
- "conversation_id": conversation_id,
131
- "message_id": message_id,
143
+ "session_id": str(session_id),
144
+ "conversation_id": str(conversation_id),
145
+ "message_id": str(message_id),
132
146
  "simulation_name": simulation_name,
133
147
  "agent_name": agent_name,
134
148
  "risk_name": risk_name,
{snowglobe-0.4.9.dist-info → snowglobe-0.4.10.dist-info}/METADATA RENAMED
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: snowglobe
3
- Version: 0.4.9
3
+ Version: 0.4.10
4
4
  Summary: client server for usage with snowglobe experiments
5
5
  Author-email: Guardrails AI <contact@guardrailsai.com>
6
6
  License: MIT License
{snowglobe-0.4.9.dist-info → snowglobe-0.4.10.dist-info}/RECORD RENAMED
@@ -6,11 +6,11 @@ snowglobe/client/src/config.py,sha256=YRx_AQEZoHaAqk6guTxynIEGV_iJ3wNNGtMmaKsYMb
6
6
  snowglobe/client/src/models.py,sha256=BX310WrDN9Fd8v68me3XGL_ic1ulvjCrZyIT2ND1eUo,866
7
7
  snowglobe/client/src/project_manager.py,sha256=Ze-qs4dQI2kIV-PmtWZ1b67hMUfsnsMHus90aT8HOow,9970
8
8
  snowglobe/client/src/stats.py,sha256=IdaXroOZBmvLVa_p9pDE6hsxsc7-fBEDnLf8O6Ch0GA,1596
9
- snowglobe/client/src/telemetry.py,sha256=N91Q37YfJaUYPa7BUAs_3x4LxjculwlETIKC5k1dbig,5045
9
+ snowglobe/client/src/telemetry.py,sha256=CLrnKtCRnq-wGraFHdMcgUqY4-YX1JqxQpnAdWy_wok,5679
10
10
  snowglobe/client/src/utils.py,sha256=hHOht0hc8fv3OuPTz2Tqs639CzSAF34JTZs5ifKV6YI,3708
11
- snowglobe-0.4.9.dist-info/licenses/LICENSE,sha256=S90V6iFU5ZeSg44JQYS1To3pa7ZEobrHc_t483qSKSI,1070
12
- snowglobe-0.4.9.dist-info/METADATA,sha256=NckGusSPGxCyboKGK79F-PiNaXstK6FOxmM0p0lOB5g,5406
13
- snowglobe-0.4.9.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
14
- snowglobe-0.4.9.dist-info/entry_points.txt,sha256=mqx4mTwFPHttjctE2ceYTYWCCIG30Ji2C89aaCYgHcM,71
15
- snowglobe-0.4.9.dist-info/top_level.txt,sha256=PoyYihnCBjRyjeIT19yBcE47JTe7i1OwRXvJ4d5EohM,10
16
- snowglobe-0.4.9.dist-info/RECORD,,
11
+ snowglobe-0.4.10.dist-info/licenses/LICENSE,sha256=S90V6iFU5ZeSg44JQYS1To3pa7ZEobrHc_t483qSKSI,1070
12
+ snowglobe-0.4.10.dist-info/METADATA,sha256=fGi-6Bq8spjaxvmS-0Xu3EtDdnoOQ8f76yzh26pyv3c,5407
13
+ snowglobe-0.4.10.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
14
+ snowglobe-0.4.10.dist-info/entry_points.txt,sha256=mqx4mTwFPHttjctE2ceYTYWCCIG30Ji2C89aaCYgHcM,71
15
+ snowglobe-0.4.10.dist-info/top_level.txt,sha256=PoyYihnCBjRyjeIT19yBcE47JTe7i1OwRXvJ4d5EohM,10
16
+ snowglobe-0.4.10.dist-info/RECORD,,