move analytics code from scripts to skyvern module (#32)

run_skyvern.sh

@@ -8,5 +8,4 @@ if [ ! -f .env ]; then
 fi
 source "$(poetry env info --path)/bin/activate"
 ./run_alembic_check.sh
-python scripts/tracking.py skyvern-oss-run-server
 poetry run python -m skyvern.forge

run_ui.sh

@@ -1,5 +1,2 @@
-
-
 source "$(poetry env info --path)/bin/activate"
-python scripts/tracking.py skyvern-oss-run-ui
 streamlit run streamlit_app/visualizer/streamlit.py -- $@

setup.sh

@@ -3,7 +3,7 @@
 # Call function to send telemetry event
 log_event() {
     if [ -n $1 ]; then
-        python scripts/tracking.py $1
+        python skyvern/analytics.py $1
     fi
 }
 
@@ -36,7 +36,7 @@ initialize_env_file() {
 
 
     # Ask for email or generate UUID
-    read -p "Please enter your email for analytics tracking (press enter to skip): " analytics_id
+    read -p "Please enter your email for analytics (press enter to skip): " analytics_id
     if [ -z "$analytics_id" ]; then
         analytics_id=$(uuidgen)
     fi
@@ -140,4 +140,4 @@ main() {
 }
 
 # Execute main function
-main
+main

skyvern/forge/__main__.py

@@ -3,12 +3,14 @@ import uvicorn
 from dotenv import load_dotenv
 
 import skyvern.forge.sdk.forge_log as forge_log
+from skyvern import analytics
 from skyvern.forge.sdk.settings_manager import SettingsManager
 
 LOG = structlog.stdlib.get_logger()
 
 
 if __name__ == "__main__":
+    analytics.capture("skyvern-oss-run-server")
     forge_log.setup_logger()
     port = SettingsManager.get_settings().PORT
     LOG.info("Agent server starting.", host="0.0.0.0", port=port)

skyvern/forge/agent.py

@@ -8,7 +8,7 @@ import requests
 import structlog
 from playwright._impl._errors import TargetClosedError
 
-from scripts import tracking
+from skyvern import analytics
 from skyvern.exceptions import (
     BrowserStateMissingPage,
     FailedToSendWebhook,
@@ -195,7 +195,7 @@ class ForgeAgent(Agent):
             await self.validate_step_execution(task, step)
             step, browser_state, detailed_output = await self._initialize_execution_state(task, step, workflow_run)
             step, detailed_output = await self.agent_step(task, step, browser_state, organization=organization)
-            tracking.capture("skyvern-oss-agent-step-status", {"status": step.status})
+            analytics.capture("skyvern-oss-agent-step-status", {"status": step.status})
             retry = False
 
             # If the step failed, mark the step as failed and retry
@@ -674,7 +674,7 @@ class ForgeAgent(Agent):
             raise TaskNotFound(task_id=task.task_id) from e
         task = refreshed_task
         # log the task status as an event
-        tracking.capture("skyvern-oss-agent-task-status", {"status": task.status})
+        analytics.capture("skyvern-oss-agent-task-status", {"status": task.status})
         # Take one last screenshot and create an artifact before closing the browser to see the final state
         browser_state: BrowserState = await app.BROWSER_MANAGER.get_or_create_for_task(task)
         page = await browser_state.get_or_create_page()

skyvern/forge/sdk/routes/agent_protocol.py

@@ -5,7 +5,7 @@ from fastapi import APIRouter, BackgroundTasks, Depends, Header, HTTPException,
 from fastapi.responses import ORJSONResponse
 from pydantic import BaseModel
 
-from scripts import tracking
+from skyvern import analytics
 from skyvern.exceptions import StepNotFound
 from skyvern.forge import app
 from skyvern.forge.sdk.artifact.models import Artifact, ArtifactType
@@ -32,7 +32,7 @@ async def webhook(
     x_skyvern_signature: Annotated[str | None, Header()] = None,
     x_skyvern_timestamp: Annotated[str | None, Header()] = None,
 ) -> Response:
-    tracking.capture("skyvern-oss-agent-webhook-received")
+    analytics.capture("skyvern-oss-agent-webhook-received")
     payload = await request.body()
 
     if not x_skyvern_signature or not x_skyvern_timestamp:
@@ -77,7 +77,7 @@ async def create_agent_task(
     x_api_key: Annotated[str | None, Header()] = None,
     x_max_steps_override: Annotated[int | None, Header()] = None,
 ) -> CreateTaskResponse:
-    tracking.capture("skyvern-oss-agent-task-create", data={"url": task.url})
+    analytics.capture("skyvern-oss-agent-task-create", data={"url": task.url})
     agent = request["agent"]
 
     created_task = await agent.create_task(task, current_org.organization_id)
@@ -111,7 +111,7 @@ async def execute_agent_task_step(
     step_id: str | None = None,
     current_org: Organization = Depends(org_auth_service.get_current_org),
 ) -> Response:
-    tracking.capture("skyvern-oss-agent-task-step-execute")
+    analytics.capture("skyvern-oss-agent-task-step-execute")
     agent = request["agent"]
     task = await app.DATABASE.get_task(task_id, organization_id=current_org.organization_id)
     if not task:
@@ -170,7 +170,7 @@ async def get_task(
     task_id: str,
     current_org: Organization = Depends(org_auth_service.get_current_org),
 ) -> TaskResponse:
-    tracking.capture("skyvern-oss-agent-task-get")
+    analytics.capture("skyvern-oss-agent-task-get")
     request["agent"]
     task_obj = await app.DATABASE.get_task(task_id, organization_id=current_org.organization_id)
     if not task_obj:
@@ -234,7 +234,7 @@ async def retry_webhook(
     current_org: Organization = Depends(org_auth_service.get_current_org),
     x_api_key: Annotated[str | None, Header()] = None,
 ) -> TaskResponse:
-    tracking.capture("skyvern-oss-agent-task-retry-webhook")
+    analytics.capture("skyvern-oss-agent-task-retry-webhook")
     agent = request["agent"]
     task_obj = await agent.db.get_task(task_id, organization_id=current_org.organization_id)
     if not task_obj:
@@ -268,7 +268,7 @@ async def get_task_internal(
     :return: List of tasks with pagination without steps populated. Steps can be populated by calling the
         get_agent_task endpoint.
     """
-    tracking.capture("skyvern-oss-agent-task-get-internal")
+    analytics.capture("skyvern-oss-agent-task-get-internal")
     task = await app.DATABASE.get_task(task_id, organization_id=current_org.organization_id)
     if not task:
         raise HTTPException(
@@ -293,7 +293,7 @@ async def get_agent_tasks(
     :return: List of tasks with pagination without steps populated. Steps can be populated by calling the
         get_agent_task endpoint.
     """
-    tracking.capture("skyvern-oss-agent-tasks-get")
+    analytics.capture("skyvern-oss-agent-tasks-get")
     request["agent"]
     tasks = await app.DATABASE.get_tasks(page, page_size, organization_id=current_org.organization_id)
     return ORJSONResponse([task.to_task_response().model_dump() for task in tasks])
@@ -314,7 +314,7 @@ async def get_agent_tasks_internal(
     :return: List of tasks with pagination without steps populated. Steps can be populated by calling the
         get_agent_task endpoint.
     """
-    tracking.capture("skyvern-oss-agent-tasks-get-internal")
+    analytics.capture("skyvern-oss-agent-tasks-get-internal")
     request["agent"]
     tasks = await app.DATABASE.get_tasks(page, page_size, organization_id=current_org.organization_id)
     return ORJSONResponse([task.model_dump() for task in tasks])
@@ -332,7 +332,7 @@ async def get_agent_task_steps(
     :param task_id:
     :return: List of steps for a task with pagination.
     """
-    tracking.capture("skyvern-oss-agent-task-steps-get")
+    analytics.capture("skyvern-oss-agent-task-steps-get")
     request["agent"]
     steps = await app.DATABASE.get_task_steps(task_id, organization_id=current_org.organization_id)
     return ORJSONResponse([step.model_dump() for step in steps])
@@ -352,7 +352,7 @@ async def get_agent_task_step_artifacts(
     :param step_id:
     :return: List of artifacts for a list of steps.
     """
-    tracking.capture("skyvern-oss-agent-task-step-artifacts-get")
+    analytics.capture("skyvern-oss-agent-task-step-artifacts-get")
     request["agent"]
     artifacts = await app.DATABASE.get_artifacts_for_task_step(
         task_id,
@@ -375,7 +375,7 @@ async def get_task_actions(
     task_id: str,
     current_org: Organization = Depends(org_auth_service.get_current_org),
 ) -> list[ActionResultTmp]:
-    tracking.capture("skyvern-oss-agent-task-actions-get")
+    analytics.capture("skyvern-oss-agent-task-actions-get")
     request["agent"]
     steps = await app.DATABASE.get_task_step_models(task_id, organization_id=current_org.organization_id)
     results: list[ActionResultTmp] = []
@@ -397,7 +397,7 @@ async def execute_workflow(
     x_api_key: Annotated[str | None, Header()] = None,
     x_max_steps_override: Annotated[int | None, Header()] = None,
 ) -> RunWorkflowResponse:
-    tracking.capture("skyvern-oss-agent-workflow-execute")
+    analytics.capture("skyvern-oss-agent-workflow-execute")
     LOG.info(
         f"Running workflow {workflow_id}",
         workflow_id=workflow_id,
@@ -434,7 +434,7 @@ async def get_workflow_run(
     workflow_run_id: str,
     current_org: Organization = Depends(org_auth_service.get_current_org),
 ) -> WorkflowRunStatusResponse:
-    tracking.capture("skyvern-oss-agent-workflow-run-get")
+    analytics.capture("skyvern-oss-agent-workflow-run-get")
     request["agent"]
     return await app.WORKFLOW_SERVICE.build_workflow_run_status_response(
         workflow_id=workflow_id, workflow_run_id=workflow_run_id, organization_id=current_org.organization_id

skyvern/forge/sdk/workflow/service.py

@@ -6,7 +6,7 @@ from datetime import datetime
 import requests
 import structlog
 
-from scripts import tracking
+from skyvern import analytics
 from skyvern.exceptions import (
     FailedToSendWebhook,
     MissingValueForParameter,
@@ -375,7 +375,7 @@ class WorkflowService:
         api_key: str | None = None,
         close_browser_on_completion: bool = True,
     ) -> None:
-        tracking.capture("skyvern-oss-agent-workflow-status", {"status": workflow_run.status})
+        analytics.capture("skyvern-oss-agent-workflow-status", {"status": workflow_run.status})
         browser_state = await app.BROWSER_MANAGER.cleanup_for_workflow_run(
             workflow_run.workflow_run_id, close_browser_on_completion
         )

streamlit_app/visualizer/streamlit.py

@@ -4,6 +4,7 @@ import clipboard
 import pandas as pd
 import streamlit as st
 
+from skyvern import analytics
 from skyvern.forge.sdk.schemas.tasks import ProxyLocation, TaskRequest
 from streamlit_app.visualizer import styles
 from streamlit_app.visualizer.api import SkyvernClient
@@ -15,6 +16,8 @@ from streamlit_app.visualizer.artifact_loader import (
 from streamlit_app.visualizer.repository import TaskRepository
 from streamlit_app.visualizer.sample_data import supported_examples
 
+analytics.capture("skyvern-oss-run-ui")
+
 # Streamlit UI Configuration
 st.set_page_config(layout="wide")