Skip to content

Commit 5c8e441

Browse files
3coins3coins
committed
WIP: Updates to match toolkit changes
1 parent 09f0de3 commit 5c8e441

File tree

2 files changed

+41
-240
lines changed

2 files changed

+41
-240
lines changed

‎jupyter_ai_tools/__init__.py‎

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
1-
from .extension import jupyter_server_extension_tools
2-
3-
__all__ = ["jupyter_server_extension_tools"]
1+
from .extension import get_git_tools, get_notebook_tools
2+
from jupyter_server_ai_tools.models import Tool, Toolkit
43

54
__version__ = "0.1.2"
65

@@ -11,3 +10,20 @@ def _jupyter_server_extension_points():
1110

1211
def _load_jupyter_server_extension(serverapp):
1312
serverapp.log.info("✅ jupyter_ai_tools extension loaded.")
13+
14+
async def _start_jupyter_server_extension(serverapp):
15+
registry = serverapp.extension_manager.extensions.get(
16+
"jupyter_server_ai_tools"
17+
)
18+
if registry:
19+
registry.register_toolkit(
20+
Toolkit(
21+
name="notebook_toolkit", tools=get_notebook_tools()
22+
)
23+
)
24+
25+
registry.register_toolkit(
26+
Toolkit(
27+
name="git_toolkit", tools=get_git_tools()
28+
)
29+
)

‎jupyter_ai_tools/extension.py‎

Lines changed: 22 additions & 237 deletions
Original file line numberDiff line numberDiff line change
@@ -1,241 +1,26 @@
1-
from jupyter_server_ai_tools.models import ToolDefinition
1+
from jupyter_server_ai_tools.models import Tool
22

33
from . import git_tools, ynotebook_tools
44

55

