vodorod/Dorod-Sky
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Add AI suggestion endpoints (#1519)

Browse Source
This commit is contained in:
Shuchang Zheng
2025-01-08 21:45:38 -08:00
committed by GitHub
parent f0ae840452
commit 5796de73d1
14 changed files with 248 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
skyvern/forge/sdk/artifact/manager.py
View File

@@ -8,6 +8,7 @@ from skyvern.forge import app
from skyvern.forge.sdk.artifact.models import Artifact, ArtifactType, LogEntityType
from skyvern.forge.sdk.db.id import generate_artifact_id
from skyvern.forge.sdk.models import Step
from skyvern.forge.sdk.schemas.ai_suggestions import AISuggestion
from skyvern.forge.sdk.schemas.observers import ObserverCruise, ObserverThought
from skyvern.forge.sdk.schemas.workflow_runs import WorkflowRunBlock
@@ -30,6 +31,7 @@ class ArtifactManager:
workflow_run_block_id: str | None = None,
observer_thought_id: str | None = None,
observer_cruise_id: str | None = None,
ai_suggestion_id: str | None = None,
organization_id: str | None = None,
data: bytes | None = None,
path: str | None = None,
@@ -49,6 +51,7 @@ class ArtifactManager:
observer_thought_id=observer_thought_id,
observer_cruise_id=observer_cruise_id,
organization_id=organization_id,
ai_suggestion_id=ai_suggestion_id,
)
if data:
# Fire and forget
@@ -173,6 +176,26 @@ class ArtifactManager:
path=path,
)
async def create_ai_suggestion_artifact(
self,
ai_suggestion: AISuggestion,
artifact_type: ArtifactType,
data: bytes | None = None,
path: str | None = None,
) -> str:
artifact_id = generate_artifact_id()
uri = app.STORAGE.build_ai_suggestion_uri(artifact_id, ai_suggestion, artifact_type)
return await self._create_artifact(
aio_task_primary_key=ai_suggestion.ai_suggestion_id,
artifact_id=artifact_id,
artifact_type=artifact_type,
uri=uri,
ai_suggestion_id=ai_suggestion.ai_suggestion_id,
organization_id=ai_suggestion.organization_id,
data=data,
path=path,
)
async def create_llm_artifact(
self,
data: bytes,
@@ -181,6 +204,7 @@ class ArtifactManager:
step: Step | None = None,
observer_thought: ObserverThought | None = None,
observer_cruise: ObserverCruise | None = None,
ai_suggestion: AISuggestion | None = None,
) -> None:
if step:
await self.create_artifact(
@@ -218,6 +242,18 @@ class ArtifactManager:
artifact_type=ArtifactType.SCREENSHOT_LLM,
data=screenshot,
)
elif ai_suggestion:
await self.create_ai_suggestion_artifact(
ai_suggestion=ai_suggestion,
artifact_type=artifact_type,
data=data,
)
for screenshot in screenshots or []:
await self.create_ai_suggestion_artifact(
ai_suggestion=ai_suggestion,
artifact_type=ArtifactType.SCREENSHOT_LLM,
data=screenshot,
)
async def update_artifact_data(
self,

1
skyvern/forge/sdk/artifact/models.py
View File

@@ -75,6 +75,7 @@ class Artifact(BaseModel):
workflow_run_block_id: str | None = None
observer_cruise_id: str | None = None
observer_thought_id: str | None = None
ai_suggestion_id: str | None = None
signed_url: str | None = None
organization_id: str | None = None

7
skyvern/forge/sdk/artifact/storage/base.py
View File

@@ -2,6 +2,7 @@ from abc import ABC, abstractmethod
from skyvern.forge.sdk.artifact.models import Artifact, ArtifactType, LogEntityType
from skyvern.forge.sdk.models import Step
from skyvern.forge.sdk.schemas.ai_suggestions import AISuggestion
from skyvern.forge.sdk.schemas.observers import ObserverCruise, ObserverThought
from skyvern.forge.sdk.schemas.workflow_runs import WorkflowRunBlock
@@ -61,6 +62,12 @@ class BaseStorage(ABC):
) -> str:
pass
@abstractmethod
def build_ai_suggestion_uri(
self, artifact_id: str, ai_suggestion: AISuggestion, artifact_type: ArtifactType
) -> str:
pass
@abstractmethod
async def store_artifact(self, artifact: Artifact, data: bytes) -> None:
pass

