langchain/tests/integration_tests/llms/test_baseten.py
Philip Kiely - Baseten a09a0e3511
Baseten integration (#5862)
This PR adds a Baseten integration. I've done my best to follow the
contributor's guidelines and add docs, an example notebook, and an
integration test modeled after similar integrations' test.

Please let me know if there is anything I can do to improve the PR. When
it is merged, please tag https://twitter.com/basetenco and
https://twitter.com/philip_kiely as contributors (the note on the PR
template said to include Twitter accounts)
2023-06-08 23:05:57 -07:00

17 lines
389 B
Python

"""Test Baseten API wrapper."""
import os
import baseten
import pytest
from langchain.llms.baseten import Baseten
@pytest.mark.requires(baseten)
def test_baseten_call() -> None:
"""Test valid call to Baseten."""
baseten.login(os.environ["BASETEN_API_KEY"])
llm = Baseten(model=os.environ["BASETEN_MODEL_ID"])
output = llm("Say foo:")
assert isinstance(output, str)