mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-13 13:17:52 +02:00
103 lines
3.4 KiB
Python
103 lines
3.4 KiB
Python
import asyncio
|
|
import logging
|
|
import random
|
|
import time
|
|
from typing import Optional
|
|
|
|
from langgraph.pregel.types import PregelExecutableTask, RetryPolicy
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def run_with_retry(
|
|
task: PregelExecutableTask,
|
|
retry_policy: Optional[RetryPolicy],
|
|
) -> None:
|
|
"""Run a task with retries."""
|
|
retry_policy = task.retry_policy or retry_policy
|
|
interval = retry_policy.initial_interval if retry_policy else 0
|
|
attempts = 0
|
|
while True:
|
|
try:
|
|
# clear any writes from previous attempts
|
|
task.writes.clear()
|
|
# run the task
|
|
task.proc.invoke(task.input, task.config)
|
|
# if successful, end
|
|
break
|
|
except Exception as exc:
|
|
if retry_policy is None:
|
|
raise
|
|
# increment attempts
|
|
attempts += 1
|
|
# check if we should retry
|
|
if callable(retry_policy.retry_on):
|
|
if not retry_policy.retry_on(exc):
|
|
raise
|
|
elif not isinstance(exc, retry_policy.retry_on):
|
|
raise
|
|
# check if we should give up
|
|
if attempts >= retry_policy.max_attempts:
|
|
raise
|
|
# sleep before retrying
|
|
interval = min(
|
|
retry_policy.max_interval,
|
|
interval * retry_policy.backoff_factor,
|
|
)
|
|
time.sleep(
|
|
interval + random.uniform(0, 1) if retry_policy.jitter else interval
|
|
)
|
|
# log the retry
|
|
logger.info(
|
|
f"Retrying task {task.name} after {interval:.2f} seconds (attempt {attempts}) after {exc.__class__.__name__} {exc}"
|
|
)
|
|
|
|
|
|
async def arun_with_retry(
|
|
task: PregelExecutableTask,
|
|
retry_policy: Optional[RetryPolicy],
|
|
stream: bool = False,
|
|
) -> None:
|
|
"""Run a task asynchronously with retries."""
|
|
retry_policy = task.retry_policy or retry_policy
|
|
interval = retry_policy.initial_interval if retry_policy else 0
|
|
attempts = 0
|
|
while True:
|
|
try:
|
|
# clear any writes from previous attempts
|
|
task.writes.clear()
|
|
# run the task
|
|
if stream:
|
|
async for _ in task.proc.astream(task.input, task.config):
|
|
pass
|
|
else:
|
|
await task.proc.ainvoke(task.input, task.config)
|
|
# if successful, end
|
|
break
|
|
except Exception as exc:
|
|
if retry_policy is None:
|
|
raise
|
|
# increment attempts
|
|
attempts += 1
|
|
# check if we should retry
|
|
if callable(retry_policy.retry_on):
|
|
if not retry_policy.retry_on(exc):
|
|
raise
|
|
elif not isinstance(exc, retry_policy.retry_on):
|
|
raise
|
|
# check if we should give up
|
|
if attempts >= retry_policy.max_attempts:
|
|
raise
|
|
# sleep before retrying
|
|
interval = min(
|
|
retry_policy.max_interval,
|
|
interval * retry_policy.backoff_factor,
|
|
)
|
|
await asyncio.sleep(
|
|
interval + random.uniform(0, 1) if retry_policy.jitter else interval
|
|
)
|
|
# log the retry
|
|
logger.info(
|
|
f"Retrying task {task.name} after {interval:.2f} seconds (attempt {attempts}) after {exc.__class__.__name__} {exc}"
|
|
)
|