llama-index

1
from typing import Dict, Type
2

3
from llama_index.legacy.llms.bedrock import Bedrock
4
from llama_index.legacy.llms.custom import CustomLLM
5
from llama_index.legacy.llms.gradient import (
6
    GradientBaseModelLLM,
7
    GradientModelAdapterLLM,
8
)
9
from llama_index.legacy.llms.huggingface import HuggingFaceLLM
10
from llama_index.legacy.llms.langchain import LangChainLLM
11
from llama_index.legacy.llms.llama_cpp import LlamaCPP
12
from llama_index.legacy.llms.llm import LLM
13
from llama_index.legacy.llms.mock import MockLLM
14
from llama_index.legacy.llms.openai import OpenAI
15
from llama_index.legacy.llms.palm import PaLM
16
from llama_index.legacy.llms.predibase import PredibaseLLM
17
from llama_index.legacy.llms.replicate import Replicate
18
from llama_index.legacy.llms.vertex import Vertex
19
from llama_index.legacy.llms.xinference import Xinference
20

21
RECOGNIZED_LLMS: Dict[str, Type[LLM]] = {
22
    MockLLM.class_name(): MockLLM,
23
    Replicate.class_name(): Replicate,
24
    HuggingFaceLLM.class_name(): HuggingFaceLLM,
25
    OpenAI.class_name(): OpenAI,
26
    Xinference.class_name(): Xinference,
27
    LlamaCPP.class_name(): LlamaCPP,
28
    LangChainLLM.class_name(): LangChainLLM,
29
    PaLM.class_name(): PaLM,
30
    PredibaseLLM.class_name(): PredibaseLLM,
31
    Bedrock.class_name(): Bedrock,
32
    CustomLLM.class_name(): CustomLLM,
33
    GradientBaseModelLLM.class_name(): GradientBaseModelLLM,
34
    GradientModelAdapterLLM.class_name(): GradientModelAdapterLLM,
35
    Vertex.class_name(): Vertex,
36
}
37

38

39
def load_llm(data: dict) -> LLM:
40
    """Load LLM by name."""
41
    if isinstance(data, LLM):
42
        return data
43
    llm_name = data.get("class_name", None)
44
    if llm_name is None:
45
        raise ValueError("LLM loading requires a class_name")
46

47
    if llm_name not in RECOGNIZED_LLMS:
48
        raise ValueError(f"Invalid LLM name: {llm_name}")
49

50
    return RECOGNIZED_LLMS[llm_name].from_dict(data)
51