llama-index

1
import os
2
from typing import Any, Awaitable, Callable, List, Optional, Sequence
3

4
from llama_index.legacy.core.llms.types import (
5
    ChatMessage,
6
    ChatResponse,
7
    ChatResponseAsyncGen,
8
    ChatResponseGen,
9
    CompletionResponse,
10
    CompletionResponseAsyncGen,
11
    CompletionResponseGen,
12
    MessageRole,
13
)
14

15

16
def messages_to_history_str(messages: Sequence[ChatMessage]) -> str:
17
    """Convert messages to a history string."""
18
    string_messages = []
19
    for message in messages:
20
        role = message.role
21
        content = message.content
22
        string_message = f"{role.value}: {content}"
23

24
        addtional_kwargs = message.additional_kwargs
25
        if addtional_kwargs:
26
            string_message += f"\n{addtional_kwargs}"
27
        string_messages.append(string_message)
28
    return "\n".join(string_messages)
29

30

31
def messages_to_prompt(messages: Sequence[ChatMessage]) -> str:
32
    """Convert messages to a prompt string."""
33
    string_messages = []
34
    for message in messages:
35
        role = message.role
36
        content = message.content
37
        string_message = f"{role.value}: {content}"
38

39
        addtional_kwargs = message.additional_kwargs
40
        if addtional_kwargs:
41
            string_message += f"\n{addtional_kwargs}"
42
        string_messages.append(string_message)
43

44
    string_messages.append(f"{MessageRole.ASSISTANT.value}: ")
45
    return "\n".join(string_messages)
46

47

48
def prompt_to_messages(prompt: str) -> List[ChatMessage]:
49
    """Convert a string prompt to a sequence of messages."""
50
    return [ChatMessage(role=MessageRole.USER, content=prompt)]
51

52

53
def completion_response_to_chat_response(
54
    completion_response: CompletionResponse,
55
) -> ChatResponse:
56
    """Convert a completion response to a chat response."""
57
    return ChatResponse(
58
        message=ChatMessage(
59
            role=MessageRole.ASSISTANT,
60
            content=completion_response.text,
61
            additional_kwargs=completion_response.additional_kwargs,
62
        ),
63
        raw=completion_response.raw,
64
    )
65

66

67
def stream_completion_response_to_chat_response(
68
    completion_response_gen: CompletionResponseGen,
69
) -> ChatResponseGen:
70
    """Convert a stream completion response to a stream chat response."""
71

72
    def gen() -> ChatResponseGen:
73
        for response in completion_response_gen:
74
            yield ChatResponse(
75
                message=ChatMessage(
76
                    role=MessageRole.ASSISTANT,
77
                    content=response.text,
78
                    additional_kwargs=response.additional_kwargs,
79
                ),
80
                delta=response.delta,
81
                raw=response.raw,
82
            )
83

84
    return gen()
85

86

87
def astream_completion_response_to_chat_response(
88
    completion_response_gen: CompletionResponseAsyncGen,
89
) -> ChatResponseAsyncGen:
90
    """Convert an async stream completion to an async stream chat response."""
91

92
    async def gen() -> ChatResponseAsyncGen:
93
        async for response in completion_response_gen:
94
            yield ChatResponse(
95
                message=ChatMessage(
96
                    role=MessageRole.ASSISTANT,
97
                    content=response.text,
98
                    additional_kwargs=response.additional_kwargs,
99
                ),
100
                delta=response.delta,
101
                raw=response.raw,
102
            )
103

104
    return gen()
105

106

107
def chat_response_to_completion_response(
108
    chat_response: ChatResponse,
109
) -> CompletionResponse:
110
    """Convert a chat response to a completion response."""
111
    return CompletionResponse(
112
        text=chat_response.message.content or "",
113
        additional_kwargs=chat_response.message.additional_kwargs,
114
        raw=chat_response.raw,
115
    )
116

117

118
def stream_chat_response_to_completion_response(
119
    chat_response_gen: ChatResponseGen,
120
) -> CompletionResponseGen:
121
    """Convert a stream chat response to a completion response."""
122

123
    def gen() -> CompletionResponseGen:
124
        for response in chat_response_gen:
125
            yield CompletionResponse(
126
                text=response.message.content or "",
127
                additional_kwargs=response.message.additional_kwargs,
128
                delta=response.delta,
129
                raw=response.raw,
130
            )
131

132
    return gen()
133

134

135
def completion_to_chat_decorator(
136
    func: Callable[..., CompletionResponse]
137
) -> Callable[..., ChatResponse]:
138
    """Convert a completion function to a chat function."""
139