7
skyvern/forge/sdk/artifact/storage/local.py
View File

@@ -11,6 +11,7 @@ from skyvern.forge.sdk.api.files import get_download_dir, get_skyvern_temp_dir
from skyvern.forge.sdk.artifact.models import Artifact, ArtifactType, LogEntityType
from skyvern.forge.sdk.artifact.storage.base import FILE_EXTENTSION_MAP, BaseStorage
from skyvern.forge.sdk.models import Step
from skyvern.forge.sdk.schemas.ai_suggestions import AISuggestion
from skyvern.forge.sdk.schemas.observers import ObserverCruise, ObserverThought
from skyvern.forge.sdk.schemas.workflow_runs import WorkflowRunBlock
@@ -47,6 +48,12 @@ class LocalStorage(BaseStorage):
file_ext = FILE_EXTENTSION_MAP[artifact_type]
return f"file://{self.artifact_path}/{settings.ENV}/workflow_runs/{workflow_run_block.workflow_run_id}/{workflow_run_block.workflow_run_block_id}/{datetime.utcnow().isoformat()}_{artifact_id}_{artifact_type}.{file_ext}"
def build_ai_suggestion_uri(
self, artifact_id: str, ai_suggestion: AISuggestion, artifact_type: ArtifactType
) -> str:
file_ext = FILE_EXTENTSION_MAP[artifact_type]
return f"file://{self.artifact_path}/{settings.ENV}/ai_suggestions/{ai_suggestion.ai_suggestion_id}/{datetime.utcnow().isoformat()}_{artifact_id}_{artifact_type}.{file_ext}"
async def store_artifact(self, artifact: Artifact, data: bytes) -> None:
file_path = None
try:

7
skyvern/forge/sdk/artifact/storage/s3.py
View File

@@ -15,6 +15,7 @@ from skyvern.forge.sdk.api.files import (
from skyvern.forge.sdk.artifact.models import Artifact, ArtifactType, LogEntityType
from skyvern.forge.sdk.artifact.storage.base import FILE_EXTENTSION_MAP, BaseStorage
from skyvern.forge.sdk.models import Step
from skyvern.forge.sdk.schemas.ai_suggestions import AISuggestion
from skyvern.forge.sdk.schemas.observers import ObserverCruise, ObserverThought
from skyvern.forge.sdk.schemas.workflow_runs import WorkflowRunBlock
@@ -50,6 +51,12 @@ class S3Storage(BaseStorage):
file_ext = FILE_EXTENTSION_MAP[artifact_type]
return f"s3://{self.bucket}/{settings.ENV}/workflow_runs/{workflow_run_block.workflow_run_id}/{workflow_run_block.workflow_run_block_id}/{datetime.utcnow().isoformat()}_{artifact_id}_{artifact_type}.{file_ext}"
def build_ai_suggestion_uri(
self, artifact_id: str, ai_suggestion: AISuggestion, artifact_type: ArtifactType
) -> str:
file_ext = FILE_EXTENTSION_MAP[artifact_type]
return f"s3://{self.bucket}/{settings.ENV}/ai_suggestions/{ai_suggestion.ai_suggestion_id}/{datetime.utcnow().isoformat()}_{artifact_id}_{artifact_type}.{file_ext}"
async def store_artifact(self, artifact: Artifact, data: bytes) -> None:
await self.async_client.upload_file(artifact.uri, data)