langchain/libs/community/tests/integration_tests/llms/test_replicate.py
fayvor 3b956b3a97
community: Update Replicate LLM and fix tests (#27655)
**Description:** 
- Fix bug in Replicate LLM class, where it was looking for parameter
names in a place where they no longer exist in pydantic 2, resulting in
the "Field required" validation error described in the issue.
- Fix Replicate LLM integration tests to:
  - Use active models on Replicate.
- Use the correct model parameter `max_new_tokens` as shown in the
[Replicate
docs](https://replicate.com/docs/guides/language-models/how-to-use#minimum-and-maximum-new-tokens).
  - Use callbacks instead of deprecated callback_manager.

**Issue:** #26937 

**Dependencies:** n/a

**Twitter handle:** n/a

---------

Signed-off-by: Fayvor Love <fayvor@gmail.com>
Co-authored-by: Chester Curme <chester.curme@gmail.com>
2024-10-30 16:07:08 +00:00

50 lines
1.7 KiB
Python

"""Test Replicate API wrapper."""
from langchain_community.llms.replicate import Replicate
from tests.unit_tests.callbacks.fake_callback_handler import FakeCallbackHandler
TEST_MODEL_HELLO = (
"replicate/hello-world:"
+ "5c7d5dc6dd8bf75c1acaa8565735e7986bc5b66206b55cca93cb72c9bf15ccaa"
)
TEST_MODEL_LANG = "meta/meta-llama-3-8b-instruct"
def test_replicate_call() -> None:
"""Test simple non-streaming call to Replicate."""
llm = Replicate(model=TEST_MODEL_HELLO)
output = llm.invoke("What is LangChain")
assert output
assert isinstance(output, str)
def test_replicate_streaming_call() -> None:
"""Test streaming call to Replicate."""
callback_handler = FakeCallbackHandler()
llm = Replicate(
streaming=True, callbacks=[callback_handler], model=TEST_MODEL_HELLO
)
output = llm.invoke("What is LangChain")
assert output
assert isinstance(output, str)
def test_replicate_model_kwargs() -> None:
"""Test simple non-streaming call to Replicate."""
llm = Replicate( # type: ignore[call-arg]
model=TEST_MODEL_LANG, model_kwargs={"max_new_tokens": 10, "temperature": 0.01}
)
long_output = llm.invoke("What is LangChain")
llm = Replicate( # type: ignore[call-arg]
model=TEST_MODEL_LANG, model_kwargs={"max_new_tokens": 5, "temperature": 0.01}
)
short_output = llm.invoke("What is LangChain")
assert len(short_output) < len(long_output)
assert llm.model_kwargs == {"max_new_tokens": 5, "temperature": 0.01}
def test_replicate_input() -> None:
llm = Replicate(model=TEST_MODEL_LANG, input={"max_new_tokens": 10})
assert llm.model_kwargs == {"max_new_tokens": 10}