feat(openai): add support for pdf files as media

Signed-off-by: Jan-Eric Harnack <janericharnack@googlemail.com>
This commit is contained in:
Jan-Eric Harnack
2025-05-18 23:43:38 +02:00
committed by Ilayaperumal Gopinathan
parent 148bf3376b
commit a33045871f
3 changed files with 63 additions and 15 deletions

View File

@@ -635,6 +635,10 @@ public class OpenAiChatModel implements ChatModel {
return new MediaContent(
new MediaContent.InputAudio(fromAudioData(media.getData()), MediaContent.InputAudio.Format.WAV));
}
if (MimeTypeUtils.parseMimeType("application/pdf").equals(mimeType)) {
return new MediaContent(new MediaContent.InputFile(media.getName(),
this.fromMediaData(media.getMimeType(), media.getData())));
}
else {
return new MediaContent(
new MediaContent.ImageUrl(this.fromMediaData(media.getMimeType(), media.getData())));

View File

@@ -1367,14 +1367,15 @@ public class OpenAiApi {
@JsonProperty("type") String type,
@JsonProperty("text") String text,
@JsonProperty("image_url") ImageUrl imageUrl,
@JsonProperty("input_audio") InputAudio inputAudio) { // @formatter:on
@JsonProperty("input_audio") InputAudio inputAudio,
@JsonProperty("file") InputFile inputFile) { // @formatter:on
/**
* Shortcut constructor for a text content.
* @param text The text content of the message.
*/
public MediaContent(String text) {
this("text", text, null, null);
this("text", text, null, null, null);
}
/**
@@ -1382,7 +1383,7 @@ public class OpenAiApi {
* @param imageUrl The image content of the message.
*/
public MediaContent(ImageUrl imageUrl) {
this("image_url", null, imageUrl, null);
this("image_url", null, imageUrl, null, null);
}
/**
@@ -1390,7 +1391,15 @@ public class OpenAiApi {
* @param inputAudio The audio content of the message.
*/
public MediaContent(InputAudio inputAudio) {
this("input_audio", null, null, inputAudio);
this("input_audio", null, null, inputAudio, null);
}
/**
* Shortcut constructor for a file content
* @param inputFile The file content of the message.
*/
public MediaContent(InputFile inputFile) {
this("file", null, null, null, inputFile);
}
/**
@@ -1428,6 +1437,18 @@ public class OpenAiApi {
}
/**
* Constructor for base64-encoded file
*
* @param filename name of the file
* @param fileData file data with format
* "data:{mimetype};base64,{base64-encoded-image-data}".
*/
public record InputFile(@JsonProperty("filename") String filename,
@JsonProperty("file_data") String fileData) {
}
}
/**

View File

@@ -18,6 +18,7 @@ package org.springframework.ai.openai.chat;
import java.net.MalformedURLException;
import java.net.URI;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;
@@ -29,6 +30,8 @@ import org.mockito.Captor;
import org.mockito.Mock;
import org.mockito.Mockito;
import org.mockito.junit.jupiter.MockitoExtension;
import org.springframework.core.io.ByteArrayResource;
import org.springframework.util.MimeType;
import reactor.core.publisher.Flux;
import org.springframework.ai.chat.messages.SystemMessage;
@@ -126,11 +129,8 @@ public class MessageTypeContentTests {
given(this.openAiApi.chatCompletionEntity(this.pomptCaptor.capture(), this.headersCaptor.capture()))
.willReturn(Mockito.mock(ResponseEntity.class));
URI mediaUri = URI.create("http://test");
this.chatModel.call(new Prompt(List.of(UserMessage.builder()
.text("test message")
.media(List.of(Media.builder().mimeType(MimeTypeUtils.IMAGE_JPEG).data(mediaUri).build()))
.build())));
this.chatModel
.call(new Prompt(List.of(UserMessage.builder().text("test message").media(this.buildMediaList()).build())));
validateComplexContent(this.pomptCaptor.getValue());
}
@@ -141,11 +141,10 @@ public class MessageTypeContentTests {
given(this.openAiApi.chatCompletionStream(this.pomptCaptor.capture(), this.headersCaptor.capture()))
.willReturn(this.fluxResponse);
URI mediaUrl = URI.create("http://test");
this.chatModel.stream(new Prompt(List.of(UserMessage.builder()
.text("test message")
.media(List.of(Media.builder().mimeType(MimeTypeUtils.IMAGE_JPEG).data(mediaUrl).build()))
.build()))).subscribe();
this.chatModel
.stream(new Prompt(
List.of(UserMessage.builder().text("test message").media(this.buildMediaList()).build())))
.subscribe();
validateComplexContent(this.pomptCaptor.getValue());
}
@@ -161,16 +160,40 @@ public class MessageTypeContentTests {
@SuppressWarnings({ "unused", "unchecked" })
List<Map<String, Object>> mediaContents = (List<Map<String, Object>>) userMessage.rawContent();
assertThat(mediaContents).hasSize(2);
assertThat(mediaContents).hasSize(3);
// Assert text content
Map<String, Object> textContent = mediaContents.get(0);
assertThat(textContent.get("type")).isEqualTo("text");
assertThat(textContent.get("text")).isEqualTo("test message");
// Assert image content
Map<String, Object> imageContent = mediaContents.get(1);
assertThat(imageContent.get("type")).isEqualTo("image_url");
assertThat(imageContent).containsKey("image_url");
// Assert file content
Map<String, Object> fileContent = mediaContents.get(2);
assertThat(fileContent.get("type")).isEqualTo("file");
assertThat(fileContent).containsKey("file");
assertThat(fileContent.get("file")).isInstanceOf(Map.class);
Map<String, Object> fileMap = (Map<String, Object>) fileContent.get("file");
assertThat(fileMap.get("file_data")).isEqualTo("data:application/pdf;base64,JVBERi0xLjc=");
}
private List<Media> buildMediaList() {
URI imageUri = URI.create("http://test");
Media imageMedia = Media.builder().mimeType(MimeTypeUtils.IMAGE_JPEG).data(imageUri).build();
byte[] pdfData = "%PDF-1.7".getBytes(StandardCharsets.UTF_8);
Media pdfMedia = Media.builder()
.mimeType(MimeType.valueOf("application/pdf"))
.data(new ByteArrayResource(pdfData))
.build();
return List.of(imageMedia, pdfMedia);
}
}