langchain/libs/community/langchain_community/chat_models/moonshot.py
Leonid Ganeline 4cb5f4c353
community[patch]: import flattening fix (#20110)
This PR should make it easier for linters to do type checking and for IDEs to jump to definition of code.

See #20050 as a template for this PR.
- As a byproduct: Added 3 missed `test_imports`.
- Added missed `SolarChat` in to __init___.py Added it into test_import
ut.
- Added `# type: ignore` to fix linting. It is not clear, why linting
errors appear after ^ changes.

---------

Co-authored-by: Eugene Yurtsev <eyurtsev@gmail.com>
2024-04-10 13:01:19 -04:00

59 lines
1.9 KiB
Python

"""Wrapper around Moonshot chat models."""
from typing import Dict
from langchain_core.pydantic_v1 import root_validator
from langchain_core.utils import get_from_dict_or_env
from langchain_community.chat_models import ChatOpenAI
from langchain_community.llms.moonshot import MOONSHOT_SERVICE_URL_BASE, MoonshotCommon
class MoonshotChat(MoonshotCommon, ChatOpenAI): # type: ignore[misc]
"""Wrapper around Moonshot large language models.
To use, you should have the ``openai`` python package installed, and the
environment variable ``MOONSHOT_API_KEY`` set with your API key.
(Moonshot's chat API is compatible with OpenAI's SDK.)
Referenced from https://platform.moonshot.cn/docs
Example:
.. code-block:: python
from langchain_community.chat_models.moonshot import MoonshotChat
moonshot = MoonshotChat(model="moonshot-v1-8k")
"""
@root_validator()
def validate_environment(cls, values: Dict) -> Dict:
"""Validate that the environment is set up correctly."""
values["moonshot_api_key"] = get_from_dict_or_env(
values, "moonshot_api_key", "MOONSHOT_API_KEY"
)
try:
import openai
except ImportError:
raise ImportError(
"Could not import openai python package. "
"Please install it with `pip install openai`."
)
client_params = {
"api_key": values["moonshot_api_key"],
"base_url": values["base_url"]
if "base_url" in values
else MOONSHOT_SERVICE_URL_BASE,
}
if not values.get("client"):
values["client"] = openai.OpenAI(**client_params).chat.completions
if not values.get("async_client"):
values["async_client"] = openai.AsyncOpenAI(
**client_params
).chat.completions
return values