Skip to content

Commit 9a48e6b

Browse files
curefatglaforge
curefat
authored and
glaforge
committed
feat: Enhance LangChain4j to support MCP tools with parametersJsonSchema
1 parent 3ba04d3 commit 9a48e6b

File tree

3 files changed

+144
-5
lines changed

3 files changed

+144
-5
lines changed

‎contrib/langchain4j/src/main/java/com/google/adk/models/langchain4j/LangChain4j.java‎

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import com.fasterxml.jackson.core.JsonProcessingException;
1919
import com.fasterxml.jackson.core.type.TypeReference;
2020
import com.fasterxml.jackson.databind.ObjectMapper;
21+
import com.google.adk.JsonBaseModel;
2122
import com.google.adk.models.BaseLlm;
2223
import com.google.adk.models.BaseLlmConnection;
2324
import com.google.adk.models.LlmRequest;
@@ -428,8 +429,24 @@ private List<ToolSpecification> toToolSpecifications(LlmRequest llmRequest) {
428429
baseTool -> {
429430
if (baseTool.declaration().isPresent()) {
430431
FunctionDeclaration functionDeclaration = baseTool.declaration().get();
431-
if (functionDeclaration.parameters().isPresent()) {
432-
Schema schema = functionDeclaration.parameters().get();
432+
Schema schema = null;
433+
if (functionDeclaration.parametersJsonSchema().isPresent()) {
434+
Object jsonSchemaObj = functionDeclaration.parametersJsonSchema().get();
435+
try {
436+
if (jsonSchemaObj instanceof Schema) {
437+
schema = (Schema) jsonSchemaObj;
438+
} else {
439+
schema = JsonBaseModel.getMapper().convertValue(jsonSchemaObj, Schema.class);
440+
}
441+
} catch (Exception e) {
442+
throw new IllegalStateException(
443+
"Failed to convert parametersJsonSchema to Schema: " + e.getMessage(), e);
444+
}
445+
} else if (functionDeclaration.parameters().isPresent()) {
446+
schema = functionDeclaration.parameters().get();
447+
}
448+
449+
if (schema != null) {
433450
ToolSpecification toolSpecification =
434451
ToolSpecification.builder()
435452
.name(baseTool.name())
@@ -438,11 +455,9 @@ private List<ToolSpecification> toToolSpecifications(LlmRequest llmRequest) {
438455
.build();
439456
toolSpecifications.add(toolSpecification);
440457
} else {
441-
// TODO exception or something else?
442458
throw new IllegalStateException("Tool lacking parameters: " + baseTool);
443459
}
444460
} else {
445-
// TODO exception or something else?
446461
throw new IllegalStateException("Tool lacking declaration: " + baseTool);
447462
}
448463
});

‎contrib/langchain4j/src/test/java/com/google/adk/models/langchain4j/LangChain4jTest.java‎

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -688,4 +688,128 @@ void testGenerateContentWithStructuredResponseJsonSchema() {
688688
final UserMessage userMessage = (UserMessage) capturedRequest.messages().get(0);
689689
assertThat(userMessage.singleText()).isEqualTo("Give me information about John Doe");
690690
}
691+
692+
@Test
693+
@DisplayName("Should handle MCP tools with parametersJsonSchema")
694+
void testGenerateContentWithMcpToolParametersJsonSchema() {
695+
// Given
696+
// Create a mock BaseTool for MCP tool
697+
final com.google.adk.tools.BaseTool mcpTool = mock(com.google.adk.tools.BaseTool.class);
698+
when(mcpTool.name()).thenReturn("mcpTool");
699+
when(mcpTool.description()).thenReturn("An MCP tool");
700+
701+
// Create a mock FunctionDeclaration
702+
final FunctionDeclaration functionDeclaration = mock(FunctionDeclaration.class);
703+
when(mcpTool.declaration()).thenReturn(Optional.of(functionDeclaration));
704+
705+
// MCP tools use parametersJsonSchema() instead of parameters()
706+
// Create a JSON schema object (Map representation)
707+
final Map<String, Object> jsonSchemaMap =
708+
Map.of(
709+
"type",
710+
"object",
711+
"properties",
712+
Map.of("city", Map.of("type", "string", "description", "City name")),
713+
"required",
714+
List.of("city"));
715+
716+
// Mock parametersJsonSchema() to return the JSON schema object
717+
when(functionDeclaration.parametersJsonSchema()).thenReturn(Optional.of(jsonSchemaMap));
718+
when(functionDeclaration.parameters()).thenReturn(Optional.empty());
719+
720+
// Create a LlmRequest with the MCP tool
721+
final LlmRequest llmRequest =
722+
LlmRequest.builder()
723+
.contents(List.of(Content.fromParts(Part.fromText("Use the MCP tool"))))
724+
.tools(Map.of("mcpTool", mcpTool))
725+
.build();
726+
727+
// Mock the AI response
728+
final AiMessage aiMessage = AiMessage.from("Tool executed successfully");
729+
730+
final ChatResponse chatResponse = mock(ChatResponse.class);
731+
when(chatResponse.aiMessage()).thenReturn(aiMessage);
732+
when(chatModel.chat(any(ChatRequest.class))).thenReturn(chatResponse);
733+
734+
// When
735+
final LlmResponse response = langChain4j.generateContent(llmRequest, false).blockingFirst();
736+
737+
// Then
738+
// Verify the response
739+
assertThat(response).isNotNull();
740+
assertThat(response.content()).isPresent();
741+
assertThat(response.content().get().text()).isEqualTo("Tool executed successfully");
742+
743+
// Verify the request was built correctly with the tool specification
744+
final ArgumentCaptor<ChatRequest> requestCaptor = ArgumentCaptor.forClass(ChatRequest.class);
745+
verify(chatModel).chat(requestCaptor.capture());
746+
final ChatRequest capturedRequest = requestCaptor.getValue();
747+
748+
// Verify tool specifications were created from parametersJsonSchema
749+
assertThat(capturedRequest.toolSpecifications()).isNotEmpty();
750+
assertThat(capturedRequest.toolSpecifications().get(0).name()).isEqualTo("mcpTool");
751+
assertThat(capturedRequest.toolSpecifications().get(0).description()).isEqualTo("An MCP tool");
752+
}
753+
754+
@Test
755+
@DisplayName("Should handle MCP tools with parametersJsonSchema when it's already a Schema")
756+
void testGenerateContentWithMcpToolParametersJsonSchemaAsSchema() {
757+
// Given
758+
// Create a mock BaseTool for MCP tool
759+
final com.google.adk.tools.BaseTool mcpTool = mock(com.google.adk.tools.BaseTool.class);
760+
when(mcpTool.name()).thenReturn("mcpTool");
761+
when(mcpTool.description()).thenReturn("An MCP tool");
762+
763+
// Create a mock FunctionDeclaration
764+
final FunctionDeclaration functionDeclaration = mock(FunctionDeclaration.class);
765+
when(mcpTool.declaration()).thenReturn(Optional.of(functionDeclaration));
766+
767+
// Create a Schema object directly (when parametersJsonSchema returns Schema)
768+
final Schema cityPropertySchema =
769+
Schema.builder().type("STRING").description("City name").build();
770+
771+
final Schema objectSchema =
772+
Schema.builder()
773+
.type("OBJECT")
774+
.properties(Map.of("city", cityPropertySchema))
775+
.required(List.of("city"))
776+
.build();
777+
778+
// Mock parametersJsonSchema() to return Schema directly
779+
when(functionDeclaration.parametersJsonSchema()).thenReturn(Optional.of(objectSchema));
780+
when(functionDeclaration.parameters()).thenReturn(Optional.empty());
781+
782+
// Create a LlmRequest with the MCP tool
783+
final LlmRequest llmRequest =
784+
LlmRequest.builder()
785+
.contents(List.of(Content.fromParts(Part.fromText("Use the MCP tool"))))
786+
.tools(Map.of("mcpTool", mcpTool))
787+
.build();
788+
789+
// Mock the AI response
790+
final AiMessage aiMessage = AiMessage.from("Tool executed successfully");
791+
792+
final ChatResponse chatResponse = mock(ChatResponse.class);
793+
when(chatResponse.aiMessage()).thenReturn(aiMessage);
794+
when(chatModel.chat(any(ChatRequest.class))).thenReturn(chatResponse);
795+
796+
// When
797+
final LlmResponse response = langChain4j.generateContent(llmRequest, false).blockingFirst();
798+
799+
// Then
800+
// Verify the response
801+
assertThat(response).isNotNull();
802+
assertThat(response.content()).isPresent();
803+
assertThat(response.content().get().text()).isEqualTo("Tool executed successfully");
804+
805+
// Verify the request was built correctly with the tool specification
806+
final ArgumentCaptor<ChatRequest> requestCaptor = ArgumentCaptor.forClass(ChatRequest.class);
807+
verify(chatModel).chat(requestCaptor.capture());
808+
final ChatRequest capturedRequest = requestCaptor.getValue();
809+
810+
// Verify tool specifications were created from parametersJsonSchema
811+
assertThat(capturedRequest.toolSpecifications()).isNotEmpty();
812+
assertThat(capturedRequest.toolSpecifications().get(0).name()).isEqualTo("mcpTool");
813+
assertThat(capturedRequest.toolSpecifications().get(0).description()).isEqualTo("An MCP tool");
814+
}
691815
}

‎core/src/main/java/com/google/adk/models/LlmRequest.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -150,7 +150,7 @@ private static Builder create() {
150150
abstract LiveConnectConfig liveConnectConfig();
151151

152152
@CanIgnoreReturnValue
153-
abstract Builder tools(Map<String, BaseTool> tools);
153+
public abstract Builder tools(Map<String, BaseTool> tools);
154154

155155
abstract Map<String, BaseTool> tools();
156156

0 commit comments

Comments
 (0)