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

snowglobe/client/src/app.py

@@ -14,6 +14,7 @@ from functools import wraps
 from logging import getLogger
 from typing import Dict
 from urllib.parse import quote_plus
+import uuid
 
 import httpx
 import uvicorn
@@ -21,6 +22,8 @@ from apscheduler import AsyncScheduler
 from apscheduler.triggers.interval import IntervalTrigger
 from fastapi import FastAPI, HTTPException, Request
 
+from snowglobe.client.src.telemetry import trace_completion_fn, trace_risk_evaluation_fn
+
 from .cli_utils import info, shutdown_manager
 from .config import config, get_api_key_or_raise
 from .models import CompletionFunctionOutputs, CompletionRequest, RiskEvaluationRequest
@@ -127,16 +130,32 @@ async def process_application_heartbeat(app_id):
     try:
         prompt = "Hello from Snowglobe!"
         test_request = CompletionRequest(messages=[{"role": "user", "content": prompt}])
-        completion_fn = apps.get(app_id, {}).get("completion_fn")
+        heartbeat_id = uuid.uuid4().hex
+        agent = apps.get(app_id, {})
+        agent_name = agent.get("name", "")
+        completion_fn = agent.get("completion_fn")
         if not completion_fn:
             LOGGER.warning(
                 f"No completion function found for application {app_id}. Skipping heartbeat."
             )
             return
-        if asyncio.iscoroutinefunction(completion_fn):
-            response = await completion_fn(test_request)
-        else:
-            response = completion_fn(test_request)
+
+        @trace_completion_fn(
+            agent_name=agent_name,
+            conversation_id=heartbeat_id,
+            message_id=heartbeat_id,
+            session_id=heartbeat_id,
+            simulation_name=f"{agent_name} Heartbeat",
+            span_type="snowglobe/heartbeat",
+        )
+        async def run_completion_fn(completion_request: CompletionRequest):
+            if asyncio.iscoroutinefunction(completion_fn):
+                response = await completion_fn(completion_request)
+            else:
+                response = completion_fn(completion_request)
+            return response
+
+        response = await run_completion_fn(test_request)
         if not isinstance(response, CompletionFunctionOutputs):
             LOGGER.error(
                 f"Completion function for application {app_id} did not return a valid response. Expected CompletionFunctionOutputs, got {type(response)}"
@@ -201,19 +220,31 @@ async def process_application_heartbeat(app_id):
     return connection_test_response.json()
 
 
-async def process_risk_evaluation(test, risk_name):
+async def process_risk_evaluation(test, risk_name, simulation_name, agent_name):
     """finds correct risk and calls the risk evaluation function and creates a risk evaluation for the test"""
     start = time.time()
 
     messages = await fetch_messages(test=test)
 
-    if asyncio.iscoroutinefunction(risks[risk_name]):
-        risk_evaluation = await risks[risk_name](
-            RiskEvaluationRequest(messages=messages)
-        )
-    else:
-        risk_evaluation = risks[risk_name](RiskEvaluationRequest(messages=messages))
+    risk_eval_req = RiskEvaluationRequest(messages=messages)
 
+    @trace_risk_evaluation_fn(
+        agent_name=agent_name,
+        conversation_id=test["conversation_id"],
+        message_id=test["id"],
+        session_id=test["conversation_id"],
+        simulation_name=simulation_name,
+        span_type=f"snowglobe/risk-evaluation/{risk_name}",
+        risk_name=risk_name,
+    )
+    async def run_risk_evaluation_fn(risk_evaluation_request: RiskEvaluationRequest):
+        if asyncio.iscoroutinefunction(risks[risk_name]):
+            risk_evaluation = await risks[risk_name](risk_evaluation_request)
+        else:
+            risk_evaluation = risks[risk_name](risk_evaluation_request)
+        return risk_evaluation
+
+    risk_evaluation = await run_risk_evaluation_fn(risk_eval_req)
     LOGGER.debug(f"Risk evaluation output: {risk_evaluation}")
 
     # Extract fields from risk_evaluation object
@@ -249,16 +280,33 @@ async def process_risk_evaluation(test, risk_name):
         raise Exception("Error posting risk evaluation, task is not healthy")
 
 
-async def process_test(test, completion_fn, app_id):
+async def process_test(test, completion_fn, app_id, simulation_name):
     """Processes a test by converting it to OpenAI style messages and calling the completion function"""
     start = time.time()
     # convert test to openai style messages
     messages = await fetch_messages(test=test)
 
-    if asyncio.iscoroutinefunction(completion_fn):
-        completionOutput = await completion_fn(CompletionRequest(messages=messages))
-    else:
-        completionOutput = completion_fn(CompletionRequest(messages=messages))
+    agent = apps.get(app_id, {})
+    agent_name = agent.get("name", "")
+
+    completion_req = CompletionRequest(messages=messages)
+
+    @trace_completion_fn(
+        agent_name=agent_name,
+        conversation_id=test["conversation_id"],
+        message_id=test["id"],
+        session_id=test["conversation_id"],
+        simulation_name=simulation_name,
+        span_type="snowglobe/completion",
+    )
+    async def run_completion_fn(completion_request: CompletionRequest):
+        if asyncio.iscoroutinefunction(completion_fn):
+            completionOutput = await completion_fn(completion_request)
+        else:
+            completionOutput = completion_fn(completion_request)
+        return completionOutput
+
+    completionOutput = await run_completion_fn(completion_req)
 
     LOGGER.debug(f"Completion output: {completionOutput}")
 
@@ -388,7 +436,11 @@ async def poll_for_completions():
                         try:
                             completion_request = await httpx.AsyncClient().post(
                                 f"{config.SNOWGLOBE_CLIENT_URL}/completion",
-                                json={"test": test, "app_id": app_id},
+                                json={
+                                    "test": test,
+                                    "app_id": app_id,
+                                    "simulation_name": experiment["name"],
+                                },
                                 timeout=30,
                             )
                         except (httpx.ConnectError, httpx.TimeoutException) as e:
@@ -543,6 +595,31 @@ async def poll_for_risk_evaluations():
                     )
                     continue
                 experiment = experiment_request.json()
