File size: 950 Bytes
ed4d993
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
30
31
32
33
34
35
36
37
"""Test Konko API wrapper.

In order to run this test, you need to have an Konko api key.
You'll then need to set KONKO_API_KEY environment variable to your api key.
"""
import pytest as pytest

from langchain_community.llms import Konko


def test_konko_call() -> None:
    """Test simple call to konko."""
    llm = Konko(
        model="mistralai/mistral-7b-v0.1",
        temperature=0.2,
        max_tokens=250,
    )
    output = llm.invoke("Say foo:")

    assert llm._llm_type == "konko"
    assert isinstance(output, str)


async def test_konko_acall() -> None:
    """Test simple call to konko."""
    llm = Konko(
        model="mistralai/mistral-7b-v0.1",
        temperature=0.2,
        max_tokens=250,
    )
    output = await llm.agenerate(["Say foo:"], stop=["bar"])

    assert llm._llm_type == "konko"
    output_text = output.generations[0][0].text
    assert isinstance(output_text, str)
    assert output_text.count("bar") <= 1