6-
def jupyter_server_extension_tools():
7-
return [
8-
ToolDefinition(
9-
callable=ynotebook_tools.delete_cell,
10-
metadata={
11-
"name": "delete_cell",
12-
"description": "Remove the cell at the specified index and return its contents.",
13-
"inputSchema": {
14-
"type": "object",
15-
"properties": {
16-
"index": {
17-
"type": "integer",
18-
"description": "The index of the cell to delete",
19-
}
20-
},
21-
"required": ["index"],
22-
},
23-
},
24-
),
25-
ToolDefinition(
26-
callable=ynotebook_tools.add_cell,
27-
metadata={
28-
"name": "add_cell",
29-
"description": "Insert a blank cell at the specified index.",
30-
"inputSchema": {
31-
"type": "object",
32-
"properties": {
33-
"index": {"type": "integer", "description": "The index to insert at"},
34-
"cell_type": {
35-
"type": "string",
36-
"description": "The type of cell: 'code' or 'markdown' ",
37-
"default": "code",
38-
},
39-
},
40-
"required": ["index"],
41-
},
42-
},
43-
),
44-
ToolDefinition(
45-
callable=ynotebook_tools.write_to_cell,
46-
metadata={
47-
"name": "write_to_cell",
48-
"description": "Overwrite the source of a cell with content at the given index "
49-
"in the notebook.",
50-
"inputSchema": {
51-
"type": "object",
52-
"properties": {
53-
"index": {"type": "integer", "description": "The index to write at"},
54-
"content": {
55-
"type": "string",
56-
"description": "The content to write into the cell, either python "
57-
"code or markdown",
58-
},
59-
},
60-
"required": ["index", "content"],
61-
},
62-
},
63-
),
64-
ToolDefinition(
65-
callable=ynotebook_tools.get_max_cell_index,
66-
metadata={
67-
"name": "get_max_cell_index",
68-
"description": "Return the highest valid cell index in the current notebook.",
69-
"inputSchema": {"type": "object", "properties": {}},
70-
},
71-
),
72-
ToolDefinition(
73-
callable=ynotebook_tools.read_cell,
74-
metadata={
75-
"name": "read_cell",
76-
"description": "Read the full content of a specific cell, including outputs, "
77-
"source, and metadata.",
78-
"inputSchema": {
79-
"type": "object",
80-
"properties": {
81-
"index": {"type": "integer", "description": "The index of the cell to read"}
82-
},
83-
"required": ["index"],
84-
},
85-
},
86-
),
87-
ToolDefinition(
88-
callable=ynotebook_tools.read_notebook,
89-
metadata={
90-
"name": "read_notebook",
91-
"description": "Return all cells in the notebook as a JSON-formatted list.",
92-
"inputSchema": {"type": "object", "properties": {}},
93-
},
94-
),
95-
ToolDefinition(
96-
callable=git_tools.git_clone,
97-
metadata={
98-
"name": "git_clone",
99-
"description": "Clone a Git repo into the specified path.",
100-
"inputSchema": {
101-
"type": "object",
102-
"properties": {
103-
"path": {"type": "string", "description": "Target path"},
104-
"url": {"type": "string", "description": "Repository URL"},
105-
},
106-
"required": ["path", "url"],
107-
},
108-
},
109-
),
110-
ToolDefinition(
111-
callable=git_tools.git_status,
112-
metadata={
113-
"name": "git_status",
114-
"description": "Get the current Git status in the specified path.",
115-
"inputSchema": {
116-
"type": "object",
117-
"properties": {
118-
"path": {
119-
"type": "string",
120-
"description": "Path to the Git repository root directory",
121-
}
122-
},
123-
"required": ["path"],
124-
},
125-
},
126-
),
127-
ToolDefinition(
128-
callable=git_tools.git_log,
129-
metadata={
130-
"name": "git_log",
131-
"description": "Get the last N Git commits.",
132-
"inputSchema": {
133-
"type": "object",
134-
"properties": {
135-
"path": {
136-
"type": "string",
137-
"description": "Path to the Git repository root directory",
138-
},
139-
"history_count": {
140-
"type": "integer",
141-
"description": "Number of commits",
142-
"default": 10,
143-
},
144-
},
145-
"required": ["path"],
146-
},
147-
},
148-
),
149-
ToolDefinition(
150-
callable=git_tools.git_pull,
151-
metadata={
152-
"name": "git_pull",
153-
"description": "Pull the latest changes from the remote.",
154-
"inputSchema": {
155-
"type": "object",
156-
"properties": {
157-
"path": {
158-
"type": "string",
159-
"description": "Path to the Git repository root directory",
160-
}
161-
},
162-
"required": ["path"],
163-
},
164-
},
165-
),
166-
ToolDefinition(
167-
callable=git_tools.git_push,
168-
metadata={
169-
"name": "git_push",
170-
"description": "Push local changes to the remote.",
171-
"inputSchema": {
172-
"type": "object",
173-
"properties": {
174-
"path": {
175-
"type": "string",
176-
"description": "Path to the Git repository root directory",
177-
},
178-
"branch": {"type": "string", "description": "Repo branch"},
179-
},
180-
"required": ["path"],
181-
},
182-
},
183-
),
184-
ToolDefinition(
185-
callable=git_tools.git_commit,
186-
metadata={
187-
"name": "git_commit",
188-
"description": "Commit staged changes with a message.",
189-
"inputSchema": {
190-
"type": "object",
191-
"properties": {
192-
"path": {
193-
"type": "string",
194-
"description": "Path to the Git repository root directory",
195-
},
196-
"message": {"type": "string", "description": "Commit message"},
197-
},
198-
"required": ["path", "message"],
199-
},
200-
},
201-
),
202-
ToolDefinition(
203-
callable=git_tools.git_add,
204-
metadata={
205-
"name": "git_add",
206-
"description": "Stage files for commit. Optionally add all files.",
207-
"inputSchema": {
208-
"type": "object",
209-
"properties": {
210-
"path": {
211-
"type": "string",
212-
"description": "Path to the Git repository root directory",
213-
},
214-
"add_all": {
215-
"type": "boolean",
216-
"default": True,
217-
"description": "Stage all files",
218-
},
219-
"filename": {
220-
"type": "string",
221-
"description": "File to add (used if add_all is false)",
222-
},
223-
},
224-
"required": ["path"],
225-
},
226-
},
227-
),
228-
ToolDefinition(
229-
callable=git_tools.git_get_repo_root,
230-
metadata={
231-
"name": "git_get_repo_root_from_notebookpath",
232-
"description": "Given the path of a file, return the path to the Repo root"
233-
" if any.",
234-
"inputSchema": {
235-
"type": "object",
236-
"properties": {"path": {"type": "string", "description": "the path to a file"}},
237-
"required": ["path"],
238-
},
239-
},
240-
),
241-
]
6+
def get_notebook_tools():
7+
return {
8+
Tool(callable=ynotebook_tools.delete_cell, delete=True),
9+
Tool(callable=ynotebook_tools.add_cell, write=True),
10+
Tool(callable=ynotebook_tools.write_to_cell, read=True, write=True),
11+
Tool(callable=ynotebook_tools.get_max_cell_index, read=True),
12+
Tool(callable=ynotebook_tools.read_cell, read=True),
13+
Tool(callable=ynotebook_tools.read_notebook, read=True)
14+
}
15+
16+
def get_git_tools():
17+
return {
18+
Tool(callable=git_tools.git_clone, write=True),
19+
Tool(callable=git_tools.git_status, read=True),
20+
Tool(callable=git_tools.git_log, read=True),
21+
Tool(callable=git_tools.git_pull, read=True, write=True),
22+
Tool(callable=git_tools.git_push, read=True, write=True),
23+
Tool(callable=git_tools.git_commit, write=True),
24+
Tool(callable=git_tools.git_add, write=True),
25+
Tool(callable=git_tools.git_get_repo_root, read=True)
26+
}

0 commit comments

Comments
 (0)