140
    def wrapper(messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
141
        # normalize input
142
        prompt = messages_to_prompt(messages)
143
        completion_response = func(prompt, **kwargs)
144
        # normalize output
145
        return completion_response_to_chat_response(completion_response)
146

147
    return wrapper
148

149

150
def stream_completion_to_chat_decorator(
151
    func: Callable[..., CompletionResponseGen]
152
) -> Callable[..., ChatResponseGen]:
153
    """Convert a completion function to a chat function."""
154

155
    def wrapper(messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponseGen:
156
        # normalize input
157
        prompt = messages_to_prompt(messages)
158
        completion_response = func(prompt, **kwargs)
159
        # normalize output
160
        return stream_completion_response_to_chat_response(completion_response)
161

162
    return wrapper
163

164

165
def chat_to_completion_decorator(
166
    func: Callable[..., ChatResponse]
167
) -> Callable[..., CompletionResponse]:
168
    """Convert a chat function to a completion function."""
169

170
    def wrapper(prompt: str, **kwargs: Any) -> CompletionResponse:
171
        # normalize input
172
        messages = prompt_to_messages(prompt)
173
        chat_response = func(messages, **kwargs)
174
        # normalize output
175
        return chat_response_to_completion_response(chat_response)
176

177
    return wrapper
178

179

180
def stream_chat_to_completion_decorator(
181
    func: Callable[..., ChatResponseGen]
182
) -> Callable[..., CompletionResponseGen]:
183
    """Convert a chat function to a completion function."""
184

185
    def wrapper(prompt: str, **kwargs: Any) -> CompletionResponseGen:
186
        # normalize input
187
        messages = prompt_to_messages(prompt)
188
        chat_response = func(messages, **kwargs)
189
        # normalize output
190
        return stream_chat_response_to_completion_response(chat_response)
191

192
    return wrapper
193

194

195
# ===== Async =====
196

197

198
def acompletion_to_chat_decorator(
199
    func: Callable[..., Awaitable[CompletionResponse]]
200
) -> Callable[..., Awaitable[ChatResponse]]:
201
    """Convert a completion function to a chat function."""
202

203
    async def wrapper(messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
204
        # normalize input
205
        prompt = messages_to_prompt(messages)
206
        completion_response = await func(prompt, **kwargs)
207
        # normalize output
208
        return completion_response_to_chat_response(completion_response)
209

210
    return wrapper
211

212

213
def achat_to_completion_decorator(
214
    func: Callable[..., Awaitable[ChatResponse]]
215
) -> Callable[..., Awaitable[CompletionResponse]]:
216
    """Convert a chat function to a completion function."""
217

218
    async def wrapper(prompt: str, **kwargs: Any) -> CompletionResponse:
219
        # normalize input
220
        messages = prompt_to_messages(prompt)
221
        chat_response = await func(messages, **kwargs)
222
        # normalize output
223
        return chat_response_to_completion_response(chat_response)
224

225
    return wrapper
226

227

228
def astream_completion_to_chat_decorator(
229
    func: Callable[..., Awaitable[CompletionResponseAsyncGen]]
230
) -> Callable[..., Awaitable[ChatResponseAsyncGen]]:
231
    """Convert a completion function to a chat function."""
232

233
    async def wrapper(
234
        messages: Sequence[ChatMessage], **kwargs: Any
235
    ) -> ChatResponseAsyncGen:
236
        # normalize input
237
        prompt = messages_to_prompt(messages)
238
        completion_response = await func(prompt, **kwargs)
239
        # normalize output
240
        return astream_completion_response_to_chat_response(completion_response)
241

242
    return wrapper
243

244

245
def astream_chat_to_completion_decorator(
246
    func: Callable[..., Awaitable[ChatResponseAsyncGen]]
247
) -> Callable[..., Awaitable[CompletionResponseAsyncGen]]:
248
    """Convert a chat function to a completion function."""
249

250
    async def wrapper(prompt: str, **kwargs: Any) -> CompletionResponseAsyncGen:
251
        # normalize input
252
        messages = prompt_to_messages(prompt)
253
        chat_response = await func(messages, **kwargs)
254
        # normalize output
255
        return astream_chat_response_to_completion_response(chat_response)
256

257
    return wrapper
258

259

260
def async_stream_completion_response_to_chat_response(
261
    completion_response_gen: CompletionResponseAsyncGen,
262
) -> ChatResponseAsyncGen:
263
    """Convert a stream completion response to a stream chat response."""
264

265
    async def gen() -> ChatResponseAsyncGen:
266
        async for response in completion_response_gen:
267
            yield ChatResponse(
268
                message=ChatMessage(
269
                    role=MessageRole.ASSISTANT,
270
                    content=response.text,
271
                    additional_kwargs=response.additional_kwargs,
272
                ),
273
                delta=response.delta,
274
                raw=response.raw,
275
            )
276

277
    return gen()
278

279

280
def astream_chat_response_to_completion_response(
281
    chat_response_gen: ChatResponseAsyncGen,
282
) -> CompletionResponseAsyncGen:
283
    """Convert a stream chat response to a completion response."""
284

285
    async def gen() -> CompletionResponseAsyncGen:
286
        async for response in chat_response_gen:
287
            yield CompletionResponse(
288
                text=response.message.content or "",
289
                additional_kwargs=response.message.additional_kwargs,
290
                delta=response.delta,
291
                raw=response.raw,
292
            )
293

294
    return gen()
295

296

297
def get_from_param_or_env(
298
    key: str,
299
    param: Optional[str] = None,
300
    env_key: Optional[str] = None,
301
    default: Optional[str] = None,
302
) -> str:
303
    """Get a value from a param or an environment variable."""
304
    if param is not None:
305
        return param
306
    elif env_key and env_key in os.environ and os.environ[env_key]:
307
        return os.environ[env_key]
308
    elif default is not None:
309
        return default
310
    else:
311
        raise ValueError(
312
            f"Did not find {key}, please add an environment variable"
313
            f" `{env_key}` which contains it, or pass"
314
            f"  `{key}` as a named parameter."
315
        )
316