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

introduce db timeout (#1222)

Browse Source
This commit is contained in:
Shuchang Zheng
2024-11-19 17:18:25 -08:00
committed by GitHub
parent 9896a70c61
commit e4c1cf65af
2 changed files with 13 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
skyvern/config.py
View File

@@ -26,6 +26,7 @@ class Settings(BaseSettings):
MAX_RETRIES_PER_STEP: int = 5 MAX_RETRIES_PER_STEP: int = 5
DEBUG_MODE: bool = False DEBUG_MODE: bool = False
DATABASE_STRING: str = "postgresql+psycopg://skyvern@localhost/skyvern" DATABASE_STRING: str = "postgresql+psycopg://skyvern@localhost/skyvern"
DATABASE_STATEMENT_TIMEOUT_MS: int = 60000
PROMPT_ACTION_HISTORY_WINDOW: int = 1 PROMPT_ACTION_HISTORY_WINDOW: int = 1
TASK_RESPONSE_ACTION_SCREENSHOT_COUNT: int = 3 TASK_RESPONSE_ACTION_SCREENSHOT_COUNT: int = 3

13
skyvern/forge/sdk/db/client.py
View File

@@ -74,12 +74,23 @@ from skyvern.webeye.actions.models import AgentStepOutput
LOG = structlog.get_logger() LOG = structlog.get_logger()
DB_CONNECT_ARGS: dict[str, Any] = {}
if "postgresql+psycopg" in settings.DATABASE_STRING:
DB_CONNECT_ARGS = {"options": f"-c statement_timeout={settings.DATABASE_STATEMENT_TIMEOUT_MS}"}
elif "postgresql+asyncpg" in settings.DATABASE_STRING:
DB_CONNECT_ARGS = {"server_settings": {"statement_timeout": str(settings.DATABASE_STATEMENT_TIMEOUT_MS)}}
class AgentDB: class AgentDB:
def __init__(self, database_string: str, debug_enabled: bool = False) -> None: def __init__(self, database_string: str, debug_enabled: bool = False) -> None:
super().__init__() super().__init__()
self.debug_enabled = debug_enabled self.debug_enabled = debug_enabled
self.engine = create_async_engine(database_string, json_serializer=_custom_json_serializer, pool_pre_ping=True) self.engine = create_async_engine(
database_string,
json_serializer=_custom_json_serializer,
connect_args=DB_CONNECT_ARGS,
)
self.Session = async_sessionmaker(bind=self.engine) self.Session = async_sessionmaker(bind=self.engine)
async def create_task( async def create_task(