community[patch]: Update YandexGPT API (#14773)

Update LLMand Chat model to use new api version --------- Co-authored-by: Dmitry Tyumentsev <[email protected]>
langchain-ai · Dec 16, 2023 · dcead81 · dcead81
1 parent eca89f8
commit dcead81
Show file tree

Hide file tree

Showing 4 changed files with 185 additions and 69 deletions.
diff --git a/docs/docs/integrations/chat/yandex.ipynb b/docs/docs/integrations/chat/yandex.ipynb
@@ -42,13 +42,20 @@
     "Next, you have two authentication options:\n",
     "- [IAM token](https://cloud.yandex.com/en/docs/iam/operations/iam-token/create-for-sa).\n",
     "    You can specify the token in a constructor parameter `iam_token` or in an environment variable `YC_IAM_TOKEN`.\n",
+    "\n",
     "- [API key](https://cloud.yandex.com/en/docs/iam/operations/api-key/create)\n",
-    "    You can specify the key in a constructor parameter `api_key` or in an environment variable `YC_API_KEY`."
+    "    You can specify the key in a constructor parameter `api_key` or in an environment variable `YC_API_KEY`.\n",
+    "\n",
+    "In the `model_uri` parameter, specify the model used, see [the documentation](https://cloud.yandex.com/en/docs/yandexgpt/concepts/models#yandexgpt-generation) for more details.\n",
+    "\n",
+    "To specify the model you can use `model_uri` parameter, see [the documentation](https://cloud.yandex.com/en/docs/yandexgpt/concepts/models#yandexgpt-generation) for more details.\n",
+    "\n",
+    "By default, the latest version of `yandexgpt-lite` is used from the folder specified in the parameter `folder_id` or `YC_FOLDER_ID` environment variable."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 1,
    "id": "eba2d63b-f871-4f61-b55f-f6092bdc297a",
    "metadata": {},
    "outputs": [],
@@ -59,7 +66,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 2,
    "id": "75905d9a-dfae-43aa-95b9-a160280e43f7",
    "metadata": {},
    "outputs": [],
@@ -69,17 +76,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 3,
    "id": "40844fe7-7fe5-4679-b6c9-1b3238807bdc",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "AIMessage(content=\"Je t'aime programmer.\")"
+       "AIMessage(content='Je adore le programmement.')"
       ]
      },
-     "execution_count": 8,
+     "execution_count": 3,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -113,7 +120,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.18"
+   "version": "3.10.13"
   }
  },
  "nbformat": 4,

diff --git a/docs/docs/integrations/llms/yandex.ipynb b/docs/docs/integrations/llms/yandex.ipynb
@@ -29,13 +29,20 @@
     "Next, you have two authentication options:\n",
     "- [IAM token](https://cloud.yandex.com/en/docs/iam/operations/iam-token/create-for-sa).\n",
     "    You can specify the token in a constructor parameter `iam_token` or in an environment variable `YC_IAM_TOKEN`.\n",
+    "\n",
     "- [API key](https://cloud.yandex.com/en/docs/iam/operations/api-key/create)\n",
-    "    You can specify the key in a constructor parameter `api_key` or in an environment variable `YC_API_KEY`."
+    "    You can specify the key in a constructor parameter `api_key` or in an environment variable `YC_API_KEY`.\n",
+    "\n",
+    "In the `model_uri` parameter, specify the model used, see [the documentation](https://cloud.yandex.com/en/docs/yandexgpt/concepts/models#yandexgpt-generation) for more details.\n",
+    "\n",
+    "To specify the model you can use `model_uri` parameter, see [the documentation](https://cloud.yandex.com/en/docs/yandexgpt/concepts/models#yandexgpt-generation) for more details.\n",
+    "\n",
+    "By default, the latest version of `yandexgpt-lite` is used from the folder specified in the parameter `folder_id` or `YC_FOLDER_ID` environment variable."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 246,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -46,7 +53,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 247,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -56,7 +63,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 248,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -65,7 +72,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 249,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -74,16 +81,16 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 250,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "'Moscow'"
+       "'The capital of Russia is Moscow.'"
       ]
      },
-     "execution_count": 250,
+     "execution_count": 5,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -111,7 +118,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.18"
+   "version": "3.10.13"
   }
  },
  "nbformat": 4,

diff --git a/libs/community/langchain_community/chat_models/yandex.py b/libs/community/langchain_community/chat_models/yandex.py
@@ -1,6 +1,6 @@
 """Wrapper around YandexGPT chat models."""
 import logging
