|
|
@ -32,18 +32,33 @@ def test_vertex_initialization(model_name: str) -> None:
|
|
|
|
"model_name",
|
|
|
|
"model_name",
|
|
|
|
model_names_to_test_with_default,
|
|
|
|
model_names_to_test_with_default,
|
|
|
|
)
|
|
|
|
)
|
|
|
|
def test_vertex_call(model_name: str) -> None:
|
|
|
|
def test_vertex_invoke(model_name: str) -> None:
|
|
|
|
llm = (
|
|
|
|
llm = (
|
|
|
|
VertexAI(model_name=model_name, temperature=0)
|
|
|
|
VertexAI(model_name=model_name, temperature=0)
|
|
|
|
if model_name
|
|
|
|
if model_name
|
|
|
|
else VertexAI(temperature=0.0)
|
|
|
|
else VertexAI(temperature=0.0)
|
|
|
|
)
|
|
|
|
)
|
|
|
|
output = llm("Say foo:")
|
|
|
|
output = llm.invoke("Say foo:")
|
|
|
|
assert isinstance(output, str)
|
|
|
|
assert isinstance(output, str)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
|
|
|
|
|
|
"model_name",
|
|
|
|
|
|
|
|
model_names_to_test_with_default,
|
|
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
def test_vertex_generate(model_name: str) -> None:
|
|
|
|
|
|
|
|
llm = (
|
|
|
|
|
|
|
|
VertexAI(model_name=model_name, temperature=0)
|
|
|
|
|
|
|
|
if model_name
|
|
|
|
|
|
|
|
else VertexAI(temperature=0.0)
|
|
|
|
|
|
|
|
)
|
|
|
|
|
|
|
|
output = llm.generate(["Say foo:"])
|
|
|
|
|
|
|
|
assert isinstance(output, LLMResult)
|
|
|
|
|
|
|
|
assert len(output.generations) == 1
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.xfail(reason="VertexAI doesn't always respect number of candidates")
|
|
|
|
@pytest.mark.xfail(reason="VertexAI doesn't always respect number of candidates")
|
|
|
|
def test_vertex_generate() -> None:
|
|
|
|
def test_vertex_generate_multiple_candidates() -> None:
|
|
|
|
llm = VertexAI(temperature=0.3, n=2, model_name="text-bison@001")
|
|
|
|
llm = VertexAI(temperature=0.3, n=2, model_name="text-bison@001")
|
|
|
|
output = llm.generate(["Say foo:"])
|
|
|
|
output = llm.generate(["Say foo:"])
|
|
|
|
assert isinstance(output, LLMResult)
|
|
|
|
assert isinstance(output, LLMResult)
|
|
|
|