langchain/libs/community/tests/integration_tests/chat_models/test_sparkllm.py
Xudong Sun 019b6ebe8d
community[minor]: Add iFlyTek Spark LLM chat model support (#13389)
- **Description:** This PR enables LangChain to access the iFlyTek's
Spark LLM via the chat_models wrapper.
  - **Dependencies:** websocket-client ^1.6.1
  - **Tag maintainer:** @baskaryan 

### SparkLLM chat model usage

Get SparkLLM's app_id, api_key and api_secret from [iFlyTek SparkLLM API
Console](https://console.xfyun.cn/services/bm3) (for more info, see
[iFlyTek SparkLLM Intro](https://xinghuo.xfyun.cn/sparkapi) ), then set
environment variables `IFLYTEK_SPARK_APP_ID`, `IFLYTEK_SPARK_API_KEY`
and `IFLYTEK_SPARK_API_SECRET` or pass parameters when using it like the
demo below:

```python3
from langchain.chat_models.sparkllm import ChatSparkLLM

client = ChatSparkLLM(
    spark_app_id="<app_id>",
    spark_api_key="<api_key>",
    spark_api_secret="<api_secret>"
)
```
2024-01-23 19:23:46 -08:00

37 lines
1.1 KiB
Python

from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage
from langchain_community.chat_models.sparkllm import ChatSparkLLM
def test_chat_spark_llm() -> None:
chat = ChatSparkLLM()
message = HumanMessage(content="Hello")
response = chat([message])
assert isinstance(response, AIMessage)
assert isinstance(response.content, str)
def test_chat_spark_llm_streaming() -> None:
chat = ChatSparkLLM(streaming=True)
for chunk in chat.stream("Hello!"):
assert isinstance(chunk, AIMessageChunk)
assert isinstance(chunk.content, str)
def test_chat_spark_llm_with_domain() -> None:
chat = ChatSparkLLM(spark_llm_domain="generalv3")
message = HumanMessage(content="Hello")
response = chat([message])
print(response)
assert isinstance(response, AIMessage)
assert isinstance(response.content, str)
def test_chat_spark_llm_with_temperature() -> None:
chat = ChatSparkLLM(temperature=0.9, top_k=2)
message = HumanMessage(content="Hello")
response = chat([message])
print(response)
assert isinstance(response, AIMessage)
assert isinstance(response.content, str)