overhual llm key override (#2677)

This commit is contained in:
Shuchang Zheng
2025-06-11 08:23:44 -07:00
committed by GitHub
parent 9a42c2ad9f
commit 9a29d966ab
5 changed files with 72 additions and 62 deletions

View File

@@ -60,7 +60,7 @@ from skyvern.forge.sdk.api.files import (
list_files_in_directory,
wait_for_download_finished,
)
from skyvern.forge.sdk.api.llm.api_handler_factory import LLMCallerManager
from skyvern.forge.sdk.api.llm.api_handler_factory import LLMAPIHandlerFactory, LLMCallerManager
from skyvern.forge.sdk.api.llm.exceptions import LLMProviderError
from skyvern.forge.sdk.core import skyvern_context
from skyvern.forge.sdk.core.aiohttp_helper import aiohttp_post
@@ -2557,7 +2557,8 @@ async def sequentially_select_from_dropdown(
select_history=json.dumps(build_sequential_select_history(select_history)),
local_datetime=datetime.now(ensure_context().tz_info).isoformat(),
)
json_response = await app.LLM_API_HANDLER(
llm_api_handler = LLMAPIHandlerFactory.get_override_llm_api_handler(task.llm_key, default=app.LLM_API_HANDLER)
json_response = await llm_api_handler(
prompt=prompt, screenshots=[screenshot], step=step, prompt_name="confirm-multi-selection-finish"
)
if json_response.get("is_mini_goal_finished", False):
@@ -2641,7 +2642,8 @@ async def select_from_emerging_elements(
task_id=task.task_id,
)
json_response = await app.LLM_API_HANDLER(prompt=prompt, step=step, prompt_name="custom-select")
llm_api_handler = LLMAPIHandlerFactory.get_override_llm_api_handler(task.llm_key, default=app.LLM_API_HANDLER)
json_response = await llm_api_handler(prompt=prompt, step=step, prompt_name="custom-select")
value: str | None = json_response.get("value", None)
LOG.info(
"LLM response for the matched element",
@@ -3385,12 +3387,12 @@ async def extract_information_for_navigation_goal(
# CUA tasks should use the default data extraction llm key
llm_key_override = None
json_response = await app.LLM_API_HANDLER(
llm_api_handler = LLMAPIHandlerFactory.get_override_llm_api_handler(llm_key_override, default=app.LLM_API_HANDLER)
json_response = await llm_api_handler(
prompt=extract_information_prompt,
step=step,
screenshots=scraped_page.screenshots,
prompt_name="extract-information",
llm_key_override=llm_key_override,
)
return ScrapeResult(