From 100b7c9963c027946dbc5195c71e883c0d55a3c9 Mon Sep 17 00:00:00 2001 From: Bagatur Date: Wed, 6 Dec 2023 11:39:53 -0800 Subject: [PATCH] nit --- libs/langchain/langchain/llms/watsonxllm.py | 2 +- .../chat_models/test_volcengine_maas.py | 9 ++------- .../tests/integration_tests/llms/test_volcengine_maas.py | 2 +- 3 files changed, 4 insertions(+), 9 deletions(-) diff --git a/libs/langchain/langchain/llms/watsonxllm.py b/libs/langchain/langchain/llms/watsonxllm.py index 36ad7fdb17127..85fd2de4edee3 100644 --- a/libs/langchain/langchain/llms/watsonxllm.py +++ b/libs/langchain/langchain/llms/watsonxllm.py @@ -5,7 +5,7 @@ from langchain.callbacks.manager import CallbackManagerForLLMRun from langchain.llms.base import BaseLLM from langchain.pydantic_v1 import Extra, SecretStr, root_validator -from langchain.schema import LLMResult +from langchain_core.outputs import LLMResult from langchain.schema.output import Generation, GenerationChunk from langchain.utils import convert_to_secret_str, get_from_dict_or_env diff --git a/libs/langchain/tests/integration_tests/chat_models/test_volcengine_maas.py b/libs/langchain/tests/integration_tests/chat_models/test_volcengine_maas.py index 31800b9a4e70e..a9ba9644fff47 100644 --- a/libs/langchain/tests/integration_tests/chat_models/test_volcengine_maas.py +++ b/libs/langchain/tests/integration_tests/chat_models/test_volcengine_maas.py @@ -2,13 +2,8 @@ from langchain.callbacks.manager import CallbackManager from langchain.chat_models.volcengine_maas import VolcEngineMaasChat -from langchain.schema import ( - AIMessage, - BaseMessage, - ChatGeneration, - HumanMessage, - LLMResult, -) +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage +from langchain_core.outputs import ChatGeneration, LLMResult from tests.unit_tests.callbacks.fake_callback_handler import FakeCallbackHandler diff --git a/libs/langchain/tests/integration_tests/llms/test_volcengine_maas.py b/libs/langchain/tests/integration_tests/llms/test_volcengine_maas.py index 22a357034d466..f88e62d65ae67 100644 --- a/libs/langchain/tests/integration_tests/llms/test_volcengine_maas.py +++ b/libs/langchain/tests/integration_tests/llms/test_volcengine_maas.py @@ -3,7 +3,7 @@ from typing import Generator from langchain.llms.volcengine_maas import VolcEngineMaasLLM -from langchain.schema import LLMResult +from langchain_core.outputs import LLMResult def test_default_call() -> None: