llama-index

1
import logging
2
from importlib.metadata import version
3
from types import ModuleType
4
from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Type
5

6
import openai
7
from packaging.version import parse
8
from tenacity import (
9
    before_sleep_log,
10
    retry,
11
    retry_if_exception_type,
12
    stop_after_attempt,
13
    wait_exponential,
14
)
15

16
from llama_index.legacy.bridge.pydantic import BaseModel
17
from llama_index.legacy.llms.generic_utils import get_from_param_or_env
18
from llama_index.legacy.llms.types import ChatMessage
19

20
DEFAULT_KONKO_API_TYPE = "open_ai"
21
DEFAULT_KONKO_API_BASE = "https://api.konko.ai/v1"
22
DEFAULT_KONKO_API_VERSION = ""
23
MISSING_API_KEY_ERROR_MESSAGE = """No Konko API key found for LLM.
24
E.g. to use konko Please set the KONKO_API_KEY environment variable or \
25
konko.api_key prior to initialization.
26
API keys can be found or created at \
27
https://www.konko.ai/
28
"""
29

30
logger = logging.getLogger(__name__)
31

32

33
def import_konko() -> ModuleType:
34
    try:
35
        import konko
36

37
        return konko
38
    except ImportError:
39
        raise ImportError(
40
            "Could not import konko python package. "
41
            "Please install it with `pip install konko`."
42
        )
43

44

45
def is_openai_v1() -> bool:
46
    try:
47
        _version = parse(version("openai"))
48
        major_version = _version.major
49
    except AttributeError:
50
        # Handle the case where version or major attribute is not present
51
        return False
52
    return bool(major_version >= 1)
53

54

55
def _create_retry_decorator(max_retries: int) -> Callable[[Any], Any]:
56
    min_seconds = 4
57
    max_seconds = 10
58
    # Wait 2^x * 1 second between each retry starting with
59
    # 4 seconds, then up to 10 seconds, then 10 seconds afterwards
60
    return retry(
61
        reraise=True,
62
        stop=stop_after_attempt(max_retries),
63
        wait=wait_exponential(multiplier=1, min=min_seconds, max=max_seconds),
64
        retry=(
65
            retry_if_exception_type(openai.APITimeoutError)
66
            | retry_if_exception_type(openai.APIError)
67
            | retry_if_exception_type(openai.APIConnectionError)
68
            | retry_if_exception_type(openai.RateLimitError)
69
            | retry_if_exception_type(openai.APIStatusError)
70
        ),
71
        before_sleep=before_sleep_log(logger, logging.WARNING),
72
    )
73

74

75
def completion_with_retry(is_chat_model: bool, max_retries: int, **kwargs: Any) -> Any:
76
    """Use tenacity to retry the completion call."""
77
    retry_decorator = _create_retry_decorator(max_retries=max_retries)
78

79
    @retry_decorator
80
    def _completion_with_retry(**kwargs: Any) -> Any:
81
        client = get_completion_endpoint(is_chat_model)
82
        return client.create(**kwargs)
83

84
    return _completion_with_retry(**kwargs)
85

86

87
def get_completion_endpoint(is_chat_model: bool) -> Any:
88
    """
89
    Get the appropriate completion endpoint based on the model type and API version.
90

91
    Args:
92
    - is_chat_model (bool): A flag indicating whether the model is a chat model.
93

94
    Returns:
95
    - The appropriate completion endpoint based on the model type and API version.
96

97
    Raises:
98
    - NotImplementedError: If the combination of is_chat_model and API version is not supported.
99
    """
100
    konko = import_konko()
101
    # For OpenAI version 1
102
    if is_openai_v1():
103
        return konko.chat.completions if is_chat_model else konko.completions
104

105
    # For other versions
106
    if not is_openai_v1():
107
        return konko.ChatCompletion if is_chat_model else konko.Completion
108

109
    # Raise error if the combination of is_chat_model and API version is not covered
110
    raise NotImplementedError(
111
        "The combination of model type and API version is not supported."
112
    )
113

114

115
def to_openai_message_dict(message: ChatMessage) -> dict:
116
    """Convert generic message to OpenAI message dict."""
117
    message_dict = {
118
        "role": message.role,
119
        "content": message.content,
120
    }
121
    message_dict.update(message.additional_kwargs)
122

123
    return message_dict
124

125

126
def to_openai_message_dicts(messages: Sequence[ChatMessage]) -> List[dict]:
127
    """Convert generic messages to OpenAI message dicts."""
128
    return [to_openai_message_dict(message) for message in messages]
129

130

131
def from_openai_message_dict(message_dict: Any) -> ChatMessage:
132
    """Convert openai message dict to generic message."""
133
    if is_openai_v1():
134
        # Handling for OpenAI version 1
135
        role = message_dict.role
136
        content = message_dict.content
137
        additional_kwargs = {
138
            attr: getattr(message_dict, attr)
139
            for attr in dir(message_dict)
140
            if not attr.startswith("_") and attr not in ["role", "content"]
141
        }
142
    else:
143
        # Handling for OpenAI version 0
144
        role = message_dict.get("role")
145
        content = message_dict.get("content", None)
146
        additional_kwargs = {
147
            key: value
148
            for key, value in message_dict.items()
149
            if key not in ["role", "content"]
150
        }
151

152
    return ChatMessage(role=role, content=content, additional_kwargs=additional_kwargs)
153

154

155
def from_openai_message_dicts(message_dicts: Sequence[dict]) -> List[ChatMessage]:
156
    """Convert openai message dicts to generic messages."""
157
    return [from_openai_message_dict(message_dict) for message_dict in message_dicts]
158

159

160
def to_openai_function(pydantic_class: Type[BaseModel]) -> Dict[str, Any]:
161
    """Convert pydantic class to OpenAI function."""
162
    schema = pydantic_class.schema()
163
    return {
164
        "name": schema["title"],
165
        "description": schema["description"],
166
        "parameters": pydantic_class.schema(),
167
    }
168

169

170
def resolve_konko_credentials(
171
    konko_api_key: Optional[str] = None,
172
    openai_api_key: Optional[str] = None,
173
    api_type: Optional[str] = None,
174
    api_base: Optional[str] = None,
175
    api_version: Optional[str] = None,
176
) -> Tuple[str, str, str, str, str]:
177
    """ "Resolve KonkoAI credentials.
178

179
    The order of precedence is:
180
    1. param
181
    2. env
182
    3. konkoai module
183
    4. default
184
    """
185
    konko = import_konko()
186
    # resolve from param or env
187
    konko_api_key = get_from_param_or_env(
188
        "konko_api_key", konko_api_key, "KONKO_API_KEY", ""
189
    )
190
    openai_api_key = get_from_param_or_env(
191
        "openai_api_key", openai_api_key, "OPENAI_API_KEY", ""
192
    )
193
    api_type = get_from_param_or_env("api_type", api_type, "KONKO_API_TYPE", "")
194
    api_base = DEFAULT_KONKO_API_BASE
195
    api_version = get_from_param_or_env(
196
        "api_version", api_version, "KONKO_API_VERSION", ""
197
    )
198

199
    # resolve from konko module or default
200
    konko_api_key = konko_api_key
201
    openai_api_key = openai_api_key
202
    api_type = api_type or DEFAULT_KONKO_API_TYPE
203
    api_base = api_base or konko.api_base or DEFAULT_KONKO_API_BASE
204
    api_version = api_version or DEFAULT_KONKO_API_VERSION
205

206
    if not konko_api_key:
207
        raise ValueError(MISSING_API_KEY_ERROR_MESSAGE)
208

209
    return konko_api_key, openai_api_key, api_type, api_base, api_version
210

211

212
async def acompletion_with_retry(
213
    is_chat_model: bool, max_retries: int, **kwargs: Any
214
) -> Any:
215
    """Use tenacity to retry the async completion call."""
216
    konko = import_konko()
217
    retry_decorator = _create_retry_decorator(max_retries=max_retries)
218

219
    @retry_decorator
220
    async def _completion_with_retry(**kwargs: Any) -> Any:
221
        if is_chat_model:
222
            if is_openai_v1():
223
                return await konko.AsyncKonko().chat.completions.create(**kwargs)
224
            else:
225
                return await konko.ChatCompletion.acreate(**kwargs)
226
        else:
227
            if is_openai_v1():
228
                return await konko.AsyncKonko().completions.create(**kwargs)
229
            else:
230
                return await konko.Completion.acreate(**kwargs)
231

232
    return await _completion_with_retry(**kwargs)
233