Skip to content

Commit

Permalink
BIGTOP-4308: Add some info tools to chatbot (#131)
Browse files Browse the repository at this point in the history
  • Loading branch information
lhpqaq authored Dec 24, 2024
1 parent ca13779 commit 928f6ce
Show file tree
Hide file tree
Showing 10 changed files with 388 additions and 88 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.SerializationFeature;
import lombok.extern.slf4j.Slf4j;

import java.io.File;
Expand All @@ -31,11 +32,15 @@
public class JsonUtils {

public static final ObjectMapper OBJECTMAPPER;
public static final ObjectMapper INDENT_MAPPER;

static {
OBJECTMAPPER = new ObjectMapper();
OBJECTMAPPER.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
OBJECTMAPPER.setSerializationInclusion(JsonInclude.Include.NON_NULL);

INDENT_MAPPER = OBJECTMAPPER.copy();
INDENT_MAPPER.enable(SerializationFeature.INDENT_OUTPUT);
}

public static <T> void writeToFile(String fileName, T obj) {
Expand Down Expand Up @@ -129,4 +134,15 @@ public static <T> String writeAsString(T obj) {
throw new RuntimeException(e);
}
}

public static <T> String indentWriteAsString(T obj) {
if (obj == null) {
return null;
}
try {
return INDENT_MAPPER.writeValueAsString(obj);
} catch (Exception e) {
throw new RuntimeException(e);
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@
and h.status = #{query.status}
</if>
</where>
group by h.id, c.name
group by h.id, c.name, c.display_name
</select>

<select id="findDetailsById" resultType="org.apache.bigtop.manager.dao.po.HostPO">
Expand All @@ -70,7 +70,7 @@
left join cluster c on h.cluster_id = c.id
left join component comp on h.id = comp.host_id
where h.id = #{id}
group by h.id, c.name
group by h.id, c.name, c.display_name
limit 1
</select>

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -93,9 +93,8 @@ public SseEmitter talk(@PathVariable Long threadId, @RequestBody ChatbotMessageR
if (command != null) {
messageReq.setMessage(
messageReq.getMessage().substring(command.getCmd().length() + 2));
return chatbotService.talk(threadId, command, messageReq.getMessage());
}
return chatbotService.talk(threadId, null, messageReq.getMessage());
return chatbotService.talk(threadId, command, messageReq.getMessage());
}

@Operation(summary = "history", description = "Get chat records")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@
import org.apache.bigtop.manager.server.model.vo.ChatThreadVO;
import org.apache.bigtop.manager.server.model.vo.TalkVO;
import org.apache.bigtop.manager.server.service.ChatbotService;
import org.apache.bigtop.manager.server.tools.AiServiceToolsProvider;
import org.apache.bigtop.manager.server.tools.provider.AIServiceToolsProvider;

import org.springframework.context.i18n.LocaleContextHolder;
import org.springframework.stereotype.Service;
Expand Down Expand Up @@ -77,6 +77,9 @@ public class ChatbotServiceImpl implements ChatbotService {
@Resource
private ChatMessageDao chatMessageDao;

@Resource
private AIServiceToolsProvider aiServiceToolsProvider;

private AIAssistantFactory aiAssistantFactory;

private static final int CHAT_THREAD_NAME_LENGTH = 100;
Expand Down Expand Up @@ -121,18 +124,12 @@ private PlatformType getPlatformType(String platformName) {

private AIAssistant buildAIAssistant(
String platformName, String model, Map<String, String> credentials, Long threadId, ChatbotCommand command) {
if (command == null) {
return getAIAssistantFactory()
.createAiService(
getPlatformType(platformName), getAIAssistantConfig(model, credentials), threadId, null);
} else {
return getAIAssistantFactory()
.createAiService(
getPlatformType(platformName),
getAIAssistantConfig(model, credentials),
threadId,
new AiServiceToolsProvider(command));
}
return getAIAssistantFactory()
.createAiService(
getPlatformType(platformName),
getAIAssistantConfig(model, credentials),
threadId,
aiServiceToolsProvider.getToolsProvide(command));
}

@Override
Expand Down Expand Up @@ -193,7 +190,7 @@ public List<ChatThreadVO> getAllChatThreads() {
return chatThreads;
}

private AIAssistant prepareTalk(Long threadId, ChatbotCommand command, String message) {
private AIAssistant prepareTalk(Long threadId, ChatbotCommand command) {
ChatThreadPO chatThreadPO = chatThreadDao.findById(threadId);
Long userId = SessionUserHolder.getUserId();
if (!Objects.equals(userId, chatThreadPO.getUserId()) || chatThreadPO.getIsDeleted()) {
Expand All @@ -219,7 +216,7 @@ private AIAssistant prepareTalk(Long threadId, ChatbotCommand command, String me

@Override
public SseEmitter talk(Long threadId, ChatbotCommand command, String message) {
AIAssistant aiAssistant = prepareTalk(threadId, command, message);
AIAssistant aiAssistant = prepareTalk(threadId, command);
Flux<String> stringFlux;
if (command == null) {
stringFlux = aiAssistant.streamAsk(message);
Expand Down

This file was deleted.

Original file line number Diff line number Diff line change
@@ -0,0 +1,101 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
package org.apache.bigtop.manager.server.tools.functions;

import org.apache.bigtop.manager.common.utils.JsonUtils;
import org.apache.bigtop.manager.server.model.vo.ClusterVO;
import org.apache.bigtop.manager.server.service.ClusterService;

import org.springframework.stereotype.Component;

import dev.langchain4j.agent.tool.JsonSchemaProperty;
import dev.langchain4j.agent.tool.ToolSpecification;
import dev.langchain4j.service.tool.ToolExecutor;
import lombok.extern.slf4j.Slf4j;

import jakarta.annotation.Resource;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Component
@Slf4j
public class ClusterFunctions {
@Resource
private ClusterService clusterService;

public Map<ToolSpecification, ToolExecutor> listCluster() {
ToolSpecification toolSpecification = ToolSpecification.builder()
.name("listCluster")
.description("Get cluster list")
.build();
ToolExecutor toolExecutor =
(toolExecutionRequest, memoryId) -> JsonUtils.indentWriteAsString(clusterService.list());

return Map.of(toolSpecification, toolExecutor);
}

public Map<ToolSpecification, ToolExecutor> getClusterById() {
ToolSpecification toolSpecification = ToolSpecification.builder()
.name("getClusterById")
.description("Get cluster information based on ID")
.addParameter("clusterId", JsonSchemaProperty.NUMBER, JsonSchemaProperty.description("cluster id"))
.build();
ToolExecutor toolExecutor = (toolExecutionRequest, memoryId) -> {
Map<String, Object> arguments = JsonUtils.readFromString(toolExecutionRequest.arguments());
Long clusterId = Long.valueOf(arguments.get("clusterId").toString());
ClusterVO clusterVO = clusterService.get(clusterId);
if (clusterVO == null) {
return "Cluster not found";
}
return JsonUtils.indentWriteAsString(clusterVO);
};

return Map.of(toolSpecification, toolExecutor);
}

public Map<ToolSpecification, ToolExecutor> getClusterByName() {
ToolSpecification toolSpecification = ToolSpecification.builder()
.name("getClusterByName")
.description("Get cluster information based on cluster name")
.addParameter("clusterName", JsonSchemaProperty.STRING, JsonSchemaProperty.description("cluster name"))
.build();
ToolExecutor toolExecutor = (toolExecutionRequest, memoryId) -> {
Map<String, Object> arguments = JsonUtils.readFromString(toolExecutionRequest.arguments());
String clusterName = arguments.get("clusterName").toString();
List<ClusterVO> clusterVOS = clusterService.list();
for (ClusterVO clusterVO : clusterVOS) {
if (clusterVO.getName().equals(clusterName)) {
return JsonUtils.indentWriteAsString(clusterVO);
}
}
return "Cluster not found";
};

return Map.of(toolSpecification, toolExecutor);
}

public Map<ToolSpecification, ToolExecutor> getAllFunctions() {
Map<ToolSpecification, ToolExecutor> functions = new HashMap<>();
functions.putAll(listCluster());
functions.putAll(getClusterById());
functions.putAll(getClusterByName());
return functions;
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,87 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
package org.apache.bigtop.manager.server.tools.functions;

import org.apache.bigtop.manager.common.utils.JsonUtils;
import org.apache.bigtop.manager.dao.query.HostQuery;
import org.apache.bigtop.manager.server.model.vo.HostVO;
import org.apache.bigtop.manager.server.model.vo.PageVO;
import org.apache.bigtop.manager.server.service.HostService;

import org.springframework.stereotype.Component;

import dev.langchain4j.agent.tool.JsonSchemaProperty;
import dev.langchain4j.agent.tool.ToolSpecification;
import dev.langchain4j.service.tool.ToolExecutor;
import lombok.extern.slf4j.Slf4j;

import jakarta.annotation.Resource;
import java.util.HashMap;
import java.util.Map;

@Component
@Slf4j
public class HostFunctions {
@Resource
private HostService hostService;

public Map<ToolSpecification, ToolExecutor> getHostById() {
ToolSpecification toolSpecification = ToolSpecification.builder()
.name("getHostById")
.description("Get host information based on ID")
.addParameter("hostId", JsonSchemaProperty.NUMBER, JsonSchemaProperty.description("host id"))
.build();
ToolExecutor toolExecutor = (toolExecutionRequest, memoryId) -> {
Map<String, Object> arguments = JsonUtils.readFromString(toolExecutionRequest.arguments());
Long hostId = Long.valueOf(arguments.get("hostId").toString());
HostVO hostVO = hostService.get(hostId);
if (hostVO == null) {
return "Host not found";
}
return JsonUtils.indentWriteAsString(hostVO);
};

return Map.of(toolSpecification, toolExecutor);
}

public Map<ToolSpecification, ToolExecutor> getHostByName() {
ToolSpecification toolSpecification = ToolSpecification.builder()
.name("getHostByName")
.description("Get host information based on cluster name")
.addParameter("hostName", JsonSchemaProperty.STRING, JsonSchemaProperty.description("host name"))
.build();
ToolExecutor toolExecutor = (toolExecutionRequest, memoryId) -> {
Map<String, Object> arguments = JsonUtils.readFromString(toolExecutionRequest.arguments());
String hostName = arguments.get("hostName").toString();
HostQuery hostQuery = new HostQuery();
hostQuery.setHostname(hostName);
PageVO<HostVO> hostVO = hostService.list(hostQuery);
return JsonUtils.indentWriteAsString(hostVO);
};

return Map.of(toolSpecification, toolExecutor);
}

public Map<ToolSpecification, ToolExecutor> getAllFunctions() {
Map<ToolSpecification, ToolExecutor> functions = new HashMap<>();
functions.putAll(getHostById());
functions.putAll(getHostByName());
return functions;
}
}
Loading

0 comments on commit 928f6ce

Please sign in to comment.