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

Initial scenario tests infra (#4118)

Browse Source
This commit is contained in:
Stanislav Novosad
2025-11-26 19:26:48 -07:00
committed by GitHub
parent 4dcc8ed725
commit 5b530cab52
6 changed files with 245 additions and 106 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
skyvern/forge/forge_app.py
View File

@@ -9,8 +9,10 @@ from openai import AsyncAzureOpenAI, AsyncOpenAI
from skyvern.config import Settings
from skyvern.forge.agent import ForgeAgent
from skyvern.forge.agent_functions import AgentFunction
from skyvern.forge.sdk.api.azure import AzureClientFactory
from skyvern.forge.sdk.api.llm.api_handler_factory import LLMAPIHandlerFactory
from skyvern.forge.sdk.api.llm.models import LLMAPIHandler
from skyvern.forge.sdk.api.real_azure import RealAzureClientFactory
from skyvern.forge.sdk.artifact.manager import ArtifactManager
from skyvern.forge.sdk.artifact.storage.base import BaseStorage
from skyvern.forge.sdk.artifact.storage.factory import StorageFactory
@@ -20,7 +22,7 @@ from skyvern.forge.sdk.cache.factory import CacheFactory
from skyvern.forge.sdk.db.client import AgentDB
from skyvern.forge.sdk.experimentation.providers import BaseExperimentationProvider, NoOpExperimentationProvider
from skyvern.forge.sdk.schemas.credentials import CredentialVaultType
from skyvern.forge.sdk.schemas.organizations import Organization
from skyvern.forge.sdk.schemas.organizations import AzureClientSecretCredential, Organization
from skyvern.forge.sdk.services.credential.azure_credential_vault_service import AzureCredentialVaultService
from skyvern.forge.sdk.services.credential.bitwarden_credential_service import BitwardenCredentialVaultService
from skyvern.forge.sdk.services.credential.credential_vault_service import CredentialVaultService
@@ -46,6 +48,7 @@ class ForgeApp:
OPENAI_CLIENT: AsyncOpenAI | AsyncAzureOpenAI
ANTHROPIC_CLIENT: AsyncAnthropic | AsyncAnthropicBedrock
UI_TARS_CLIENT: AsyncOpenAI | None
AZURE_CLIENT_FACTORY: AzureClientFactory
SECONDARY_LLM_API_HANDLER: LLMAPIHandler
SELECT_AGENT_LLM_API_HANDLER: LLMAPIHandler
NORMAL_SELECT_AGENT_LLM_API_HANDLER: LLMAPIHandler
@@ -166,12 +169,17 @@ def create_forge_app() -> ForgeApp:
app.AGENT_FUNCTION = AgentFunction()
app.PERSISTENT_SESSIONS_MANAGER = PersistentSessionsManager(database=app.DATABASE)
app.AZURE_CLIENT_FACTORY = RealAzureClientFactory()
app.BITWARDEN_CREDENTIAL_VAULT_SERVICE = BitwardenCredentialVaultService()
app.AZURE_CREDENTIAL_VAULT_SERVICE = (
AzureCredentialVaultService(
tenant_id=settings.AZURE_TENANT_ID, # type: ignore[arg-type]
client_id=settings.AZURE_CLIENT_ID, # type: ignore[arg-type]
client_secret=settings.AZURE_CLIENT_SECRET, # type: ignore[arg-type]
app.AZURE_CLIENT_FACTORY.create_from_client_secret(
AzureClientSecretCredential(
tenant_id=settings.AZURE_TENANT_ID, # type: ignore
client_id=settings.AZURE_CLIENT_ID, # type: ignore
client_secret=settings.AZURE_CLIENT_SECRET, # type: ignore
)
),
vault_name=settings.AZURE_CREDENTIAL_VAULT, # type: ignore[arg-type]
)
if settings.AZURE_CREDENTIAL_VAULT

187
skyvern/forge/sdk/api/azure.py
View File

@@ -1,115 +1,126 @@
from typing import Self
import structlog
from azure.identity.aio import ClientSecretCredential, DefaultAzureCredential
from azure.keyvault.secrets.aio import SecretClient
from azure.storage.blob.aio import BlobServiceClient
from typing import Protocol, Self
from skyvern.forge.sdk.schemas.organizations import AzureClientSecretCredential
LOG = structlog.get_logger()
class AsyncAzureVaultClient(Protocol):
"""Protocol defining the interface for Azure Vault clients.
class AsyncAzureVaultClient:
def __init__(self, credential: ClientSecretCredential | DefaultAzureCredential) -> None:
self.credential = credential
This client provides methods to interact with Azure Key Vault for secret management.
"""
async def __aenter__(self) -> Self:
return self
"""Enter async context manager."""
...
async def __aexit__(
self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: object
) -> None:
await self.credential.close()
"""Exit async context manager and cleanup resources."""
...
async def get_secret(self, secret_name: str, vault_name: str) -> str | None:
secret_client = await self._get_secret_client(vault_name)
try:
secret = await secret_client.get_secret(secret_name)
return secret.value
except Exception as e:
LOG.exception("Failed to get secret from Azure Key Vault.", secret_name=secret_name, error=e)
return None
finally:
await secret_client.close()
"""Retrieve a secret from Azure Key Vault.
Args:
secret_name: The name of the secret to retrieve
vault_name: The name of the Azure Key Vault
Returns:
The secret value as a string, or None if the secret doesn't exist or an error occurs
"""
...
async def create_or_update_secret(self, secret_name: str, secret_value: str, vault_name: str) -> str:
secret_client = await self._get_secret_client(vault_name)
try:
secret = await secret_client.set_secret(secret_name, secret_value)
return secret.name
except Exception as e:
LOG.exception("Failed to create secret from Azure Key Vault.", secret_name=secret_name, error=e)
raise e
finally:
await secret_client.close()
"""Create or update a secret in Azure Key Vault.
Args:
secret_name: The name of the secret to create or update
secret_value: The value to store
vault_name: The name of the Azure Key Vault
Returns:
The name of the created/updated secret
Raises:
Exception: If the operation fails
"""
...
async def delete_secret(self, secret_name: str, vault_name: str) -> str:
secret_client = await self._get_secret_client(vault_name)
try:
secret = await secret_client.delete_secret(secret_name)
return secret.name
except Exception as e:
LOG.exception("Failed to delete secret from Azure Key Vault.", secret_name=secret_name, error=e)
raise e
finally:
await secret_client.close()
"""Delete a secret from Azure Key Vault.
async def _get_secret_client(self, vault_name: str) -> SecretClient:
# Azure Key Vault URL format: https://<your-key-vault-name>.vault.azure.net
# Assuming the secret_name is actually the Key Vault URL and the secret name
# This needs to be clarified or passed as separate parameters
# For now, let's assume secret_name is the actual secret name and Key Vault URL is in settings.
key_vault_url = f"https://{vault_name}.vault.azure.net" # Placeholder, adjust as needed
return SecretClient(vault_url=key_vault_url, credential=self.credential)
Args:
secret_name: The name of the secret to delete
vault_name: The name of the Azure Key Vault
Returns:
The name of the deleted secret
Raises:
Exception: If the operation fails
"""
...
async def close(self) -> None:
await self.credential.close()
@classmethod
def create_default(cls) -> "AsyncAzureVaultClient":
return cls(DefaultAzureCredential())
@classmethod
def create_from_client_secret(
cls,
credential: AzureClientSecretCredential,
) -> "AsyncAzureVaultClient":
cred = ClientSecretCredential(
tenant_id=credential.tenant_id,
client_id=credential.client_id,
client_secret=credential.client_secret,
)
return cls(cred)
"""Close the client and release all resources."""
...
class AsyncAzureStorageClient:
def __init__(self, storage_account_name: str, storage_account_key: str):
self.blob_service_client = BlobServiceClient(
account_url=f"https://{storage_account_name}.blob.core.windows.net",
credential=storage_account_key,
)
class AsyncAzureStorageClient(Protocol):
"""Protocol defining the interface for Azure Storage clients.
This client provides methods to interact with Azure Blob Storage for file operations.
"""
async def upload_file_from_path(self, container_name: str, blob_name: str, file_path: str) -> None:
try:
container_client = self.blob_service_client.get_container_client(container_name)
# Create the container if it doesn't exist
try:
await container_client.create_container()
except Exception as e:
LOG.info("Azure container already exists or failed to create", container_name=container_name, error=e)
"""Upload a file from the local filesystem to Azure Blob Storage.
with open(file_path, "rb") as data:
await container_client.upload_blob(name=blob_name, data=data, overwrite=True)
LOG.info("File uploaded to Azure Blob Storage", container_name=container_name, blob_name=blob_name)
except Exception as e:
LOG.error(
"Failed to upload file to Azure Blob Storage",
container_name=container_name,
blob_name=blob_name,
error=e,
)
raise e
Args:
container_name: The name of the Azure Blob container
blob_name: The name to give the blob in storage
file_path: The local path to the file to upload
Raises:
Exception: If the upload fails
"""
...
async def close(self) -> None:
await self.blob_service_client.close()
"""Close the storage client and release resources."""
...
class AzureClientFactory(Protocol):
"""Protocol defining the interface for creating Azure Vault and Storage clients."""
def create_default(self) -> "AsyncAzureVaultClient":
"""Create an Azure Vault client using default credentials.
Returns:
An AsyncAzureVaultClient instance using DefaultAzureCredential
"""
...
def create_from_client_secret(self, credential: AzureClientSecretCredential) -> "AsyncAzureVaultClient":
"""Create an Azure Vault client using client secret credentials.
Args:
credential: Azure client secret credentials containing tenant_id, client_id, and client_secret
Returns:
An AsyncAzureVaultClient instance
"""
...
def create_storage_client(self, storage_account_name: str, storage_account_key: str) -> "AsyncAzureStorageClient":
"""Create an Azure Storage client with the provided credentials.
Args:
storage_account_name: The name of the Azure storage account
storage_account_key: The access key for the storage account
Returns:
An AsyncAzureStorageClient instance
"""
...

127
skyvern/forge/sdk/api/real_azure.py Normal file
View File

@@ -0,0 +1,127 @@
"""Real implementations of Azure clients (Vault and Storage) and their factories."""
from typing import Self
import structlog
from azure.identity.aio import ClientSecretCredential, DefaultAzureCredential
from azure.keyvault.secrets.aio import SecretClient
from azure.storage.blob.aio import BlobServiceClient
from skyvern.forge.sdk.api.azure import AsyncAzureStorageClient, AsyncAzureVaultClient, AzureClientFactory
from skyvern.forge.sdk.schemas.organizations import AzureClientSecretCredential
LOG = structlog.get_logger()
class RealAsyncAzureVaultClient(AsyncAzureVaultClient):
"""Real implementation of Azure Vault client using Azure SDK."""
def __init__(self, credential: ClientSecretCredential | DefaultAzureCredential) -> None:
self.credential = credential
async def __aenter__(self) -> Self:
return self
async def __aexit__(
self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: object
) -> None:
await self.credential.close()
async def get_secret(self, secret_name: str, vault_name: str) -> str | None:
secret_client = await self._get_secret_client(vault_name)
try:
secret = await secret_client.get_secret(secret_name)
return secret.value
except Exception as e:
LOG.exception("Failed to get secret from Azure Key Vault.", secret_name=secret_name, error=e)
return None
finally:
await secret_client.close()
async def create_or_update_secret(self, secret_name: str, secret_value: str, vault_name: str) -> str:
secret_client = await self._get_secret_client(vault_name)
try:
secret = await secret_client.set_secret(secret_name, secret_value)
return secret.name
except Exception as e:
LOG.exception("Failed to create secret from Azure Key Vault.", secret_name=secret_name, error=e)
raise e
finally:
await secret_client.close()
async def delete_secret(self, secret_name: str, vault_name: str) -> str:
secret_client = await self._get_secret_client(vault_name)
try:
secret = await secret_client.delete_secret(secret_name)
return secret.name
except Exception as e:
LOG.exception("Failed to delete secret from Azure Key Vault.", secret_name=secret_name, error=e)
raise e
finally:
await secret_client.close()
async def _get_secret_client(self, vault_name: str) -> SecretClient:
# Azure Key Vault URL format: https://<your-key-vault-name>.vault.azure.net
# Assuming the secret_name is actually the Key Vault URL and the secret name
# This needs to be clarified or passed as separate parameters
# For now, let's assume secret_name is the actual secret name and Key Vault URL is in settings.
key_vault_url = f"https://{vault_name}.vault.azure.net" # Placeholder, adjust as needed
return SecretClient(vault_url=key_vault_url, credential=self.credential)
async def close(self) -> None:
await self.credential.close()
class RealAsyncAzureStorageClient(AsyncAzureStorageClient):
"""Real implementation of Azure Storage client using Azure SDK."""
def __init__(self, storage_account_name: str, storage_account_key: str):
self.blob_service_client = BlobServiceClient(
account_url=f"https://{storage_account_name}.blob.core.windows.net",
credential=storage_account_key,
)
async def upload_file_from_path(self, container_name: str, blob_name: str, file_path: str) -> None:
try:
container_client = self.blob_service_client.get_container_client(container_name)
# Create the container if it doesn't exist
try:
await container_client.create_container()
except Exception as e:
LOG.info("Azure container already exists or failed to create", container_name=container_name, error=e)
with open(file_path, "rb") as data:
await container_client.upload_blob(name=blob_name, data=data, overwrite=True)
LOG.info("File uploaded to Azure Blob Storage", container_name=container_name, blob_name=blob_name)
except Exception as e:
LOG.error(
"Failed to upload file to Azure Blob Storage",
container_name=container_name,
blob_name=blob_name,
error=e,
)
raise e
async def close(self) -> None:
await self.blob_service_client.close()
class RealAzureClientFactory(AzureClientFactory):
"""Factory for creating real Azure Vault and Storage clients."""
def create_default(self) -> AsyncAzureVaultClient:
"""Create an Azure Vault client using DefaultAzureCredential."""
return RealAsyncAzureVaultClient(DefaultAzureCredential())
def create_from_client_secret(self, credential: AzureClientSecretCredential) -> AsyncAzureVaultClient:
"""Create an Azure Vault client using client secret credentials."""
cred = ClientSecretCredential(
tenant_id=credential.tenant_id,
client_id=credential.client_id,
client_secret=credential.client_secret,
)
return RealAsyncAzureVaultClient(cred)
def create_storage_client(self, storage_account_name: str, storage_account_key: str) -> AsyncAzureStorageClient:
"""Create an Azure Storage client with the provided credentials."""
return RealAsyncAzureStorageClient(storage_account_name, storage_account_key)

11
skyvern/forge/sdk/services/credential/azure_credential_vault_service.py
View File

@@ -2,7 +2,6 @@ import uuid
from typing import Annotated, Literal, Union
import structlog
from azure.identity.aio import ClientSecretCredential
from pydantic import BaseModel, Field, TypeAdapter
from skyvern.forge import app
@@ -41,14 +40,8 @@ class AzureCredentialVaultService(CredentialVaultService):
Union[_PasswordCredentialDataImage, _CreditCardCredentialDataImage], Field(discriminator="type")
]
def __init__(self, tenant_id: str, client_id: str, client_secret: str, vault_name: str):
self._client = AsyncAzureVaultClient(
ClientSecretCredential(
tenant_id=tenant_id,
client_id=client_id,
client_secret=client_secret,
)
)
def __init__(self, client: AsyncAzureVaultClient, vault_name: str):
self._client = client
self._vault_name = vault_name
async def create_credential(self, organization_id: str, data: CreateCredentialRequest) -> Credential:

7
skyvern/forge/sdk/workflow/context_manager.py
View File

@@ -387,7 +387,8 @@ class WorkflowRunContext:
# If the parameter is an Azure secret, fetch the secret value and store it in the secrets dict
# The value of the parameter will be the random secret id with format `secret_<uuid>`.
# We'll replace the random secret id with the actual secret value when we need to use it.
async with AsyncAzureVaultClient.create_default() as azure_vault_client:
azure_vault_client = app.AZURE_CLIENT_FACTORY.create_default()
async with azure_vault_client:
secret_value = await azure_vault_client.get_secret(parameter.azure_key, vault_name)
if secret_value is not None:
random_secret_id = self.generate_random_secret_id()
@@ -989,10 +990,10 @@ class WorkflowRunContext:
organization.organization_id, OrganizationAuthTokenType.azure_client_secret_credential.value
)
if org_auth_token:
azure_vault_client = AsyncAzureVaultClient.create_from_client_secret(org_auth_token.credential)
azure_vault_client = app.AZURE_CLIENT_FACTORY.create_from_client_secret(org_auth_token.credential)
else:
# Use the DefaultAzureCredential if not configured on organization level
azure_vault_client = AsyncAzureVaultClient.create_default()
azure_vault_client = app.AZURE_CLIENT_FACTORY.create_default()
return azure_vault_client
def _add_secret_parameter_value(self, parameter: Parameter, key: str, value: str) -> None:

3
skyvern/forge/sdk/workflow/models/block.py
View File

@@ -49,7 +49,6 @@ from skyvern.forge import app
from skyvern.forge.prompts import prompt_engine
from skyvern.forge.sdk.api import email
from skyvern.forge.sdk.api.aws import AsyncAWSClient
from skyvern.forge.sdk.api.azure import AsyncAzureStorageClient
from skyvern.forge.sdk.api.files import (
calculate_sha256_for_file,
create_named_temporary_file,
@@ -2322,7 +2321,7 @@ class FileUploadBlock(Block):
if actual_azure_storage_account_name is None or actual_azure_storage_account_key is None:
raise AzureConfigurationError("Azure Storage is not configured")
azure_client = AsyncAzureStorageClient(
azure_client = app.AZURE_CLIENT_FACTORY.create_storage_client(
storage_account_name=actual_azure_storage_account_name,
storage_account_key=actual_azure_storage_account_key,
)