llama-index

1
from typing import Any, Dict, List, Optional, Sequence, Tuple
2

3
from llama_index.legacy.core.llms.types import ChatMessage, MessageRole
4
from llama_index.legacy.llms.generic_utils import get_from_param_or_env
5

6
DEFAULT_ANYSCALE_API_BASE = "https://api.endpoints.anyscale.com/v1"
7
DEFAULT_ANYSCALE_API_VERSION = ""
8

9
LLAMA_MODELS = {
10
    "meta-llama/Llama-2-7b-chat-hf": 4096,
11
    "meta-llama/Llama-2-13b-chat-hf": 4096,
12
    "meta-llama/Llama-2-70b-chat-hf": 4096,
13
    "codellama/CodeLlama-34b-Instruct-hf": 16384,
14
    "Meta-Llama/Llama-Guard-7b": 4096,
15
}
16

17
MISTRAL_MODELS = {
18
    "mistralai/Mistral-7B-Instruct-v0.1": 16384,
19
    "Open-Orca/Mistral-7B-OpenOrca": 8192,
20
    "mistralai/Mixtral-8x7B-Instruct-v0.1": 32768,
21
}
22

23
ZEPHYR_MODELS = {
24
    "HuggingFaceH4/zephyr-7b-beta": 16384,
25
}
26

27
ALL_AVAILABLE_MODELS = {
28
    **LLAMA_MODELS,
29
    **MISTRAL_MODELS,
30
    **ZEPHYR_MODELS,
31
}
32

33
DISCONTINUED_MODELS: Dict[str, int] = {}
34

35

36
def anyscale_modelname_to_contextsize(modelname: str) -> int:
37
    """
38
    Calculate the maximum number of tokens possible to generate for a model.
39

40
    Args:
41
        modelname: The modelname we want to know the context size for.
42

43
    Returns:
44
        The maximum context size
45

46
    Example:
47
        .. code-block:: python
48

49
            max_tokens = anyscale_modelname_to_contextsize(model_name)
50
    """
51
    # handling finetuned models
52
    # TO BE FILLED
53

54
    if modelname in DISCONTINUED_MODELS:
55
        raise ValueError(
56
            f"Anyscale hosted model {modelname} has been discontinued. "
57
            "Please choose another model."
58
        )
59

60
    context_size = ALL_AVAILABLE_MODELS.get(modelname, None)
61

62
    if context_size is None:
63
        raise ValueError(
64
            f"Unknown model: {modelname}. Please provide a valid Anyscale model name."
65
            "Known models are: " + ", ".join(ALL_AVAILABLE_MODELS.keys())
66
        )
67

68
    return context_size
69

70

71
def _message_to_anyscale_prompt(message: ChatMessage) -> Dict[str, Any]:
72
    if message.role == MessageRole.USER:
73
        prompt = {"role": "user", "content": message.content}
74
    elif message.role == MessageRole.ASSISTANT:
75
        prompt = {"role": "assistant", "content": message.content}
76
    elif message.role == MessageRole.SYSTEM:
77
        prompt = {"role": "system", "content": message.content}
78
    elif message.role == MessageRole.FUNCTION:
79
        raise ValueError(f"Message role {MessageRole.FUNCTION} is not supported.")
80
    else:
81
        raise ValueError(f"Unknown message role: {message.role}")
82

83
    return prompt
84

85

86
def messages_to_anyscale_prompt(messages: Sequence[ChatMessage]) -> List[Dict]:
87
    if len(messages) == 0:
88
        raise ValueError("Got empty list of messages.")
89

90
    return [_message_to_anyscale_prompt(message) for message in messages]
91

92

93
def resolve_anyscale_credentials(
94
    api_key: Optional[str] = None,
95
    api_base: Optional[str] = None,
96
    api_version: Optional[str] = None,
97
) -> Tuple[Optional[str], str, str]:
98
    """
99
    "Resolve OpenAI credentials.
100

101
    The order of precedence is:
102
    1. param
103
    2. env
104
    3. openai module
105
    4. default
106
    """
107
    # resolve from param or env
108
    api_key = get_from_param_or_env("api_key", api_key, "ANYSCALE_API_KEY", "")
109
    api_base = get_from_param_or_env("api_base", api_base, "ANYSCALE_API_BASE", "")
110
    api_version = get_from_param_or_env(
111
        "api_version", api_version, "ANYSCALE_API_VERSION", ""
112
    )
113

114
    # resolve from openai module or default
115
    final_api_key = api_key or ""
116
    final_api_base = api_base or DEFAULT_ANYSCALE_API_BASE
117
    final_api_version = api_version or DEFAULT_ANYSCALE_API_VERSION
118

119
    return final_api_key, str(final_api_base), final_api_version
120