From e15e2bd3b47f67e85a34eb9a1e48cde412780ddd Mon Sep 17 00:00:00 2001 From: Zach Sailer Date: Mon, 18 Aug 2025 08:24:38 -0700 Subject: [PATCH 1/8] first attempt at a message template --- README.md | 19 ++- src/jupyter_ai_claude_code/persona.py | 181 +++++++++----------- src/jupyter_ai_claude_code/templates.py | 214 ++++++++++++++++++++++++ 3 files changed, 308 insertions(+), 106 deletions(-) create mode 100644 src/jupyter_ai_claude_code/templates.py diff --git a/README.md b/README.md index 78572bf..8b19a3d 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,13 @@ # Jupyter AI Claude Code -Jupyter AI integration with Claude Code. +Jupyter AI integration with Claude Code persona for enhanced development capabilities. + +## Features + +- **Claude Code Integration**: Full Claude Code persona for Jupyter AI +- **Development Tools**: Access to Claude Code's built-in development tools +- **Seamless Integration**: Works with existing Jupyter AI workflow +- **Template Management**: Interactive task progress tracking and updates ## Setup @@ -39,7 +46,14 @@ This will: pixi run start ``` -This will start JupyterLab with the Jupyter AI extension and this package available. +This will start JupyterLab with the Jupyter AI extension and Claude Code persona available. + +### Using Claude Code Persona + +1. Open JupyterLab +2. Open the Jupyter AI chat panel +3. Select "Claude" persona +4. Interact with Claude Code's development tools ### Build the Package @@ -57,6 +71,7 @@ The package source code is located in `src/jupyter_ai_claude_code/`. - **JupyterLab**: Latest stable version from conda-forge - **Jupyter AI**: Version 3.0.0b5 from PyPI +- **Claude Code SDK**: For Claude Code integration - **Python**: >=3.8 ## License diff --git a/src/jupyter_ai_claude_code/persona.py b/src/jupyter_ai_claude_code/persona.py index 7305a18..e3a8d37 100644 --- a/src/jupyter_ai_claude_code/persona.py +++ b/src/jupyter_ai_claude_code/persona.py @@ -1,92 +1,20 @@ from typing import Dict, Any, List, Optional, AsyncIterator +import datetime from jupyter_ai.personas.base_persona import BasePersona, PersonaDefaults -from jupyterlab_chat.models import Message +from jupyterlab_chat.models import Message, NewMessage from claude_code_sdk import ( - query, ClaudeCodeOptions, - Message, SystemMessage, AssistantMessage, ResultMessage, - TextBlock, ToolUseBlock + query, ClaudeCodeOptions, AssistantMessage, TextBlock, ToolUseBlock, + UserMessage, SystemMessage ) +from .templates import ( + TemplateManager, claude_message_to_str +) -OMIT_INPUT_ARGS = ['content'] - -TOOL_PARAM_MAPPING = { - 'Task': 'description', - 'Bash': 'command', - 'Glob': 'pattern', - 'Grep': 'pattern', - 'LS': 'path', - 'Read': 'file_path', - 'Edit': 'file_path', - 'MultiEdit': 'file_path', - 'Write': 'file_path', - 'NotebookRead': 'notebook_path', - 'NotebookWrite': 'notebook_path', - 'WebFetch': 'url', - 'WebSearch': 'query', -} - -PROMPT_TEMPLATE = """ -{{body}} - -The user has selected the following files as attachements: - - -""" - -def input_dict_to_str(d: Dict[str, Any]) -> str: - """Convert input dictionary to string representation, omitting specified args.""" - args = [] - for k, v in d.items(): - if k not in OMIT_INPUT_ARGS: - args.append(f"{k}={v}") - return ', '.join(args) - - -def tool_to_str(block: ToolUseBlock, persona_instance=None) -> str: - """Convert a ToolUseBlock to its string representation.""" - results = [] - - if block.name == 'TodoWrite': - block_id = block.id if hasattr(block, 'id') else str(hash(str(block.input))) - - if persona_instance and block_id in persona_instance._printed_todowrite_blocks: - return "" - - if persona_instance: - persona_instance._printed_todowrite_blocks.add(block_id) - - todos = block.input.get('todos', []) - results.append('TodoWrite()') - for todo in todos: - content = todo.get('content') - if content: - results.append(f"* {content}") - elif block.name in TOOL_PARAM_MAPPING: - param_key = TOOL_PARAM_MAPPING[block.name] - param_value = block.input.get(param_key, '') - results.append(f"🛠️ {block.name}({param_value})") - else: - results.append(f"🛠️ {block.name}({input_dict_to_str(block.input)})") - - return '\n'.join(results) -def claude_message_to_str(message, persona_instance=None) -> Optional[str]: - """Convert a Claude Message to a string by extracting text content.""" - text_parts = [] - for block in message.content: - if isinstance(block, TextBlock): - text_parts.append(block.text) - elif isinstance(block, ToolUseBlock): - tool_str = tool_to_str(block, persona_instance) - if tool_str: - text_parts.append(tool_str) - else: - text_parts.append(str(block)) - return '\n'.join(text_parts) if text_parts else None class ClaudeCodePersona(BasePersona): @@ -94,7 +22,7 @@ class ClaudeCodePersona(BasePersona): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self._printed_todowrite_blocks = set() + self.template_mgr = TemplateManager(self) @property def defaults(self) -> PersonaDefaults: @@ -102,18 +30,33 @@ def defaults(self) -> PersonaDefaults: return PersonaDefaults( name="Claude", avatar_path="/files/.jupyter/claude.svg", - description="Claude Code", - system_prompt="...", + description="Claude Code persona", + system_prompt="I am Claude Code, an AI assistant with access to development tools.", ) async def _process_response_message(self, message_iterator) -> AsyncIterator[str]: - """Process response messages from Claude Code SDK.""" - async for response_message in message_iterator: - self.log.info(str(response_message)) - if isinstance(response_message, AssistantMessage): - msg_str = claude_message_to_str(response_message, self) - if msg_str is not None: - yield msg_str + '\n\n' + """Process response messages with template updates.""" + has_content = False + template_was_used = False + + async for message in message_iterator: + self.log.info(str(message)) + if isinstance(message, AssistantMessage): + result = await claude_message_to_str(message, self.template_mgr) + if result is not None: # Only yield if we got actual content + has_content = True + yield result + '\n\n' + elif self.template_mgr.active: + # Template handled this message + template_was_used = True + + # Complete template if active + if self.template_mgr.active: + await self.template_mgr.complete() + + # Only yield empty string if no content was produced and template wasn't used + if not has_content and not template_was_used: + yield "" # Ensure stream completes for empty responses def _generate_prompt(self, message: Message) -> str: attachment_ids = message.attachments @@ -133,24 +76,54 @@ def _generate_prompt(self, message: Message) -> str: self.log.info(prompt) return prompt + def _get_system_prompt(self): + """Get the system prompt for Claude Code options.""" + return "I am Claude Code, an AI assistant with access to development tools." + async def process_message(self, message: Message) -> None: """Process incoming message and stream Claude Code response.""" - self._printed_todowrite_blocks.clear() - async_gen = None - prompt = self._generate_prompt(message) + # Always set writing state at the start + self.awareness.set_local_state_field("isWriting", True) + + self.template_mgr.reset() + try: - async_gen = query( - prompt=prompt, - options=ClaudeCodeOptions( - max_turns=20, - cwd=self.get_workspace_dir(), - permission_mode='bypassPermissions' - ) - ) + # Configure Claude Code - use workspace dir for better working directory detection + chat_dir = self.get_chat_dir() + workspace_dir = self.get_workspace_dir() + + # Prefer workspace dir if available, fallback to chat dir + working_dir = chat_dir if chat_dir else workspace_dir + + self.log.info(f"Chat directory: {chat_dir}") + self.log.info(f"Workspace directory: {workspace_dir}") + self.log.info(f"Using working directory: {working_dir}") + + options = { + 'max_turns': 20, + 'cwd': working_dir, + 'permission_mode': 'bypassPermissions', + 'system_prompt': self._get_system_prompt() + } + + # Generate prompt from current message + user_prompt = self._generate_prompt(message) + + # Stream response directly with prompt + async_gen = query(prompt=user_prompt, options=ClaudeCodeOptions(**options)) + + # Use stream_message to handle the streaming await self.stream_message(self._process_response_message(async_gen)) + except Exception as e: - self.log.error(f"Error in process_message: {e}") - await self.send_message(f"Sorry, I have had an internal error while working on that: {e}") + self.log.error(f"Error: {e}") + if self.template_mgr.active: + await self.template_mgr.complete() + + try: + await self.send_message(f"Sorry, error: {e}") + except TypeError: + self.send_message(f"Sorry, error: {e}") finally: - if async_gen is not None: - await async_gen.aclose() + # Always clear writing state when done + self.awareness.set_local_state_field("isWriting", False) diff --git a/src/jupyter_ai_claude_code/templates.py b/src/jupyter_ai_claude_code/templates.py new file mode 100644 index 0000000..73aa3af --- /dev/null +++ b/src/jupyter_ai_claude_code/templates.py @@ -0,0 +1,214 @@ +"""Message template management for Claude Code persona.""" + +from typing import Dict, Any, List, Optional +import datetime +from jinja2 import Template +from jupyterlab_chat.models import Message, NewMessage +from claude_code_sdk import TextBlock, ToolUseBlock + + +# Template for rendering todo lists with progress +TODO_TEMPLATE = Template(""" +{%- if todos %} +#### Task Progress + +{%- for todo in todos %} +{%- if todo.status == 'completed' %} +- [x] ~~{{ todo.content }}~~ +{%- elif todo.status == 'in_progress' %} +- [ ] **{{ todo.content }}** *(in progress)* +{%- else %} +- [ ] {{ todo.content }} +{%- endif %} +{%- endfor %} + +{%- endif %} + +{%- if current_action %} +#### Current Action +{{ current_action }} + +{%- endif %} + +{%- if previous_actions %} +
+Previous Actions ({{ previous_actions|length }}) + +{%- for action in previous_actions %} +{{ loop.index }}. {{ action }} +{%- endfor %} + +
+ +{%- endif %} + +{%- if response_text %} +{{ response_text }} +{%- endif %} +""".strip()) + + +class TemplateManager: + """Manages in-place template message updates.""" + + def __init__(self, persona): + self.persona = persona + self.message_id = None + self.todos = [] + self.action = None + self.previous_actions = [] + self.text_parts = [] + self.active = False + + def _same_todo_list(self, new_todos): + """Check if this is the same todo list (just status updates).""" + if not self.todos: + return False + old_ids = {t['id'] for t in self.todos} + new_ids = {t['id'] for t in new_todos} + return old_ids == new_ids + + async def update_todos(self, todos): + """Update todo list, creating or updating message as needed.""" + if self._same_todo_list(todos) and self.active: + self.todos = todos + await self._update_message() + else: + # New todo list - create new message + self.todos = todos + self.active = True + await self._create_message() + return "" + + async def update_action(self, action): + """Update current action in template.""" + if self.active: + # Move current action to previous actions if it exists + if self.action and self.action != action: + self.previous_actions.append(self.action) + + self.action = action + await self._update_message() + return "" + return action + + async def update_text(self, text): + """Add text to template response.""" + if self.active: + self.text_parts.append(text) + await self._update_message() + return "" + return text + + async def _create_message(self): + """Create new template message.""" + # Don't override main persona's writing state - it's already set + content = self._render_template() + new_msg = NewMessage(body=content, sender=self.persona.id) + self.message_id = self.persona.ychat.add_message(new_msg) + + async def _update_message(self): + """Update existing template message.""" + if not self.message_id: + return + + # Update awareness to show writing to specific message + self.persona.awareness.set_local_state_field("isWriting", self.message_id) + content = self._render_template() + + msg = Message( + id=self.message_id, + time=datetime.datetime.now().timestamp(), + body=content, + sender=self.persona.id + ) + self.persona.ychat.update_message(msg, append=False) + + def _render_template(self): + """Render current template state.""" + return TODO_TEMPLATE.render( + todos=self.todos, + current_action=self.action, + previous_actions=self.previous_actions if len(self.previous_actions) > 0 else None, + response_text='\n'.join(self.text_parts) if self.text_parts else None + ) + + async def complete(self): + """Complete template - move current action to history and clear it.""" + if self.active and self.message_id: + # Move current action to previous actions if it exists + if self.action: + self.previous_actions.append(self.action) + self.action = None # Clear current action + + # Do final template update to show completed state + await self._update_message() + + # Don't reset here - preserve the completed state + # Just mark as inactive and clear message ID + self.message_id = None + self.active = False + + def reset(self): + """Reset for new conversation.""" + self.message_id = None + self.todos = [] + self.action = None + self.previous_actions = [] + self.text_parts = [] + self.active = False + + +# Tool parameter mapping for display formatting +TOOL_PARAM_MAPPING = { + 'Task': 'description', 'Bash': 'command', 'Glob': 'pattern', 'Grep': 'pattern', + 'LS': 'path', 'Read': 'file_path', 'Edit': 'file_path', 'MultiEdit': 'file_path', + 'Write': 'file_path', 'NotebookRead': 'notebook_path', 'NotebookWrite': 'notebook_path', + 'WebFetch': 'url', 'WebSearch': 'query' +} + + +def format_tool_input(tool_name, tool_input): + """Format tool input for display.""" + if tool_name in TOOL_PARAM_MAPPING: + key = TOOL_PARAM_MAPPING[tool_name] + return tool_input.get(key, '') + + # Format remaining args (excluding content) + args = [f"{k}={v}" for k, v in tool_input.items() if k != 'content'] + return ', '.join(args) + + +async def process_message_block(block, template_mgr=None): + """Process a single message block (text or tool).""" + if isinstance(block, TextBlock): + if template_mgr and template_mgr.active: + await template_mgr.update_text(block.text) + return None # Template handles display, don't stream + return block.text + + elif isinstance(block, ToolUseBlock): + if block.name == 'TodoWrite' and template_mgr: + todos = block.input.get('todos', []) + await template_mgr.update_todos(todos) + return None # Template handles display, don't stream + + # Regular tool display + tool_display = f"🛠️ {block.name}({format_tool_input(block.name, block.input)})" + + if template_mgr and template_mgr.active: + await template_mgr.update_action(tool_display) + return None # Template handles display, don't stream + return tool_display + + return str(block) + + +async def claude_message_to_str(message, template_mgr=None) -> Optional[str]: + """Convert Claude Message to string, handling template updates.""" + text_parts = [] + for block in message.content: + result = await process_message_block(block, template_mgr) + if result is not None: # Only add non-None results + text_parts.append(result) + return '\n'.join(text_parts) if text_parts else None \ No newline at end of file From b495baf17e5955d787ee508a4722ee1edaef826e Mon Sep 17 00:00:00 2001 From: Zach Sailer Date: Mon, 18 Aug 2025 09:03:43 -0700 Subject: [PATCH 2/8] since the message template is pretty specific to claude code, rename it to reflec tthis --- src/jupyter_ai_claude_code/persona.py | 6 +- src/jupyter_ai_claude_code/templates.py | 193 ++++++++++++++---------- 2 files changed, 114 insertions(+), 85 deletions(-) diff --git a/src/jupyter_ai_claude_code/persona.py b/src/jupyter_ai_claude_code/persona.py index e3a8d37..ffc0aa4 100644 --- a/src/jupyter_ai_claude_code/persona.py +++ b/src/jupyter_ai_claude_code/persona.py @@ -10,7 +10,7 @@ ) from .templates import ( - TemplateManager, claude_message_to_str + ClaudeCodeTemplateManager ) @@ -22,7 +22,7 @@ class ClaudeCodePersona(BasePersona): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self.template_mgr = TemplateManager(self) + self.template_mgr = ClaudeCodeTemplateManager(self) @property def defaults(self) -> PersonaDefaults: @@ -42,7 +42,7 @@ async def _process_response_message(self, message_iterator) -> AsyncIterator[str async for message in message_iterator: self.log.info(str(message)) if isinstance(message, AssistantMessage): - result = await claude_message_to_str(message, self.template_mgr) + result = await self.template_mgr.claude_message_to_str(message) if result is not None: # Only yield if we got actual content has_content = True yield result + '\n\n' diff --git a/src/jupyter_ai_claude_code/templates.py b/src/jupyter_ai_claude_code/templates.py index 73aa3af..4007030 100644 --- a/src/jupyter_ai_claude_code/templates.py +++ b/src/jupyter_ai_claude_code/templates.py @@ -1,10 +1,10 @@ -"""Message template management for Claude Code persona.""" +"""Claude Code message template management for Jupyter AI persona.""" from typing import Dict, Any, List, Optional import datetime from jinja2 import Template from jupyterlab_chat.models import Message, NewMessage -from claude_code_sdk import TextBlock, ToolUseBlock +from claude_code_sdk import TextBlock, ToolUseBlock, ToolResultBlock, AssistantMessage # Template for rendering todo lists with progress @@ -24,39 +24,52 @@ {%- endif %} -{%- if current_action %} -#### Current Action -{{ current_action }} - -{%- endif %} - -{%- if previous_actions %} +{%- if completed_actions %}
-Previous Actions ({{ previous_actions|length }}) +Completed Actions ({{ completed_actions|length }}) + +{%- for action in completed_actions %} +{{ loop.index }}. **{{ action.tool_call }}** +``` +{{ action.result }} +``` -{%- for action in previous_actions %} -{{ loop.index }}. {{ action }} {%- endfor %}
{%- endif %} +{%- if current_action %} +#### Current Action +{{ current_action }} + +{%- endif %} + {%- if response_text %} {{ response_text }} {%- endif %} """.strip()) -class TemplateManager: - """Manages in-place template message updates.""" +class ClaudeCodeTemplateManager: + """Manages in-place template message updates for Claude Code SDK messages.""" + + # Tool parameter mapping for display formatting + TOOL_PARAM_MAPPING = { + 'Task': 'description', 'Bash': 'command', 'Glob': 'pattern', 'Grep': 'pattern', + 'LS': 'path', 'Read': 'file_path', 'Edit': 'file_path', 'MultiEdit': 'file_path', + 'Write': 'file_path', 'NotebookRead': 'notebook_path', 'NotebookWrite': 'notebook_path', + 'WebFetch': 'url', 'WebSearch': 'query' + } def __init__(self, persona): self.persona = persona self.message_id = None self.todos = [] - self.action = None - self.previous_actions = [] + self.current_action = None + self.current_action_result = None + self.completed_actions = [] # List of {'tool_call': str, 'result': str} self.text_parts = [] self.active = False @@ -83,15 +96,27 @@ async def update_todos(self, todos): async def update_action(self, action): """Update current action in template.""" if self.active: - # Move current action to previous actions if it exists - if self.action and self.action != action: - self.previous_actions.append(self.action) + # Complete previous action if it exists + if self.current_action: + self.completed_actions.append({ + 'tool_call': self.current_action, + 'result': self.current_action_result or 'No result captured' + }) - self.action = action + self.current_action = action + self.current_action_result = None # Reset for new action await self._update_message() return "" return action + async def update_action_result(self, result): + """Update the result of the current action.""" + if self.active and self.current_action: + self.current_action_result = result + await self._update_message() + return "" + return result + async def update_text(self, text): """Add text to template response.""" if self.active: @@ -128,18 +153,22 @@ def _render_template(self): """Render current template state.""" return TODO_TEMPLATE.render( todos=self.todos, - current_action=self.action, - previous_actions=self.previous_actions if len(self.previous_actions) > 0 else None, + current_action=self.current_action, + completed_actions=self.completed_actions if len(self.completed_actions) > 0 else None, response_text='\n'.join(self.text_parts) if self.text_parts else None ) async def complete(self): - """Complete template - move current action to history and clear it.""" + """Complete template - move current action to completed actions and clear it.""" if self.active and self.message_id: - # Move current action to previous actions if it exists - if self.action: - self.previous_actions.append(self.action) - self.action = None # Clear current action + # Move current action to completed actions if it exists + if self.current_action: + self.completed_actions.append({ + 'tool_call': self.current_action, + 'result': self.current_action_result or 'No result captured' + }) + self.current_action = None # Clear current action + self.current_action_result = None # Do final template update to show completed state await self._update_message() @@ -149,66 +178,66 @@ async def complete(self): self.message_id = None self.active = False + def format_tool_input(self, tool_name, tool_input): + """Format tool input for display.""" + if tool_name in self.TOOL_PARAM_MAPPING: + key = self.TOOL_PARAM_MAPPING[tool_name] + return tool_input.get(key, '') + + # Format remaining args (excluding content) + args = [f"{k}={v}" for k, v in tool_input.items() if k != 'content'] + return ', '.join(args) + + async def process_message_block(self, block): + """Process a single Claude SDK message block (text or tool).""" + if isinstance(block, TextBlock): + if self.active: + await self.update_text(block.text) + return None # Template handles display, don't stream + return block.text + + elif isinstance(block, ToolUseBlock): + if block.name == 'TodoWrite': + todos = block.input.get('todos', []) + await self.update_todos(todos) + return None # Template handles display, don't stream + + # Regular tool display + tool_display = f"🛠️ {block.name}({self.format_tool_input(block.name, block.input)})" + + if self.active: + await self.update_action(tool_display) + return None # Template handles display, don't stream + return tool_display + + elif isinstance(block, ToolResultBlock): + # Handle tool result + if self.active: + # Format the result for display + result_text = str(block.content) if hasattr(block, 'content') else str(block) + await self.update_action_result(result_text) + return None # Template handles display, don't stream + return f"Result: {block.content if hasattr(block, 'content') else block}" + + return str(block) + + async def claude_message_to_str(self, message) -> Optional[str]: + """Convert Claude SDK Message to string, handling template updates.""" + text_parts = [] + for block in message.content: + result = await self.process_message_block(block) + if result is not None: # Only add non-None results + text_parts.append(result) + return '\n'.join(text_parts) if text_parts else None + def reset(self): """Reset for new conversation.""" self.message_id = None self.todos = [] - self.action = None - self.previous_actions = [] + self.current_action = None + self.current_action_result = None + self.completed_actions = [] self.text_parts = [] self.active = False -# Tool parameter mapping for display formatting -TOOL_PARAM_MAPPING = { - 'Task': 'description', 'Bash': 'command', 'Glob': 'pattern', 'Grep': 'pattern', - 'LS': 'path', 'Read': 'file_path', 'Edit': 'file_path', 'MultiEdit': 'file_path', - 'Write': 'file_path', 'NotebookRead': 'notebook_path', 'NotebookWrite': 'notebook_path', - 'WebFetch': 'url', 'WebSearch': 'query' -} - - -def format_tool_input(tool_name, tool_input): - """Format tool input for display.""" - if tool_name in TOOL_PARAM_MAPPING: - key = TOOL_PARAM_MAPPING[tool_name] - return tool_input.get(key, '') - - # Format remaining args (excluding content) - args = [f"{k}={v}" for k, v in tool_input.items() if k != 'content'] - return ', '.join(args) - - -async def process_message_block(block, template_mgr=None): - """Process a single message block (text or tool).""" - if isinstance(block, TextBlock): - if template_mgr and template_mgr.active: - await template_mgr.update_text(block.text) - return None # Template handles display, don't stream - return block.text - - elif isinstance(block, ToolUseBlock): - if block.name == 'TodoWrite' and template_mgr: - todos = block.input.get('todos', []) - await template_mgr.update_todos(todos) - return None # Template handles display, don't stream - - # Regular tool display - tool_display = f"🛠️ {block.name}({format_tool_input(block.name, block.input)})" - - if template_mgr and template_mgr.active: - await template_mgr.update_action(tool_display) - return None # Template handles display, don't stream - return tool_display - - return str(block) - - -async def claude_message_to_str(message, template_mgr=None) -> Optional[str]: - """Convert Claude Message to string, handling template updates.""" - text_parts = [] - for block in message.content: - result = await process_message_block(block, template_mgr) - if result is not None: # Only add non-None results - text_parts.append(result) - return '\n'.join(text_parts) if text_parts else None \ No newline at end of file From d75affb046243b12fcc2f906becb223e2aea6202 Mon Sep 17 00:00:00 2001 From: Zach Sailer Date: Mon, 18 Aug 2025 09:42:56 -0700 Subject: [PATCH 3/8] format tool headers --- src/jupyter_ai_claude_code/persona.py | 12 +- src/jupyter_ai_claude_code/templates.py | 145 ++++++++++++++---------- 2 files changed, 94 insertions(+), 63 deletions(-) diff --git a/src/jupyter_ai_claude_code/persona.py b/src/jupyter_ai_claude_code/persona.py index ffc0aa4..3730d1e 100644 --- a/src/jupyter_ai_claude_code/persona.py +++ b/src/jupyter_ai_claude_code/persona.py @@ -43,16 +43,18 @@ async def _process_response_message(self, message_iterator) -> AsyncIterator[str self.log.info(str(message)) if isinstance(message, AssistantMessage): result = await self.template_mgr.claude_message_to_str(message) - if result is not None: # Only yield if we got actual content + # Template now handles everything - never stream individual components + if self.template_mgr.active: + template_was_used = True + elif result is not None: + # Only for messages without any tool usage (rare) has_content = True yield result + '\n\n' - elif self.template_mgr.active: - # Template handled this message - template_was_used = True # Complete template if active if self.template_mgr.active: await self.template_mgr.complete() + template_was_used = True # Only yield empty string if no content was produced and template wasn't used if not has_content and not template_was_used: @@ -78,7 +80,7 @@ def _generate_prompt(self, message: Message) -> str: def _get_system_prompt(self): """Get the system prompt for Claude Code options.""" - return "I am Claude Code, an AI assistant with access to development tools." + return "..." async def process_message(self, message: Message) -> None: """Process incoming message and stream Claude Code response.""" diff --git a/src/jupyter_ai_claude_code/templates.py b/src/jupyter_ai_claude_code/templates.py index 4007030..05bf5df 100644 --- a/src/jupyter_ai_claude_code/templates.py +++ b/src/jupyter_ai_claude_code/templates.py @@ -7,7 +7,7 @@ from claude_code_sdk import TextBlock, ToolUseBlock, ToolResultBlock, AssistantMessage -# Template for rendering todo lists with progress +# Template for rendering consolidated actions and final response TODO_TEMPLATE = Template(""" {%- if todos %} #### Task Progress @@ -24,30 +24,35 @@ {%- endif %} +{%- if initial_text %} +{{ initial_text }} + +{%- endif %} + +{%- if current_action %} + +{{ current_action.tool_call }} +⎿ {{ current_action.result }} + +{%- endif %} + {%- if completed_actions %}
-Completed Actions ({{ completed_actions|length }}) +Actions Taken ({{ completed_actions|length }}) {%- for action in completed_actions %} -{{ loop.index }}. **{{ action.tool_call }}** -``` -{{ action.result }} -``` +{{ action.tool_call }} +⎿ {{ action.result | replace('\n', '\n ') }} {%- endfor %}
{%- endif %} -{%- if current_action %} -#### Current Action -{{ current_action }} +{%- if final_text %} -{%- endif %} - -{%- if response_text %} -{{ response_text }} +{{ final_text }} {%- endif %} """.strip()) @@ -67,11 +72,13 @@ def __init__(self, persona): self.persona = persona self.message_id = None self.todos = [] - self.current_action = None - self.current_action_result = None + self.current_action = None # Currently executing action self.completed_actions = [] # List of {'tool_call': str, 'result': str} - self.text_parts = [] + self.initial_text_parts = [] # Text before any tool calls + self.final_text_parts = [] # Text after tool calls self.active = False + self.turn_active = False # Track if we're in an active Claude turn + self.has_actions = False # Track if we've seen any tool calls def _same_todo_list(self, new_todos): """Check if this is the same todo list (just status updates).""" @@ -94,17 +101,25 @@ async def update_todos(self, todos): return "" async def update_action(self, action): - """Update current action in template.""" + """Start a new action - show it as current action.""" if self.active: # Complete previous action if it exists if self.current_action: - self.completed_actions.append({ - 'tool_call': self.current_action, - 'result': self.current_action_result or 'No result captured' - }) + self.completed_actions.append(self.current_action) + + # Mark that we've seen actions + self.has_actions = True - self.current_action = action - self.current_action_result = None # Reset for new action + # Always start template if not active yet + if not self.turn_active: + self.turn_active = True + await self._create_message() + + # Set new current action + self.current_action = { + 'tool_call': action, + 'result': 'Executing...' + } await self._update_message() return "" return action @@ -112,7 +127,8 @@ async def update_action(self, action): async def update_action_result(self, result): """Update the result of the current action.""" if self.active and self.current_action: - self.current_action_result = result + # Update the current action's result + self.current_action['result'] = result await self._update_message() return "" return result @@ -120,7 +136,12 @@ async def update_action_result(self, result): async def update_text(self, text): """Add text to template response.""" if self.active: - self.text_parts.append(text) + # If we haven't seen actions yet, add to initial text + # If we have seen actions, add to final text + if not self.has_actions: + self.initial_text_parts.append(text) + else: + self.final_text_parts.append(text) await self._update_message() return "" return text @@ -155,46 +176,58 @@ def _render_template(self): todos=self.todos, current_action=self.current_action, completed_actions=self.completed_actions if len(self.completed_actions) > 0 else None, - response_text='\n'.join(self.text_parts) if self.text_parts else None + initial_text='\n'.join(self.initial_text_parts) if self.initial_text_parts else None, + final_text='\n'.join(self.final_text_parts) if self.final_text_parts else None ) async def complete(self): - """Complete template - move current action to completed actions and clear it.""" + """Complete template - move current action to completed actions.""" if self.active and self.message_id: # Move current action to completed actions if it exists if self.current_action: - self.completed_actions.append({ - 'tool_call': self.current_action, - 'result': self.current_action_result or 'No result captured' - }) - self.current_action = None # Clear current action - self.current_action_result = None + self.completed_actions.append(self.current_action) + self.current_action = None # Do final template update to show completed state await self._update_message() - # Don't reset here - preserve the completed state - # Just mark as inactive and clear message ID + # Mark as inactive and clear message ID but preserve completed actions self.message_id = None self.active = False + self.turn_active = False def format_tool_input(self, tool_name, tool_input): - """Format tool input for display.""" + """Format tool input for Claude Code CLI style display.""" if tool_name in self.TOOL_PARAM_MAPPING: key = self.TOOL_PARAM_MAPPING[tool_name] - return tool_input.get(key, '') + value = tool_input.get(key, '') + # For long values, truncate with ellipsis + if len(str(value)) > 60: + return str(value)[:60] + '…' + return str(value) # Format remaining args (excluding content) - args = [f"{k}={v}" for k, v in tool_input.items() if k != 'content'] + args = [] + for k, v in tool_input.items(): + if k != 'content': + val_str = str(v) + if len(val_str) > 30: + val_str = val_str[:30] + '…' + args.append(f"{k}={val_str}") return ', '.join(args) async def process_message_block(self, block): """Process a single Claude SDK message block (text or tool).""" if isinstance(block, TextBlock): - if self.active: + # Always capture text in template during active turn + if not self.active: + # Start template on first content + self.active = True await self.update_text(block.text) - return None # Template handles display, don't stream - return block.text + return None + else: + await self.update_text(block.text) + return None # Template handles all display elif isinstance(block, ToolUseBlock): if block.name == 'TodoWrite': @@ -202,24 +235,18 @@ async def process_message_block(self, block): await self.update_todos(todos) return None # Template handles display, don't stream - # Regular tool display - tool_display = f"🛠️ {block.name}({self.format_tool_input(block.name, block.input)})" - - if self.active: - await self.update_action(tool_display) - return None # Template handles display, don't stream - return tool_display + # Regular tool display - always capture in template + tool_display = f"{block.name}({self.format_tool_input(block.name, block.input)})" + await self.update_action(tool_display) + return None # Template handles all display elif isinstance(block, ToolResultBlock): - # Handle tool result - if self.active: - # Format the result for display - result_text = str(block.content) if hasattr(block, 'content') else str(block) - await self.update_action_result(result_text) - return None # Template handles display, don't stream - return f"Result: {block.content if hasattr(block, 'content') else block}" + # Handle tool result - always capture in template + result_text = str(block.content) if hasattr(block, 'content') else str(block) + await self.update_action_result(result_text) + return None # Template handles all display - return str(block) + return None # Don't stream anything - template handles all async def claude_message_to_str(self, message) -> Optional[str]: """Convert Claude SDK Message to string, handling template updates.""" @@ -235,9 +262,11 @@ def reset(self): self.message_id = None self.todos = [] self.current_action = None - self.current_action_result = None self.completed_actions = [] - self.text_parts = [] + self.initial_text_parts = [] + self.final_text_parts = [] self.active = False + self.turn_active = False + self.has_actions = False From 4c2d7049adc88da651b01c035e9640df4d9c2dd0 Mon Sep 17 00:00:00 2001 From: Zach Sailer Date: Mon, 18 Aug 2025 09:58:24 -0700 Subject: [PATCH 4/8] better formatted tools --- src/jupyter_ai_claude_code/persona.py | 6 ++++-- src/jupyter_ai_claude_code/templates.py | 7 ++++++- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/src/jupyter_ai_claude_code/persona.py b/src/jupyter_ai_claude_code/persona.py index 3730d1e..4089931 100644 --- a/src/jupyter_ai_claude_code/persona.py +++ b/src/jupyter_ai_claude_code/persona.py @@ -27,11 +27,13 @@ def __init__(self, *args, **kwargs): @property def defaults(self) -> PersonaDefaults: """Return default configuration for the Claude Code persona.""" + system_prompt = "I am Claude Code, an AI assistant with access to development tools. When formatting responses, I use **bold text** for emphasis and section headers instead of markdown headings (# ## ###). I keep formatting clean and readable without large headers." + return PersonaDefaults( name="Claude", avatar_path="/files/.jupyter/claude.svg", description="Claude Code persona", - system_prompt="I am Claude Code, an AI assistant with access to development tools.", + system_prompt=system_prompt, ) async def _process_response_message(self, message_iterator) -> AsyncIterator[str]: @@ -80,7 +82,7 @@ def _generate_prompt(self, message: Message) -> str: def _get_system_prompt(self): """Get the system prompt for Claude Code options.""" - return "..." + return "I am Claude Code, an AI assistant with access to development tools. When formatting responses, I use **bold text** for emphasis and section headers instead of markdown headings (# ## ###). I keep formatting clean and readable without large headers." async def process_message(self, message: Message) -> None: """Process incoming message and stream Claude Code response.""" diff --git a/src/jupyter_ai_claude_code/templates.py b/src/jupyter_ai_claude_code/templates.py index 05bf5df..d8a7cdb 100644 --- a/src/jupyter_ai_claude_code/templates.py +++ b/src/jupyter_ai_claude_code/templates.py @@ -31,14 +31,17 @@ {%- if current_action %} +**Current Action:** {{ current_action.tool_call }} ⎿ {{ current_action.result }} {%- endif %} {%- if completed_actions %} + +**Actions Taken:**
-Actions Taken ({{ completed_actions|length }}) +See details ({{ completed_actions|length }}) {%- for action in completed_actions %} @@ -52,6 +55,8 @@ {%- if final_text %} +--- + {{ final_text }} {%- endif %} """.strip()) From 88fc22d9d1ff2ab5d15051c41860ba453174703f Mon Sep 17 00:00:00 2001 From: Zach Sailer Date: Mon, 18 Aug 2025 13:59:54 -0700 Subject: [PATCH 5/8] final tuning of formating --- src/jupyter_ai_claude_code/persona.py | 10 +- src/jupyter_ai_claude_code/templates.py | 137 +++++++++++++++++------- 2 files changed, 107 insertions(+), 40 deletions(-) diff --git a/src/jupyter_ai_claude_code/persona.py b/src/jupyter_ai_claude_code/persona.py index 4089931..4bb91a0 100644 --- a/src/jupyter_ai_claude_code/persona.py +++ b/src/jupyter_ai_claude_code/persona.py @@ -27,7 +27,10 @@ def __init__(self, *args, **kwargs): @property def defaults(self) -> PersonaDefaults: """Return default configuration for the Claude Code persona.""" - system_prompt = "I am Claude Code, an AI assistant with access to development tools. When formatting responses, I use **bold text** for emphasis and section headers instead of markdown headings (# ## ###). I keep formatting clean and readable without large headers." + system_prompt = ("I am Claude Code, an AI assistant with access to development tools. " + "When formatting responses, I use **bold text** for emphasis and section headers instead of markdown headings (# ## ###). " + "I keep formatting clean and readable without large headers. " + "For complex tasks requiring multiple steps (3+ actions), I proactively create a todo list using TodoWrite to track progress and keep the user informed of my plan.") return PersonaDefaults( name="Claude", @@ -82,7 +85,10 @@ def _generate_prompt(self, message: Message) -> str: def _get_system_prompt(self): """Get the system prompt for Claude Code options.""" - return "I am Claude Code, an AI assistant with access to development tools. When formatting responses, I use **bold text** for emphasis and section headers instead of markdown headings (# ## ###). I keep formatting clean and readable without large headers." + return ("I am Claude Code, an AI assistant with access to development tools. " + "When formatting responses, I use **bold text** for emphasis and section headers instead of markdown headings (# ## ###). " + "I keep formatting clean and readable without large headers. " + "For complex tasks requiring multiple steps (3+ actions), I proactively create a todo list using TodoWrite to track progress and keep the user informed of my plan.") async def process_message(self, message: Message) -> None: """Process incoming message and stream Claude Code response.""" diff --git a/src/jupyter_ai_claude_code/templates.py b/src/jupyter_ai_claude_code/templates.py index d8a7cdb..c61ccc9 100644 --- a/src/jupyter_ai_claude_code/templates.py +++ b/src/jupyter_ai_claude_code/templates.py @@ -9,8 +9,13 @@ # Template for rendering consolidated actions and final response TODO_TEMPLATE = Template(""" -{%- if todos %} -#### Task Progress +{%- if initial_text %} +{{ initial_text }} + +{%- endif %} + +{% if todos %} +**Task Progress:** {%- for todo in todos %} {%- if todo.status == 'completed' %} @@ -24,41 +29,39 @@ {%- endif %} -{%- if initial_text %} -{{ initial_text }} - -{%- endif %} - {%- if current_action %} -**Current Action:** +**Current Tool:** {{ current_action.tool_call }} -⎿ {{ current_action.result }} +⎿ Executing... {%- endif %} -{%- if completed_actions %} +{% if completed_actions %} -**Actions Taken:** +**Tools Called:**
See details ({{ completed_actions|length }}) {%- for action in completed_actions %} -{{ action.tool_call }} -⎿ {{ action.result | replace('\n', '\n ') }} +{{ action }} +⎿ Completed {%- endfor %}
- -{%- endif %} - -{%- if final_text %} - ---- - +
+{% endif %} + +{% if current_result or final_text %} +**Response:** +{% if current_result %} +{{ current_result }} +{% endif %} +{% if final_text %} {{ final_text }} -{%- endif %} +{% endif %} +{% endif %} """.strip()) @@ -78,12 +81,14 @@ def __init__(self, persona): self.message_id = None self.todos = [] self.current_action = None # Currently executing action - self.completed_actions = [] # List of {'tool_call': str, 'result': str} + self.completed_actions = [] # List of tool calls (without results) + self.current_result = None # Current result message (replaces previous) self.initial_text_parts = [] # Text before any tool calls self.final_text_parts = [] # Text after tool calls self.active = False self.turn_active = False # Track if we're in an active Claude turn self.has_actions = False # Track if we've seen any tool calls + self.in_final_phase = False # Track if we're in final summary phase def _same_todo_list(self, new_todos): """Check if this is the same todo list (just status updates).""" @@ -95,22 +100,22 @@ def _same_todo_list(self, new_todos): async def update_todos(self, todos): """Update todo list, creating or updating message as needed.""" - if self._same_todo_list(todos) and self.active: - self.todos = todos - await self._update_message() - else: - # New todo list - create new message - self.todos = todos + self.todos = todos + + # Always ensure template is active when we have todos + if not self.active: self.active = True await self._create_message() + else: + await self._update_message() return "" async def update_action(self, action): """Start a new action - show it as current action.""" if self.active: - # Complete previous action if it exists + # Complete previous action if it exists (just add tool call to completed) if self.current_action: - self.completed_actions.append(self.current_action) + self.completed_actions.append(self.current_action['tool_call']) # Mark that we've seen actions self.has_actions = True @@ -125,6 +130,10 @@ async def update_action(self, action): 'tool_call': action, 'result': 'Executing...' } + + # Reset final phase when we start a new action + self.in_final_phase = False + await self._update_message() return "" return action @@ -133,7 +142,17 @@ async def update_action_result(self, result): """Update the result of the current action.""" if self.active and self.current_action: # Update the current action's result - self.current_action['result'] = result + self.current_action['result'] = self._escape_markdown(result) + + # Set this as the current result (just the result text, no tool call) + self.current_result = self.current_action['result'] + + # After receiving a result, the action is complete + # Move just the tool call to completed actions + self.completed_actions.append(self.current_action['tool_call']) + self.current_action = None + self.in_final_phase = True # Now any subsequent text should go to final_text + await self._update_message() return "" return result @@ -142,11 +161,24 @@ async def update_text(self, text): """Add text to template response.""" if self.active: # If we haven't seen actions yet, add to initial text - # If we have seen actions, add to final text if not self.has_actions: self.initial_text_parts.append(text) + elif self.current_action and not self.in_final_phase: + # If we have a current action and not in final phase, treat this text as its result + # This handles cases where tool results come as text blocks + if self.current_action['result'] == 'Executing...': + self.current_action['result'] = text + # Update current result display (just the result text) + self.current_result = text + else: + # Append to existing result if already has content + self.current_action['result'] += '\\n' + text + self.current_result = self.current_action['result'] else: + # No current action or we're in final phase - this is final summary text + # This should appear after the horizontal rule self.final_text_parts.append(text) + self.in_final_phase = True # Mark that we're now in final phase await self._update_message() return "" return text @@ -181,6 +213,7 @@ def _render_template(self): todos=self.todos, current_action=self.current_action, completed_actions=self.completed_actions if len(self.completed_actions) > 0 else None, + current_result=self.current_result, initial_text='\n'.join(self.initial_text_parts) if self.initial_text_parts else None, final_text='\n'.join(self.final_text_parts) if self.final_text_parts else None ) @@ -193,14 +226,40 @@ async def complete(self): self.completed_actions.append(self.current_action) self.current_action = None + # Mark that we're now in final phase - any subsequent text should go to final_text + self.in_final_phase = True + # Do final template update to show completed state await self._update_message() + elif self.active: + # Still mark final phase even without message_id + self.in_final_phase = True - # Mark as inactive and clear message ID but preserve completed actions - self.message_id = None - self.active = False + # Keep template active but mark turn as inactive + # This allows final text to still be processed self.turn_active = False + def _escape_markdown(self, text): + """Escape markdown characters in text to prevent formatting issues.""" + # Escape common markdown characters that could cause formatting problems + escapes = { + '*': '\\*', + '_': '\\_', + '`': '\\`', + '#': '\\#', + '[': '\\[', + ']': '\\]', + '(': '\\(', + ')': '\\)', + '{': '\\{', + '}': '\\}', + '\\': '\\\\' + } + result = str(text) + for char, escape in escapes.items(): + result = result.replace(char, escape) + return result + def format_tool_input(self, tool_name, tool_input): """Format tool input for Claude Code CLI style display.""" if tool_name in self.TOOL_PARAM_MAPPING: @@ -208,8 +267,8 @@ def format_tool_input(self, tool_name, tool_input): value = tool_input.get(key, '') # For long values, truncate with ellipsis if len(str(value)) > 60: - return str(value)[:60] + '…' - return str(value) + return self._escape_markdown(str(value)[:60] + '…') + return self._escape_markdown(str(value)) # Format remaining args (excluding content) args = [] @@ -218,7 +277,7 @@ def format_tool_input(self, tool_name, tool_input): val_str = str(v) if len(val_str) > 30: val_str = val_str[:30] + '…' - args.append(f"{k}={val_str}") + args.append(f"{k}={self._escape_markdown(val_str)}") return ', '.join(args) async def process_message_block(self, block): @@ -268,10 +327,12 @@ def reset(self): self.todos = [] self.current_action = None self.completed_actions = [] + self.current_result = None self.initial_text_parts = [] self.final_text_parts = [] self.active = False self.turn_active = False self.has_actions = False + self.in_final_phase = False From 0c75168635a88ca2af69c845d5d21bd63402d4dd Mon Sep 17 00:00:00 2001 From: Zach Sailer Date: Mon, 18 Aug 2025 14:22:16 -0700 Subject: [PATCH 6/8] move data to a dataclass --- src/jupyter_ai_claude_code/templates.py | 131 ++++++++++++++---------- 1 file changed, 76 insertions(+), 55 deletions(-) diff --git a/src/jupyter_ai_claude_code/templates.py b/src/jupyter_ai_claude_code/templates.py index c61ccc9..b0e98e8 100644 --- a/src/jupyter_ai_claude_code/templates.py +++ b/src/jupyter_ai_claude_code/templates.py @@ -2,22 +2,41 @@ from typing import Dict, Any, List, Optional import datetime +from dataclasses import dataclass from jinja2 import Template from jupyterlab_chat.models import Message, NewMessage from claude_code_sdk import TextBlock, ToolUseBlock, ToolResultBlock, AssistantMessage +@dataclass +class MessageData: + """Dataclass containing all data needed for message template rendering.""" + todos: List[Dict[str, Any]] = None + current_action: Dict[str, str] = None + completed_actions: List[str] = None + current_result: str = None + initial_text: str = None + final_text: str = None + + def __post_init__(self): + """Initialize mutable defaults.""" + if self.todos is None: + self.todos = [] + if self.completed_actions is None: + self.completed_actions = [] + + # Template for rendering consolidated actions and final response TODO_TEMPLATE = Template(""" -{%- if initial_text %} -{{ initial_text }} +{%- if data.initial_text %} +{{ data.initial_text }} {%- endif %} -{% if todos %} +{% if data.todos %} **Task Progress:** -{%- for todo in todos %} +{%- for todo in data.todos %} {%- if todo.status == 'completed' %} - [x] ~~{{ todo.content }}~~ {%- elif todo.status == 'in_progress' %} @@ -29,21 +48,21 @@ {%- endif %} -{%- if current_action %} +{%- if data.current_action %} **Current Tool:** -{{ current_action.tool_call }} +{{ data.current_action.tool_call }} ⎿ Executing... {%- endif %} -{% if completed_actions %} +{% if data.completed_actions %} **Tools Called:**
-See details ({{ completed_actions|length }}) +See details ({{ data.completed_actions|length }}) -{%- for action in completed_actions %} +{%- for action in data.completed_actions %} {{ action }} ⎿ Completed @@ -53,13 +72,13 @@
{% endif %} -{% if current_result or final_text %} +{% if data.current_result or data.final_text %} **Response:** -{% if current_result %} -{{ current_result }} +{% if data.current_result %} +{{ data.current_result }} {% endif %} -{% if final_text %} -{{ final_text }} +{% if data.final_text %} +{{ data.final_text }} {% endif %} {% endif %} """.strip()) @@ -79,12 +98,7 @@ class ClaudeCodeTemplateManager: def __init__(self, persona): self.persona = persona self.message_id = None - self.todos = [] - self.current_action = None # Currently executing action - self.completed_actions = [] # List of tool calls (without results) - self.current_result = None # Current result message (replaces previous) - self.initial_text_parts = [] # Text before any tool calls - self.final_text_parts = [] # Text after tool calls + self.message_data = MessageData() # Single dataclass for all template data self.active = False self.turn_active = False # Track if we're in an active Claude turn self.has_actions = False # Track if we've seen any tool calls @@ -92,15 +106,15 @@ def __init__(self, persona): def _same_todo_list(self, new_todos): """Check if this is the same todo list (just status updates).""" - if not self.todos: + if not self.message_data.todos: return False - old_ids = {t['id'] for t in self.todos} + old_ids = {t['id'] for t in self.message_data.todos} new_ids = {t['id'] for t in new_todos} return old_ids == new_ids async def update_todos(self, todos): """Update todo list, creating or updating message as needed.""" - self.todos = todos + self.message_data.todos = todos # Always ensure template is active when we have todos if not self.active: @@ -114,8 +128,8 @@ async def update_action(self, action): """Start a new action - show it as current action.""" if self.active: # Complete previous action if it exists (just add tool call to completed) - if self.current_action: - self.completed_actions.append(self.current_action['tool_call']) + if self.message_data.current_action: + self.message_data.completed_actions.append(self.message_data.current_action['tool_call']) # Mark that we've seen actions self.has_actions = True @@ -126,7 +140,7 @@ async def update_action(self, action): await self._create_message() # Set new current action - self.current_action = { + self.message_data.current_action = { 'tool_call': action, 'result': 'Executing...' } @@ -140,17 +154,17 @@ async def update_action(self, action): async def update_action_result(self, result): """Update the result of the current action.""" - if self.active and self.current_action: + if self.active and self.message_data.current_action: # Update the current action's result - self.current_action['result'] = self._escape_markdown(result) + self.message_data.current_action['result'] = self._escape_markdown(result) # Set this as the current result (just the result text, no tool call) - self.current_result = self.current_action['result'] + self.message_data.current_result = self.message_data.current_action['result'] # After receiving a result, the action is complete # Move just the tool call to completed actions - self.completed_actions.append(self.current_action['tool_call']) - self.current_action = None + self.message_data.completed_actions.append(self.message_data.current_action['tool_call']) + self.message_data.current_action = None self.in_final_phase = True # Now any subsequent text should go to final_text await self._update_message() @@ -162,22 +176,28 @@ async def update_text(self, text): if self.active: # If we haven't seen actions yet, add to initial text if not self.has_actions: - self.initial_text_parts.append(text) - elif self.current_action and not self.in_final_phase: + if self.message_data.initial_text: + self.message_data.initial_text += '\n' + text + else: + self.message_data.initial_text = text + elif self.message_data.current_action and not self.in_final_phase: # If we have a current action and not in final phase, treat this text as its result # This handles cases where tool results come as text blocks - if self.current_action['result'] == 'Executing...': - self.current_action['result'] = text + if self.message_data.current_action['result'] == 'Executing...': + self.message_data.current_action['result'] = text # Update current result display (just the result text) - self.current_result = text + self.message_data.current_result = text else: # Append to existing result if already has content - self.current_action['result'] += '\\n' + text - self.current_result = self.current_action['result'] + self.message_data.current_action['result'] += '\\n' + text + self.message_data.current_result = self.message_data.current_action['result'] else: # No current action or we're in final phase - this is final summary text # This should appear after the horizontal rule - self.final_text_parts.append(text) + if self.message_data.final_text: + self.message_data.final_text += '\n' + text + else: + self.message_data.final_text = text self.in_final_phase = True # Mark that we're now in final phase await self._update_message() return "" @@ -209,22 +229,28 @@ async def _update_message(self): def _render_template(self): """Render current template state.""" - return TODO_TEMPLATE.render( - todos=self.todos, - current_action=self.current_action, - completed_actions=self.completed_actions if len(self.completed_actions) > 0 else None, - current_result=self.current_result, - initial_text='\n'.join(self.initial_text_parts) if self.initial_text_parts else None, - final_text='\n'.join(self.final_text_parts) if self.final_text_parts else None + # Prepare completed_actions for template (only if non-empty) + completed_actions = self.message_data.completed_actions if len(self.message_data.completed_actions) > 0 else None + + # Create a copy of message_data with the processed completed_actions + data = MessageData( + todos=self.message_data.todos, + current_action=self.message_data.current_action, + completed_actions=completed_actions, + current_result=self.message_data.current_result, + initial_text=self.message_data.initial_text, + final_text=self.message_data.final_text ) + + return TODO_TEMPLATE.render(data=data) async def complete(self): """Complete template - move current action to completed actions.""" if self.active and self.message_id: # Move current action to completed actions if it exists - if self.current_action: - self.completed_actions.append(self.current_action) - self.current_action = None + if self.message_data.current_action: + self.message_data.completed_actions.append(self.message_data.current_action) + self.message_data.current_action = None # Mark that we're now in final phase - any subsequent text should go to final_text self.in_final_phase = True @@ -324,12 +350,7 @@ async def claude_message_to_str(self, message) -> Optional[str]: def reset(self): """Reset for new conversation.""" self.message_id = None - self.todos = [] - self.current_action = None - self.completed_actions = [] - self.current_result = None - self.initial_text_parts = [] - self.final_text_parts = [] + self.message_data = MessageData() # Reset to new dataclass instance self.active = False self.turn_active = False self.has_actions = False From 115e42d497cbd5372aa5e9f86c1b6916c76eb6d0 Mon Sep 17 00:00:00 2001 From: Zach Sailer Date: Mon, 18 Aug 2025 16:34:26 -0700 Subject: [PATCH 7/8] Make file links clickable --- src/jupyter_ai_claude_code/persona.py | 6 +- src/jupyter_ai_claude_code/templates.py | 118 +++++++++++++++--------- 2 files changed, 80 insertions(+), 44 deletions(-) diff --git a/src/jupyter_ai_claude_code/persona.py b/src/jupyter_ai_claude_code/persona.py index 4bb91a0..888e575 100644 --- a/src/jupyter_ai_claude_code/persona.py +++ b/src/jupyter_ai_claude_code/persona.py @@ -61,8 +61,10 @@ async def _process_response_message(self, message_iterator) -> AsyncIterator[str await self.template_mgr.complete() template_was_used = True - # Only yield empty string if no content was produced and template wasn't used - if not has_content and not template_was_used: + # Always yield something to complete the stream + if template_was_used: + yield "" # Empty yield to signal completion when template handled everything + elif not has_content: yield "" # Ensure stream completes for empty responses def _generate_prompt(self, message: Message) -> str: diff --git a/src/jupyter_ai_claude_code/templates.py b/src/jupyter_ai_claude_code/templates.py index b0e98e8..4d75801 100644 --- a/src/jupyter_ai_claude_code/templates.py +++ b/src/jupyter_ai_claude_code/templates.py @@ -2,7 +2,9 @@ from typing import Dict, Any, List, Optional import datetime -from dataclasses import dataclass +import os +import re +from dataclasses import dataclass, field from jinja2 import Template from jupyterlab_chat.models import Message, NewMessage from claude_code_sdk import TextBlock, ToolUseBlock, ToolResultBlock, AssistantMessage @@ -11,19 +13,12 @@ @dataclass class MessageData: """Dataclass containing all data needed for message template rendering.""" - todos: List[Dict[str, Any]] = None + todos: List[Dict[str, Any]] = field(default_factory=list) current_action: Dict[str, str] = None - completed_actions: List[str] = None + completed_actions: List[str] = field(default_factory=list) current_result: str = None initial_text: str = None final_text: str = None - - def __post_init__(self): - """Initialize mutable defaults.""" - if self.todos is None: - self.todos = [] - if self.completed_actions is None: - self.completed_actions = [] # Template for rendering consolidated actions and final response @@ -56,7 +51,7 @@ def __post_init__(self): {%- endif %} -{% if data.completed_actions %} +{% if data.completed_actions and data.completed_actions|length > 0 %} **Tools Called:**
@@ -87,6 +82,10 @@ def __post_init__(self): class ClaudeCodeTemplateManager: """Manages in-place template message updates for Claude Code SDK messages.""" + # Constants + MAX_TOOL_VALUE_LENGTH = 60 + MAX_ARG_LENGTH = 30 + # Tool parameter mapping for display formatting TOOL_PARAM_MAPPING = { 'Task': 'description', 'Bash': 'command', 'Glob': 'pattern', 'Grep': 'pattern', @@ -95,6 +94,9 @@ class ClaudeCodeTemplateManager: 'WebFetch': 'url', 'WebSearch': 'query' } + # Tools that should have clickable file links + FILE_LINK_TOOLS = {'Read', 'Edit', 'MultiEdit', 'Write'} + def __init__(self, persona): self.persona = persona self.message_id = None @@ -112,6 +114,13 @@ def _same_todo_list(self, new_todos): new_ids = {t['id'] for t in new_todos} return old_ids == new_ids + async def _ensure_message_exists(self): + """Ensure a message exists, creating one if needed.""" + if not self.message_id: + await self._create_message() + else: + await self._update_message() + async def update_todos(self, todos): """Update todo list, creating or updating message as needed.""" self.message_data.todos = todos @@ -119,9 +128,8 @@ async def update_todos(self, todos): # Always ensure template is active when we have todos if not self.active: self.active = True - await self._create_message() - else: - await self._update_message() + + await self._ensure_message_exists() return "" async def update_action(self, action): @@ -137,7 +145,8 @@ async def update_action(self, action): # Always start template if not active yet if not self.turn_active: self.turn_active = True - await self._create_message() + if not self.message_id: + await self._create_message() # Set new current action self.message_data.current_action = { @@ -148,14 +157,14 @@ async def update_action(self, action): # Reset final phase when we start a new action self.in_final_phase = False - await self._update_message() + await self._ensure_message_exists() return "" return action async def update_action_result(self, result): """Update the result of the current action.""" if self.active and self.message_data.current_action: - # Update the current action's result + # Update the current action's result (just escape markdown, no aggressive path linking) self.message_data.current_action['result'] = self._escape_markdown(result) # Set this as the current result (just the result text, no tool call) @@ -167,7 +176,7 @@ async def update_action_result(self, result): self.message_data.current_action = None self.in_final_phase = True # Now any subsequent text should go to final_text - await self._update_message() + await self._ensure_message_exists() return "" return result @@ -189,7 +198,7 @@ async def update_text(self, text): self.message_data.current_result = text else: # Append to existing result if already has content - self.message_data.current_action['result'] += '\\n' + text + self.message_data.current_action['result'] += '\n' + text self.message_data.current_result = self.message_data.current_action['result'] else: # No current action or we're in final phase - this is final summary text @@ -199,7 +208,9 @@ async def update_text(self, text): else: self.message_data.final_text = text self.in_final_phase = True # Mark that we're now in final phase - await self._update_message() + + # Create or update message + await self._ensure_message_exists() return "" return text @@ -229,34 +240,21 @@ async def _update_message(self): def _render_template(self): """Render current template state.""" - # Prepare completed_actions for template (only if non-empty) - completed_actions = self.message_data.completed_actions if len(self.message_data.completed_actions) > 0 else None - - # Create a copy of message_data with the processed completed_actions - data = MessageData( - todos=self.message_data.todos, - current_action=self.message_data.current_action, - completed_actions=completed_actions, - current_result=self.message_data.current_result, - initial_text=self.message_data.initial_text, - final_text=self.message_data.final_text - ) - - return TODO_TEMPLATE.render(data=data) + return TODO_TEMPLATE.render(data=self.message_data) async def complete(self): """Complete template - move current action to completed actions.""" if self.active and self.message_id: # Move current action to completed actions if it exists if self.message_data.current_action: - self.message_data.completed_actions.append(self.message_data.current_action) + self.message_data.completed_actions.append(self.message_data.current_action['tool_call']) self.message_data.current_action = None # Mark that we're now in final phase - any subsequent text should go to final_text self.in_final_phase = True # Do final template update to show completed state - await self._update_message() + await self._ensure_message_exists() elif self.active: # Still mark final phase even without message_id self.in_final_phase = True @@ -286,23 +284,59 @@ def _escape_markdown(self, text): result = result.replace(char, escape) return result + def _make_jupyter_file_link(self, file_path): + """Convert file path to clickable JupyterLab file link.""" + # Get server root reference for path resolution + server_root_reference = self._get_server_root_reference() + relative_path = self._resolve_relative_path(file_path, server_root_reference) + return f"[{file_path}](/files/{relative_path})" + + def _get_server_root_reference(self): + """Get server root directory reference from persona.""" + try: + workspace_dir = getattr(self.persona, 'get_workspace_dir', lambda: None)() + chat_dir = getattr(self.persona, 'get_chat_dir', lambda: None)() + return workspace_dir or chat_dir + except Exception: + return None + + def _resolve_relative_path(self, file_path, server_root_reference): + """Resolve file path to be relative to server root.""" + if not file_path.startswith('/') or not server_root_reference: + return file_path.lstrip('/') + + try: + relative_path = os.path.relpath(file_path, start=server_root_reference) + # If path goes outside server root, use basename only + return os.path.basename(file_path) if relative_path.startswith('..') else relative_path + except (ValueError, OSError): + return os.path.basename(file_path) + def format_tool_input(self, tool_name, tool_input): """Format tool input for Claude Code CLI style display.""" if tool_name in self.TOOL_PARAM_MAPPING: key = self.TOOL_PARAM_MAPPING[tool_name] value = tool_input.get(key, '') - # For long values, truncate with ellipsis - if len(str(value)) > 60: - return self._escape_markdown(str(value)[:60] + '…') - return self._escape_markdown(str(value)) + + # Make file paths clickable for file-related tools + if tool_name in self.FILE_LINK_TOOLS and value: + if len(str(value)) > self.MAX_TOOL_VALUE_LENGTH: + truncated = str(value)[:self.MAX_TOOL_VALUE_LENGTH] + '…' + return self._make_jupyter_file_link(truncated) + return self._make_jupyter_file_link(str(value)) + else: + # For other tools, just escape markdown + if len(str(value)) > self.MAX_TOOL_VALUE_LENGTH: + return self._escape_markdown(str(value)[:self.MAX_TOOL_VALUE_LENGTH] + '…') + return self._escape_markdown(str(value)) # Format remaining args (excluding content) args = [] for k, v in tool_input.items(): if k != 'content': val_str = str(v) - if len(val_str) > 30: - val_str = val_str[:30] + '…' + if len(val_str) > self.MAX_ARG_LENGTH: + val_str = val_str[:self.MAX_ARG_LENGTH] + '…' args.append(f"{k}={self._escape_markdown(val_str)}") return ', '.join(args) From c41ff43e76ecbfe304413db4ddd7bbe3f8cc4b3e Mon Sep 17 00:00:00 2001 From: Zach Sailer Date: Mon, 18 Aug 2025 16:58:53 -0700 Subject: [PATCH 8/8] ensure relative links work --- src/jupyter_ai_claude_code/persona.py | 7 +--- src/jupyter_ai_claude_code/templates.py | 43 +++++++++++++++++++++---- 2 files changed, 38 insertions(+), 12 deletions(-) diff --git a/src/jupyter_ai_claude_code/persona.py b/src/jupyter_ai_claude_code/persona.py index 888e575..1e2cdcf 100644 --- a/src/jupyter_ai_claude_code/persona.py +++ b/src/jupyter_ai_claude_code/persona.py @@ -27,16 +27,11 @@ def __init__(self, *args, **kwargs): @property def defaults(self) -> PersonaDefaults: """Return default configuration for the Claude Code persona.""" - system_prompt = ("I am Claude Code, an AI assistant with access to development tools. " - "When formatting responses, I use **bold text** for emphasis and section headers instead of markdown headings (# ## ###). " - "I keep formatting clean and readable without large headers. " - "For complex tasks requiring multiple steps (3+ actions), I proactively create a todo list using TodoWrite to track progress and keep the user informed of my plan.") - return PersonaDefaults( name="Claude", avatar_path="/files/.jupyter/claude.svg", description="Claude Code persona", - system_prompt=system_prompt, + system_prompt="...", ) async def _process_response_message(self, message_iterator) -> AsyncIterator[str]: diff --git a/src/jupyter_ai_claude_code/templates.py b/src/jupyter_ai_claude_code/templates.py index 4d75801..312eae6 100644 --- a/src/jupyter_ai_claude_code/templates.py +++ b/src/jupyter_ai_claude_code/templates.py @@ -284,12 +284,23 @@ def _escape_markdown(self, text): result = result.replace(char, escape) return result - def _make_jupyter_file_link(self, file_path): - """Convert file path to clickable JupyterLab file link.""" - # Get server root reference for path resolution + def _make_jupyter_file_link(self, file_path, tool_name=None): + """Convert file path to clickable JupyterLab file link if path exists or will be created.""" server_root_reference = self._get_server_root_reference() relative_path = self._resolve_relative_path(file_path, server_root_reference) - return f"[{file_path}](/files/{relative_path})" + + # Always create links for Write tools (file will be created) + # For other tools, only create link if file exists + should_create_link = ( + tool_name == 'Write' or + self._path_exists_on_server(relative_path, server_root_reference) + ) + + if should_create_link: + return f"[{file_path}](/files/{relative_path})" + else: + # Return plain text if path doesn't exist and it's not a Write operation + return str(file_path) def _get_server_root_reference(self): """Get server root directory reference from persona.""" @@ -311,6 +322,26 @@ def _resolve_relative_path(self, file_path, server_root_reference): return os.path.basename(file_path) if relative_path.startswith('..') else relative_path except (ValueError, OSError): return os.path.basename(file_path) + + def _path_exists_on_server(self, relative_path, server_root_reference): + """Check if the relative path exists on the server.""" + if not server_root_reference or not relative_path: + return False + + try: + # Construct full path from server root + full_path = os.path.join(server_root_reference, relative_path) + # Normalize path to handle any .. or . components + normalized_path = os.path.normpath(full_path) + + # Security check: ensure normalized path is still within server root + if not normalized_path.startswith(os.path.normpath(server_root_reference)): + return False + + # Check if file exists + return os.path.exists(normalized_path) + except Exception: + return False def format_tool_input(self, tool_name, tool_input): """Format tool input for Claude Code CLI style display.""" @@ -322,8 +353,8 @@ def format_tool_input(self, tool_name, tool_input): if tool_name in self.FILE_LINK_TOOLS and value: if len(str(value)) > self.MAX_TOOL_VALUE_LENGTH: truncated = str(value)[:self.MAX_TOOL_VALUE_LENGTH] + '…' - return self._make_jupyter_file_link(truncated) - return self._make_jupyter_file_link(str(value)) + return self._make_jupyter_file_link(truncated, tool_name) + return self._make_jupyter_file_link(str(value), tool_name) else: # For other tools, just escape markdown if len(str(value)) > self.MAX_TOOL_VALUE_LENGTH: