mirror of
https://github.com/hwchase17/langchain.git
synced 2025-08-12 22:28:03 +00:00
Default On Retry (#9230)
Base callbacks don't have a default on retry event Fix #8542 --------- Co-authored-by: landonsilla <landon.silla@stepstone.com>
This commit is contained in:
parent
d0a0d560ad
commit
c478fc208e
@ -4,6 +4,8 @@ from __future__ import annotations
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Union
|
||||
from uuid import UUID
|
||||
|
||||
from tenacity import RetryCallState
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langchain.schema.agent import AgentAction, AgentFinish
|
||||
from langchain.schema.document import Document
|
||||
@ -222,6 +224,16 @@ class RunManagerMixin:
|
||||
) -> Any:
|
||||
"""Run on arbitrary text."""
|
||||
|
||||
def on_retry(
|
||||
self,
|
||||
retry_state: RetryCallState,
|
||||
*,
|
||||
run_id: UUID,
|
||||
parent_run_id: Optional[UUID] = None,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
"""Run on a retry event."""
|
||||
|
||||
|
||||
class BaseCallbackHandler(
|
||||
LLMManagerMixin,
|
||||
@ -414,6 +426,16 @@ class AsyncCallbackHandler(BaseCallbackHandler):
|
||||
) -> None:
|
||||
"""Run on arbitrary text."""
|
||||
|
||||
async def on_retry(
|
||||
self,
|
||||
retry_state: RetryCallState,
|
||||
*,
|
||||
run_id: UUID,
|
||||
parent_run_id: Optional[UUID] = None,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
"""Run on a retry event."""
|
||||
|
||||
async def on_agent_action(
|
||||
self,
|
||||
action: AgentAction,
|
||||
|
@ -1,3 +1,6 @@
|
||||
from unittest.mock import MagicMock
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
|
||||
from langchain.callbacks import OpenAICallbackHandler
|
||||
@ -124,3 +127,7 @@ def test_on_llm_end_no_cost_invalid_model(
|
||||
)
|
||||
handler.on_llm_end(response)
|
||||
assert handler.total_cost == 0
|
||||
|
||||
|
||||
def test_on_retry_works(handler: OpenAICallbackHandler) -> None:
|
||||
handler.on_retry(MagicMock(), run_id=uuid4())
|
||||
|
Loading…
Reference in New Issue
Block a user