llama-index

1
from typing import TYPE_CHECKING, Optional, Union
2

3
if TYPE_CHECKING:
4
    from langchain.base_language import BaseLanguageModel
5

6
from llama_index.legacy.llms.llama_cpp import LlamaCPP
7
from llama_index.legacy.llms.llama_utils import completion_to_prompt, messages_to_prompt
8
from llama_index.legacy.llms.llm import LLM
9
from llama_index.legacy.llms.mock import MockLLM
10
from llama_index.legacy.llms.openai import OpenAI
11
from llama_index.legacy.llms.openai_utils import validate_openai_api_key
12

13
LLMType = Union[str, LLM, "BaseLanguageModel"]
14

15

16
def resolve_llm(llm: Optional[LLMType] = None) -> LLM:
17
    """Resolve LLM from string or LLM instance."""
18
    try:
19
        from langchain.base_language import BaseLanguageModel
20

21
        from llama_index.legacy.llms.langchain import LangChainLLM
22
    except ImportError:
23
        BaseLanguageModel = None  # type: ignore
24

25
    if llm == "default":
26
        # return default OpenAI model. If it fails, return LlamaCPP
27
        try:
28
            llm = OpenAI()
29
            validate_openai_api_key(llm.api_key)
30
        except ValueError as e:
31
            raise ValueError(
32
                "\n******\n"
33
                "Could not load OpenAI model. "
34
                "If you intended to use OpenAI, please check your OPENAI_API_KEY.\n"
35
                "Original error:\n"
36
                f"{e!s}"
37
                "\nTo disable the LLM entirely, set llm=None."
38
                "\n******"
39
            )
40

41
    if isinstance(llm, str):
42
        splits = llm.split(":", 1)
43
        is_local = splits[0]
44
        model_path = splits[1] if len(splits) > 1 else None
45
        if is_local != "local":
46
            raise ValueError(
47
                "llm must start with str 'local' or of type LLM or BaseLanguageModel"
48
            )
49
        llm = LlamaCPP(
50
            model_path=model_path,
51
            messages_to_prompt=messages_to_prompt,
52
            completion_to_prompt=completion_to_prompt,
53
            model_kwargs={"n_gpu_layers": 1},
54
        )
55
    elif BaseLanguageModel is not None and isinstance(llm, BaseLanguageModel):
56
        # NOTE: if it's a langchain model, wrap it in a LangChainLLM
57
        llm = LangChainLLM(llm=llm)
58
    elif llm is None:
59
        print("LLM is explicitly disabled. Using MockLLM.")
60
        llm = MockLLM()
61

62
    return llm
63