mirror of
https://github.com/hwchase17/langchain
synced 2024-11-10 01:10:59 +00:00
481d3855dc
- `llm(prompt)` -> `llm.invoke(prompt)`
- `llm(prompt=prompt` -> `llm.invoke(prompt)` (same with `messages=`)
- `llm(prompt, callbacks=callbacks)` -> `llm.invoke(prompt,
config={"callbacks": callbacks})`
- `llm(prompt, **kwargs)` -> `llm.invoke(prompt, **kwargs)`
49 lines
1.6 KiB
Python
49 lines
1.6 KiB
Python
"""Test GradientAI API wrapper.
|
|
|
|
In order to run this test, you need to have an GradientAI api key.
|
|
You can get it by registering for free at https://gradient.ai/.
|
|
|
|
You'll then need to set:
|
|
- `GRADIENT_ACCESS_TOKEN` environment variable to your api key.
|
|
- `GRADIENT_WORKSPACE_ID` environment variable to your workspace id.
|
|
- `GRADIENT_MODEL` environment variable to your workspace id.
|
|
"""
|
|
import os
|
|
|
|
from langchain_community.llms import GradientLLM
|
|
|
|
|
|
def test_gradient_acall() -> None:
|
|
"""Test simple call to gradient.ai."""
|
|
model = os.environ["GRADIENT_MODEL"]
|
|
gradient_access_token = os.environ["GRADIENT_ACCESS_TOKEN"]
|
|
gradient_workspace_id = os.environ["GRADIENT_WORKSPACE_ID"]
|
|
llm = GradientLLM(
|
|
model=model,
|
|
gradient_access_token=gradient_access_token,
|
|
gradient_workspace_id=gradient_workspace_id,
|
|
)
|
|
output = llm.invoke("Say hello:", temperature=0.2, max_tokens=250)
|
|
|
|
assert llm._llm_type == "gradient"
|
|
|
|
assert isinstance(output, str)
|
|
assert len(output)
|
|
|
|
|
|
async def test_gradientai_acall() -> None:
|
|
"""Test async call to gradient.ai."""
|
|
model = os.environ["GRADIENT_MODEL"]
|
|
gradient_access_token = os.environ["GRADIENT_ACCESS_TOKEN"]
|
|
gradient_workspace_id = os.environ["GRADIENT_WORKSPACE_ID"]
|
|
llm = GradientLLM(
|
|
model=model,
|
|
gradient_access_token=gradient_access_token,
|
|
gradient_workspace_id=gradient_workspace_id,
|
|
)
|
|
output = await llm.agenerate(["Say hello:"], temperature=0.2, max_tokens=250)
|
|
assert llm._llm_type == "gradient"
|
|
|
|
assert isinstance(output, str)
|
|
assert len(output)
|