langchain/libs/community/tests/integration_tests/llms/test_baseten.py
Philip Kiely - Baseten 6342da333a
community: refactor Baseten integration with new API endpoints & docs (#15017)
- **Description:** In response to user feedback, this PR refactors the
Baseten integration with updated model endpoints, as well as updates
relevant documentation. This PR has been tested by end users in
production and works as expected.
  - **Issue:** N/A
- **Dependencies:** This PR actually removes the dependency on the
`baseten` package!
  - **Twitter handle:** https://twitter.com/basetenco
2023-12-22 12:46:24 -08:00

14 lines
396 B
Python

"""Test Baseten API wrapper."""
import os
from langchain_community.llms.baseten import Baseten
# This test requires valid BASETEN_MODEL_ID and BASETEN_API_KEY environment variables
def test_baseten_call() -> None:
"""Test valid call to Baseten."""
llm = Baseten(model=os.environ["BASETEN_MODEL_ID"])
output = llm("Test prompt, please respond.")
assert isinstance(output, str)