fi.evolver.ai.spring.provider.openai.OpenAiAssistant Maven / Gradle / Ivy

Go to download
package fi.evolver.ai.spring.provider.openai;

import java.time.Duration;
import java.util.ArrayList;
import java.util.List;

import fi.evolver.ai.spring.AssistantContentSubscriber;
import fi.evolver.ai.spring.assistant.Assistant;
import fi.evolver.ai.spring.assistant.AssistantPrompt;
import fi.evolver.ai.spring.assistant.AssistantResponse;
import fi.evolver.ai.spring.chat.prompt.Message;
import fi.evolver.ai.spring.file.AiFile;
import fi.evolver.ai.spring.provider.openai.OpenAiAssistantResponse.AssistantContent;
import fi.evolver.ai.spring.provider.openai.response.assistants.OTool;
import fi.evolver.ai.spring.provider.openai.response.files.OFile;


public class OpenAiAssistant implements Assistant {
	private final OpenAiService openAiService;
	private final String provider;
	private final AssistantPrompt prompt;
	private final Duration timeout;
	private final String id;

	private String beforeId;
	private String threadId;
	private String vectorStoreId;
	private final List fileIds = new ArrayList<>();
	private final boolean hasFileSearchTool;

	OpenAiAssistant(
			OpenAiService openAiService,
			String provider,
			AssistantPrompt prompt,
			String assistantId,
			Duration timeout) {

		this.openAiService = openAiService;
		this.provider = provider;
		this.prompt = prompt;
		this.id = assistantId;
		this.timeout = timeout;
		this.hasFileSearchTool = prompt.tools().stream()
				.map(OTool::type)
				.anyMatch("file_search"::equals);
	}

	@Override
	public AssistantResponse ask(Message message) {
		if (threadId == null)
			threadId = openAiService.createThread(provider).id();

		openAiService.createMessage(provider, threadId, message, timeout);

		NextMessagesPrompt nextMessagesPrompt = NextMessagesPrompt.builder(threadId, id)
				.setProvider(provider)
				.setBeforeId(beforeId)
				.setParameter(OpenAiRequestParameters.STREAM, prompt.getBooleanProperty(OpenAiRequestParameters.STREAM).orElse(false))
				.build();

		AssistantResponse response = openAiService.getNextThreadMessages(nextMessagesPrompt, timeout);

		if (response instanceof OpenAiAssistantResponse)
			response.addSubscriber(new AssistantResponseConsumer(((OpenAiAssistantResponse)response), this));

		return response;
	}

	@Override
	public String addFile(AiFile file) {
		OFile uploadFileResponse = openAiService.uploadFile(provider, file);
		fileIds.add(uploadFileResponse.id());
		addVectorStoreFile(uploadFileResponse.id());
		openAiService.modifyAssistant(prompt.withFilesAndStores(fileIds, vectorStoreId != null ? List.of(vectorStoreId) : List.of()), id);
		return uploadFileResponse.id();
	}

	private void addVectorStoreFile(String fileId) {
		if (!hasFileSearchTool)
			return;

		if (vectorStoreId == null)
			vectorStoreId = openAiService.createVectorStore(provider, fileIds).id();
		else
			openAiService.createVectorStoreFile(provider, vectorStoreId, fileId);
	}

	@Override
	public void deleteFile(String fileId) {
		if (fileIds.remove(fileId)) {
			openAiService.deleteFile(provider, fileId);
			if (vectorStoreId != null)
				openAiService.deleteVectorStoreFile(provider, vectorStoreId, fileId);

			AssistantPrompt updatedPrompt = prompt.withFilesAndStores(fileIds, vectorStoreId != null ? List.of(vectorStoreId) : List.of());
			openAiService.modifyAssistant(updatedPrompt, id);
		}
	}

	@Override
	public void close() {
		if (threadId != null)
			openAiService.deleteThread(provider, threadId);
		openAiService.deleteAssistant(provider, prompt.model(), id);
		for (String fileId: fileIds)
			openAiService.deleteFile(provider, fileId);
		if (vectorStoreId != null)
			openAiService.deleteVectorStore(provider, vectorStoreId);
	}

	@Override
	public String getProvider() {
		return provider;
	}

	private static class AssistantResponseConsumer implements AssistantContentSubscriber {
		private final OpenAiAssistantResponse response;
		private final OpenAiAssistant assistant;

		public AssistantResponseConsumer(OpenAiAssistantResponse response, OpenAiAssistant assistant) {
			this.response = response;
			this.assistant = assistant;
		}

		@Override
		public void onComplete(String status) {
			assistant.beforeId = response.getMessageId();
		}

		@Override
		public void onContent(AssistantContent content) { /* NOOP */ }

	}

}