community[patch]: fix agenerate return value (#14815)

Fixed: - `_agenerate` return value in the YandexGPT Chat Model - duplicate line in the documentation Co-authored-by: Dmitry Tyumentsev <[email protected]>
langchain-ai · Dec 17, 2023 · 78ae276 · 78ae276
1 parent f1d3f29
commit 78ae276
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 13 deletions.
diff --git a/docs/docs/integrations/chat/yandex.ipynb b/docs/docs/integrations/chat/yandex.ipynb
@@ -46,8 +46,6 @@
     "- [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`.\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."

diff --git a/docs/docs/integrations/llms/yandex.ipynb b/docs/docs/integrations/llms/yandex.ipynb
@@ -33,8 +33,6 @@
     "- [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`.\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."

diff --git a/libs/community/langchain_community/chat_models/yandex.py b/libs/community/langchain_community/chat_models/yandex.py
@@ -192,9 +192,9 @@ async def _agenerate(
                         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
+            completion_response = CompletionResponse()
+            operation.response.Unpack(completion_response)
+            text = 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)])
diff --git a/libs/community/langchain_community/llms/yandex.py b/libs/community/langchain_community/llms/yandex.py
@@ -232,9 +232,9 @@ async def _acall(
                         operation_request, metadata=self._grpc_metadata
                     )
 
-            instruct_response = CompletionResponse()
-            operation.response.Unpack(instruct_response)
-            text = instruct_response.alternatives[0].message.text
+            completion_response = CompletionResponse()
+            operation.response.Unpack(completion_response)
+            text = completion_response.alternatives[0].message.text
             if stop is not None:
                 text = enforce_stop_tokens(text, stop)
             return text