lm-lib first commit;

git-svn-id: svn://home0/lm-lib@3603 71dcc7fd-b1f2-4e5d-bb8e-274146e81173

src/main/java/de/extio/lmlib/client/textcompletion/ModelNameLoader.java

@@ -0,0 +1,65 @@
+package de.extio.lmlib.client.textcompletion;
+
+import java.time.Duration;
+import java.util.function.Consumer;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.beans.factory.annotation.Qualifier;
+import org.springframework.http.HttpMethod;
+import org.springframework.web.reactive.function.client.WebClient;
+
+import de.extio.lmlib.client.textcompletion.dto.Model;
+import de.extio.lmlib.client.textcompletion.dto.ModelsResponse;
+
+import reactor.util.retry.Retry;
+
+class ModelNameLoader implements Runnable {
+	
+	private static final Logger LOGGER = LoggerFactory.getLogger(ModelNameLoader.class);
+	
+	private final WebClient webClient;
+	
+	private final Consumer<String> consumer;
+	
+	public ModelNameLoader(@Qualifier("textCompletionWebClientBuilder") final WebClient webClient, final Consumer<String> consumer) {
+		this.consumer = consumer;
+		this.webClient = webClient;
+	}
+	
+	@Override
+	public void run() {
+		String modelName = null;
+		do {
+			try {
+				LOGGER.info("Querying model name");
+				
+				modelName = this.webClient
+						.method(HttpMethod.GET)
+						.uri(uriBuilder -> uriBuilder.path("/v1/models").build())
+						.retrieve()
+						.bodyToMono(ModelsResponse.class)
+						.retryWhen(Retry.backoff(99, Duration.ofSeconds(1)).maxBackoff(Duration.ofSeconds(10)).jitter(0.25d).doAfterRetry(rs -> LOGGER.warn("Request failed", rs.failure())))
+						.block()
+						.getData()
+						.stream()
+						.map(Model::getId)
+						.findFirst()
+						.orElseThrow();
+			}
+			catch (final Exception e) {
+				LOGGER.error(e.getMessage(), e);
+				try {
+					Thread.sleep(1000l);
+				}
+				catch (final InterruptedException e1) {
+					throw new RuntimeException("Interrupted", e1);
+				}
+			}
+		} while (modelName == null);
+		
+		LOGGER.info("Model name: {}", modelName);
+		this.consumer.accept(modelName);
+	}
+	
+}

src/main/java/de/extio/lmlib/client/textcompletion/ModelNameSupplier.java

@@ -0,0 +1,61 @@
+package de.extio.lmlib.client.textcompletion;
+
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CountDownLatch;
+
+import org.springframework.beans.factory.InitializingBean;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.beans.factory.annotation.Qualifier;
+import org.springframework.stereotype.Component;
+import org.springframework.web.reactive.function.client.WebClient;
+
+import de.extio.lmlib.client.profile.ModelProfile.ModelProvider;
+import de.extio.lmlib.client.profile.ModelProfileService;
+
+@Component
+class ModelNameSupplier implements InitializingBean {
+	
+	@Autowired
+	@Qualifier("lmLibWebClientBuilder")
+	private WebClient.Builder webClientBuilder;
+	
+	private final Map<String, CompletableFuture<String>> modelNames = new ConcurrentHashMap<>();
+	
+	private final Map<String, CountDownLatch> countDownLatches = new ConcurrentHashMap<>();
+	
+	@Autowired
+	private ModelProfileService modelProfileService;
+	
+	@Override
+	public void afterPropertiesSet() throws Exception {
+		final var urls = this.modelProfileService.getModelProfileUrls(ModelProvider.OAI_TEXT_COMPLETION);
+		for (final var url : urls) {
+			this.countDownLatches.put(url, new CountDownLatch(1));
+			this.fetchModelName(url);
+		}
+	}
+	
+	private void fetchModelName(final String url) {
+		final var webClient = this.webClientBuilder.baseUrl(url).build();
+		CompletableFuture.runAsync(() -> {
+			final var modelNameLoader = new ModelNameLoader(webClient, modelName -> {
+				this.modelNames.put(url, CompletableFuture.completedFuture(modelName));
+				this.countDownLatches.get(url).countDown();
+			});
+			modelNameLoader.run();
+		});
+	}
+	
+	public String getModelName(final String url) {
+		try {
+			this.countDownLatches.get(url).await();
+		}
+		catch (final InterruptedException e) {
+			throw new RuntimeException("Interrupted", e);
+		}
+		return this.modelNames.get(url).join();
+	}
+	
+}

