import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

@Slf4j
@RestController
@RequestMapping("/base")
public class OllamaModelController {

    @Resource
    private OllamaChatModel ollamaChatModel;

    /**
     * 直接输出全部
     *
     * @return
     */
    @RequestMapping(value = "/chat/simple")
    public String simpleChat() {
        return ollamaChatModel.call("你是谁");
    }

    /**
     * 流式输出
     *
     * @return
     */
    @RequestMapping(value = "/chat/stream", produces = "text/stream;charset=UTF-8")
    public Flux<String> chatStream() {

        return ollamaChatModel.stream("你是谁");
    }

    /**
     * 流式输出
     *
     * @return
     */
    @RequestMapping(value = "/chat/stream2", produces = "text/stream;charset=UTF-8")
    public Flux<String> chatStream2() {
        Flux<ChatResponse> responseEntityFlux = ollamaChatModel.stream(new Prompt("武汉今天的天气?"));
        return responseEntityFlux.map(item -> item.getResult().getOutput().getText());
    }
}