package com.pablotj.ia.chat.boot.infraestructure.llm;

import de.kherud.llama.InferenceParameters;
import de.kherud.llama.LlamaOutput;
import org.springframework.stereotype.Component;

@Component
public class LlmModelClient {

    private final LlmModelLoader modelLoader;

    public LlmModelClient(LlmModelLoader modelLoader) {
        this.modelLoader = modelLoader;
    }

    public String generate(String prompt) {
        InferenceParameters inf = new InferenceParameters(prompt)
                .setNPredict(1024)
                .setTemperature(0.7f)
                .setTopP(0.9f)
                .setTopK(40)
                .setUseChatTemplate(false);

        StringBuilder sb = new StringBuilder();
        for (LlamaOutput out : modelLoader.getModel().generate(inf)) {
            sb.append(out.text);
        }
        return sb.toString().replace("<|end_of_turn|>", "");
    }
}
