Pass base parameters to acompletion (#343)

Skyvern-AI · May 17, 2024 · 6908ca5 · 6908ca5
1 parent 9c19442
commit 6908ca5
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/skyvern/forge/sdk/api/llm/api_handler_factory.py b/skyvern/forge/sdk/api/llm/api_handler_factory.py
@@ -133,7 +133,7 @@ async def llm_api_handler_with_router_and_fallback(
         return llm_api_handler_with_router_and_fallback
 
     @staticmethod
-    def get_llm_api_handler(llm_key: str) -> LLMAPIHandler:
+    def get_llm_api_handler(llm_key: str, base_parameters: dict[str, Any] | None = None) -> LLMAPIHandler:
         llm_config = LLMConfigRegistry.get_config(llm_key)
 
         if LLMConfigRegistry.is_router_config(llm_key):
@@ -145,9 +145,12 @@ async def llm_api_handler(
             screenshots: list[bytes] | None = None,
             parameters: dict[str, Any] | None = None,
         ) -> dict[str, Any]:
+            active_parameters = base_parameters or {}
             if parameters is None:
                 parameters = LLMAPIHandlerFactory.get_api_parameters()
 
+            active_parameters.update(parameters)
+
             if step:
                 await app.ARTIFACT_MANAGER.create_artifact(
                     step=step,
@@ -174,6 +177,7 @@ async def llm_api_handler(
                         {
                             "model": llm_config.model_name,
                             "messages": messages,
+                            # we're not using active_parameters here because it may contain sensitive information
                             **parameters,
                         }
                     ).encode("utf-8"),
@@ -185,7 +189,7 @@ async def llm_api_handler(
                 response = await litellm.acompletion(
                     model=llm_config.model_name,
                     messages=messages,
-                    **parameters,
+                    **active_parameters,
                 )
             except openai.OpenAIError as e:
                 raise LLMProviderError(llm_key) from e

diff --git a/skyvern/forge/sdk/experimentation/providers.py b/skyvern/forge/sdk/experimentation/providers.py
@@ -17,7 +17,11 @@ def is_feature_enabled_cached(self, feature_name: str, distinct_id: str, propert
         if feature_name not in self.result_map:
             self.result_map[feature_name] = {}
         if distinct_id not in self.result_map[feature_name]:
-            self.result_map[feature_name][distinct_id] = self.is_feature_enabled(feature_name, distinct_id, properties)
+            feature_flag_value = self.is_feature_enabled(feature_name, distinct_id, properties)
+            self.result_map[feature_name][distinct_id] = feature_flag_value
+            if feature_flag_value:
+                LOG.info("Feature flag is enabled", flag=feature_name, distinct_id=distinct_id)
+
         return self.result_map[feature_name][distinct_id]