From 59f07d373afa1969c0645ff0d25b4b48e5466714 Mon Sep 17 00:00:00 2001
From: ykeremy <keremyilmazcs@gmail.com>
Date: Thu, 2 May 2024 22:03:07 +0000
Subject: [PATCH] =?UTF-8?q?=F0=9F=94=84=20synced=20local=20'skyvern/'=20wi?=
 =?UTF-8?q?th=20remote=20'skyvern/'?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 skyvern/config.py                             |  1 +
 skyvern/forge/app.py                          |  2 ++
 skyvern/forge/sdk/api/llm/config_registry.py  | 24 ++++++++++++++++++-
 skyvern/forge/sdk/api/llm/models.py           |  2 +-
 skyvern/forge/sdk/experimentation/__init__.py |  0
 .../forge/sdk/experimentation/providers.py    | 16 +++++++++++++
 6 files changed, 43 insertions(+), 2 deletions(-)
 create mode 100644 skyvern/forge/sdk/experimentation/__init__.py
 create mode 100644 skyvern/forge/sdk/experimentation/providers.py

diff --git a/skyvern/config.py b/skyvern/config.py
index 3edcf0745..bec0fd798 100644
--- a/skyvern/config.py
+++ b/skyvern/config.py
@@ -69,6 +69,7 @@ class Settings(BaseSettings):
     ENABLE_OPENAI: bool = True
     ENABLE_ANTHROPIC: bool = False
     ENABLE_AZURE: bool = False
+    ENABLE_BEDROCK: bool = False
     # OPENAI
     OPENAI_API_KEY: str | None = None
     # ANTHROPIC
diff --git a/skyvern/forge/app.py b/skyvern/forge/app.py
index c0c006a41..4d5525ef0 100644
--- a/skyvern/forge/app.py
+++ b/skyvern/forge/app.py
@@ -11,6 +11,7 @@
 from skyvern.forge.sdk.artifact.manager import ArtifactManager
 from skyvern.forge.sdk.artifact.storage.factory import StorageFactory
 from skyvern.forge.sdk.db.client import AgentDB
+from skyvern.forge.sdk.experimentation.providers import BaseExperimentationProvider, NoOpExperimentationProvider
 from skyvern.forge.sdk.forge_log import setup_logger
 from skyvern.forge.sdk.models import Organization
 from skyvern.forge.sdk.schemas.tasks import Task
@@ -35,6 +36,7 @@
 STORAGE = StorageFactory.get_storage()
 ARTIFACT_MANAGER = ArtifactManager()
 BROWSER_MANAGER = BrowserManager()
+EXPERIMENTATION_PROVIDER: BaseExperimentationProvider = NoOpExperimentationProvider()
 LLM_API_HANDLER = LLMAPIHandlerFactory.get_llm_api_handler(SettingsManager.get_settings().LLM_KEY)
 WORKFLOW_CONTEXT_MANAGER = WorkflowContextManager()
 WORKFLOW_SERVICE = WorkflowService()
diff --git a/skyvern/forge/sdk/api/llm/config_registry.py b/skyvern/forge/sdk/api/llm/config_registry.py
index eb191f84d..24ffec428 100644
--- a/skyvern/forge/sdk/api/llm/config_registry.py
+++ b/skyvern/forge/sdk/api/llm/config_registry.py
@@ -60,7 +60,29 @@ def get_config(cls, llm_key: str) -> LLMRouterConfig | LLMConfig:
 
 if SettingsManager.get_settings().ENABLE_ANTHROPIC:
     LLMConfigRegistry.register_config(
-        "ANTHROPIC_CLAUDE3", LLMConfig("anthropic/claude-3-opus-20240229", ["ANTHROPIC_API_KEY"], True)
+        "ANTHROPIC_CLAUDE3_OPUS", LLMConfig("anthropic/claude-3-opus-20240229", ["ANTHROPIC_API_KEY"], True)
+    )
+    LLMConfigRegistry.register_config(
+        "ANTHROPIC_CLAUDE3_SONNET", LLMConfig("anthropic/claude-3-sonnet-20240229", ["ANTHROPIC_API_KEY"], True)
+    )
+
+if SettingsManager.get_settings().ENABLE_BEDROCK:
+    # Supported through AWS IAM authentication
+    LLMConfigRegistry.register_config(
+        "BEDROCK_ANTHROPIC_CLAUDE3_OPUS",
+        LLMConfig(
+            "bedrock/anthropic.claude-3-opus-20240229-v1:0",
+            ["AWS_REGION"],
+            True,
+        ),
+    )
+    LLMConfigRegistry.register_config(
+        "BEDROCK_ANTHROPIC_CLAUDE3_SONNET",
+        LLMConfig(
+            "bedrock/anthropic.claude-3-sonnet-20240229-v1:0",
+            ["AWS_REGION"],
+            True,
+        ),
     )
 
 if SettingsManager.get_settings().ENABLE_AZURE:
diff --git a/skyvern/forge/sdk/api/llm/models.py b/skyvern/forge/sdk/api/llm/models.py
index 2f65d81c2..41c8d1522 100644
--- a/skyvern/forge/sdk/api/llm/models.py
+++ b/skyvern/forge/sdk/api/llm/models.py
@@ -46,7 +46,7 @@ class LLMRouterConfig(LLMConfig):
         "usage-based-routing-v2",
         "latency-based-routing",
     ] = "usage-based-routing"
-    num_retries: int = 2
+    num_retries: int = 1
     retry_delay_seconds: int = 15
     set_verbose: bool = False
 
diff --git a/skyvern/forge/sdk/experimentation/__init__.py b/skyvern/forge/sdk/experimentation/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/skyvern/forge/sdk/experimentation/providers.py b/skyvern/forge/sdk/experimentation/providers.py
new file mode 100644
index 000000000..312576270
--- /dev/null
+++ b/skyvern/forge/sdk/experimentation/providers.py
@@ -0,0 +1,16 @@
+from abc import ABC, abstractmethod
+
+import structlog
+
+LOG = structlog.get_logger()
+
+
+class BaseExperimentationProvider(ABC):
+    @abstractmethod
+    def is_feature_enabled(self, feature_name: str, distinct_id: str) -> bool:
+        """Check if a specific feature is enabled."""
+
+
+class NoOpExperimentationProvider(BaseExperimentationProvider):
+    def is_feature_enabled(self, feature_name: str, distinct_id: str) -> bool:
+        return False