src/main/java/de/extio/lmlib/client/textcompletion/TextCompletionClient.java

@@ -0,0 +1,304 @@
+package de.extio.lmlib.client.textcompletion;
+
+import java.time.Duration;
+import java.time.LocalDateTime;
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.Deque;
+import java.util.List;
+import java.util.stream.Collectors;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.beans.factory.DisposableBean;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.beans.factory.annotation.Qualifier;
+import org.springframework.beans.factory.annotation.Value;
+import org.springframework.http.HttpMethod;
+import org.springframework.stereotype.Component;
+import org.springframework.web.reactive.function.client.WebClient;
+
+import de.extio.lmlib.client.Client;
+import de.extio.lmlib.client.Completion;
+import de.extio.lmlib.client.Conversation;
+import de.extio.lmlib.client.Conversation.Turn;
+import de.extio.lmlib.client.profile.ModelCategory;
+import de.extio.lmlib.client.profile.ModelProfile;
+import de.extio.lmlib.client.profile.ModelProfile.ModelProvider;
+import de.extio.lmlib.client.profile.ModelProfileService;
+import de.extio.lmlib.client.prompt.PromptStrategy;
+import de.extio.lmlib.client.prompt.PromptStrategyFactory;
+import de.extio.lmlib.client.textcompletion.dto.Choice;
+import de.extio.lmlib.client.textcompletion.dto.CompletionRequest;
+import de.extio.lmlib.client.textcompletion.dto.CompletionResponse;
+import de.extio.lmlib.client.token.Tokenizer;
+
+import reactor.util.retry.Retry;
+
+@Component
+public class TextCompletionClient implements Client, DisposableBean {
+	
+	private static final Logger LOGGER = LoggerFactory.getLogger(TextCompletionClient.class);
+	
+	@Autowired
+	@Qualifier("lmLibWebClientBuilder")
+	private WebClient.Builder webClientBuilder;
+	
+	@Autowired
+	private PromptStrategyFactory promptStrategyFactory;
+	
+	@Autowired
+	private ModelNameSupplier modelNameSupplier;
+	
+	@Autowired
+	private Tokenizer tokenizer;
+	
+	@Autowired
+	private ModelProfileService modelProfileService;
+	
+	@Value("${client.collectStatistics:false}")
+	private boolean collectStatistics;
+	
+	private final CompletionStatistics totalStatistics = new CompletionStatistics();
+	
+	@Override
+	public ModelProvider getModelProvider() {
+		return ModelProvider.OAI_TEXT_COMPLETION;
+	}
+	
+	@Override
+	public Completion completion(final ModelCategory modelCategory, final String system, final String question, final String fullText) {
+		String text = "";
+		if (!question.isEmpty() && !fullText.isEmpty()) {
+			text = question + "\n" + fullText;
+		}
+		else if (!question.isEmpty()) {
+			text = question;
+		}
+		else if (!fullText.isEmpty()) {
+			text = fullText;
+		}
+		else {
+			throw new IllegalArgumentException("Either question or fullText must be set");
+		}
+		
+		return this.conversation(modelCategory, Conversation.create(system, text));
+	}
+	
+	@Override
+	public Completion conversation(ModelCategory modelCategory, final Conversation conversation) {
+		if (modelCategory == null) {
+			modelCategory = ModelCategory.MEDIUM;
+		}
+		final var modelProfile = this.modelProfileService.getModelProfile(modelCategory.getModelProfile());
+		if (modelProfile == null || modelProfile.modelProvider() != ModelProvider.OAI_TEXT_COMPLETION) {
+			throw new IllegalArgumentException("ModelProfile " + modelCategory.getModelProfile());
+		}
+		
+		final PromptStrategy promptStrategy = this.promptStrategyFactory.getStrategy(modelProfile.prompt());
+		if (promptStrategy == null) {
+			throw new IllegalArgumentException("Prompt strategy not found: " + modelProfile.prompt());
+		}
+		
+		final List<StringBuilder> prompts = this.createPrompts(conversation, modelCategory, modelProfile, promptStrategy);
+		
+		final List<StringBuilder> answers = new ArrayList<>();
+		final CompletionStatistics statistics = new CompletionStatistics();
+		for (final StringBuilder prompt : prompts) {
+			final StringBuilder answer = new StringBuilder();
+			this.requestCompletionsContinuations(prompt, answer, statistics, modelProfile, promptStrategy);
+			answers.add(answer);
+		}
+		
+		if (answers.size() > 1) {
+			this.summarize(answers, modelProfile, promptStrategy, statistics);
+		}
+		
+		return new Completion(
+				answers.stream().collect(Collectors.joining()),
+				statistics.requests,
+				statistics.duration,
+				statistics.inTokens,
+				statistics.outTokens,
+				false);
+	}
+	
+	private List<StringBuilder> createPrompts(final Conversation conversation, final ModelCategory modelCategory, final ModelProfile modelProfile, final PromptStrategy promptStrategy) {
+		StringBuilder prompt = this.formatConversation(conversation, promptStrategy);
+		
+		List<Long> tokenized = this.tokenizer.tokenize(prompt.toString(), modelProfile);
+		final int maxLength = modelProfile.maxContextLength() - modelProfile.maxTokens() - 15; // 15 is a margin for EOT and other special tokens
+		if (tokenized.size() > maxLength) {
+			// TODO: Cut smarter, e.g. split conversation into parts while keeping system prompt
+			LOGGER.warn("Prompt too long: {} tokens / {} max. Prompt will be cut!", tokenized.size(), maxLength);
+			tokenized = tokenized.subList(0, maxLength);
+			prompt = new StringBuilder(this.tokenizer.detokenize(tokenized, modelProfile));
+		}
+		
+		LOGGER.info("Completion request for {}. Input tokens: {}", modelCategory, tokenized.size());
+		
+		return List.of(prompt);
+	}
+	
+	private StringBuilder formatConversation(final Conversation conversation, final PromptStrategy promptStrategy) {
+		final StringBuilder prompt = new StringBuilder();
+		
+		final Deque<Turn> turns = new ArrayDeque<>(conversation.getConversation());
+		Turn systemTurn = null;
+		Turn userTurn = null;
+		Turn assistantTurn = null;
+		boolean first = true;
+		
+		while (!turns.isEmpty()) {
+			final Turn turn = turns.pop();
+			switch (turn.type()) {
+				case SYSTEM:
+					systemTurn = turn;
+					break;
+				case USER:
+					userTurn = turn;
+					break;
+				case ASSISTANT:
+					assistantTurn = turn;
+					break;
+				default:
+					throw new IllegalArgumentException("Unknown turn type: " + turn.type());
+			}
+			
+			if (first && userTurn != null) {
+				prompt.append(promptStrategy.start(systemTurn != null ? systemTurn.text() : "", "", userTurn.text()));
+				userTurn = null;
+				first = false;
+			}
+			else if (!first && assistantTurn != null && userTurn != null) {
+				promptStrategy.next(prompt, assistantTurn.text(), userTurn.text());
+				userTurn = null;
+				assistantTurn = null;
+			}
+		}
+		
+		return prompt;
+	}
+	
+	private void requestCompletionsContinuations(final StringBuilder prompt, final StringBuilder answer, final CompletionStatistics statistics,
+			final ModelProfile modelProfile, final PromptStrategy promptStrategy) {
+		for (int continuation = 0; continuation < modelProfile.maxContinuations(); continuation++) {
+			final LocalDateTime start = LocalDateTime.now();
+			
+			final String promptStr = prompt.toString();
+			
+			final var response = this.requestCompletions(promptStr, modelProfile);
+			
+			String content = null;
+			var complete = false;
+			if (response.getContent() != null) {
+				content = promptStrategy.removeEOT(response.getContent());
+				complete = response.isStoppedEos();
+			}
+			else {
+				final var choice = response.getChoices().getFirst();
+				content = promptStrategy.removeEOT(choice.getText());
+				complete = !Choice.FINISH_REASON_LENGTH.equals(choice.getFinishReason());
+			}
+			
+			if (this.collectStatistics) {
+				final var dur = Duration.between(start, LocalDateTime.now());
+				final var in = this.tokenizer.count(promptStr, modelProfile);
+				final var out = this.tokenizer.count(content, modelProfile);
+				LOGGER.debug("Request duration: " + dur + "; in tokens: " + in + "; out tokens: " + out);
+				statistics.add(dur, in, out);
+				this.totalStatistics.add(dur, in, out);
+			}
+			
+			answer.append(content);
+			
+			if (complete) {
+				LOGGER.debug("Prompt response complete");
+				break;
+			}
+			
+			answer.append('\n');
+			promptStrategy.continue_(prompt, content);
+		}
+	}
+	
+	private CompletionResponse requestCompletions(final String prompt, final ModelProfile modelProfile) {
+		LOGGER.debug("Requesting completion at {}", modelProfile.url());
+		
+		final var request = new CompletionRequest();
+		request.setModel(this.modelNameSupplier.getModelName(modelProfile.url()));
+		request.setPrompt(prompt);
+		request.setMaxTokens(modelProfile.maxTokens());
+		request.setTemperature(modelProfile.temperature());
+		request.setTopP(modelProfile.topP());
+		request.setStream(false);
+		
+		final var webClient = this.webClientBuilder.baseUrl(modelProfile.url()).build();
+		return webClient
+				.method(HttpMethod.POST)
+				.uri(uriBuilder -> uriBuilder.path("/v1/completions").build())
+				.bodyValue(request)
+				.retrieve()
+				.bodyToMono(CompletionResponse.class)
+				.retryWhen(Retry.backoff(99, Duration.ofSeconds(1)).maxBackoff(Duration.ofSeconds(10)).jitter(0.25d).doAfterRetry(rs -> LOGGER.warn("Request failed", rs.failure())))
+				.block();
+	}
+	
+	private void summarize(final List<StringBuilder> answers, final ModelProfile modelProfile, final PromptStrategy promptStrategy, final CompletionStatistics statistics) {
+		// TODO: Reactivate when multiple (split) prompts have been implemented again
+		//		LOGGER.debug("Summarizing answers");
+		//		
+		//		final StringBuilder summary = new StringBuilder();
+		//		
+		//		final StringBuilder prompt = promptStrategy.start(instruction, question, "");
+		//		answers.forEach(answer -> promptStrategy.next(prompt, answer.toString(), "Continue"));
+		//		promptStrategy.next(prompt, "", "Generate now a full summary");
+		//		
+		//		this.requestCompletionsContinuations(prompt, summary, statistics, modelProfile, promptStrategy);
+		//		
+		//		answers.clear();
+		//		answers.add(summary);
+	}
+	
+	@Override
+	public void destroy() throws Exception {
+		if (this.collectStatistics) {
+			LOGGER.info("Client Statistics: {}", this.totalStatistics);
+		}
+	}
+	
+	static class CompletionStatistics {
+		
+		volatile int requests;
+		
+		volatile Duration duration = Duration.ofMillis(0l);
+		
+		volatile long inTokens;
+		
+		volatile long outTokens;
+		
+		synchronized void add(final Duration duration, final int inTokens, final int outTokens) {
+			this.requests++;
+			this.duration = this.duration.plus(duration);
+			this.inTokens += inTokens;
+			this.outTokens += outTokens;
+		}
+		
+		@Override
+		public String toString() {
+			final StringBuilder builder = new StringBuilder();
+			builder.append("CompletionStatistics [requests=");
+			builder.append(this.requests);
+			builder.append(", duration=");
+			builder.append(this.duration);
+			builder.append(", inTokens=");
+			builder.append(this.inTokens);
+			builder.append(", outTokens=");
+			builder.append(this.outTokens);
+			builder.append("]");
+			return builder.toString();
+		}
+		
+	}
+	
+}