-from typing import Any, Dict, List, Optional, Tuple, cast
+from typing import Any, Dict, List, Optional, cast
 
 from langchain_core.callbacks import (
     AsyncCallbackManagerForLLMRun,
@@ -25,23 +25,22 @@ def _parse_message(role: str, text: str) -> Dict:
     return {"role": role, "text": text}
 
 
-def _parse_chat_history(history: List[BaseMessage]) -> Tuple[List[Dict[str, str]], str]:
+def _parse_chat_history(history: List[BaseMessage]) -> List[Dict[str, str]]:
     """Parse a sequence of messages into history.
 
     Returns:
-        A tuple of a list of parsed messages and an instruction message for the model.
+        A list of parsed messages.
     """
     chat_history = []
-    instruction = ""
     for message in history:
         content = cast(str, message.content)
         if isinstance(message, HumanMessage):
             chat_history.append(_parse_message("user", content))
         if isinstance(message, AIMessage):
             chat_history.append(_parse_message("assistant", content))
         if isinstance(message, SystemMessage):
-            instruction = content
-    return chat_history, instruction
+            chat_history.append(_parse_message("system", content))
+    return chat_history
 
 
 class ChatYandexGPT(_BaseYandexGPT, BaseChatModel):
@@ -84,9 +83,14 @@ def _generate(
         try:
             import grpc
             from google.protobuf.wrappers_pb2 import DoubleValue, Int64Value
-            from yandex.cloud.ai.llm.v1alpha.llm_pb2 import GenerationOptions, Message
-            from yandex.cloud.ai.llm.v1alpha.llm_service_pb2 import ChatRequest
-            from yandex.cloud.ai.llm.v1alpha.llm_service_pb2_grpc import (
+            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:
@@ -97,25 +101,20 @@ def _generate(
             raise ValueError(
                 "You should provide at least one message to start the chat!"
             )
-        message_history, instruction = _parse_chat_history(messages)
+        message_history = _parse_chat_history(messages)
         channel_credentials = grpc.ssl_channel_credentials()
         channel = grpc.secure_channel(self.url, channel_credentials)
-        request = ChatRequest(
-            model=self.model_name,
-            generation_options=GenerationOptions(
+        request = CompletionRequest(
+            model_uri=self.model_uri,
+            completion_options=CompletionOptions(
                 temperature=DoubleValue(value=self.temperature),
                 max_tokens=Int64Value(value=self.max_tokens),
             ),
-            instruction_text=instruction,
             messages=[Message(**message) for message in message_history],
         )
         stub = TextGenerationServiceStub(channel)
-        if self.iam_token:
-            metadata = (("authorization", f"Bearer {self.iam_token}"),)
-        else:
-            metadata = (("authorization", f"Api-Key {self.api_key}"),)
-        res = stub.Chat(request, metadata=metadata)
-        text = list(res)[0].message.text
+        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)])
@@ -127,6 +126,75 @@ async def _agenerate(
         run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
         **kwargs: Any,
     ) -> ChatResult:
-        raise NotImplementedError(
-            """YandexGPT doesn't support async requests at the moment."""
-        )
+        """Async method to generate next turn in the conversation.
+
+        Args:
+            messages: The history of the conversation as a list of messages.
+            stop: The list of stop words (optional).
+            run_manager: The CallbackManager for LLM run, it's not used at the moment.
+
+        Returns:
+            The ChatResult that contains outputs generated by the model.
+
+        Raises:
+            ValueError: if the last message in the list is not from human.
+        """
+        try:
+            import asyncio
+
+            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,
+                CompletionResponse,
+            )
+            from yandex.cloud.ai.foundation_models.v1.foundation_models_service_pb2_grpc import (  # noqa: E501
+                TextGenerationAsyncServiceStub,
+            )
+            from yandex.cloud.operation.operation_service_pb2 import GetOperationRequest
+            from yandex.cloud.operation.operation_service_pb2_grpc import (
+                OperationServiceStub,
+            )
+        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)
+        operation_api_url = "operation.api.cloud.yandex.net:443"
+        channel_credentials = grpc.ssl_channel_credentials()
+        async with grpc.aio.secure_channel(self.url, channel_credentials) as channel:
+            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 = TextGenerationAsyncServiceStub(channel)
+            operation = await stub.Completion(request, metadata=self._grpc_metadata)
+            async with grpc.aio.secure_channel(
+                operation_api_url, channel_credentials
+            ) as operation_channel:
+                operation_stub = OperationServiceStub(operation_channel)
+                while not operation.done:
+                    await asyncio.sleep(1)
+                    operation_request = GetOperationRequest(operation_id=operation.id)
+                    operation = await operation_stub.Get(
+                        operation_request, metadata=self._grpc_metadata
+                    )
+
+            instruct_response = CompletionResponse()
+            operation.response.Unpack(instruct_response)
+            text = instruct_response.alternatives[0].message.text
+            if stop is not None:
+                text = enforce_stop_tokens(text, stop)
+            return text