This is an automated email from the ASF dual-hosted git repository.
voidmatcha pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/zeppelin.git
The following commit(s) were added to refs/heads/master by this push:
new 2d14bb99bd [ZEPPELIN-6759] Add get_paragraph tool to the AI assistant
2d14bb99bd is described below
commit 2d14bb99bdf2091b4ff36edd5a623e12a93ead1c
Author: SeungYoung Oh <[email protected]>
AuthorDate: Mon Oct 26 21:55:34 2026 +0900
[ZEPPELIN-6759] Add get_paragraph tool to the AI assistant
### What is this PR for?
Add a get_paragraph assistant tool to retrieve a paragraph’s code,
execution status, and output.
### What type of PR is it?
Feature
### Todos
### What is the Jira issue?
[[ZEPPELIN-6759]](https://issues.apache.org/jira/browse/ZEPPELIN-6759)
### How should this be tested?
Run GetParagraphToolTest, ListParagraphsToolTest, and AssistantTest.
### Screenshots (if appropriate)
### Questions:
* Does the license files need to update? N
* Is there breaking changes for older versions? N
* Does this needs documentation? N
Closes #5586 from seung-00/feature/ZEPPELIN-6759.
Signed-off-by: YONGJAE LEE <[email protected]>
---
.../zeppelin/service/assistant/AssistantImpl.java | 3 +-
...stParagraphsTool.java => GetParagraphTool.java} | 65 +++++++++++++---------
.../service/assistant/ListParagraphsTool.java | 24 +++-----
...aphsToolTest.java => GetParagraphToolTest.java} | 53 +++++++++---------
.../service/assistant/ListParagraphsToolTest.java | 5 +-
5 files changed, 78 insertions(+), 72 deletions(-)
diff --git
a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/AssistantImpl.java
b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/AssistantImpl.java
index 9020c36625..fd405311ae 100644
---
a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/AssistantImpl.java
+++
b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/AssistantImpl.java
@@ -82,7 +82,8 @@ public class AssistantImpl implements Assistant {
this.conversationRepository = conversationRepository;
this.toolExecutor = new ToolExecutor(
List.of(
- new ListParagraphsTool(notebookService)
+ new ListParagraphsTool(notebookService),
+ new GetParagraphTool(notebookService)
)
);
}
diff --git
a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ListParagraphsTool.java
b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/GetParagraphTool.java
similarity index 52%
copy from
zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ListParagraphsTool.java
copy to
zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/GetParagraphTool.java
index 9104ae45f3..8ac66c9cbf 100644
---
a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ListParagraphsTool.java
+++
b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/GetParagraphTool.java
@@ -18,67 +18,80 @@
package org.apache.zeppelin.service.assistant;
import java.io.IOException;
-import java.util.ArrayList;
-import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.Optional;
import java.util.Set;
-import org.apache.commons.lang3.StringUtils;
-import org.apache.zeppelin.user.AuthenticationInfo;
-import org.apache.zeppelin.notebook.Paragraph;
+import org.apache.commons.lang3.StringUtils;
+import org.apache.zeppelin.notebook.Note;
+import org.apache.zeppelin.rest.exception.ParagraphNotFoundException;
import org.apache.zeppelin.service.NotebookService;
import org.apache.zeppelin.service.ServiceContext;
import org.apache.zeppelin.service.SimpleServiceCallback;
+import org.apache.zeppelin.user.AuthenticationInfo;
-public class ListParagraphsTool implements Tool {
+public class GetParagraphTool implements Tool {
+
+ private static final int MAX_RESULT_LENGTH = 10_000;
private final NotebookService notebookService;
- public ListParagraphsTool(NotebookService notebookService) {
+ public GetParagraphTool(NotebookService notebookService) {
this.notebookService = notebookService;
}
@Override
public String name() {
- return "list_paragraphs";
+ return "get_paragraph";
}
@Override
public String description() {
- return "List notebook paragraphs.";
+ return "Get a paragraph's input, status, and latest result (truncated at
10000 chars).";
}
@Override
public Map<String, Object> parameters() {
- return Map.of("type", "object", "properties", Map.of());
+ return Map.of(
+ "type", "object",
+ "properties", Map.of("paragraphId", Map.of("type", "string")),
+ "required", List.of("paragraphId")
+ );
}
@Override
- public List<Map<String, Object>> call(
+ public Map<String, Object> call(
String noteId,
Map<String, Object> args,
AuthenticationInfo authInfo,
Set<String> userAndRoles
) throws IOException {
- ServiceContext context = new ServiceContext(authInfo, userAndRoles);
- return notebookService.getNote(noteId, context, callback(), note -> {
- List<Map<String, Object>> result = new ArrayList<>();
- for (int i = 0; i < note.getParagraphCount(); i++) {
- Paragraph p = note.getParagraph(i);
- Map<String, Object> info = new HashMap<>();
- info.put("id", p.getId());
- info.put("title", p.getTitle());
- info.put("text", StringUtils.abbreviate(p.getText(), 500));
- info.put("index", i);
- result.add(info);
- }
- return result;
+ Object value = args.get("paragraphId");
+ if (!(value instanceof String) || ((String) value).isBlank()) {
+ throw new IllegalArgumentException("paragraphId must be a non-blank
string");
+ }
+ var paragraphId = (String) value;
+ var context = new ServiceContext(authInfo, userAndRoles);
+
+ return notebookService.getNote(noteId, context, throwOnFailure(), note -> {
+ var paragraph = Optional.ofNullable(note.getParagraph(paragraphId))
+ .orElseThrow(() -> new ParagraphNotFoundException(paragraphId));
+
+ return Map.of(
+ "id", paragraph.getId(),
+ "title", StringUtils.defaultString(paragraph.getTitle()),
+ "input", StringUtils.defaultString(paragraph.getText()),
+ "status", paragraph.getStatus().name(),
+ "result", Optional.ofNullable(paragraph.getReturn())
+ .map(r -> StringUtils.abbreviate(r.toString(),
MAX_RESULT_LENGTH))
+ .orElse("")
+ );
});
}
- private static <T> SimpleServiceCallback<T> callback() {
- return new SimpleServiceCallback<T>() {
+ private static SimpleServiceCallback<Note> throwOnFailure() {
+ return new SimpleServiceCallback<Note>() {
@Override
public void onFailure(Exception ex, ServiceContext context) throws
IOException {
if (ex instanceof RuntimeException) {
diff --git
a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ListParagraphsTool.java
b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ListParagraphsTool.java
index 9104ae45f3..a3bdde7d3c 100644
---
a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ListParagraphsTool.java
+++
b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ListParagraphsTool.java
@@ -18,15 +18,13 @@
package org.apache.zeppelin.service.assistant;
import java.io.IOException;
-import java.util.ArrayList;
-import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
+import java.util.stream.Collectors;
import org.apache.commons.lang3.StringUtils;
import org.apache.zeppelin.user.AuthenticationInfo;
-import org.apache.zeppelin.notebook.Paragraph;
import org.apache.zeppelin.service.NotebookService;
import org.apache.zeppelin.service.ServiceContext;
import org.apache.zeppelin.service.SimpleServiceCallback;
@@ -62,19 +60,13 @@ public class ListParagraphsTool implements Tool {
Set<String> userAndRoles
) throws IOException {
ServiceContext context = new ServiceContext(authInfo, userAndRoles);
- return notebookService.getNote(noteId, context, callback(), note -> {
- List<Map<String, Object>> result = new ArrayList<>();
- for (int i = 0; i < note.getParagraphCount(); i++) {
- Paragraph p = note.getParagraph(i);
- Map<String, Object> info = new HashMap<>();
- info.put("id", p.getId());
- info.put("title", p.getTitle());
- info.put("text", StringUtils.abbreviate(p.getText(), 500));
- info.put("index", i);
- result.add(info);
- }
- return result;
- });
+ return notebookService.getNote(noteId, context, callback(), note ->
+ note.getParagraphs().stream()
+ .map(p -> Map.<String, Object>of(
+ "id", p.getId(),
+ "title", StringUtils.defaultString(p.getTitle()),
+ "input",
StringUtils.abbreviate(StringUtils.defaultString(p.getText()), 500)))
+ .collect(Collectors.toList()));
}
private static <T> SimpleServiceCallback<T> callback() {
diff --git
a/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/ListParagraphsToolTest.java
b/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/GetParagraphToolTest.java
similarity index 61%
copy from
zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/ListParagraphsToolTest.java
copy to
zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/GetParagraphToolTest.java
index 4cdab27b63..35887dc32a 100644
---
a/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/ListParagraphsToolTest.java
+++
b/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/GetParagraphToolTest.java
@@ -23,22 +23,25 @@ import static org.mockito.Mockito.*;
import java.util.Map;
import java.util.Set;
+
import org.apache.zeppelin.interpreter.InterpreterFactory;
+import org.apache.zeppelin.interpreter.InterpreterResult;
import org.apache.zeppelin.notebook.AuthorizationService;
import org.apache.zeppelin.notebook.Note;
import org.apache.zeppelin.notebook.Notebook;
import org.apache.zeppelin.notebook.Notebook.NoteProcessor;
import org.apache.zeppelin.rest.exception.ForbiddenException;
+import org.apache.zeppelin.scheduler.Job.Status;
import org.apache.zeppelin.service.NotebookService;
import org.apache.zeppelin.user.AuthenticationInfo;
import org.junit.jupiter.api.Test;
-class ListParagraphsToolTest {
+class GetParagraphToolTest {
private final AuthenticationInfo authInfo = AuthenticationInfo.ANONYMOUS;
private final Set<String> userAndRoles = Set.of("user");
@Test
- void listsParagraphs() throws Exception {
+ void getsParagraph() throws Exception {
var note = new Note();
note.setInterpreterFactory(mock(InterpreterFactory.class));
var notebook = mock(Notebook.class);
@@ -46,42 +49,40 @@ class ListParagraphsToolTest {
((NoteProcessor<?>) invocation.getArgument(2)).process(note));
var authorization = mock(AuthorizationService.class);
when(authorization.isReader("noteId", userAndRoles)).thenReturn(true);
- var notebookService = new NotebookService(notebook, authorization, null,
null);
- var sut = new ListParagraphsTool(notebookService);
+ var sut = new GetParagraphTool(new NotebookService(notebook,
authorization, null, null));
- var shortParagraph = note.addNewParagraph(AuthenticationInfo.ANONYMOUS);
- shortParagraph.setText("%md hello");
- var longParagraph = note.addNewParagraph(AuthenticationInfo.ANONYMOUS);
- var longText = "%python\n" + "x".repeat(500);
- longParagraph.setText(longText);
+ var paragraph = note.addNewParagraph(authInfo);
+ var text = "%python\n" + "x".repeat(600);
+ paragraph.setText(text);
+ paragraph.setTitle("Example");
+ paragraph.setStatus(Status.FINISHED);
+ paragraph.setResult(new InterpreterResult(
+ InterpreterResult.Code.SUCCESS, InterpreterResult.Type.TEXT,
"x".repeat(20_000)));
- var paragraphs = sut.call("noteId", Map.of(), authInfo, userAndRoles);
- assertEquals(2, paragraphs.size());
- assertEquals(shortParagraph.getId(), paragraphs.get(0).get("id"));
- assertEquals(longParagraph.getId(), paragraphs.get(1).get("id"));
- assertEquals("%md hello", paragraphs.get(0).get("text"));
- assertEquals(longText.substring(0, 497) + "...",
paragraphs.get(1).get("text"));
- assertFalse(paragraphs.get(1).containsKey("interpreter"));
- assertEquals(longText, longParagraph.getText());
- assertEquals(1, paragraphs.get(1).get("index"));
+ var actual = sut.call(
+ "noteId",
+ Map.of("paragraphId", paragraph.getId()),
+ authInfo,
+ userAndRoles
+ );
+ assertEquals(paragraph.getId(), actual.get("id"));
+ assertEquals("Example", actual.get("title"));
+ assertEquals(text, actual.get("input"));
+ assertEquals("FINISHED", actual.get("status"));
}
@Test
void rejectsNonReaders() throws Exception {
var notebook = mock(Notebook.class);
- when(notebook.processNote(eq("noteId"), eq(false), any()))
- .thenAnswer(invocation -> ((NoteProcessor<?>)
invocation.getArgument(2))
- .process(new Note()));
+ when(notebook.processNote(eq("noteId"), eq(false),
any())).thenAnswer(invocation ->
+ ((NoteProcessor<?>) invocation.getArgument(2)).process(new Note()));
var authorization = mock(AuthorizationService.class);
when(authorization.isReader("noteId", userAndRoles)).thenReturn(false);
-
- var notebookService = new NotebookService(notebook, authorization, null,
null);
-
- var sut = new ListParagraphsTool(notebookService);
+ var sut = new GetParagraphTool(new NotebookService(notebook,
authorization, null, null));
assertThrows(
ForbiddenException.class,
- () -> sut.call("noteId", Map.of(), authInfo, userAndRoles)
+ () -> sut.call("noteId", Map.of("paragraphId", "any"), authInfo,
userAndRoles)
);
}
}
diff --git
a/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/ListParagraphsToolTest.java
b/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/ListParagraphsToolTest.java
index 4cdab27b63..a6a31ba893 100644
---
a/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/ListParagraphsToolTest.java
+++
b/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/ListParagraphsToolTest.java
@@ -59,11 +59,10 @@ class ListParagraphsToolTest {
assertEquals(2, paragraphs.size());
assertEquals(shortParagraph.getId(), paragraphs.get(0).get("id"));
assertEquals(longParagraph.getId(), paragraphs.get(1).get("id"));
- assertEquals("%md hello", paragraphs.get(0).get("text"));
- assertEquals(longText.substring(0, 497) + "...",
paragraphs.get(1).get("text"));
+ assertEquals("%md hello", paragraphs.get(0).get("input"));
+ assertEquals(longText.substring(0, 497) + "...",
paragraphs.get(1).get("input"));
assertFalse(paragraphs.get(1).containsKey("interpreter"));
assertEquals(longText, longParagraph.getText());
- assertEquals(1, paragraphs.get(1).get("index"));
}
@Test