src/main/java/de/extio/lmlib/client/textcompletion/dto/Choice.java

@@ -0,0 +1,59 @@
+package de.extio.lmlib.client.textcompletion.dto;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+public final class Choice {
+	
+	public static final String FINISH_REASON_STOP = "stop";
+	
+	public static final String FINISH_REASON_LENGTH = "length";
+	
+	public static final String FINISH_REASON_CONTENT_FILTER = "content_filter";
+	
+	private String text;
+	
+	private Double score;
+	
+	private String finishReason;
+	
+	@JsonProperty("text")
+	public String getText() {
+		return this.text;
+	}
+	
+	public void setText(final String text) {
+		this.text = text;
+	}
+	
+	@JsonProperty("score")
+	public Double getScore() {
+		return this.score;
+	}
+	
+	public void setScore(final Double score) {
+		this.score = score;
+	}
+	
+	@JsonProperty("finish_reason")
+	public String getFinishReason() {
+		return this.finishReason;
+	}
+	
+	public void setFinishReason(final String finishReason) {
+		this.finishReason = finishReason;
+	}
+	
+	@Override
+	public String toString() {
+		final StringBuilder builder = new StringBuilder();
+		builder.append("Choice [text=");
+		builder.append(this.text);
+		builder.append(", score=");
+		builder.append(this.score);
+		builder.append(", finishReason=");
+		builder.append(this.finishReason);
+		builder.append("]");
+		return builder.toString();
+	}
+	
+}

