langchain/libs/community/tests/integration_tests/llms/test_gradient_ai.py
Bagatur a0c2281540
infra: update mypy 1.10, ruff 0.5 (#23721)
```python
"""python scripts/update_mypy_ruff.py"""
import glob
import tomllib
from pathlib import Path

import toml
import subprocess
import re

ROOT_DIR = Path(__file__).parents[1]


def main():
    for path in glob.glob(str(ROOT_DIR / "libs/**/pyproject.toml"), recursive=True):
        print(path)
        with open(path, "rb") as f:
            pyproject = tomllib.load(f)
        try:
            pyproject["tool"]["poetry"]["group"]["typing"]["dependencies"]["mypy"] = (
                "^1.10"
            )
            pyproject["tool"]["poetry"]["group"]["lint"]["dependencies"]["ruff"] = (
                "^0.5"
            )
        except KeyError:
            continue
        with open(path, "w") as f:
            toml.dump(pyproject, f)
        cwd = "/".join(path.split("/")[:-1])
        completed = subprocess.run(
            "poetry lock --no-update; poetry install --with typing; poetry run mypy . --no-color",
            cwd=cwd,
            shell=True,
            capture_output=True,
            text=True,
        )
        logs = completed.stdout.split("\n")

        to_ignore = {}
        for l in logs:
            if re.match("^(.*)\:(\d+)\: error:.*\[(.*)\]", l):
                path, line_no, error_type = re.match(
                    "^(.*)\:(\d+)\: error:.*\[(.*)\]", l
                ).groups()
                if (path, line_no) in to_ignore:
                    to_ignore[(path, line_no)].append(error_type)
                else:
                    to_ignore[(path, line_no)] = [error_type]
        print(len(to_ignore))
        for (error_path, line_no), error_types in to_ignore.items():
            all_errors = ", ".join(error_types)
            full_path = f"{cwd}/{error_path}"
            try:
                with open(full_path, "r") as f:
                    file_lines = f.readlines()
            except FileNotFoundError:
                continue
            file_lines[int(line_no) - 1] = (
                file_lines[int(line_no) - 1][:-1] + f"  # type: ignore[{all_errors}]\n"
            )
            with open(full_path, "w") as f:
                f.write("".join(file_lines))

        subprocess.run(
            "poetry run ruff format .; poetry run ruff --select I --fix .",
            cwd=cwd,
            shell=True,
            capture_output=True,
            text=True,
        )


if __name__ == "__main__":
    main()

```
2024-07-03 10:33:27 -07:00

50 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)