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

snowglobe/client/src/app.py

@@ -26,7 +26,7 @@ from snowglobe.client.src.telemetry import trace_completion_fn, trace_risk_evalu
 
 from .cli_utils import info, shutdown_manager
 from .config import config, get_api_key_or_raise
-from .models import CompletionFunctionOutputs, CompletionRequest, RiskEvaluationRequest
+from .models import CompletionFunctionOutputs, CompletionRequest, RiskEvaluationRequest, SnowglobeData, SnowglobeMessage
 from .stats import initialize_stats, track_batch_completion
 from .utils import fetch_experiments, fetch_messages
 
@@ -129,7 +129,14 @@ async def process_application_heartbeat(app_id):
     }
     try:
         prompt = "Hello from Snowglobe!"
-        test_request = CompletionRequest(messages=[{"role": "user", "content": prompt}])
+        test_request = CompletionRequest(
+            messages=[SnowglobeMessage(
+                    role="user",
+                    content=prompt,
+                    snowglobe_data=SnowglobeData(
+                        conversation_id="test", test_id="test"
+                    ),
+                )])
         heartbeat_id = uuid.uuid4().hex
         agent = apps.get(app_id, {})
         agent_name = agent.get("name", "")
@@ -193,6 +200,10 @@ async def process_application_heartbeat(app_id):
         connection_test_payload["error"] = f"{str(e)}\n{traceback.format_exc()}"
         connection_test_payload["app_id"] = app_id
         connection_test_payload["applicationId"] = app_id