+
+                try:
+                    app_request = await client.get(
+                        f"{config.CONTROL_PLANE_URL}/api/applications/{experiment['app_id']}",
+                        headers={"x-api-key": get_api_key_or_raise()},
+                    )
+                except (httpx.ConnectError, httpx.TimeoutException) as e:
+                    if shutdown_manager.is_shutdown_requested():
+                        LOGGER.debug(f"HTTP error during shutdown (expected): {e}")
+                        return
+                    else:
+                        LOGGER.error(
+                            f"Connection error fetching application {experiment['app_id']}: {e}"
+                        )
+                        continue
+
+                app_name = experiment["app_id"]
+                if not app_request.is_success:
+                    LOGGER.error(
+                        f"Error fetching application {experiment['app_id']}: {app_request.text}"
+                    )
+                else:
+                    application = app_request.json()
+                    app_name = application["name"]
+
                 risk_eval_count = 0
 
                 for risk_name in risks.keys():
@@ -608,7 +685,12 @@ async def poll_for_risk_evaluations():
                                 try:
                                     risk_eval_response = await httpx.AsyncClient().post(
                                         f"{config.SNOWGLOBE_CLIENT_URL}/risk-evaluation",
-                                        json={"test": test, "risk_name": risk_name},
+                                        json={
+                                            "test": test,
+                                            "risk_name": risk_name,
+                                            "simulation_name": experiment["name"],
+                                            "agent_name": app_name,
+                                        },
                                         timeout=30,
                                     )
                                 except (
@@ -686,7 +768,7 @@ async def lifespan(app: FastAPI):
                             "agent_wrapper", agent_file_path
                         )
                         agent_module = importlib.util.module_from_spec(spec)
-                       
+
                         # Add current directory to path
                         sys_path_backup = sys.path.copy()
                         current_dir = os.getcwd()
@@ -848,6 +930,7 @@ def create_client():
         completion_body = await request.json()
         test = completion_body.get("test")
         app_id = completion_body.get("app_id")