src/main/java/de/extio/lmlib/client/textcompletion/dto/CompletionRequest.java

@@ -0,0 +1,72 @@
+package de.extio.lmlib.client.textcompletion.dto;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+public final class CompletionRequest {
+	
+	private String model;
+	
+	private String prompt;
+	
+	private Integer maxTokens;
+	
+	private Double temperature;
+	
+	private Double topP;
+	
+	private Boolean stream;
+	
+	@JsonProperty("model")
+	public String getModel() {
+		return this.model;
+	}
+	
+	public void setModel(final String model) {
+		this.model = model;
+	}
+	
+	@JsonProperty("prompt")
+	public String getPrompt() {
+		return this.prompt;
+	}
+	
+	public void setPrompt(final String prompt) {
+		this.prompt = prompt;
+	}
+	
+	@JsonProperty("max_tokens")
+	public Integer getMaxTokens() {
+		return this.maxTokens;
+	}
+	
+	public void setMaxTokens(final Integer maxTokens) {
+		this.maxTokens = maxTokens;
+	}
+	
+	@JsonProperty("temperature")
+	public Double getTemperature() {
+		return this.temperature;
+	}
+	
+	public void setTemperature(final Double temperature) {
+		this.temperature = temperature;
+	}
+	
+	@JsonProperty("top_p")
+	public Double getTopP() {
+		return this.topP;
+	}
+	
+	public void setTopP(final Double topP) {
+		this.topP = topP;
+	}
+	
+	@JsonProperty("stream")
+	public Boolean isStream() {
+		return this.stream;
+	}
+	
+	public void setStream(final Boolean stream) {
+		this.stream = stream;
+	}
+}

