diff --git a/skyvern/forge/agent.py b/skyvern/forge/agent.py index 10420c68..b64208fd 100644 --- a/skyvern/forge/agent.py +++ b/skyvern/forge/agent.py @@ -83,7 +83,6 @@ from skyvern.forge.sdk.schemas.organizations import Organization from skyvern.forge.sdk.schemas.tasks import Task, TaskRequest, TaskResponse, TaskStatus from skyvern.forge.sdk.schemas.totp_codes import OTPType from skyvern.forge.sdk.trace import TraceManager -from skyvern.forge.sdk.trace.experiment_utils import collect_experiment_metadata_safely from skyvern.forge.sdk.workflow.context_manager import WorkflowRunContext from skyvern.forge.sdk.workflow.models.block import ActionBlock, BaseTaskBlock, ValidationBlock from skyvern.forge.sdk.workflow.models.workflow import Workflow, WorkflowRun, WorkflowRunStatus @@ -315,11 +314,6 @@ class ForgeAgent: close_browser_on_completion and browser_session_id is None and not task.browser_address ) - # Collect and add experiment metadata to the trace - experiment_metadata = await collect_experiment_metadata_safely(app.EXPERIMENTATION_PROVIDER) - if experiment_metadata: - TraceManager.add_experiment_metadata(experiment_metadata) - workflow_run: WorkflowRun | None = None if task.workflow_run_id: workflow_run = await app.DATABASE.get_workflow_run( diff --git a/skyvern/forge/sdk/trace/__init__.py b/skyvern/forge/sdk/trace/__init__.py index 8efdd7e0..94bdc600 100644 --- a/skyvern/forge/sdk/trace/__init__.py +++ b/skyvern/forge/sdk/trace/__init__.py @@ -1,9 +1,11 @@ from functools import wraps from typing import Any, Awaitable, Callable, ParamSpec, TypeVar +from skyvern.forge import app from skyvern.forge.sdk.core import skyvern_context from skyvern.forge.sdk.settings_manager import SettingsManager from skyvern.forge.sdk.trace.base import BaseTrace, NoOpTrace +from skyvern.forge.sdk.trace.experiment_utils import collect_experiment_metadata_safely P = ParamSpec("P") R = TypeVar("R") @@ -37,6 +39,11 @@ class TraceManager: new_metadata["organization_name"] = context.organization_name user_id = context.run_id + # Collect experiment metadata and include it in the span metadata + experiment_metadata = await collect_experiment_metadata_safely(app.EXPERIMENTATION_PROVIDER) + if experiment_metadata: + new_metadata.update(experiment_metadata) + new_tags: list[str] = tags or [] new_tags.append(SettingsManager.get_settings().ENV) diff --git a/skyvern/forge/sdk/trace/lmnr.py b/skyvern/forge/sdk/trace/lmnr.py index 4e142eb4..17f21962 100644 --- a/skyvern/forge/sdk/trace/lmnr.py +++ b/skyvern/forge/sdk/trace/lmnr.py @@ -42,7 +42,7 @@ class LaminarTrace(BaseTrace): # Map status to appropriate tag status_tag_map = { - "completed": "COMPLETION", + "completed": "COMPLETED", "failed": "FAILURE", "timed_out": "TIMEOUT", "canceled": "CANCELED", diff --git a/skyvern/forge/sdk/workflow/service.py b/skyvern/forge/sdk/workflow/service.py index 92db2214..a452e9b7 100644 --- a/skyvern/forge/sdk/workflow/service.py +++ b/skyvern/forge/sdk/workflow/service.py @@ -41,7 +41,6 @@ from skyvern.forge.sdk.schemas.organizations import Organization from skyvern.forge.sdk.schemas.tasks import Task from skyvern.forge.sdk.schemas.workflow_runs import WorkflowRunBlock, WorkflowRunTimeline, WorkflowRunTimelineType from skyvern.forge.sdk.trace import TraceManager -from skyvern.forge.sdk.trace.experiment_utils import collect_experiment_metadata_safely from skyvern.forge.sdk.workflow.exceptions import ( ContextParameterSourceNotDefined, InvalidWaitBlockTime, @@ -347,11 +346,6 @@ class WorkflowService: """Execute a workflow.""" organization_id = organization.organization_id - # Collect and add experiment metadata to the trace - experiment_metadata = await collect_experiment_metadata_safely(app.EXPERIMENTATION_PROVIDER) - if experiment_metadata: - TraceManager.add_experiment_metadata(experiment_metadata) - LOG.info( "Executing workflow", workflow_run_id=workflow_run_id, diff --git a/skyvern/services/task_v2_service.py b/skyvern/services/task_v2_service.py index 3311036b..432d921e 100644 --- a/skyvern/services/task_v2_service.py +++ b/skyvern/services/task_v2_service.py @@ -29,7 +29,6 @@ from skyvern.forge.sdk.schemas.organizations import Organization from skyvern.forge.sdk.schemas.task_v2 import TaskV2, TaskV2Metadata, TaskV2Status, ThoughtScenario, ThoughtType from skyvern.forge.sdk.schemas.workflow_runs import WorkflowRunTimeline, WorkflowRunTimelineType from skyvern.forge.sdk.trace import TraceManager -from skyvern.forge.sdk.trace.experiment_utils import collect_experiment_metadata_safely from skyvern.forge.sdk.workflow.models.block import ( BlockTypeVar, ExtractionBlock, @@ -327,11 +326,6 @@ async def run_task_v2( LOG.error("Task v2 not found", task_v2_id=task_v2_id, organization_id=organization_id) raise TaskV2NotFound(task_v2_id=task_v2_id) - # Collect and add experiment metadata to the trace - experiment_metadata = await collect_experiment_metadata_safely(app.EXPERIMENTATION_PROVIDER) - if experiment_metadata: - TraceManager.add_experiment_metadata(experiment_metadata) - workflow, workflow_run = None, None try: workflow, workflow_run, task_v2 = await run_task_v2_helper(