File size: 898 Bytes
58d33f0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
"""Test GooseAI API wrapper."""

from langchain.llms.gooseai import GooseAI


def test_gooseai_call() -> None:
    """Test valid call to gooseai."""
    llm = GooseAI(max_tokens=10)
    output = llm("Say foo:")
    assert isinstance(output, str)


def test_gooseai_call_fairseq() -> None:
    """Test valid call to gooseai with fairseq model."""
    llm = GooseAI(model_name="fairseq-1-3b", max_tokens=10)
    output = llm("Say foo:")
    assert isinstance(output, str)


def test_gooseai_stop_valid() -> None:
    """Test gooseai stop logic on valid configuration."""
    query = "write an ordered list of five items"
    first_llm = GooseAI(stop="3", temperature=0)
    first_output = first_llm(query)
    second_llm = GooseAI(temperature=0)
    second_output = second_llm(query, stop=["3"])
    # Because it stops on new lines, shouldn't return anything
    assert first_output == second_output