src/main/java/de/extio/lmlib/client/textcompletion/dto/CompletionResponse.java

@@ -0,0 +1,93 @@
+package de.extio.lmlib.client.textcompletion.dto;
+
+import java.util.List;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+public final class CompletionResponse {
+	
+	private String id;
+	
+	private String object;
+	
+	private Integer created;
+	
+	private List<Choice> choices;
+	
+	private String content;
+	
+	private boolean isStoppedEos;
+	
+	@JsonProperty("id")
+	public String getId() {
+		return this.id;
+	}
+	
+	public void setId(final String id) {
+		this.id = id;
+	}
+	
+	@JsonProperty("object")
+	public String getObject() {
+		return this.object;
+	}
+	
+	public void setObject(final String object) {
+		this.object = object;
+	}
+	
+	@JsonProperty("created")
+	public Integer getCreated() {
+		return this.created;
+	}
+	
+	public void setCreated(final Integer created) {
+		this.created = created;
+	}
+	
+	@JsonProperty("choices")
+	public List<Choice> getChoices() {
+		return this.choices;
+	}
+	
+	public void setChoices(final List<Choice> choices) {
+		this.choices = choices;
+	}
+	
+	@JsonProperty("content")
+	public String getContent() {
+		return this.content;
+	}
+	
+	public void setContent(final String content) {
+		this.content = content;
+	}
+	
+	@JsonProperty("stopped_eos")
+	public boolean isStoppedEos() {
+		return this.isStoppedEos;
+	}
+	
+	public void setStoppedEos(final boolean isStoppedEos) {
+		this.isStoppedEos = isStoppedEos;
+	}
+	
+	@Override
+	public String toString() {
+		final StringBuilder builder = new StringBuilder();
+		builder.append("CompletionResponse [id=");
+		builder.append(this.id);
+		builder.append(", object=");
+		builder.append(this.object);
+		builder.append(", created=");
+		builder.append(this.created);
+		builder.append(", choices=");
+		builder.append(this.choices);
+		builder.append(", content=");
+		builder.append(this.content);
+		builder.append(", isStoppedEos=");
+		builder.append(this.isStoppedEos);
+		builder.append("]");
+		return builder.toString();
+	}
+}