+        LOGGER.error(
+            f"Error processing heartbeat for application {app_id}: {connection_test_payload['error']}"
+        )
+        LOGGER.error(traceback.format_exc())
 
     connection_test_url = (
         f"{config.CONTROL_PLANE_URL}/api/successful-code-connection-tests"

snowglobe/client/src/telemetry.py

@@ -8,6 +8,7 @@ from snowglobe.client.src.models import CompletionRequest, RiskEvaluationRequest
 try:
     import mlflow
     import mlflow.tracing
+    from databricks.sdk import WorkspaceClient
 
     mlflow.tracing.enable()
 except ImportError:
@@ -28,11 +29,17 @@ def trace_completion_fn(
     def trace_decorator(completion_fn: Callable):
         disable_mlflow = os.getenv("SNOWGLOBE_DISABLE_MLFLOW_TRACING") or ""
         if mlflow and disable_mlflow.lower() != "true":
+            w = WorkspaceClient()
+            current_user = w.current_user.me()
+
+            formatted_sim_name = simulation_name.lower().replace(" ", "_")
+            default_experiment_name = f"/Users/{current_user.user_name}/{formatted_sim_name}"
+
             mlflow_experiment_name = (
-                os.getenv("MLFLOW_EXPERIMENT_NAME") or simulation_name
+                os.getenv("MLFLOW_EXPERIMENT_NAME") or default_experiment_name
             )
             mlflow.set_experiment(mlflow_experiment_name)
-            
+
             mlflow_active_model_id = os.getenv("MLFLOW_ACTIVE_MODEL_ID")
             if mlflow_active_model_id:
                 mlflow.set_active_model(model_id=mlflow_active_model_id)
@@ -42,9 +49,9 @@ def trace_completion_fn(
             span_attributes = {
                 "snowglobe.version": SNOWGLOBE_VERSION,
                 "type": span_type,
-                "session_id": session_id,
-                "conversation_id": conversation_id,
-                "message_id": message_id,
+                "session_id": str(session_id),
+                "conversation_id": str(conversation_id),
+                "message_id": str(message_id),
                 "simulation_name": simulation_name,
                 "agent_name": agent_name,
             }
@@ -58,11 +65,11 @@ def trace_completion_fn(
             async def completion_fn_wrapper(test_request: CompletionRequest):
                 try:
                     mlflow.update_current_trace(
-                        metadata={"mlflow.trace.session": session_id},
+                        metadata={"mlflow.trace.session": str(session_id)},
                         tags={
-                            "session_id": session_id,
-                            "conversation_id": conversation_id,
-                            "message_id": message_id,
+                            "session_id": str(session_id),
+                            "conversation_id": str(conversation_id),
+                            "message_id": str(message_id),
                             "simulation_name": simulation_name,
                             "agent_name": agent_name,
                         },
@@ -92,8 +99,14 @@ def trace_risk_evaluation_fn(
     def trace_decorator(risk_evaluation_fn: Callable):
         disable_mlflow = os.getenv("SNOWGLOBE_DISABLE_MLFLOW_TRACING") or ""
         if mlflow and disable_mlflow.lower() != "true":
+            w = WorkspaceClient()
+            current_user = w.current_user.me()
+
+            formatted_sim_name = simulation_name.lower().replace(" ", "_")
+            default_experiment_name = f"/Users/{current_user.user_name}/{formatted_sim_name}"
+        
             mlflow_experiment_name = (
-                os.getenv("MLFLOW_EXPERIMENT_NAME") or simulation_name
+                os.getenv("MLFLOW_EXPERIMENT_NAME") or default_experiment_name
             )
             mlflow.set_experiment(mlflow_experiment_name)
 
@@ -102,12 +115,13 @@ def trace_risk_evaluation_fn(
                 mlflow.set_active_model(model_id=mlflow_active_model_id)
             else:
                 mlflow.set_active_model(name=agent_name)
+
             span_attributes = {
                 "snowglobe.version": SNOWGLOBE_VERSION,
                 "type": span_type,
-                "session_id": session_id,
-                "conversation_id": conversation_id,
-                "message_id": message_id,
+                "session_id": str(session_id),
+                "conversation_id": str(conversation_id),
+                "message_id": str(message_id),
                 "simulation_name": simulation_name,
                 "agent_name": agent_name,
                 "risk_name": risk_name,
@@ -124,11 +138,11 @@ def trace_risk_evaluation_fn(
             ):
                 try:
                     mlflow.update_current_trace(
-                        metadata={"mlflow.trace.session": session_id},
+                        metadata={"mlflow.trace.session": str(session_id)},
                         tags={
-                            "session_id": session_id,
-                            "conversation_id": conversation_id,
-                            "message_id": message_id,
+                            "session_id": str(session_id),
+                            "conversation_id": str(conversation_id),
+                            "message_id": str(message_id),
                             "simulation_name": simulation_name,
                             "agent_name": agent_name,
                             "risk_name": risk_name,

snowglobe/client/src/utils.py

@@ -16,13 +16,22 @@ async def fetch_experiments(app_id: str = None) -> list[dict]:
         list[dict]: A list of experiments.
     """
     async with httpx.AsyncClient() as client:
-        experiments_url = f"{config.CONTROL_PLANE_URL}/api/experiments?&evaluated=false"
+        experiments_url = f"{config.CONTROL_PLANE_URL}/api/experiments?evaluated=false"
+
         if app_id:
             experiments_url += f"&appId={config.APPLICATION_ID}"
-        experiments_response = await client.get(
+        try:
+            # get elapsed time for this request
+            import time
+            start_time = time.monotonic()
+            experiments_response = await client.get(
             experiments_url,
             headers={"x-api-key": get_api_key_or_raise()},
-        )
+            timeout=60.0,  # Set timeout to 60 seconds
+            )
+        except httpx.ConnectTimeout:
+            elapsed_time = time.monotonic() - start_time
+            raise Exception(f"Warning: Connection timed out while fetching experiments. Elapsed time: {elapsed_time:.2f} seconds. Polling will continue. If this persists please contact Snowglobe support.")
 
     if not experiments_response.status_code == 200:
         try:

{snowglobe-0.4.9.dist-info → snowglobe-0.4.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowglobe
-Version: 0.4.9
+Version: 0.4.11
 Summary: client server for usage with snowglobe experiments
 Author-email: Guardrails AI <contact@guardrailsai.com>
 License: MIT License

snowglobe-0.4.11.dist-info/RECORD

@@ -0,0 +1,16 @@
+snowglobe/client/__init__.py,sha256=kzp9wPUUYBXqDSKZbfmD4vrAQvrWSW5HOvtpFlEJWfs,353
+snowglobe/client/src/app.py,sha256=QOqwK4134siIDnrlY2vpRJ6DiwDVKVmrJszzp5z_RE0,43117
+snowglobe/client/src/cli.py,sha256=I3LVWJmyvUzOQlV0gv_AeMuEazfy8GsYdH6svo7cZOU,28544
+snowglobe/client/src/cli_utils.py,sha256=6C7J5gow8xveQYF4w6ewtQJKI7VvlLTx7FS_7gl7RwI,17227
+snowglobe/client/src/config.py,sha256=YRx_AQEZoHaAqk6guTxynIEGV_iJ3wNNGtMmaKsYMbc,10488
+snowglobe/client/src/models.py,sha256=BX310WrDN9Fd8v68me3XGL_ic1ulvjCrZyIT2ND1eUo,866
+snowglobe/client/src/project_manager.py,sha256=Ze-qs4dQI2kIV-PmtWZ1b67hMUfsnsMHus90aT8HOow,9970
+snowglobe/client/src/stats.py,sha256=IdaXroOZBmvLVa_p9pDE6hsxsc7-fBEDnLf8O6Ch0GA,1596
+snowglobe/client/src/telemetry.py,sha256=CLrnKtCRnq-wGraFHdMcgUqY4-YX1JqxQpnAdWy_wok,5679
+snowglobe/client/src/utils.py,sha256=F86gku68IyjXzUleZXc9QyzFfWV4PBa9LKukAo9gJl8,4198
+snowglobe-0.4.11.dist-info/licenses/LICENSE,sha256=S90V6iFU5ZeSg44JQYS1To3pa7ZEobrHc_t483qSKSI,1070
+snowglobe-0.4.11.dist-info/METADATA,sha256=wExQyglcFQS8zPMvSy81EEin2t7i37XvAb5Ln4XtZIU,5407
+snowglobe-0.4.11.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+snowglobe-0.4.11.dist-info/entry_points.txt,sha256=mqx4mTwFPHttjctE2ceYTYWCCIG30Ji2C89aaCYgHcM,71
+snowglobe-0.4.11.dist-info/top_level.txt,sha256=PoyYihnCBjRyjeIT19yBcE47JTe7i1OwRXvJ4d5EohM,10
+snowglobe-0.4.11.dist-info/RECORD,,

snowglobe-0.4.9.dist-info/RECORD

@@ -1,16 +0,0 @@
-snowglobe/client/__init__.py,sha256=kzp9wPUUYBXqDSKZbfmD4vrAQvrWSW5HOvtpFlEJWfs,353
-snowglobe/client/src/app.py,sha256=CaDtbMn6ZXooQJuaPAHOXs2r9FkFqDxTrgqoW3Rl_2I,42686
-snowglobe/client/src/cli.py,sha256=I3LVWJmyvUzOQlV0gv_AeMuEazfy8GsYdH6svo7cZOU,28544
-snowglobe/client/src/cli_utils.py,sha256=6C7J5gow8xveQYF4w6ewtQJKI7VvlLTx7FS_7gl7RwI,17227
-snowglobe/client/src/config.py,sha256=YRx_AQEZoHaAqk6guTxynIEGV_iJ3wNNGtMmaKsYMbc,10488
-snowglobe/client/src/models.py,sha256=BX310WrDN9Fd8v68me3XGL_ic1ulvjCrZyIT2ND1eUo,866
-snowglobe/client/src/project_manager.py,sha256=Ze-qs4dQI2kIV-PmtWZ1b67hMUfsnsMHus90aT8HOow,9970
-snowglobe/client/src/stats.py,sha256=IdaXroOZBmvLVa_p9pDE6hsxsc7-fBEDnLf8O6Ch0GA,1596
-snowglobe/client/src/telemetry.py,sha256=N91Q37YfJaUYPa7BUAs_3x4LxjculwlETIKC5k1dbig,5045
-snowglobe/client/src/utils.py,sha256=hHOht0hc8fv3OuPTz2Tqs639CzSAF34JTZs5ifKV6YI,3708
-snowglobe-0.4.9.dist-info/licenses/LICENSE,sha256=S90V6iFU5ZeSg44JQYS1To3pa7ZEobrHc_t483qSKSI,1070
-snowglobe-0.4.9.dist-info/METADATA,sha256=NckGusSPGxCyboKGK79F-PiNaXstK6FOxmM0p0lOB5g,5406
-snowglobe-0.4.9.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-snowglobe-0.4.9.dist-info/entry_points.txt,sha256=mqx4mTwFPHttjctE2ceYTYWCCIG30Ji2C89aaCYgHcM,71
-snowglobe-0.4.9.dist-info/top_level.txt,sha256=PoyYihnCBjRyjeIT19yBcE47JTe7i1OwRXvJ4d5EohM,10
-snowglobe-0.4.9.dist-info/RECORD,,