@RestController
public class ChatRobotController {
    @Autowired
    private ChatGPTService chatGPTService;

    @PostMapping("/chat")
    public Mono<String> chat(@RequestBody String content) {
        return chatGPTService.chat(ChatCompletionRequest.of(content))
                .map(ChatCompletionResponse::getReplyText);
    }

    @GetMapping("/stream-chat")
    public Flux<String> streamChat(@RequestParam String content) {
        return chatGPTService.stream(ChatCompletionRequest.of(content))
                .map(ChatCompletionResponse::getReplyText);
    }
}

@GPTExchange
public interface GPTHelloService {

    @ChatCompletion("You are a language translator, please translate the below text to Chinese.\n")
    Mono<String> translateIntoChinese(String text);

    @ChatCompletion("You are a language translator, please translate the below text from {0} to {1}.\n {2}")
    Mono<String> translate(String sourceLanguage, String targetLanguage, String text);
  
    @Completion("please complete poem: {0}")
  	Mono<String> completePoem(String text);
    
}

import jakarta.annotation.Nonnull;

@Component
public class GPTFunctions {

    public record SendEmailRequest(
            @Nonnull @Parameter("Recipients of email") List<String> recipients,
            @Nonnull @Parameter("Subject of email") String subject,
            @Parameter("Content of email") String content) {
    }

    @GPTFunction(name = "send_email", value = "Send email to receiver")
    public String sendEmail(SendEmailRequest request) {
        System.out.println("Recipients: " + String.join(",", request.recipients));
        System.out.println("Subject: " + request.subject);
        System.out.println("Content:\n" + request.content);
        return "Email sent to " + String.join(",", request.recipients) + " successfully!";
    }

    public record SQLQueryRequest(
            @Parameter(required = true, value = "SQL to query") String sql) {
    }

    @GPTFunction(name = "execute_sql_query", value = "Execute SQL query and return the result set")
    public String executeSQLQuery(SQLQueryRequest request) {
        System.out.println("Execute SQL: " + request.sql);
        return "id, name, salary\n1,Jackie,8000\n2,Libing,78000\n3,Sam,7500";
    }
}

public class ChatGPTServiceImplTest {
    @Test
    public void testChatWithFunctions() throws Exception {
        final String prompt = "Hi Jackie, could you write an email to Libing([email protected]) and Sam([email protected]) and invite them to join Mike's birthday party at 4 pm tomorrow? Thanks!";
        final ChatCompletionRequest request = ChatCompletionRequest.functions(prompt, List.of("send_email"));
        final ChatCompletionResponse response = chatGPTService.chat(request).block();
        // display reply combined text with function call
        System.out.println(response.getReplyCombinedText());
        // call function manually
        for (ChatMessage chatMessage : response.getReply()) {
            final FunctionCall functionCall = chatMessage.getFunctionCall();
            if (functionCall != null) {
                final Object result = functionCall.getFunctionStub().call();
                System.out.println(result);
            }
        }
    }

    @Test
    public void testExecuteSQLQuery() {
        String context = "You are SQL developer. Write SQL according to requirements, and execute it in MySQL database.";
        final String prompt = "Query all employees whose salary is greater than the average.";
        final ChatCompletionRequest request = ChatCompletionRequest.functions(prompt, List.of("execute_sql_query"));
        // add prompt context as system message
        request.addMessage(ChatMessage.systemMessage(context));
        final ChatCompletionResponse response = chatGPTService.chat(request).block();
        System.out.println(response.getReplyCombinedText());
    }
}

public class PromptLambdaTest {
    @Test
    public void testPromptAsFunction() {
        Function<String, Mono<String>> translateIntoChineseFunction = chatGPTService.promptAsLambda("translate-into-chinese");
        Function<String, Mono<String>> sendEmailFunction = chatGPTService.promptAsLambda("send-email");
        String result = Mono.just("Hi Jackie, could you write an email to Libing([email protected]) and Sam([email protected]) and invite them to join Mike's birthday party at 4 pm tomorrow? Thanks!")
                .flatMap(translateIntoChineseFunction)
                .flatMap(sendEmailFunction)
                .block();
        System.out.println(result);
    }
}

public class PromptTest {
    public record TranslateRequest(String from, String to, String text) {
    }

    @Test
    public void testLambdaWithRecord() {
        Function<TranslateRequest, Mono<String>> translateFunction = chatGPTService.promptAsLambda("translate");
        String result = Mono.just(new TranslateRequest("Chinese", "English", "你好！"))
                .flatMap(translateFunction)
                .block();
        System.out.println(result);
    }
}

@RestController
public class OpenAIProxyController {
    @Autowired
    private OpenAIChatAPI openAIChatAPI;

    @PostMapping("/v1/chat/completions")
    public Publisher<ChatCompletionResponse> completions(@RequestBody ChatCompletionRequest request) {
        return openAIChatAPI.proxy(request);
    }
}