+        simulation_name = completion_body.get("simulation_name")
         # both are required non empty strings
         if not test or not app_id:
             raise HTTPException(
@@ -862,7 +945,7 @@ def create_client():
         completion_fn = apps.get(app_id, {}).get("completion_fn")
         LOGGER.debug(f"Received test: {test['id']}")
 
-        await process_test(test, completion_fn, app_id)
+        await process_test(test, completion_fn, app_id, simulation_name)
         return {"status": "processed"}
 
     @app.post("/heartbeat")
@@ -902,10 +985,12 @@ def create_client():
         body = await request.json()
         test = body.get("test")
         risk_name = body.get("risk_name")
+        simulation_name = body.get("simulation_name")
+        agent_name = body.get("agent_name")
         LOGGER.debug(f"Received risk evaluation for test: {test['id']}")
 
         # For now, just simulate processing
-        await process_risk_evaluation(test, risk_name)
+        await process_risk_evaluation(test, risk_name, simulation_name, agent_name)
         return {"status": "risk evaluation processed"}
 
     return app

snowglobe/client/src/cli.py

@@ -6,6 +6,7 @@ import signal
 import sys
 import threading
 import time
+import uuid
 import webbrowser
 from importlib.metadata import version
 from typing import Optional, Tuple
@@ -15,6 +16,8 @@ import uvicorn
 from fastapi import FastAPI, Request
 from fastapi.middleware.cors import CORSMiddleware
 
+from snowglobe.client.src.telemetry import trace_completion_fn
+
 # Import start_client lazily inside the start command to avoid config initialization
 from .cli_utils import (
     check_auth_status,
@@ -198,7 +201,7 @@ def test(
         else:
             info("Check your implementation and try again.")
             docs_link(
-                "Troubleshooting guide", "https://docs.snowglobe.so/troubleshooting"
+                "Troubleshooting guide", "https://snowglobe.so/docs/troubleshooting"
             )
         raise typer.Exit(1)
 
@@ -231,7 +234,7 @@ def init(
     if not is_auth:
         error("Authentication required to initialize agents")
         info("Please run 'snowglobe-connect auth' first to set up authentication")
-        docs_link("Setup guide", "https://docs.snowglobe.so/setup")
+        docs_link("Setup guide", "https://snowglobe.so/docs/setup")
         raise typer.Exit(1)
 
     success("Authenticated successfully")
@@ -256,9 +259,7 @@ def init(
         raise typer.Exit(0)
     elif selected == "new":
         info("Creating new application not yet implemented in init command")
-        info(
-            "Please visit https://snowglobe.guardrails-ai.com/applications/create to create a new app"
-        )
+        info("Please visit https://snowglobe.guardrailsai.com/app to create a new app")
         info("Then run this command again to select it")
         raise typer.Exit(0)
 
@@ -388,10 +389,24 @@ def test_agent_wrapper(filename: str, app_id: str, app_name: str) -> Tuple[bool,
             ]
         )
 
-        if asyncio.iscoroutinefunction(process_scenario):
-            response = asyncio.run(process_scenario(test_request))
-        else:
-            response = process_scenario(test_request)
+        test_id = uuid.uuid4()
+
+        @trace_completion_fn(
+            agent_name=app_name,
+            conversation_id=test_id,
+            message_id=test_id,
+            session_id=test_id,
+            simulation_name=f"{app_name} CLI Test",
+            span_type="snowglobe/cli-test",
+        )
+        async def run_process_scenario(completion_request: CompletionRequest):
+            if asyncio.iscoroutinefunction(process_scenario):
+                response = asyncio.run(process_scenario(completion_request))
+            else:
+                response = process_scenario(completion_request)
+            return response
+
+        response = asyncio.run(run_process_scenario(test_request))
 
         if hasattr(response, "response") and isinstance(response.response, str):
             if response.response == "Your response here":
@@ -402,6 +417,7 @@ def test_agent_wrapper(filename: str, app_id: str, app_name: str) -> Tuple[bool,
 
     except Exception as e:
         import traceback
+
         traceback.print_exc()
         return False, f"Error: {str(e)}"
 
@@ -434,7 +450,7 @@ def enhanced_error_handler(status_code: int, operation: str = "operation") -> No
         error("Authentication failed")
         info("Your API key may be invalid or expired")
         info("Run 'snowglobe-connect auth' to set up authentication")
-        docs_link("Authentication help", "https://docs.snowglobe.so/auth")
+        docs_link("Authentication help", "https://snowglobe.so/docs/auth")
     elif status_code == 403:
         error("Access forbidden")
         info("You don't have permission for this operation")
@@ -606,6 +622,7 @@ def _create_auth_server(config_key: str, rc_path: str) -> FastAPI:
             return {"written": True}
         except Exception as e:
             import traceback
+
             traceback.print_exc()
             error(f"Failed to process key configuration: {e}")
             return {"error": "Failed to process key configuration request"}
@@ -624,7 +641,7 @@ def _show_auth_success_next_steps() -> None:
     console.print("3. Start the client:")
     console.print("   [bold green]snowglobe-connect start[/bold green]")
     console.print()
-    docs_link("Getting started guide", "https://docs.snowglobe.so/getting-started")
+    docs_link("Getting started guide", "https://snowglobe.so/docs/getting-started")
 
 
 def _poll_for_api_key(rc_path: str, timeout: int = 300) -> bool:
@@ -792,6 +809,7 @@ def start(
                     console.print()
     except Exception:
         import traceback
+
         traceback.print_exc()
         # Do not block startup if we cannot load agents mapping
         pass

snowglobe/client/src/telemetry.py

@@ -0,0 +1,146 @@
+import os
+from importlib.metadata import version as importlib_version
+from typing import Callable
+from functools import wraps
+
+from snowglobe.client.src.models import CompletionRequest, RiskEvaluationRequest
+
+try:
+    import mlflow
+    import mlflow.tracing
+
+    mlflow.tracing.enable()
+except ImportError:
+    mlflow = None
+
+SNOWGLOBE_VERSION = importlib_version("snowglobe")
+
+
+def trace_completion_fn(
+    *,
+    session_id: str,
+    conversation_id: str,
+    message_id: str,
+    simulation_name: str,
+    agent_name: str,
+    span_type: str,
+):
+    def trace_decorator(completion_fn: Callable):
+        disable_mlflow = os.getenv("SNOWGLOBE_DISABLE_MLFLOW_TRACING") or ""
+        if mlflow and disable_mlflow.lower() != "true":
+            mlflow_experiment_name = (
+                os.getenv("MLFLOW_EXPERIMENT_NAME") or simulation_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)
+            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,
+                "simulation_name": simulation_name,
+                "agent_name": agent_name,
+            }
+
+            @mlflow.trace(
+                name=span_type,
+                span_type=span_type,
+                attributes=span_attributes,
+            )
+            @wraps(completion_fn)
+            async def completion_fn_wrapper(test_request: CompletionRequest):
+                try:
+                    mlflow.update_current_trace(
+                        metadata={"mlflow.trace.session": session_id},
+                        tags={
+                            "session_id": session_id,
+                            "conversation_id": conversation_id,
+                            "message_id": message_id,
+                            "simulation_name": simulation_name,
+                            "agent_name": agent_name,
+                        },
+                    )
+                    response = await completion_fn(test_request)
+                    return response
+                except Exception as e:
+                    raise e
+
+            return completion_fn_wrapper
+        else:
+            return completion_fn
+
+    return trace_decorator
+
+
+def trace_risk_evaluation_fn(
+    *,
+    session_id: str,
+    conversation_id: str,
+    message_id: str,
+    simulation_name: str,
+    agent_name: str,
+    span_type: str,
+    risk_name,
+):
+    def trace_decorator(risk_evaluation_fn: Callable):
+        disable_mlflow = os.getenv("SNOWGLOBE_DISABLE_MLFLOW_TRACING") or ""
+        if mlflow and disable_mlflow.lower() != "true":
+            mlflow_experiment_name = (
+                os.getenv("MLFLOW_EXPERIMENT_NAME") or simulation_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)
+            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,
+                "simulation_name": simulation_name,
+                "agent_name": agent_name,
+                "risk_name": risk_name,
+            }
+
+            @mlflow.trace(
+                name=span_type,
+                span_type=span_type,
+                attributes=span_attributes,
+            )
+            @wraps(risk_evaluation_fn)
+            async def risk_evaluation_fn_wrapper(
+                risk_evaluation_request: RiskEvaluationRequest,
+            ):
+                try:
+                    mlflow.update_current_trace(
+                        metadata={"mlflow.trace.session": session_id},
+                        tags={
+                            "session_id": session_id,
+                            "conversation_id": conversation_id,
+                            "message_id": message_id,
+                            "simulation_name": simulation_name,
+                            "agent_name": agent_name,
+                            "risk_name": risk_name,
+                        },
+                    )
+                    response = await risk_evaluation_fn(risk_evaluation_request)
+                    return response
+                except Exception as e:
+                    raise e
+
+            return risk_evaluation_fn_wrapper
+        else:
+            return risk_evaluation_fn
+
+    return trace_decorator

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: snowglobe
-Version: 0.4.8
+Version: 0.4.9
 Summary: client server for usage with snowglobe experiments
 Author-email: Guardrails AI <contact@guardrailsai.com>
 License: MIT License
@@ -119,3 +119,17 @@ def process_scenario(request: CompletionRequest) -> CompletionFunctionOutputs:
     )
     return CompletionFunctionOutputs(response=response.choices[0].message.content)
 ```
+
+## Tracing with MLflow
+The Snowglobe Connect SDK has MLflow tracing built in! Simply `pip install mlflow` and the sdk will take care of the rest.  Read more about MLflow's tracing capability for GenAI Apps [here](https://mlflow.org/docs/latest/genai/tracing/app-instrumentation/).
+
+### Enhancing Snowglobe Connect SDK's Traces with Autologging
+You can turn on mlflow autologging in your app to add additional context to the traces the Snowglobe Connect SDK captures.  In you app's entry point simply call the appropriate autolog method for the LLM provider you're using.  The below example shows how to enable this for LiteLLM:
+```py
+import mlflow
+
+mlflow.litellm.autolog()
+```
+
+### Disable Snowglobe Connect SDK's MLflow Tracing
+If you already use MLflow and don't want the Snowglobe Connect SDK to capture additional traces, you can disable this feature by setting the `SNOWGLOBE_DISABLE_MLFLOW_TRACING` environment variable to `true`.

{snowglobe-0.4.8.dist-info → snowglobe-0.4.9.dist-info}/RECORD

@@ -1,15 +1,16 @@
 snowglobe/client/__init__.py,sha256=kzp9wPUUYBXqDSKZbfmD4vrAQvrWSW5HOvtpFlEJWfs,353
-snowglobe/client/src/app.py,sha256=Ij7nB88VeGtz4iEKc9iwCQnG-AdsYaIDscwcU7SzDXg,39112
-snowglobe/client/src/cli.py,sha256=DIOI4E0ZeEljXDiEZb9rkbTWdHi3m4owWvyKtAC5Ro0,27999
+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.8.dist-info/licenses/LICENSE,sha256=S90V6iFU5ZeSg44JQYS1To3pa7ZEobrHc_t483qSKSI,1070
-snowglobe-0.4.8.dist-info/METADATA,sha256=158Z9QgzGaU4HIBqe020xSQMI7evAz0bZdlpl29yPZI,4467
-snowglobe-0.4.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-snowglobe-0.4.8.dist-info/entry_points.txt,sha256=mqx4mTwFPHttjctE2ceYTYWCCIG30Ji2C89aaCYgHcM,71
-snowglobe-0.4.8.dist-info/top_level.txt,sha256=PoyYihnCBjRyjeIT19yBcE47JTe7i1OwRXvJ4d5EohM,10
-snowglobe-0.4.8.dist-info/RECORD,,
+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,,