File size: 2,807 Bytes
ad33df7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
from pathlib import Path
from unittest.mock import patch

from kotaemon.base.schema import LLMInterface
from kotaemon.llms import AzureOpenAI, LlamaCpp, OpenAI

try:
    from langchain_openai import AzureOpenAI as AzureOpenAILC
    from langchain_openai import OpenAI as OpenAILC
except ImportError:
    from langchain.llms import AzureOpenAI as AzureOpenAILC
    from langchain.llms import OpenAI as OpenAILC

from openai.types.completion import Completion

from .conftest import skip_llama_cpp_not_installed, skip_openai_lc_wrapper_test

_openai_completion_response = Completion.parse_obj(
    {
        "id": "cmpl-7qyNoIo6gRSCJR0hi8o3ZKBH4RkJ0",
        "object": "text_completion",
        "created": 1392751226,
        "model": "gpt-35-turbo",
        "system_fingerprint": None,
        "choices": [
            {
                "text": "completion",
                "index": 0,
                "finish_reason": "length",
                "logprobs": None,
            }
        ],
        "usage": {"completion_tokens": 20, "prompt_tokens": 2, "total_tokens": 22},
    }
)


@skip_openai_lc_wrapper_test
@patch(
    "openai.resources.completions.Completions.create",
    side_effect=lambda *args, **kwargs: _openai_completion_response,
)
def test_azureopenai_model(openai_completion):
    model = AzureOpenAI(
        azure_endpoint="https://test.openai.azure.com/",
        openai_api_key="some-key",
        openai_api_version="2023-03-15-preview",
        deployment_name="gpt35turbo",
        temperature=0,
        request_timeout=60,
    )
    assert isinstance(
        model.to_langchain_format(), AzureOpenAILC
    ), "Agent not wrapped in Langchain's AzureOpenAI"

    output = model("hello world")
    assert isinstance(
        output, LLMInterface
    ), "Output for single text is not LLMInterface"


@patch(
    "openai.resources.completions.Completions.create",
    side_effect=lambda *args, **kwargs: _openai_completion_response,
)
def test_openai_model(openai_completion):
    model = OpenAI(
        openai_api_base="https://test.openai.azure.com/",
        openai_api_key="some-key",
        openai_api_version="2023-03-15-preview",
        deployment_name="gpt35turbo",
        temperature=0,
        request_timeout=60,
    )
    assert isinstance(
        model.to_langchain_format(), OpenAILC
    ), "Agent is not wrapped in Langchain's OpenAI"

    output = model("hello world")
    assert isinstance(
        output, LLMInterface
    ), "Output for single text is not LLMInterface"


@skip_llama_cpp_not_installed
def test_llamacpp_model():
    weight_path = Path(__file__).parent / "resources" / "ggml-vocab-llama.gguf"

    # test initialization
    model = LlamaCpp(model_path=str(weight_path), vocab_only=True)
    assert isinstance(model._obj, model._get_lc_class())