community[patch]: Add retry logic to Yandex GPT API Calls (#14907)

**Description:** Added logic for re-calling the YandexGPT API in case of
an error

---------

Co-authored-by: Dmitry Tyumentsev <dmitry.tyumentsev@raftds.com>

libs/community/langchain_community/chat_models/yandex.py

@@ -1,6 +1,8 @@
 """Wrapper around YandexGPT chat models."""
+from __future__ import annotations
+
 import logging
-from typing import Any, Dict, List, Optional, cast
+from typing import Any, Callable, Dict, List, Optional, cast
 
 from langchain_core.callbacks import (
     AsyncCallbackManagerForLLMRun,
@@ -14,6 +16,13 @@ from langchain_core.messages import (
     SystemMessage,
 )
 from langchain_core.outputs import ChatGeneration, ChatResult
+from tenacity import (
+    before_sleep_log,
+    retry,
+    retry_if_exception_type,
+    stop_after_attempt,
+    wait_exponential,
+)
 
 from langchain_community.llms.utils import enforce_stop_tokens
 from langchain_community.llms.yandex import _BaseYandexGPT
@@ -80,41 +89,7 @@ class ChatYandexGPT(_BaseYandexGPT, BaseChatModel):
         Raises:
             ValueError: if the last message in the list is not from human.
         """
-        try:
+        text = completion_with_retry(self, messages=messages)
-            import grpc
-            from google.protobuf.wrappers_pb2 import DoubleValue, Int64Value
-            from yandex.cloud.ai.foundation_models.v1.foundation_models_pb2 import (
-                CompletionOptions,
-                Message,
-            )
-            from yandex.cloud.ai.foundation_models.v1.foundation_models_service_pb2 import (  # noqa: E501
-                CompletionRequest,
-            )
-            from yandex.cloud.ai.foundation_models.v1.foundation_models_service_pb2_grpc import (  # noqa: E501
-                TextGenerationServiceStub,
-            )
-        except ImportError as e:
-            raise ImportError(
-                "Please install YandexCloud SDK" " with `pip install yandexcloud`."
-            ) from e
-        if not messages:
-            raise ValueError(
-                "You should provide at least one message to start the chat!"
-            )
-        message_history = _parse_chat_history(messages)
-        channel_credentials = grpc.ssl_channel_credentials()
-        channel = grpc.secure_channel(self.url, channel_credentials)
-        request = CompletionRequest(
-            model_uri=self.model_uri,
-            completion_options=CompletionOptions(
-                temperature=DoubleValue(value=self.temperature),
-                max_tokens=Int64Value(value=self.max_tokens),
-            ),
-            messages=[Message(**message) for message in message_history],
-        )
-        stub = TextGenerationServiceStub(channel)
-        res = stub.Completion(request, metadata=self._grpc_metadata)
-        text = list(res)[0].alternatives[0].message.text
         text = text if stop is None else enforce_stop_tokens(text, stop)
         message = AIMessage(content=text)
         return ChatResult(generations=[ChatGeneration(message=message)])
@@ -139,6 +114,52 @@ class ChatYandexGPT(_BaseYandexGPT, BaseChatModel):
         Raises:
             ValueError: if the last message in the list is not from human.
         """
+        text = await acompletion_with_retry(self, messages=messages)
+        text = text if stop is None else enforce_stop_tokens(text, stop)
+        message = AIMessage(content=text)
+        return ChatResult(generations=[ChatGeneration(message=message)])
+
+
+def _make_request(
+    self: ChatYandexGPT,
+    messages: List[BaseMessage],
+) -> str:
+    try:
+        import grpc
+        from google.protobuf.wrappers_pb2 import DoubleValue, Int64Value
+        from yandex.cloud.ai.foundation_models.v1.foundation_models_pb2 import (
+            CompletionOptions,
+            Message,
+        )
+        from yandex.cloud.ai.foundation_models.v1.foundation_models_service_pb2 import (  # noqa: E501
+            CompletionRequest,
+        )
+        from yandex.cloud.ai.foundation_models.v1.foundation_models_service_pb2_grpc import (  # noqa: E501
+            TextGenerationServiceStub,
+        )
+    except ImportError as e:
+        raise ImportError(
+            "Please install YandexCloud SDK" " with `pip install yandexcloud`."
+        ) from e
+    if not messages:
+        raise ValueError("You should provide at least one message to start the chat!")
+    message_history = _parse_chat_history(messages)
+    channel_credentials = grpc.ssl_channel_credentials()
+    channel = grpc.secure_channel(self.url, channel_credentials)
+    request = CompletionRequest(
+        model_uri=self.model_uri,
+        completion_options=CompletionOptions(
+            temperature=DoubleValue(value=self.temperature),
+            max_tokens=Int64Value(value=self.max_tokens),
+        ),
+        messages=[Message(**message) for message in message_history],
+    )
+    stub = TextGenerationServiceStub(channel)
+    res = stub.Completion(request, metadata=self._grpc_metadata)
+    return list(res)[0].alternatives[0].message.text
+
+
+async def _amake_request(self: ChatYandexGPT, messages: List[BaseMessage]) -> str:
     try:
         import asyncio
 
@@ -164,9 +185,7 @@ class ChatYandexGPT(_BaseYandexGPT, BaseChatModel):
             "Please install YandexCloud SDK" " with `pip install yandexcloud`."
         ) from e
     if not messages:
-            raise ValueError(
+        raise ValueError("You should provide at least one message to start the chat!")
-                "You should provide at least one message to start the chat!"
-            )
     message_history = _parse_chat_history(messages)
     operation_api_url = "operation.api.cloud.yandex.net:443"
     channel_credentials = grpc.ssl_channel_credentials()
@@ -194,7 +213,40 @@ class ChatYandexGPT(_BaseYandexGPT, BaseChatModel):
 
         completion_response = CompletionResponse()
         operation.response.Unpack(completion_response)
-            text = completion_response.alternatives[0].message.text
+        return completion_response.alternatives[0].message.text
-            text = text if stop is None else enforce_stop_tokens(text, stop)
+
-            message = AIMessage(content=text)
+
-            return ChatResult(generations=[ChatGeneration(message=message)])
+def _create_retry_decorator(llm: ChatYandexGPT) -> Callable[[Any], Any]:
+    from grpc import RpcError
+
+    min_seconds = 1
+    max_seconds = 60
+    return retry(
+        reraise=True,
+        stop=stop_after_attempt(llm.max_retries),
+        wait=wait_exponential(multiplier=1, min=min_seconds, max=max_seconds),
+        retry=(retry_if_exception_type((RpcError))),
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+
+
+def completion_with_retry(llm: ChatYandexGPT, **kwargs: Any) -> Any:
+    """Use tenacity to retry the completion call."""
+    retry_decorator = _create_retry_decorator(llm)
+
+    @retry_decorator
+    def _completion_with_retry(**_kwargs: Any) -> Any:
+        return _make_request(llm, **_kwargs)
+
+    return _completion_with_retry(**kwargs)
+
+
+async def acompletion_with_retry(llm: ChatYandexGPT, **kwargs: Any) -> Any:
+    """Use tenacity to retry the async completion call."""
+    retry_decorator = _create_retry_decorator(llm)
+
+    @retry_decorator
+    async def _completion_with_retry(**_kwargs: Any) -> Any:
+        return await _amake_request(llm, **_kwargs)
+
+    return await _completion_with_retry(**kwargs)

libs/community/langchain_community/llms/yandex.py

@@ -1,4 +1,7 @@
-from typing import Any, Dict, List, Mapping, Optional
+from __future__ import annotations
+
+import logging
+from typing import Any, Callable, Dict, List, Mapping, Optional
 
 from langchain_core.callbacks import (
     AsyncCallbackManagerForLLMRun,
@@ -8,9 +11,18 @@ from langchain_core.language_models.llms import LLM
 from langchain_core.load.serializable import Serializable
 from langchain_core.pydantic_v1 import root_validator
 from langchain_core.utils import get_from_dict_or_env
+from tenacity import (
+    before_sleep_log,
+    retry,
+    retry_if_exception_type,
+    stop_after_attempt,
+    wait_exponential,
+)
 
 from langchain_community.llms.utils import enforce_stop_tokens
 
+logger = logging.getLogger(__name__)
+
 
 class _BaseYandexGPT(Serializable):
     iam_token: str = ""
@@ -38,11 +50,24 @@ class _BaseYandexGPT(Serializable):
     """Sequences when completion generation will stop."""
     url: str = "llm.api.cloud.yandex.net:443"
     """The url of the API."""
+    max_retries: int = 6
+    """Maximum number of retries to make when generating."""
 
     @property
     def _llm_type(self) -> str:
         return "yandex_gpt"
 
+    @property
+    def _identifying_params(self) -> Mapping[str, Any]:
+        """Get the identifying parameters."""
+        return {
+            "model_uri": self.model_uri,
+            "temperature": self.temperature,
+            "max_tokens": self.max_tokens,
+            "stop": self.stop,
+            "max_retries": self.max_retries,
+        }
+
     @root_validator()
     def validate_environment(cls, values: Dict) -> Dict:
         """Validate that iam token exists in environment."""
@@ -99,16 +124,6 @@ class YandexGPT(_BaseYandexGPT, LLM):
             yandex_gpt = YandexGPT(iam_token="t1.9eu...", folder_id="b1g...")
     """
 
-    @property
-    def _identifying_params(self) -> Mapping[str, Any]:
-        """Get the identifying parameters."""
-        return {
-            "model_uri": self.model_uri,
-            "temperature": self.temperature,
-            "max_tokens": self.max_tokens,
-            "stop": self.stop,
-        }
-
     def _call(
         self,
         prompt: str,
@@ -130,6 +145,37 @@ class YandexGPT(_BaseYandexGPT, LLM):
 
                 response = YandexGPT("Tell me a joke.")
         """
+        text = completion_with_retry(self, prompt=prompt)
+        if stop is not None:
+            text = enforce_stop_tokens(text, stop)
+        return text
+
+    async def _acall(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> str:
+        """Async call the Yandex GPT model and return the output.
+
+        Args:
+            prompt: The prompt to pass into the model.
+            stop: Optional list of stop words to use when generating.
+
+        Returns:
+            The string generated by the model.
+        """
+        text = await acompletion_with_retry(self, prompt=prompt)
+        if stop is not None:
+            text = enforce_stop_tokens(text, stop)
+        return text
+
+
+def _make_request(
+    self: YandexGPT,
+    prompt: str,
+) -> str:
     try:
         import grpc
         from google.protobuf.wrappers_pb2 import DoubleValue, Int64Value
@@ -158,32 +204,11 @@ class YandexGPT(_BaseYandexGPT, LLM):
         messages=[Message(role="user", text=prompt)],
     )
     stub = TextGenerationServiceStub(channel)
-        if self.iam_token:
+    res = stub.Completion(request, metadata=self._grpc_metadata)
-            metadata = (("authorization", f"Bearer {self.iam_token}"),)
+    return list(res)[0].alternatives[0].message.text
-        else:
-            metadata = (("authorization", f"Api-Key {self.api_key}"),)
-        res = stub.Completion(request, metadata=metadata)
-        text = list(res)[0].alternatives[0].message.text
-        if stop is not None:
-            text = enforce_stop_tokens(text, stop)
-        return text
 
-    async def _acall(
-        self,
-        prompt: str,
-        stop: Optional[List[str]] = None,
-        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
-        **kwargs: Any,
-    ) -> str:
-        """Async call the Yandex GPT model and return the output.
 
-        Args:
+async def _amake_request(self: YandexGPT, prompt: str) -> str:
-            prompt: The prompt to pass into the model.
-            stop: Optional list of stop words to use when generating.
-
-        Returns:
-            The string generated by the model.
-        """
     try:
         import asyncio
 
@@ -234,7 +259,40 @@ class YandexGPT(_BaseYandexGPT, LLM):
 
         completion_response = CompletionResponse()
         operation.response.Unpack(completion_response)
-            text = completion_response.alternatives[0].message.text
+        return completion_response.alternatives[0].message.text
-            if stop is not None:
+
-                text = enforce_stop_tokens(text, stop)
+
-            return text
+def _create_retry_decorator(llm: YandexGPT) -> Callable[[Any], Any]:
+    from grpc import RpcError
+
+    min_seconds = 1
+    max_seconds = 60
+    return retry(
+        reraise=True,
+        stop=stop_after_attempt(llm.max_retries),
+        wait=wait_exponential(multiplier=1, min=min_seconds, max=max_seconds),
+        retry=(retry_if_exception_type((RpcError))),
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+    )
+
+
+def completion_with_retry(llm: YandexGPT, **kwargs: Any) -> Any:
+    """Use tenacity to retry the completion call."""
+    retry_decorator = _create_retry_decorator(llm)
+
+    @retry_decorator
+    def _completion_with_retry(**_kwargs: Any) -> Any:
+        return _make_request(llm, **_kwargs)
+
+    return _completion_with_retry(**kwargs)
+
+
+async def acompletion_with_retry(llm: YandexGPT, **kwargs: Any) -> Any:
+    """Use tenacity to retry the async completion call."""
+    retry_decorator = _create_retry_decorator(llm)
+
+    @retry_decorator
+    async def _completion_with_retry(**_kwargs: Any) -> Any:
+        return await _amake_request(llm, **_kwargs)
+
+    return await _completion_with_retry(**kwargs)