llama-index

1
from typing import Dict, Sequence
2

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

5
HUMAN_PREFIX = "\n\nHuman:"
6
ASSISTANT_PREFIX = "\n\nAssistant:"
7

8

9
CLAUDE_MODELS: Dict[str, int] = {
10
    "claude-instant-1": 100000,
11
    "claude-instant-1.2": 100000,
12
    "claude-2": 100000,
13
    "claude-2.0": 100000,
14
    "claude-2.1": 200000,
15
}
16

17

18
def anthropic_modelname_to_contextsize(modelname: str) -> int:
19
    if modelname not in CLAUDE_MODELS:
20
        raise ValueError(
21
            f"Unknown model: {modelname}. Please provide a valid Anthropic model name."
22
            "Known models are: " + ", ".join(CLAUDE_MODELS.keys())
23
        )
24

25
    return CLAUDE_MODELS[modelname]
26

27

28
def _message_to_anthropic_prompt(message: ChatMessage) -> str:
29
    if message.role == MessageRole.USER:
30
        prompt = f"{HUMAN_PREFIX} {message.content}"
31
    elif message.role == MessageRole.ASSISTANT:
32
        prompt = f"{ASSISTANT_PREFIX} {message.content}"
33
    elif message.role == MessageRole.SYSTEM:
34
        prompt = f"{HUMAN_PREFIX} <system>{message.content}</system>"
35
    elif message.role == MessageRole.FUNCTION:
36
        raise ValueError(f"Message role {MessageRole.FUNCTION} is not supported.")
37
    else:
38
        raise ValueError(f"Unknown message role: {message.role}")
39

40
    return prompt
41

42

43
def messages_to_anthropic_prompt(messages: Sequence[ChatMessage]) -> str:
44
    if len(messages) == 0:
45
        raise ValueError("Got empty list of messages.")
46

47
    # NOTE: make sure the prompt ends with the assistant prefix
48
    if messages[-1].role != MessageRole.ASSISTANT:
49
        messages = [
50
            *list(messages),
51
            ChatMessage(role=MessageRole.ASSISTANT, content=""),
52
        ]
53

54
    str_list = [_message_to_anthropic_prompt(message) for message in messages]
55
    return "".join(str_list)
56