/
MistralAiApi.java
800 lines (719 loc) · 28.2 KB
/
MistralAiApi.java
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
/*
* Copyright 2023 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.mistralai.api;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Consumer;
import java.util.function.Predicate;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonInclude.Include;
import com.fasterxml.jackson.annotation.JsonProperty;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.ai.retry.RetryUtils;
import org.springframework.boot.context.properties.bind.ConstructorBinding;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import org.springframework.web.client.ResponseErrorHandler;
import org.springframework.web.client.RestClient;
import org.springframework.web.reactive.function.client.WebClient;
/**
* Single-class, Java Client library for Mistral AI platform. Provides implementation for
* the <a href="https://docs.mistral.ai/api/#operation/createEmbedding">MistralAI
* Embedding API</a> and the
* <a href="https://docs.mistral.ai/api/#operation/createChatCompletion">Chat
* Completion</a> APIs.
* <p>
* Implements <b>Synchronous</b> and <b>Streaming</b> chat completion and supports latest
* <b>Function Calling</b> features.
* </p>
*
* @author Ricken Bazolo
* @author Christian Tzolov
* @since 0.8.1
*/
public class MistralAiApi {
private static final String DEFAULT_BASE_URL = "https://api.mistral.ai";
private static final Predicate<String> SSE_DONE_PREDICATE = "[DONE]"::equals;
private final RestClient restClient;
private WebClient webClient;
/**
* Create a new client api with DEFAULT_BASE_URL
* @param mistralAiApiKey Mistral api Key.
*/
public MistralAiApi(String mistralAiApiKey) {
this(DEFAULT_BASE_URL, mistralAiApiKey);
}
/**
* Create a new client api.
* @param baseUrl api base URL.
* @param mistralAiApiKey Mistral api Key.
*/
public MistralAiApi(String baseUrl, String mistralAiApiKey) {
this(baseUrl, mistralAiApiKey, RestClient.builder(), RetryUtils.DEFAULT_RESPONSE_ERROR_HANDLER);
}
/**
* Create a new client api.
* @param baseUrl api base URL.
* @param mistralAiApiKey Mistral api Key.
* @param restClientBuilder RestClient builder.
* @param responseErrorHandler Response error handler.
*/
public MistralAiApi(String baseUrl, String mistralAiApiKey, RestClient.Builder restClientBuilder,
ResponseErrorHandler responseErrorHandler) {
Consumer<HttpHeaders> jsonContentHeaders = headers -> {
headers.setBearerAuth(mistralAiApiKey);
headers.setContentType(MediaType.APPLICATION_JSON);
};
this.restClient = restClientBuilder.baseUrl(baseUrl)
.defaultHeaders(jsonContentHeaders)
.defaultStatusHandler(responseErrorHandler)
.build();
this.webClient = WebClient.builder().baseUrl(baseUrl).defaultHeaders(jsonContentHeaders).build();
}
/**
* Represents a tool the model may call. Currently, only functions are supported as a
* tool.
*
* @param type The type of the tool. Currently, only 'function' is supported.
* @param function The function definition.
*/
@JsonInclude(Include.NON_NULL)
public record FunctionTool(@JsonProperty("type") Type type, @JsonProperty("function") Function function) {
/**
* Create a tool of type 'function' and the given function definition.
* @param function function definition.
*/
@ConstructorBinding
public FunctionTool(Function function) {
this(Type.FUNCTION, function);
}
/**
* Create a tool of type 'function' and the given function definition.
*/
public enum Type {
/**
* Function tool type.
*/
@JsonProperty("function")
FUNCTION
}
/**
* Function definition.
*
* @param description A description of what the function does, used by the model
* to choose when and how to call the function.
* @param name The name of the function to be called. Must be a-z, A-Z, 0-9, or
* contain underscores and dashes, with a maximum length of 64.
* @param parameters The parameters the functions accepts, described as a JSON
* Schema object. To describe a function that accepts no parameters, provide the
* value {"type": "object", "properties": {}}.
*/
public record Function(@JsonProperty("description") String description, @JsonProperty("name") String name,
@JsonProperty("parameters") Map<String, Object> parameters) {
/**
* Create tool function definition.
* @param description tool function description.
* @param name tool function name.
* @param jsonSchema tool function schema as json.
*/
@ConstructorBinding
public Function(String description, String name, String jsonSchema) {
this(description, name, ModelOptionsUtils.jsonToMap(jsonSchema));
}
}
}
/**
* Usage statistics.
*
* @param promptTokens Number of tokens in the prompt.
* @param totalTokens Total number of tokens used in the request (prompt +
* completion).
* @param completionTokens Number of tokens in the generated completion. Only
* applicable for completion requests.
*/
@JsonInclude(Include.NON_NULL)
public record Usage(
// @formatter:off
@JsonProperty("prompt_tokens") Integer promptTokens,
@JsonProperty("total_tokens") Integer totalTokens,
@JsonProperty("completion_tokens") Integer completionTokens) {
// @formatter:on
}
/**
* Represents an embedding vector returned by embedding endpoint.
*
* @param index The index of the embedding in the list of embeddings.
* @param embedding The embedding vector, which is a list of floats. The length of
* vector depends on the model.
* @param object The object type, which is always 'embedding'.
*/
@JsonInclude(Include.NON_NULL)
public record Embedding(
// @formatter:off
@JsonProperty("index") Integer index,
@JsonProperty("embedding") List<Double> embedding,
@JsonProperty("object") String object) {
// @formatter:on
/**
* Create an embedding with the given index, embedding and object type set to
* 'embedding'.
* @param index The index of the embedding in the list of embeddings.
* @param embedding The embedding vector, which is a list of floats. The length of
* vector depends on the model.
*/
public Embedding(Integer index, List<Double> embedding) {
this(index, embedding, "embedding");
}
}
/**
* Creates an embedding vector representing the input text.
*
* @param input Input text to embed, encoded as a string or array of tokens
* @param model ID of the model to use.
* @param encodingFormat The format to return the embeddings in. Can be either float
* or base64.
*/
@JsonInclude(Include.NON_NULL)
public record EmbeddingRequest<T>(
// @formatter:off
@JsonProperty("input") T input,
@JsonProperty("model") String model,
@JsonProperty("encoding_format") String encodingFormat) {
// @formatter:on
/**
* Create an embedding request with the given input, model and encoding format set
* to float.
* @param input Input text to embed.
* @param model ID of the model to use.
*/
public EmbeddingRequest(T input, String model) {
this(input, model, "float");
}
/**
* Create an embedding request with the given input. Encoding format is set to
* float and user is null and the model is set to 'mistral-embed'.
* @param input Input text to embed.
*/
public EmbeddingRequest(T input) {
this(input, EmbeddingModel.EMBED.getValue());
}
}
/**
* List of multiple embedding responses.
*
* @param <T> Type of the entities in the data list.
* @param object Must have value "list".
* @param data List of entities.
* @param model ID of the model to use.
* @param usage Usage statistics for the completion request.
*/
@JsonInclude(Include.NON_NULL)
public record EmbeddingList<T>(
// @formatter:off
@JsonProperty("object") String object,
@JsonProperty("data") List<T> data,
@JsonProperty("model") String model,
@JsonProperty("usage") Usage usage) {
// @formatter:on
}
/**
* Creates an embedding vector representing the input text or token array.
* @param embeddingRequest The embedding request.
* @return Returns list of {@link Embedding} wrapped in {@link EmbeddingList}.
* @param <T> Type of the entity in the data list. Can be a {@link String} or
* {@link List} of tokens (e.g. Integers). For embedding multiple inputs in a single
* request, You can pass a {@link List} of {@link String} or {@link List} of
* {@link List} of tokens. For example:
*
* <pre>{@code List.of("text1", "text2", "text3") or List.of(List.of(1, 2, 3), List.of(3, 4, 5))} </pre>
*/
public <T> ResponseEntity<EmbeddingList<Embedding>> embeddings(EmbeddingRequest<T> embeddingRequest) {
Assert.notNull(embeddingRequest, "The request body can not be null.");
// Input text to embed, encoded as a string or array of tokens. To embed multiple
// inputs in a single
// request, pass an array of strings or array of token arrays.
Assert.notNull(embeddingRequest.input(), "The input can not be null.");
Assert.isTrue(embeddingRequest.input() instanceof String || embeddingRequest.input() instanceof List,
"The input must be either a String, or a List of Strings or List of List of integers.");
// The input must not an empty string, and any array must be 1024 dimensions or
// less.
if (embeddingRequest.input() instanceof List list) {
Assert.isTrue(!CollectionUtils.isEmpty(list), "The input list can not be empty.");
Assert.isTrue(list.size() <= 1024, "The list must be 1024 dimensions or less");
Assert.isTrue(
list.get(0) instanceof String || list.get(0) instanceof Integer || list.get(0) instanceof List,
"The input must be either a String, or a List of Strings or list of list of integers.");
}
return this.restClient.post()
.uri("/v1/embeddings")
.body(embeddingRequest)
.retrieve()
.toEntity(new ParameterizedTypeReference<>() {
});
}
/**
* Creates a model request for chat conversation.
*
* @param model ID of the model to use.
* @param messages The prompt(s) to generate completions for, encoded as a list of
* dict with role and content. The first prompt role should be user or system.
* @param tools A list of tools the model may call. Currently, only functions are
* supported as a tool. Use this to provide a list of functions the model may generate
* JSON inputs for.
* @param toolChoice Controls which (if any) function is called by the model. none
* means the model will not call a function and instead generates a message. auto
* means the model can pick between generating a message or calling a function. Any
* means the model must call a function.
* @param temperature What sampling temperature to use, between 0.0 and 1.0. Higher
* values like 0.8 will make the output more random, while lower values like 0.2 will
* make it more focused and deterministic. We generally recommend altering this or
* top_p but not both.
* @param topP Nucleus sampling, where the model considers the results of the tokens
* with top_p probability mass. So 0.1 means only the tokens comprising the top 10%
* probability mass are considered. We generally recommend altering this or
* temperature but not both.
* @param maxTokens The maximum number of tokens to generate in the completion. The
* token count of your prompt plus max_tokens cannot exceed the model's context
* length.
* @param stream Whether to stream back partial progress. If set, tokens will be sent
* as data-only server-sent events as they become available, with the stream
* terminated by a data: [DONE] message. Otherwise, the server will hold the request
* open until the timeout or until completion, with the response containing the full
* result as JSON.
* @param safePrompt Whether to inject a safety prompt before all conversations.
* @param randomSeed The seed to use for random sampling. If set, different calls will
* generate deterministic results.
* @param responseFormat An object specifying the format that the model must output.
* Setting to { "type": "json_object" } enables JSON mode, which guarantees the
* message the model generates is valid JSON.
*/
@JsonInclude(Include.NON_NULL)
public record ChatCompletionRequest(
// @formatter:off
@JsonProperty("model") String model,
@JsonProperty("messages") List<ChatCompletionMessage> messages,
@JsonProperty("tools") List<FunctionTool> tools,
@JsonProperty("tool_choice") ToolChoice toolChoice,
@JsonProperty("temperature") Float temperature,
@JsonProperty("top_p") Float topP,
@JsonProperty("max_tokens") Integer maxTokens,
@JsonProperty("stream") Boolean stream,
@JsonProperty("safe_prompt") Boolean safePrompt,
@JsonProperty("random_seed") Integer randomSeed,
@JsonProperty("response_format") ResponseFormat responseFormat) {
// @formatter:on
/**
* Shortcut constructor for a chat completion request with the given messages and
* model.
* @param messages The prompt(s) to generate completions for, encoded as a list of
* dict with role and content. The first prompt role should be user or system.
* @param model ID of the model to use.
*/
public ChatCompletionRequest(List<ChatCompletionMessage> messages, String model) {
this(model, messages, null, null, 0.7f, 1f, null, false, false, null, null);
}
/**
* Shortcut constructor for a chat completion request with the given messages,
* model and temperature.
* @param messages The prompt(s) to generate completions for, encoded as a list of
* dict with role and content. The first prompt role should be user or system.
* @param model ID of the model to use.
* @param temperature What sampling temperature to use, between 0.0 and 1.0.
* @param stream Whether to stream back partial progress. If set, tokens will be
* sent
*/
public ChatCompletionRequest(List<ChatCompletionMessage> messages, String model, Float temperature,
boolean stream) {
this(model, messages, null, null, temperature, 1f, null, stream, false, null, null);
}
/**
* Shortcut constructor for a chat completion request with the given messages,
* model and temperature.
* @param messages The prompt(s) to generate completions for, encoded as a list of
* dict with role and content. The first prompt role should be user or system.
* @param model ID of the model to use.
* @param temperature What sampling temperature to use, between 0.0 and 1.0.
*
*/
public ChatCompletionRequest(List<ChatCompletionMessage> messages, String model, Float temperature) {
this(model, messages, null, null, temperature, 1f, null, false, false, null, null);
}
/**
* Shortcut constructor for a chat completion request with the given messages,
* model, tools and tool choice. Streaming is set to false, temperature to 0.8 and
* all other parameters are null.
* @param messages A list of messages comprising the conversation so far.
* @param model ID of the model to use.
* @param tools A list of tools the model may call. Currently, only functions are
* supported as a tool.
* @param toolChoice Controls which (if any) function is called by the model.
*/
public ChatCompletionRequest(List<ChatCompletionMessage> messages, String model, List<FunctionTool> tools,
ToolChoice toolChoice) {
this(model, messages, tools, toolChoice, null, 1f, null, false, false, null, null);
}
/**
* Shortcut constructor for a chat completion request with the given messages and
* stream.
*/
public ChatCompletionRequest(List<ChatCompletionMessage> messages, Boolean stream) {
this(null, messages, null, null, 0.7f, 1f, null, stream, false, null, null);
}
/**
* Specifies a tool the model should use. Use to force the model to call a
* specific function.
*
*/
public enum ToolChoice {
// @formatter:off
@JsonProperty("auto") AUTO,
@JsonProperty("any") ANY,
@JsonProperty("none") NONE
// @formatter:on
}
/**
* An object specifying the format that the model must output.
*
* @param type Must be one of 'text' or 'json_object'.
*/
@JsonInclude(Include.NON_NULL)
public record ResponseFormat(@JsonProperty("type") String type) {
}
}
/**
* Message comprising the conversation.
*
* @param content The contents of the message.
* @param role The role of the messages author. Could be one of the {@link Role}
* types.
* @param toolCalls The tool calls generated by the model, such as function calls.
* Applicable only for {@link Role#ASSISTANT} role and null otherwise.
*/
@JsonInclude(Include.NON_NULL)
public record ChatCompletionMessage(
// @formatter:off
@JsonProperty("content") String content,
@JsonProperty("role") Role role,
@JsonProperty("name") String name,
@JsonProperty("tool_calls") List<ToolCall> toolCalls) {
// @formatter:on
/**
* Create a chat completion message with the given content and role. All other
* fields are null.
* @param content The contents of the message.
* @param role The role of the author of this message.
*/
public ChatCompletionMessage(String content, Role role) {
this(content, role, null, null);
}
/**
* The role of the author of this message.
*
* NOTE: Mistral expects the system message to be before the user message or will
* fail with 400 error.
*/
public enum Role {
// @formatter:off
@JsonProperty("system") SYSTEM,
@JsonProperty("user") USER,
@JsonProperty("assistant") ASSISTANT,
@JsonProperty("tool") TOOL
// @formatter:on
}
/**
* The relevant tool call.
*
* @param id The ID of the tool call. This ID must be referenced when you submit
* the tool outputs in using the Submit tool outputs to run endpoint.
* @param type The type of tool call the output is required for. For now, this is
* always function.
* @param function The function definition.
*/
@JsonInclude(Include.NON_NULL)
public record ToolCall(@JsonProperty("id") String id, @JsonProperty("type") String type,
@JsonProperty("function") ChatCompletionFunction function) {
}
/**
* The function definition.
*
* @param name The name of the function.
* @param arguments The arguments that the model expects you to pass to the
* function.
*/
@JsonInclude(Include.NON_NULL)
public record ChatCompletionFunction(@JsonProperty("name") String name,
@JsonProperty("arguments") String arguments) {
}
}
/**
* The reason the model stopped generating tokens.
*/
public enum ChatCompletionFinishReason {
// @formatter:off
/**
* The model hit a natural stop point or a provided stop sequence.
*/
@JsonProperty("stop") STOP,
/**
* The maximum number of tokens specified in the request was reached.
*/
@JsonProperty("length") LENGTH,
/**
* The content was omitted due to a flag from our content filters.
*/
@JsonProperty("model_length") MODEL_LENGTH,
/**
*
*/
@JsonProperty("error") ERROR,
/**
* The model requested a tool call.
*/
@JsonProperty("tool_calls") TOOL_CALLS
// @formatter:on
}
/**
* Represents a chat completion response returned by model, based on the provided
* input.
*
* @param id A unique identifier for the chat completion.
* @param object The object type, which is always chat.completion.
* @param created The Unix timestamp (in seconds) of when the chat completion was
* created.
* @param model The model used for the chat completion.
* @param choices A list of chat completion choices.
* @param usage Usage statistics for the completion request.
*/
@JsonInclude(Include.NON_NULL)
public record ChatCompletion(
// @formatter:off
@JsonProperty("id") String id,
@JsonProperty("object") String object,
@JsonProperty("created") Long created,
@JsonProperty("model") String model,
@JsonProperty("choices") List<Choice> choices,
@JsonProperty("usage") Usage usage) {
// @formatter:on
/**
* Chat completion choice.
*
* @param index The index of the choice in the list of choices.
* @param message A chat completion message generated by the model.
* @param finishReason The reason the model stopped generating tokens.
* @param logprobs Log probability information for the choice.
*/
@JsonInclude(Include.NON_NULL)
public record Choice(
// @formatter:off
@JsonProperty("index") Integer index,
@JsonProperty("message") ChatCompletionMessage message,
@JsonProperty("finish_reason") ChatCompletionFinishReason finishReason,
@JsonProperty("logprobs") LogProbs logprobs) {
// @formatter:on
}
}
/**
*
* Log probability information for the choice. anticipation of future changes.
*
* @param content A list of message content tokens with log probability information.
*/
@JsonInclude(Include.NON_NULL)
public record LogProbs(@JsonProperty("content") List<Content> content) {
/**
* Message content tokens with log probability information.
*
* @param token The token.
* @param logprob The log probability of the token.
* @param probBytes A list of integers representing the UTF-8 bytes representation
* of the token. Useful in instances where characters are represented by multiple
* tokens and their byte representations must be combined to generate the correct
* text representation. Can be null if there is no bytes representation for the
* token.
* @param topLogprobs List of the most likely tokens and their log probability, at
* this token position. In rare cases, there may be fewer than the number of
* requested top_logprobs returned.
*/
@JsonInclude(Include.NON_NULL)
public record Content(@JsonProperty("token") String token, @JsonProperty("logprob") Float logprob,
@JsonProperty("bytes") List<Integer> probBytes,
@JsonProperty("top_logprobs") List<TopLogProbs> topLogprobs) {
/**
* The most likely tokens and their log probability, at this token position.
*
* @param token The token.
* @param logprob The log probability of the token.
* @param probBytes A list of integers representing the UTF-8 bytes
* representation of the token. Useful in instances where characters are
* represented by multiple tokens and their byte representations must be
* combined to generate the correct text representation. Can be null if there
* is no bytes representation for the token.
*/
@JsonInclude(Include.NON_NULL)
public record TopLogProbs(@JsonProperty("token") String token, @JsonProperty("logprob") Float logprob,
@JsonProperty("bytes") List<Integer> probBytes) {
}
}
}
/**
* Represents a streamed chunk of a chat completion response returned by model, based
* on the provided input.
*
* @param id A unique identifier for the chat completion. Each chunk has the same ID.
* @param object The object type, which is always 'chat.completion.chunk'.
* @param created The Unix timestamp (in seconds) of when the chat completion was
* created. Each chunk has the same timestamp.
* @param model The model used for the chat completion.
* @param choices A list of chat completion choices. Can be more than one if n is
* greater than 1.
*/
@JsonInclude(Include.NON_NULL)
public record ChatCompletionChunk(
// @formatter:off
@JsonProperty("id") String id,
@JsonProperty("object") String object,
@JsonProperty("created") Long created,
@JsonProperty("model") String model,
@JsonProperty("choices") List<ChunkChoice> choices) {
// @formatter:on
/**
* Chat completion choice.
*
* @param index The index of the choice in the list of choices.
* @param delta A chat completion delta generated by streamed model responses.
* @param finishReason The reason the model stopped generating tokens.
* @param logprobs Log probability information for the choice.
*/
@JsonInclude(Include.NON_NULL)
public record ChunkChoice(
// @formatter:off
@JsonProperty("index") Integer index,
@JsonProperty("delta") ChatCompletionMessage delta,
@JsonProperty("finish_reason") ChatCompletionFinishReason finishReason,
@JsonProperty("logprobs") LogProbs logprobs) {
// @formatter:on
}
}
/**
* List of well-known Mistral chat models.
* https://docs.mistral.ai/platform/endpoints/#mistral-ai-generative-models
*
* <p>
* Mistral AI provides five API endpoints featuring five leading Large Language
* Models:
* </p>
* <ul>
* <li><b>TINY</b> - open-mistral-7b (aka mistral-tiny-2312)</li>
* <li><b>MIXTRAL</b> - open-mixtral-8x7b (aka mistral-small-2312)</li>
* <li><b>SMALL_LATEST</b> - mistral-small-latest (aka mistral-small-2402)</li>
* <li><b>MEDIUM</b> - mistral-medium-latest (aka mistral-medium-2312)</li>
* <li><b>LARGE</b> - mistral-large-latest (aka mistral-large-2402)</li>
* </ul>
*/
public enum ChatModel {
// @formatter:off
TINY("open-mistral-7b"),
MIXTRAL("open-mixtral-8x7b"),
SMALL("mistral-small-latest"),
MEDIUM("mistral-medium-latest"),
LARGE("mistral-large-latest");
// @formatter:on
private final String value;
ChatModel(String value) {
this.value = value;
}
public String getValue() {
return this.value;
}
}
/**
* List of well-known Mistral embedding models.
* https://docs.mistral.ai/platform/endpoints/#mistral-ai-embedding-model
*/
public enum EmbeddingModel {
// @formatter:off
@JsonProperty("mistral-embed") EMBED("mistral-embed");
// @formatter:on
private final String value;
EmbeddingModel(String value) {
this.value = value;
}
public String getValue() {
return this.value;
}
}
/**
* Creates a model response for the given chat conversation.
* @param chatRequest The chat completion request.
* @return Entity response with {@link ChatCompletion} as a body and HTTP status code
* and headers.
*/
public ResponseEntity<ChatCompletion> chatCompletionEntity(ChatCompletionRequest chatRequest) {
Assert.notNull(chatRequest, "The request body can not be null.");
Assert.isTrue(!chatRequest.stream(), "Request must set the steam property to false.");
return this.restClient.post()
.uri("/v1/chat/completions")
.body(chatRequest)
.retrieve()
.toEntity(ChatCompletion.class);
}
private MistralAiStreamFunctionCallingHelper chunkMerger = new MistralAiStreamFunctionCallingHelper();
/**
* Creates a streaming chat response for the given chat conversation.
* @param chatRequest The chat completion request. Must have the stream property set
* to true.
* @return Returns a {@link Flux} stream from chat completion chunks.
*/
public Flux<ChatCompletionChunk> chatCompletionStream(ChatCompletionRequest chatRequest) {
Assert.notNull(chatRequest, "The request body can not be null.");
Assert.isTrue(chatRequest.stream(), "Request must set the steam property to true.");
AtomicBoolean isInsideTool = new AtomicBoolean(false);
return this.webClient.post()
.uri("/v1/chat/completions")
.body(Mono.just(chatRequest), ChatCompletionRequest.class)
.retrieve()
.bodyToFlux(String.class)
.takeUntil(SSE_DONE_PREDICATE)
.filter(SSE_DONE_PREDICATE.negate())
.map(content -> ModelOptionsUtils.jsonToObject(content, ChatCompletionChunk.class))
.map(chunk -> {
if (this.chunkMerger.isStreamingToolFunctionCall(chunk)) {
isInsideTool.set(true);
}
return chunk;
})
.windowUntil(chunk -> {
if (isInsideTool.get() && this.chunkMerger.isStreamingToolFunctionCallFinish(chunk)) {
isInsideTool.set(false);
return true;
}
return !isInsideTool.get();
})
.concatMapIterable(window -> {
Mono<ChatCompletionChunk> mono1 = window.reduce(new ChatCompletionChunk(null, null, null, null, null),
(previous, current) -> this.chunkMerger.merge(previous, current));
return List.of(mono1);
})
.flatMap(mono -> mono);
}
}