Add Anthropic default request timeout (#3540)

thanks @hitflame!

---------

Co-authored-by: Wenqiang Zhao <hitzhaowenqiang@sina.com>
Co-authored-by: delta@com <delta@com>
This commit is contained in:
Davis Chase
2023-04-25 11:40:41 -07:00
committed by vowelparrot
parent b7f4a410a3
commit 6f514361be
2 changed files with 10 additions and 4 deletions

View File

@@ -1,6 +1,6 @@
"""Wrapper around Anthropic APIs."""
import re
from typing import Any, Callable, Dict, Generator, List, Mapping, Optional
from typing import Any, Callable, Dict, Generator, List, Mapping, Optional, Tuple, Union
from pydantic import BaseModel, Extra, root_validator
@@ -28,6 +28,9 @@ class _AnthropicCommon(BaseModel):
streaming: bool = False
"""Whether to stream the results."""
default_request_timeout: Optional[Union[float, Tuple[float, float]]] = None
"""Timeout for requests to Anthropic Completion API. Default is 600 seconds."""
anthropic_api_key: Optional[str] = None
HUMAN_PROMPT: Optional[str] = None
@@ -43,7 +46,10 @@ class _AnthropicCommon(BaseModel):
try:
import anthropic
values["client"] = anthropic.Client(anthropic_api_key)
values["client"] = anthropic.Client(
api_key=anthropic_api_key,
default_request_timeout=values["default_request_timeout"],
)
values["HUMAN_PROMPT"] = anthropic.HUMAN_PROMPT
values["AI_PROMPT"] = anthropic.AI_PROMPT
values["count_tokens"] = anthropic.count_tokens

View File

@@ -11,14 +11,14 @@ from tests.unit_tests.callbacks.fake_callback_handler import FakeCallbackHandler
def test_anthropic_call() -> None:
"""Test valid call to anthropic."""
llm = Anthropic(model="bare-nano-0")
llm = Anthropic(model="test")
output = llm("Say foo:")
assert isinstance(output, str)
def test_anthropic_streaming() -> None:
"""Test streaming tokens from anthropic."""
llm = Anthropic(model="bare-nano-0")
llm = Anthropic(model="test")
generator = llm.stream("I'm Pickle Rick")
assert isinstance(generator, Generator)