src/main/java/de/extio/lmlib/client/textcompletion/dto/Model.java

@@ -0,0 +1,74 @@
+package de.extio.lmlib.client.textcompletion.dto;
+
+import java.util.List;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+public final class Model {
+	
+	private String id;
+	
+	private String object;
+	
+	private int created;
+	
+	private String ownedBy;
+	
+	private List<String> permission;
+	
+	private String root;
+	
+	@JsonProperty("id")
+	public String getId() {
+		return this.id;
+	}
+	
+	public void setId(final String id) {
+		this.id = id;
+	}
+	
+	@JsonProperty("object")
+	public String getObject() {
+		return this.object;
+	}
+	
+	public void setObject(final String object) {
+		this.object = object;
+	}
+	
+	@JsonProperty("created")
+	public int getCreated() {
+		return this.created;
+	}
+	
+	public void setCreated(final int created) {
+		this.created = created;
+	}
+	
+	@JsonProperty("owned_by")
+	public String getOwned_by() {
+		return this.ownedBy;
+	}
+	
+	public void setOwned_by(final String owned_by) {
+		this.ownedBy = owned_by;
+	}
+	
+	@JsonProperty("permission")
+	public List<String> getPermission() {
+		return this.permission;
+	}
+	
+	public void setPermission(final List<String> permission) {
+		this.permission = permission;
+	}
+	
+	@JsonProperty("root")
+	public String getRoot() {
+		return this.root;
+	}
+	
+	public void setRoot(final String root) {
+		this.root = root;
+	}
+}

src/main/java/de/extio/lmlib/client/textcompletion/dto/ModelsResponse.java

@@ -0,0 +1,42 @@
+package de.extio.lmlib.client.textcompletion.dto;
+
+import java.util.List;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+public final class ModelsResponse {
+	
+	private String object;
+	
+	private List<Model> data;
+	
+	@JsonProperty("object")
+	public String getObject() {
+		return this.object;
+	}
+	
+	public void setObject(final String object) {
+		this.object = object;
+	}
+	
+	@JsonProperty("data")
+	public List<Model> getData() {
+		return this.data;
+	}
+	
+	public void setData(final List<Model> data) {
+		this.data = data;
+	}
+	
+	@Override
+	public String toString() {
+		final StringBuilder builder = new StringBuilder();
+		builder.append("ModelsResponse [object=");
+		builder.append(this.object);
+		builder.append(", data=");
+		builder.append(this.data);
+		builder.append("]");
+		return builder.toString();
+	}
+	
+}