├── .envrc ├── assistant_ui_anywidget ├── static │ └── .gitkeep ├── ai │ ├── __init__.py │ ├── prompt_config.py │ ├── logger.py │ └── mock.py ├── __init__.py ├── global_agent.py ├── module_inspector.py └── simple_handlers.py ├── tests ├── __init__.py ├── test_widget.ipynb ├── test_widget.py ├── conftest.py ├── test_widget_simple.py ├── test_notebook.py ├── test_kernel_error_handling.py ├── test_prompt_config_regression.py ├── test_langgraph_approval.py ├── test_imported_modules_context.py ├── test_global_agent_auto_detection.py ├── test_widget_basic.py ├── test_ai_service_regression.py ├── test_git_native_tools.py ├── test_kernel_interface.py └── test_message_handlers.py ├── frontend ├── tsconfig.node.json ├── vitest.config.ts ├── vite.config.test.ts ├── tsconfig.json ├── scripts │ ├── serve-demo.cjs │ └── screenshot.cjs ├── eslint.config.js ├── src │ ├── test │ │ ├── basic.test.ts │ │ └── setup.ts │ ├── kernelApi.ts │ ├── types.ts │ ├── VariableExplorer.test.tsx │ ├── kernelApi.test.ts │ └── VariableExplorer.tsx ├── vite.config.ts ├── package.json ├── README.md └── demo.html ├── .prettierignore ├── .prettierrc ├── package.json ├── .vscode ├── extensions.json └── settings.json ├── .env.example ├── .github ├── dependabot.yml └── workflows │ └── ci.yml ├── .pre-commit-config.yaml ├── .gitignore ├── take_screenshot.py ├── examples ├── langgraph_approval_demo.ipynb ├── demo_global_agent.ipynb └── demo_kernel_assistant.ipynb ├── CLAUDE.md └── pyproject.toml /.envrc: -------------------------------------------------------------------------------- 1 | source .venv/bin/activate 2 | -------------------------------------------------------------------------------- /assistant_ui_anywidget/static/.gitkeep: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- 1 | """Test suite for assistant-ui-anywidget.""" 2 | -------------------------------------------------------------------------------- /frontend/tsconfig.node.json: -------------------------------------------------------------------------------- 1 | { 2 | "compilerOptions": { 3 | "composite": true, 4 | "skipLibCheck": true, 5 | "module": "ESNext", 6 | "moduleResolution": "bundler", 7 | "allowSyntheticDefaultImports": true, 8 | "strict": true 9 | }, 10 | "include": ["vite.config.ts"] 11 | } 12 | -------------------------------------------------------------------------------- /.prettierignore: -------------------------------------------------------------------------------- 1 | # Dependencies 2 | node_modules/ 3 | frontend/node_modules/ 4 | 5 | # Build outputs 6 | frontend/dist/ 7 | *.min.js 8 | *.min.css 9 | 10 | # Generated files 11 | .pytest_cache/ 12 | .mypy_cache/ 13 | *.pyc 14 | __pycache__/ 15 | 16 | # IDE files 17 | .vscode/ 18 | .idea/ 19 | 20 | # Other 21 | *.ipynb 22 | -------------------------------------------------------------------------------- /.prettierrc: -------------------------------------------------------------------------------- 1 | { 2 | "semi": true, 3 | "trailingComma": "es5", 4 | "singleQuote": false, 5 | "printWidth": 100, 6 | "tabWidth": 2, 7 | "useTabs": false, 8 | "quoteProps": "as-needed", 9 | "bracketSpacing": true, 10 | "bracketSameLine": false, 11 | "arrowParens": "avoid", 12 | "endOfLine": "lf" 13 | } 14 | -------------------------------------------------------------------------------- /package.json: -------------------------------------------------------------------------------- 1 | { 2 | "devDependencies": { 3 | "@testing-library/jest-dom": "^6.6.3", 4 | "@testing-library/react": "^16.3.0", 5 | "@testing-library/user-event": "^14.6.1", 6 | "@vitest/coverage-v8": "^3.2.4", 7 | "@vitest/ui": "^3.2.4", 8 | "jsdom": "^26.1.0", 9 | "puppeteer": "^24.14.0", 10 | "vitest": "^3.2.4" 11 | } 12 | } 13 | -------------------------------------------------------------------------------- /.vscode/extensions.json: -------------------------------------------------------------------------------- 1 | { 2 | "recommendations": [ 3 | "esbenp.prettier-vscode", 4 | "ms-python.python", 5 | "ms-python.pylint", 6 | "charliermarsh.ruff", 7 | "dbaeumer.vscode-eslint", 8 | "bradlc.vscode-tailwindcss", 9 | "ms-toolsai.jupyter", 10 | "ms-vscode.vscode-typescript-next", 11 | "ms-vscode.vscode-json" 12 | ] 13 | } 14 | -------------------------------------------------------------------------------- /assistant_ui_anywidget/ai/__init__.py: -------------------------------------------------------------------------------- 1 | """AI integration module for the assistant widget.""" 2 | 3 | from .langgraph_service import LangGraphAIService, ChatResult 4 | from .logger import ConversationLogger 5 | 6 | # Use LangGraph as the only AI service 7 | AIService = LangGraphAIService 8 | 9 | __all__ = [ 10 | "AIService", 11 | "LangGraphAIService", 12 | "ChatResult", 13 | "ConversationLogger", 14 | ] 15 | -------------------------------------------------------------------------------- /frontend/vitest.config.ts: -------------------------------------------------------------------------------- 1 | import { defineConfig } from "vitest/config"; 2 | import react from "@vitejs/plugin-react"; 3 | 4 | export default defineConfig({ 5 | plugins: [react()], 6 | test: { 7 | environment: "jsdom", 8 | globals: true, 9 | setupFiles: "./src/test/setup.ts", 10 | coverage: { 11 | reporter: ["text", "json", "html"], 12 | exclude: ["node_modules/", "src/test/", "dist/", "*.config.ts", "*.config.js"], 13 | }, 14 | }, 15 | }); 16 | -------------------------------------------------------------------------------- /.env.example: -------------------------------------------------------------------------------- 1 | # AI Provider API Keys 2 | # Uncomment and set the API key(s) you want to use 3 | # The widget will automatically detect and use the first available provider 4 | 5 | # OpenAI (GPT-4, GPT-3.5, etc.) 6 | # OPENAI_API_KEY=sk-... 7 | 8 | # Anthropic (Claude 3 Opus, Sonnet, Haiku) 9 | # ANTHROPIC_API_KEY=sk-ant-... 10 | 11 | # Google AI (Gemini Pro, Gemini Ultra) 12 | # GOOGLE_API_KEY=... 13 | 14 | # Optional: Override default model selection 15 | # AI_MODEL=gpt-4-turbo-preview 16 | # AI_PROVIDER=openai 17 | -------------------------------------------------------------------------------- /frontend/vite.config.test.ts: -------------------------------------------------------------------------------- 1 | /// 2 | import { defineConfig } from "vite"; 3 | import react from "@vitejs/plugin-react"; 4 | 5 | export default defineConfig({ 6 | plugins: [react()], 7 | test: { 8 | globals: true, 9 | environment: "jsdom", 10 | setupFiles: ["./src/test/setup.ts"], 11 | css: false, 12 | coverage: { 13 | provider: "v8", 14 | reporter: ["text", "json", "html"], 15 | exclude: [ 16 | "node_modules/", 17 | "src/test/", 18 | "**/*.d.ts", 19 | "**/*.test.{ts,tsx}", 20 | "**/*.config.{ts,js}", 21 | ], 22 | }, 23 | }, 24 | }); 25 | -------------------------------------------------------------------------------- /frontend/tsconfig.json: -------------------------------------------------------------------------------- 1 | { 2 | "compilerOptions": { 3 | "target": "ES2020", 4 | "useDefineForClassFields": true, 5 | "lib": ["ES2020", "DOM", "DOM.Iterable"], 6 | "module": "ESNext", 7 | "skipLibCheck": true, 8 | "moduleResolution": "bundler", 9 | "allowImportingTsExtensions": true, 10 | "resolveJsonModule": true, 11 | "isolatedModules": true, 12 | "noEmit": true, 13 | "jsx": "react-jsx", 14 | "strict": true, 15 | "noUnusedLocals": true, 16 | "noUnusedParameters": true, 17 | "noFallthroughCasesInSwitch": true 18 | }, 19 | "include": ["src"], 20 | "references": [{ "path": "./tsconfig.node.json" }] 21 | } 22 | -------------------------------------------------------------------------------- /assistant_ui_anywidget/__init__.py: -------------------------------------------------------------------------------- 1 | """Assistant UI AnyWidget package.""" 2 | 3 | from .agent_widget import AgentWidget 4 | 5 | # Global agent interface for notebook convenience 6 | from .global_agent import ( 7 | get_agent, 8 | reset_agent, 9 | ) 10 | from .kernel_interface import ExecutionResult, KernelInterface, StackFrame, VariableInfo 11 | from .simple_handlers import SimpleHandlers 12 | 13 | __all__ = [ 14 | # Core widget classes 15 | "AgentWidget", 16 | # Kernel and messaging interfaces 17 | "KernelInterface", 18 | "VariableInfo", 19 | "ExecutionResult", 20 | "StackFrame", 21 | "SimpleHandlers", 22 | # Global agent interface (recommended for notebooks) 23 | "get_agent", 24 | "reset_agent", 25 | ] 26 | -------------------------------------------------------------------------------- /frontend/scripts/serve-demo.cjs: -------------------------------------------------------------------------------- 1 | const http = require('http'); 2 | const fs = require('fs'); 3 | const path = require('path'); 4 | 5 | const PORT = 3000; 6 | const HTML_FILE = path.join(__dirname, '..', 'demo.html'); 7 | 8 | const server = http.createServer((req, res) => { 9 | if (req.url === '/' || req.url === '/index.html') { 10 | fs.readFile(HTML_FILE, 'utf8', (err, content) => { 11 | if (err) { 12 | res.writeHead(500); 13 | res.end('Error loading demo page'); 14 | return; 15 | } 16 | res.writeHead(200, { 'Content-Type': 'text/html' }); 17 | res.end(content); 18 | }); 19 | } else { 20 | res.writeHead(404); 21 | res.end('Not found'); 22 | } 23 | }); 24 | 25 | server.listen(PORT, () => { 26 | console.log(`Demo server running at http://localhost:${PORT}`); 27 | }); 28 | -------------------------------------------------------------------------------- /.github/dependabot.yml: -------------------------------------------------------------------------------- 1 | version: 2 2 | updates: 3 | # Enable version updates for npm 4 | - package-ecosystem: "npm" 5 | directory: "/frontend" 6 | schedule: 7 | interval: "weekly" 8 | commit-message: 9 | prefix: "npm" 10 | include: "scope" 11 | reviewers: 12 | - "basnijholt" 13 | open-pull-requests-limit: 5 14 | 15 | # Enable version updates for Python pip 16 | - package-ecosystem: "pip" 17 | directory: "/" 18 | schedule: 19 | interval: "weekly" 20 | commit-message: 21 | prefix: "pip" 22 | include: "scope" 23 | reviewers: 24 | - "basnijholt" 25 | open-pull-requests-limit: 5 26 | 27 | # Enable version updates for GitHub Actions 28 | - package-ecosystem: "github-actions" 29 | directory: "/" 30 | schedule: 31 | interval: "weekly" 32 | commit-message: 33 | prefix: "github-actions" 34 | include: "scope" 35 | reviewers: 36 | - "basnijholt" 37 | open-pull-requests-limit: 3 38 | -------------------------------------------------------------------------------- /tests/test_widget.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": null, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "from agent_widget import AgentWidget" 10 | ] 11 | }, 12 | { 13 | "cell_type": "code", 14 | "execution_count": null, 15 | "metadata": {}, 16 | "outputs": [], 17 | "source": [ 18 | "widget = AgentWidget(show_help=False)\n", 19 | "widget" 20 | ] 21 | }, 22 | { 23 | "cell_type": "code", 24 | "execution_count": null, 25 | "metadata": {}, 26 | "outputs": [], 27 | "source": [] 28 | } 29 | ], 30 | "metadata": { 31 | "kernelspec": { 32 | "display_name": "Python 3 (ipykernel)", 33 | "language": "python", 34 | "name": "python3" 35 | }, 36 | "language_info": { 37 | "codemirror_mode": { 38 | "name": "ipython", 39 | "version": 3 40 | }, 41 | "file_extension": ".py", 42 | "mimetype": "text/x-python", 43 | "name": "python", 44 | "nbconvert_exporter": "python", 45 | "pygments_lexer": "ipython3", 46 | "version": "3.13.0" 47 | } 48 | }, 49 | "nbformat": 4, 50 | "nbformat_minor": 4 51 | } 52 | -------------------------------------------------------------------------------- /tests/test_widget.py: -------------------------------------------------------------------------------- 1 | """Test script for the AgentWidget.""" 2 | 3 | import pathlib 4 | 5 | from assistant_ui_anywidget.agent_widget import AgentWidget 6 | 7 | 8 | def test_widget_creation() -> None: 9 | """Test that we can create a widget instance.""" 10 | widget = AgentWidget(show_help=False) 11 | print("✓ Widget created successfully") 12 | print(f" - Initial message: '{widget.message}'") 13 | 14 | # Debug the ESM path 15 | esm_path = pathlib.Path(__file__).parent.parent / "frontend" / "dist" / "index.js" 16 | print(f" - Expected ESM path: {esm_path}") 17 | print(f" - ESM path exists: {esm_path.exists()}") 18 | print(f" - Widget ESM path: {widget._esm}") 19 | 20 | # Assert the widget was created successfully 21 | assert widget is not None 22 | assert hasattr(widget, "message") 23 | 24 | 25 | if __name__ == "__main__": 26 | print("Testing AgentWidget...") 27 | test_widget_creation() 28 | print("\nWidget ready! In a Jupyter notebook, you can use:") 29 | print(" from assistant_ui_anywidget import AgentWidget") 30 | print(" widget = AgentWidget(show_help=False)") 31 | print(" widget # Display the widget") 32 | -------------------------------------------------------------------------------- /tests/conftest.py: -------------------------------------------------------------------------------- 1 | """Pytest configuration and fixtures for widget tests.""" 2 | 3 | from collections.abc import Callable 4 | 5 | import pytest 6 | 7 | from assistant_ui_anywidget.agent_widget import AgentWidget 8 | 9 | 10 | @pytest.fixture # type: ignore[misc] 11 | def widget() -> AgentWidget: 12 | """Create a fresh widget instance for each test.""" 13 | return AgentWidget(enable_ai=False, show_help=False) 14 | 15 | 16 | @pytest.fixture # type: ignore[misc] 17 | def sample_messages() -> list[dict[str, str]]: 18 | """Sample message data for testing.""" 19 | return [ 20 | {"role": "user", "content": "Hello world"}, 21 | {"role": "assistant", "content": "Hello! How can I help you?"}, 22 | {"role": "user", "content": "What is the weather like?"}, 23 | { 24 | "role": "assistant", 25 | "content": "I don't have access to current weather data.", 26 | }, 27 | ] 28 | 29 | 30 | @pytest.fixture # type: ignore[misc] 31 | def ui_message_factory() -> Callable[[str, str], dict[str, str]]: 32 | """Factory for creating UI message objects.""" 33 | 34 | def _create_ui_message( 35 | text: str, message_type: str = "user_message" 36 | ) -> dict[str, str]: 37 | return {"type": message_type, "text": text} 38 | 39 | return _create_ui_message 40 | -------------------------------------------------------------------------------- /.vscode/settings.json: -------------------------------------------------------------------------------- 1 | { 2 | "editor.formatOnSave": true, 3 | "editor.defaultFormatter": "esbenp.prettier-vscode", 4 | "editor.codeActionsOnSave": { 5 | "source.fixAll.eslint": "explicit" 6 | }, 7 | "python.defaultInterpreterPath": ".venv/bin/python", 8 | "python.linting.enabled": true, 9 | "python.linting.pylintEnabled": false, 10 | "python.linting.flake8Enabled": false, 11 | "python.formatting.provider": "none", 12 | "python.analysis.typeCheckingMode": "strict", 13 | "[python]": { 14 | "editor.defaultFormatter": "charliermarsh.ruff", 15 | "editor.formatOnSave": true, 16 | "editor.codeActionsOnSave": { 17 | "source.organizeImports": "explicit" 18 | } 19 | }, 20 | "[typescript]": { 21 | "editor.defaultFormatter": "esbenp.prettier-vscode" 22 | }, 23 | "[typescriptreact]": { 24 | "editor.defaultFormatter": "esbenp.prettier-vscode" 25 | }, 26 | "[javascript]": { 27 | "editor.defaultFormatter": "esbenp.prettier-vscode" 28 | }, 29 | "[javascriptreact]": { 30 | "editor.defaultFormatter": "esbenp.prettier-vscode" 31 | }, 32 | "[json]": { 33 | "editor.defaultFormatter": "esbenp.prettier-vscode" 34 | }, 35 | "[markdown]": { 36 | "editor.defaultFormatter": "esbenp.prettier-vscode" 37 | }, 38 | "files.associations": { 39 | "*.ipynb": "jupyter-notebook" 40 | }, 41 | "jupyter.askForKernelRestart": false, 42 | "typescript.preferences.importModuleSpecifier": "relative" 43 | } 44 | -------------------------------------------------------------------------------- /frontend/eslint.config.js: -------------------------------------------------------------------------------- 1 | import js from "@eslint/js"; 2 | import globals from "globals"; 3 | import reactHooks from "eslint-plugin-react-hooks"; 4 | import reactRefresh from "eslint-plugin-react-refresh"; 5 | import tseslint from "@typescript-eslint/eslint-plugin"; 6 | import tsparser from "@typescript-eslint/parser"; 7 | import react from "eslint-plugin-react"; 8 | 9 | export default [ 10 | { 11 | ignores: ["dist", "node_modules"], 12 | }, 13 | { 14 | files: ["**/*.{ts,tsx}"], 15 | languageOptions: { 16 | ecmaVersion: 2020, 17 | globals: globals.browser, 18 | parser: tsparser, 19 | parserOptions: { 20 | ecmaVersion: "latest", 21 | ecmaFeatures: { jsx: true }, 22 | sourceType: "module", 23 | }, 24 | }, 25 | settings: { 26 | react: { version: "18.2" }, 27 | }, 28 | plugins: { 29 | "@typescript-eslint": tseslint, 30 | react: react, 31 | "react-hooks": reactHooks, 32 | "react-refresh": reactRefresh, 33 | }, 34 | rules: { 35 | ...js.configs.recommended.rules, 36 | ...tseslint.configs.recommended.rules, 37 | ...reactHooks.configs.recommended.rules, 38 | "react-refresh/only-export-components": ["warn", { allowConstantExport: true }], 39 | "@typescript-eslint/no-unused-vars": ["error", { argsIgnorePattern: "^_" }], 40 | "@typescript-eslint/no-explicit-any": "warn", 41 | "prefer-const": "error", 42 | "no-var": "error", 43 | }, 44 | }, 45 | ]; 46 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | repos: 2 | - repo: https://github.com/pre-commit/pre-commit-hooks 3 | rev: v5.0.0 4 | hooks: 5 | - id: trailing-whitespace 6 | - id: end-of-file-fixer 7 | - id: check-docstring-first 8 | - id: check-yaml 9 | - id: debug-statements 10 | - id: check-ast 11 | - repo: https://github.com/astral-sh/ruff-pre-commit 12 | rev: "v0.12.4" 13 | hooks: 14 | - id: ruff 15 | exclude: docs/source/conf.py|ipynb_filter.py 16 | args: ["--fix"] 17 | - id: ruff-format 18 | exclude: docs/source/conf.py|ipynb_filter.py 19 | - repo: https://github.com/pre-commit/mirrors-mypy 20 | rev: "v1.17.0" 21 | hooks: 22 | - id: mypy 23 | exclude: ipynb_filter.py|docs/source/conf.py 24 | additional_dependencies: 25 | - types-setuptools 26 | - pandas-stubs 27 | - types-psutil 28 | - types-PyYAML 29 | - repo: https://github.com/kynan/nbstripout 30 | rev: 0.8.1 31 | hooks: 32 | - id: nbstripout 33 | - repo: https://github.com/pre-commit/mirrors-prettier 34 | rev: "v4.0.0-alpha.8" 35 | hooks: 36 | - id: prettier 37 | files: \.(js|jsx|ts|tsx|json|css|scss|md|yaml|yml)$ 38 | exclude: | 39 | (?x)^( 40 | frontend/dist/.*| 41 | frontend/node_modules/.*| 42 | .*\.min\.(js|css)$ 43 | )$ 44 | - repo: local 45 | hooks: 46 | - id: eslint 47 | name: eslint 48 | entry: bash -c 'cd frontend && npm run lint' 49 | language: system 50 | files: ^frontend/src/.*\.(ts|tsx)$ 51 | pass_filenames: false 52 | -------------------------------------------------------------------------------- /frontend/src/test/basic.test.ts: -------------------------------------------------------------------------------- 1 | import { describe, it, expect } from "vitest"; 2 | 3 | describe("Basic Test Suite", () => { 4 | it("should run basic tests", () => { 5 | expect(2 + 2).toBe(4); 6 | }); 7 | 8 | it("should handle string operations", () => { 9 | expect("hello world".toUpperCase()).toBe("HELLO WORLD"); 10 | }); 11 | 12 | it("should work with arrays", () => { 13 | const arr = [1, 2, 3]; 14 | expect(arr).toHaveLength(3); 15 | expect(arr).toContain(2); 16 | }); 17 | 18 | it("should work with objects", () => { 19 | const obj = { name: "test", value: 42 }; 20 | expect(obj).toHaveProperty("name"); 21 | expect(obj.name).toBe("test"); 22 | }); 23 | }); 24 | 25 | describe("Widget Utilities", () => { 26 | it("should validate message format", () => { 27 | const isValidMessage = (msg: { role: string; content: string }) => { 28 | return ( 29 | Boolean(msg.role) && 30 | Boolean(msg.content) && 31 | typeof msg.role === "string" && 32 | typeof msg.content === "string" 33 | ); 34 | }; 35 | 36 | expect(isValidMessage({ role: "user", content: "hello" })).toBe(true); 37 | expect(isValidMessage({ role: "", content: "hello" })).toBe(false); 38 | expect(isValidMessage({ role: "user", content: "" })).toBe(false); 39 | }); 40 | 41 | it("should handle button configurations", () => { 42 | const normalizeButton = (button: string | { text: string; color?: string }) => { 43 | return typeof button === "string" ? { text: button } : button; 44 | }; 45 | 46 | expect(normalizeButton("Click me")).toEqual({ text: "Click me" }); 47 | expect(normalizeButton({ text: "Submit", color: "blue" })).toEqual({ 48 | text: "Submit", 49 | color: "blue", 50 | }); 51 | }); 52 | }); 53 | -------------------------------------------------------------------------------- /assistant_ui_anywidget/global_agent.py: -------------------------------------------------------------------------------- 1 | """Global agent instance management for notebook convenience.""" 2 | 3 | from typing import Optional 4 | 5 | from .agent_widget import AgentWidget 6 | 7 | # Global agent instance 8 | _AGENT: Optional[AgentWidget] = None 9 | 10 | 11 | def get_agent( 12 | model: Optional[str] = None, 13 | provider: Optional[str] = "auto", 14 | temperature: float = 0.7, 15 | max_tokens: int = 2000, 16 | system_prompt: str = "You are a helpful AI assistant with access to the Jupyter kernel...", 17 | require_approval: bool = False, 18 | reset: bool = False, 19 | show_help: bool = True, 20 | ) -> AgentWidget: 21 | """Get or create the global agent instance. 22 | 23 | Args: 24 | model: AI model name (auto-detected if None) 25 | provider: AI provider ('openai', 'anthropic', 'google_genai', or 'auto') 26 | temperature: Response randomness (0.0-1.0) 27 | max_tokens: Maximum response length 28 | system_prompt: System prompt for the AI 29 | require_approval: Whether to require approval for code execution 30 | reset: If True, creates a new agent instance 31 | show_help: Whether to show the welcome message 32 | 33 | Returns: 34 | AgentWidget: The global agent instance 35 | """ 36 | global _AGENT 37 | 38 | if _AGENT is None or reset: 39 | # Create new agent with provided parameters 40 | _AGENT = AgentWidget( 41 | model=model, 42 | provider=provider, 43 | temperature=temperature, 44 | max_tokens=max_tokens, 45 | system_prompt=system_prompt, 46 | require_approval=require_approval, 47 | show_help=show_help, 48 | ) 49 | 50 | return _AGENT 51 | 52 | 53 | def reset_agent() -> None: 54 | """Reset the global agent instance.""" 55 | global _AGENT 56 | _AGENT = None 57 | -------------------------------------------------------------------------------- /frontend/src/test/setup.ts: -------------------------------------------------------------------------------- 1 | import "@testing-library/jest-dom"; 2 | import { vi } from "vitest"; 3 | 4 | // Mock AnyWidget to avoid import errors 5 | vi.mock("@anywidget/react", () => ({ 6 | createRender: vi.fn(() => vi.fn()), 7 | useModelState: vi.fn(() => [[], vi.fn()]), 8 | useModel: vi.fn(() => ({ 9 | send: vi.fn(), 10 | save_changes: vi.fn(), 11 | on: vi.fn(), 12 | off: vi.fn(), 13 | })), 14 | })); 15 | 16 | // Mock react-markdown to avoid complex dependencies 17 | vi.mock("react-markdown", () => ({ 18 | default: vi.fn(({ children }) => children), 19 | })); 20 | 21 | // Mock syntax highlighter 22 | vi.mock("react-syntax-highlighter", () => ({ 23 | Prism: vi.fn(({ children }) => children), 24 | })); 25 | 26 | vi.mock("react-syntax-highlighter/dist/esm/styles/prism", () => ({ 27 | vscDarkPlus: {}, 28 | })); 29 | 30 | vi.mock("remark-gfm", () => ({ 31 | default: vi.fn(), 32 | })); 33 | 34 | // Mock global objects that might be needed 35 | Object.defineProperty(window, "matchMedia", { 36 | writable: true, 37 | value: vi.fn().mockImplementation(query => ({ 38 | matches: false, 39 | media: query, 40 | onchange: null, 41 | addListener: vi.fn(), // deprecated 42 | removeListener: vi.fn(), // deprecated 43 | addEventListener: vi.fn(), 44 | removeEventListener: vi.fn(), 45 | dispatchEvent: vi.fn(), 46 | })), 47 | }); 48 | 49 | // Mock ResizeObserver 50 | globalThis.ResizeObserver = vi.fn().mockImplementation(() => ({ 51 | observe: vi.fn(), 52 | unobserve: vi.fn(), 53 | disconnect: vi.fn(), 54 | })); 55 | 56 | // Mock IntersectionObserver 57 | globalThis.IntersectionObserver = vi.fn().mockImplementation(() => ({ 58 | observe: vi.fn(), 59 | unobserve: vi.fn(), 60 | disconnect: vi.fn(), 61 | })); 62 | 63 | // Mock clipboard API 64 | Object.assign(navigator, { 65 | clipboard: { 66 | writeText: vi.fn(() => Promise.resolve()), 67 | readText: vi.fn(() => Promise.resolve("")), 68 | }, 69 | }); 70 | -------------------------------------------------------------------------------- /frontend/vite.config.ts: -------------------------------------------------------------------------------- 1 | /// 2 | import { defineConfig } from "vite"; 3 | import react from "@vitejs/plugin-react"; 4 | import { visualizer } from "rollup-plugin-visualizer"; 5 | 6 | export default defineConfig({ 7 | plugins: [ 8 | react(), 9 | visualizer({ 10 | filename: "dist/stats.html", 11 | open: false, // Don't automatically open the stats page 12 | gzipSize: true, 13 | brotliSize: true, 14 | }), 15 | ], 16 | define: { 17 | "process.env.NODE_ENV": JSON.stringify("production"), 18 | global: "globalThis", 19 | }, 20 | build: { 21 | lib: { 22 | entry: "src/index.tsx", 23 | formats: ["es"], 24 | fileName: "index", 25 | }, 26 | rollupOptions: { 27 | // Don't externalize React - bundle it instead 28 | external: [], 29 | output: { 30 | globals: {}, 31 | manualChunks: undefined, // Keep everything in a single bundle 32 | }, 33 | }, 34 | // Output directly to the static directory 35 | outDir: "../assistant_ui_anywidget/static", 36 | emptyOutDir: false, // Don't delete existing files in static 37 | // Optimize build 38 | minify: "terser", 39 | sourcemap: false, 40 | target: "esnext", 41 | reportCompressedSize: true, 42 | }, 43 | // Development optimizations 44 | server: { 45 | hmr: true, 46 | }, 47 | optimizeDeps: { 48 | include: ["react", "react-dom", "@anywidget/react", "@assistant-ui/react", "react-markdown"], 49 | }, 50 | // Test configuration 51 | test: { 52 | globals: true, 53 | environment: "jsdom", 54 | setupFiles: ["./src/test/setup.ts"], 55 | css: false, 56 | coverage: { 57 | provider: "v8", 58 | reporter: ["text", "json", "html"], 59 | exclude: [ 60 | "node_modules/", 61 | "src/test/", 62 | "**/*.d.ts", 63 | "**/*.test.{ts,tsx}", 64 | "**/*.config.{ts,js}", 65 | ], 66 | }, 67 | }, 68 | }); 69 | -------------------------------------------------------------------------------- /tests/test_widget_simple.py: -------------------------------------------------------------------------------- 1 | """Simple test to verify the widget works.""" 2 | 3 | import os 4 | import sys 5 | 6 | sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "python")) # noqa: PTH118, PTH120 7 | 8 | from assistant_ui_anywidget.agent_widget import AgentWidget 9 | 10 | 11 | def test_widget() -> None: 12 | """Test the widget thoroughly.""" 13 | print("=== Widget Test ===") 14 | 15 | # Test widget creation 16 | print("1. Creating widget...") 17 | widget = AgentWidget(show_help=False) 18 | print(f" ✓ Widget created: {widget.__class__.__name__}") 19 | print(f" ✓ Widget ID: {widget.model_id}") 20 | 21 | # Test JavaScript bundle 22 | print("\n2. Checking JavaScript bundle...") 23 | if hasattr(widget, "_esm"): 24 | if isinstance(widget._esm, str) and widget._esm.startswith("var"): # type: ignore[unreachable] 25 | print(" ✓ JavaScript is bundled (no external imports)") # type: ignore[unreachable] 26 | else: 27 | print(f" ✗ Unexpected ESM content: {widget._esm[:50]}...") # type: ignore[unused-ignore,index] 28 | else: 29 | print(" ✗ No _esm attribute found") 30 | 31 | # Test message handling 32 | print("\n3. Testing message handling...") 33 | try: 34 | # Simulate receiving a message 35 | widget._handle_message(None, {"type": "user_message", "text": "Hello"}) 36 | print(" ✓ Message handling works") 37 | except Exception as e: # noqa: BLE001 38 | print(f" ✗ Message handling failed: {e}") 39 | 40 | print("\n=== Widget is ready! ===") 41 | print("\nTo use in Jupyter:") 42 | print("1. uv run jupyter notebook") 43 | print("2. Create new notebook") 44 | print("3. Run:") 45 | print(" import sys") 46 | print(" sys.path.insert(0, 'python')") 47 | print(" from assistant_ui_anywidget.agent_widget import AgentWidget") 48 | print(" widget = AgentWidget(show_help=False)") 49 | print(" widget") 50 | 51 | 52 | if __name__ == "__main__": 53 | test_widget() 54 | -------------------------------------------------------------------------------- /assistant_ui_anywidget/ai/prompt_config.py: -------------------------------------------------------------------------------- 1 | """System prompt configuration using pydantic-settings.""" 2 | 3 | from pathlib import Path 4 | from typing import Tuple, Type 5 | 6 | from pydantic import ConfigDict 7 | from pydantic_settings import ( 8 | BaseSettings, 9 | PydanticBaseSettingsSource, 10 | YamlConfigSettingsSource, 11 | ) 12 | 13 | 14 | class SystemPromptConfig(BaseSettings): 15 | """System prompt configuration loaded from YAML file.""" 16 | 17 | approval_note: str 18 | main_prompt: str 19 | slash_commands: str 20 | examples_of_proactive_behavior: str 21 | scientific_computing_awareness: str 22 | final_reminder: str 23 | 24 | def get_full_prompt(self, require_approval: bool = True) -> str: 25 | """Combine all sections into the complete system prompt.""" 26 | prompt_parts = [self.main_prompt] 27 | if require_approval: 28 | prompt_parts.append(self.approval_note) 29 | prompt_parts.extend( 30 | [ 31 | self.slash_commands, 32 | self.examples_of_proactive_behavior, 33 | self.scientific_computing_awareness, 34 | self.final_reminder, 35 | ] 36 | ) 37 | return "\n\n".join(prompt_parts) 38 | 39 | model_config = ConfigDict(extra="forbid") 40 | 41 | @classmethod 42 | def settings_customise_sources( 43 | cls, 44 | settings_cls: Type[BaseSettings], 45 | init_settings: PydanticBaseSettingsSource, 46 | env_settings: PydanticBaseSettingsSource, 47 | dotenv_settings: PydanticBaseSettingsSource, 48 | file_secret_settings: PydanticBaseSettingsSource, 49 | ) -> Tuple[PydanticBaseSettingsSource, ...]: 50 | """Configure pydantic-settings to load from YAML file.""" 51 | yaml_file = Path(__file__).parent / "system_prompt.yaml" 52 | return ( 53 | init_settings, 54 | YamlConfigSettingsSource(settings_cls, yaml_file=yaml_file), 55 | env_settings, 56 | file_secret_settings, 57 | ) 58 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | *.egg-info/ 24 | .installed.cfg 25 | *.egg 26 | MANIFEST 27 | 28 | # PyInstaller 29 | # Usually these files are written by a python script from a template 30 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 31 | *.manifest 32 | *.spec 33 | 34 | # Installer logs 35 | pip-log.txt 36 | pip-delete-this-directory.txt 37 | 38 | # Unit test / coverage reports 39 | htmlcov/ 40 | .tox/ 41 | .coverage 42 | .coverage.* 43 | .cache 44 | nosetests.xml 45 | coverage.xml 46 | *.cover 47 | .hypothesis/ 48 | .pytest_cache/ 49 | 50 | # Translations 51 | *.mo 52 | *.pot 53 | 54 | # Django stuff: 55 | *.log 56 | local_settings.py 57 | db.sqlite3 58 | 59 | # Flask stuff: 60 | instance/ 61 | .webassets-cache 62 | 63 | # Scrapy stuff: 64 | .scrapy 65 | 66 | # Sphinx documentation 67 | docs/_build/ 68 | 69 | # PyBuilder 70 | target/ 71 | 72 | # Jupyter Notebook 73 | .ipynb_checkpoints 74 | 75 | # pyenv 76 | .python-version 77 | 78 | # celery beat schedule file 79 | celerybeat-schedule 80 | 81 | # SageMath parsed files 82 | *.sage.py 83 | 84 | # Environments 85 | .env 86 | .venv 87 | env/ 88 | venv/ 89 | ENV/ 90 | env.bak/ 91 | venv.bak/ 92 | 93 | # Spyder project settings 94 | .spyderproject 95 | .spyproject 96 | 97 | # Rope project settings 98 | .ropeproject 99 | 100 | # mkdocs documentation 101 | /site 102 | 103 | # mypy 104 | .mypy_cache/ 105 | 106 | # MacOS files 107 | *.DS_Store 108 | 109 | # scheduler 110 | *.sbatch 111 | *.out 112 | *.batch 113 | 114 | # JavaScript 115 | # Ignore node_modules in any directory 116 | **/node_modules/ 117 | 118 | src/graphviz_anywidget/static/ 119 | 120 | frontend/screenshots/* 121 | 122 | # Build artifacts 123 | assistant_ui_anywidget/static/* 124 | !assistant_ui_anywidget/static/.gitkeep 125 | 126 | # AI conversation logs 127 | ai_conversation_logs/ 128 | -------------------------------------------------------------------------------- /frontend/package.json: -------------------------------------------------------------------------------- 1 | { 2 | "name": "assistant-ui-widget", 3 | "version": "0.1.0", 4 | "type": "module", 5 | "scripts": { 6 | "build": "vite build", 7 | "build:analyze": "vite build --analyze", 8 | "dev": "vite build --watch", 9 | "preview": "vite preview", 10 | "lint": "eslint . --ext ts,tsx", 11 | "lint:fix": "eslint . --ext ts,tsx --fix", 12 | "format": "prettier --write .", 13 | "format:check": "prettier --check .", 14 | "type-check": "tsc --noEmit", 15 | "test": "vitest", 16 | "test:ui": "vitest --ui", 17 | "test:run": "vitest run", 18 | "test:coverage": "vitest run --coverage", 19 | "clean": "rm -rf dist node_modules/.vite coverage", 20 | "preinstall": "npx only-allow npm", 21 | "postinstall": "npm run type-check", 22 | "demo": "node scripts/serve-demo.cjs", 23 | "screenshot": "node scripts/screenshot.cjs", 24 | "screenshot:demo": "npm run demo & sleep 5 && npm run screenshot && kill $!" 25 | }, 26 | "devDependencies": { 27 | "@eslint/js": "^9.31.0", 28 | "@testing-library/jest-dom": "^6.6.3", 29 | "@testing-library/react": "^16.3.0", 30 | "@types/react": "^18.2.43", 31 | "@types/react-dom": "^18.2.17", 32 | "@typescript-eslint/eslint-plugin": "^8.38.0", 33 | "@typescript-eslint/parser": "^8.38.0", 34 | "@vitejs/plugin-react": "^4.2.1", 35 | "@vitest/ui": "^3.2.4", 36 | "eslint": "^9.31.0", 37 | "eslint-plugin-jsx-a11y": "^6.10.2", 38 | "eslint-plugin-react": "^7.37.5", 39 | "eslint-plugin-react-hooks": "^5.2.0", 40 | "eslint-plugin-react-refresh": "^0.4.20", 41 | "globals": "^16.3.0", 42 | "jsdom": "^26.1.0", 43 | "prettier": "^3.6.2", 44 | "rollup-plugin-visualizer": "^6.0.3", 45 | "terser": "^5.43.1", 46 | "typescript": "^5.2.2", 47 | "vite": "^7.0.5", 48 | "vitest": "^3.2.4" 49 | }, 50 | "dependencies": { 51 | "@anywidget/react": "^0.2.0", 52 | "@assistant-ui/react": "^0.10.26", 53 | "@assistant-ui/react-markdown": "^0.10.6", 54 | "@types/react-syntax-highlighter": "^15.5.13", 55 | "react": "^18.2.0", 56 | "react-dom": "^18.2.0", 57 | "react-markdown": "^10.1.0", 58 | "react-syntax-highlighter": "^15.6.1", 59 | "remark-gfm": "^4.0.1" 60 | } 61 | } 62 | -------------------------------------------------------------------------------- /frontend/README.md: -------------------------------------------------------------------------------- 1 | # Frontend Structure Documentation 2 | 3 | This directory contains the TypeScript/React implementation of the Assistant UI anywidget. 4 | 5 | ## Module Structure 6 | 7 | ### Core Widget 8 | 9 | - **`src/index.tsx`** - Main widget component that integrates with anywidget 10 | - Exports the widget using `createRender` from `@anywidget/react` 11 | - Contains the `ChatWidget` component with full chat functionality 12 | - Handles state synchronization with Python backend via `useModelState` and `useModel` 13 | 14 | ### Testing 15 | 16 | - **`src/test/setup.ts`** - Test environment setup with mocked dependencies 17 | - **`src/test/basic.test.ts`** - Basic unit tests for widget functionality 18 | - **`vitest.config.ts`** - Vitest test runner configuration 19 | 20 | ### Build Configuration 21 | 22 | - **`vite.config.ts`** - Main build configuration 23 | - Builds the widget as a single ES module (`dist/index.js`) 24 | - Bundles all dependencies including React 25 | - Optimized for integration with Python anywidget 26 | 27 | ### Demo & Screenshots 28 | 29 | - **`demo.html`** - Standalone HTML demo page with inline React widget 30 | - Simple, self-contained demo that doesn't require the full anywidget stack 31 | - Used for visual testing and screenshots 32 | - **`scripts/serve-demo.cjs`** - Simple HTTP server for the demo page 33 | - **`scripts/screenshot.cjs`** - Puppeteer script for automated screenshots 34 | 35 | ## How It Works 36 | 37 | 1. **In Jupyter**: The Python `AgentWidget` class loads `dist/index.js` and the widget communicates with Python via the anywidget protocol 38 | 39 | 2. **For Screenshots**: The `demo.html` file contains a simplified version that renders the same UI without requiring Jupyter or Python 40 | 41 | ## Development 42 | 43 | ```bash 44 | # Install dependencies 45 | npm install 46 | 47 | # Build the widget 48 | npm run build 49 | 50 | # Run tests 51 | npm test 52 | 53 | # Start demo server for screenshots 54 | npm run demo 55 | 56 | # Take screenshots 57 | npm run screenshot 58 | ``` 59 | 60 | ## Screenshot System 61 | 62 | The screenshot system allows viewing the widget without Jupyter: 63 | 64 | 1. Run `npm run demo` to start a simple server 65 | 2. Run `npm run screenshot` to capture the widget 66 | 3. Screenshots are saved in `screenshots/` 67 | 68 | This is useful for: 69 | 70 | - Visual regression testing 71 | - Documentation 72 | - Sharing the widget appearance without running Jupyter 73 | -------------------------------------------------------------------------------- /take_screenshot.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | """ 3 | Script to take screenshots of the Assistant UI Widget using Puppeteer. 4 | This allows Claude to view the widget's appearance. 5 | """ 6 | 7 | import os 8 | import subprocess 9 | import sys 10 | import time 11 | from pathlib import Path 12 | 13 | 14 | def start_demo_server() -> subprocess.Popen[bytes]: 15 | """Start the demo server in the background.""" 16 | print("Starting demo server...") 17 | process = subprocess.Popen( 18 | ["npm", "run", "demo"], 19 | cwd=Path(__file__).parent / "frontend", 20 | stdout=subprocess.PIPE, 21 | stderr=subprocess.PIPE, 22 | ) 23 | 24 | # Wait for server to start 25 | time.sleep(8) 26 | return process 27 | 28 | 29 | def take_screenshot(add_demo_messages: bool = True, interactive: bool = False) -> bool: 30 | """Take a screenshot of the widget using Puppeteer.""" 31 | env = { 32 | "DEMO_URL": "http://localhost:3000", 33 | "ADD_DEMO_MESSAGES": "true" if add_demo_messages else "false", 34 | "INTERACTIVE_SCREENSHOT": "true" if interactive else "false", 35 | } 36 | 37 | print("Taking screenshot...") 38 | result = subprocess.run( 39 | ["npm", "run", "screenshot"], 40 | cwd=Path(__file__).parent / "frontend", 41 | env={**os.environ, **env}, 42 | capture_output=True, 43 | text=True, 44 | ) 45 | 46 | if result.returncode != 0: 47 | print(f"Error taking screenshot: {result.stderr}") 48 | return False 49 | 50 | print(result.stdout) 51 | return True 52 | 53 | 54 | def main() -> None: 55 | """Main function to coordinate the screenshot process.""" 56 | server_process = None 57 | 58 | try: 59 | # Start the demo server 60 | server_process = start_demo_server() 61 | 62 | # Take screenshots 63 | success = take_screenshot(add_demo_messages=True, interactive=True) 64 | 65 | if success: 66 | print("\nScreenshots saved to frontend/screenshots/") 67 | print("You can now share these with Claude to view the widget!") 68 | else: 69 | print("\nFailed to take screenshots.") 70 | sys.exit(1) 71 | 72 | except KeyboardInterrupt: 73 | print("\nInterrupted by user.") 74 | except Exception as e: 75 | print(f"\nError: {e}") 76 | sys.exit(1) 77 | finally: 78 | # Clean up: stop the server 79 | if server_process: 80 | print("\nStopping demo server...") 81 | server_process.terminate() 82 | server_process.wait() 83 | 84 | 85 | if __name__ == "__main__": 86 | main() 87 | -------------------------------------------------------------------------------- /tests/test_notebook.py: -------------------------------------------------------------------------------- 1 | """Test the widget in a Jupyter notebook environment.""" 2 | 3 | import os 4 | import sys 5 | 6 | sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "python")) # noqa: PTH118, PTH120 7 | 8 | from assistant_ui_anywidget.agent_widget import AgentWidget 9 | 10 | 11 | def test_widget_in_notebook() -> AgentWidget | None: 12 | """Test widget creation and display.""" 13 | print("Creating AgentWidget...") 14 | widget = AgentWidget(show_help=False) 15 | 16 | print("Widget created successfully!") 17 | print(f"Widget ID: {widget.model_id}") 18 | print(f"Widget class: {widget.__class__.__name__}") 19 | 20 | # Display the widget 21 | print("\nDisplaying widget...") 22 | if __name__ == "__main__": 23 | return widget 24 | # For pytest, we don't return anything 25 | return None 26 | 27 | 28 | if __name__ == "__main__": 29 | widget = test_widget_in_notebook() 30 | 31 | print("\nTo use in Jupyter:") 32 | print("1. Start Jupyter: uv run jupyter notebook") 33 | print("2. Create new notebook") 34 | print("3. Run: from python.agent_widget import AgentWidget") 35 | print("4. Run: widget = AgentWidget(show_help=False)") 36 | print("5. Run: widget") 37 | 38 | # For testing, we can also create a simple notebook file 39 | # Note: This function is used both as a test and as a standalone script 40 | notebook_content = """ 41 | { 42 | "cells": [ 43 | { 44 | "cell_type": "code", 45 | "execution_count": null, 46 | "metadata": {}, 47 | "outputs": [], 48 | "source": [ 49 | "import sys\\n", 50 | "sys.path.insert(0, 'python')\\n", 51 | "from assistant_ui_anywidget.agent_widget import AgentWidget" 52 | ] 53 | }, 54 | { 55 | "cell_type": "code", 56 | "execution_count": null, 57 | "metadata": {}, 58 | "outputs": [], 59 | "source": [ 60 | "widget = AgentWidget(show_help=False)\\n", 61 | "widget" 62 | ] 63 | } 64 | ], 65 | "metadata": { 66 | "kernelspec": { 67 | "display_name": "Python 3", 68 | "language": "python", 69 | "name": "python3" 70 | }, 71 | "language_info": { 72 | "codemirror_mode": { 73 | "name": "ipython", 74 | "version": 3 75 | }, 76 | "file_extension": ".py", 77 | "mimetype": "text/x-python", 78 | "name": "python", 79 | "nbconvert_exporter": "python", 80 | "pygments_lexer": "ipython3", 81 | "version": "3.8.0" 82 | } 83 | }, 84 | "nbformat": 4, 85 | "nbformat_minor": 4 86 | } 87 | """ 88 | 89 | with open("test_widget.ipynb", "w") as f: # noqa: PTH123 90 | f.write(notebook_content) 91 | 92 | print("\nCreated test_widget.ipynb for testing in Jupyter!") 93 | -------------------------------------------------------------------------------- /tests/test_kernel_error_handling.py: -------------------------------------------------------------------------------- 1 | """Test kernel error handling, especially get_last_error edge cases.""" 2 | 3 | from unittest.mock import Mock 4 | 5 | from assistant_ui_anywidget.kernel_interface import KernelContext, KernelInterface 6 | 7 | 8 | def test_get_last_error_no_exception() -> None: 9 | """Test get_last_error when no exception exists.""" 10 | kernel = KernelInterface() 11 | 12 | # Mock IPython shell 13 | kernel.shell = Mock() 14 | 15 | # Simulate _get_exc_info raising ValueError when no exception 16 | kernel.shell._get_exc_info.side_effect = ValueError("No exception to find") 17 | 18 | # Should return None instead of raising 19 | result = kernel.get_last_error() 20 | assert result is None 21 | 22 | 23 | def test_get_last_error_attribute_error() -> None: 24 | """Test get_last_error with different IPython versions.""" 25 | kernel = KernelInterface() 26 | 27 | # Mock IPython shell without _get_exc_info 28 | kernel.shell = Mock() 29 | del kernel.shell._get_exc_info # Simulate missing method 30 | 31 | # Should return None instead of raising 32 | result = kernel.get_last_error() 33 | assert result is None 34 | 35 | 36 | def test_get_last_error_no_kernel() -> None: 37 | """Test get_last_error when kernel is not available.""" 38 | kernel = KernelInterface() 39 | kernel.shell = None 40 | 41 | result = kernel.get_last_error() 42 | assert result is None 43 | 44 | 45 | def test_get_last_error_with_exception() -> None: 46 | """Test get_last_error when there is an actual exception.""" 47 | kernel = KernelInterface() 48 | 49 | # Mock IPython shell 50 | kernel.shell = Mock() 51 | 52 | # Mock exception info 53 | class MockException(Exception): 54 | pass 55 | 56 | mock_exception = MockException("Test error") 57 | kernel.shell._get_exc_info.return_value = (MockException, mock_exception, None) 58 | 59 | # Mock _format_traceback 60 | kernel._format_traceback = Mock(return_value=["line 1", "line 2"]) # type: ignore 61 | 62 | result = kernel.get_last_error() 63 | assert result is not None 64 | assert result["type"] == "MockException" 65 | assert result["message"] == "Test error" 66 | assert result["traceback"] == ["line 1", "line 2"] 67 | 68 | 69 | def test_agent_widget_context_no_kernel() -> None: 70 | """Test that _get_kernel_context handles missing kernel gracefully.""" 71 | from assistant_ui_anywidget import AgentWidget 72 | 73 | widget = AgentWidget(show_help=False) 74 | 75 | # Should not raise even when kernel operations fail 76 | context = widget._get_kernel_context() 77 | assert isinstance(context, KernelContext) 78 | assert context.kernel_info is not None 79 | assert context.variables is not None 80 | -------------------------------------------------------------------------------- /.github/workflows/ci.yml: -------------------------------------------------------------------------------- 1 | name: CI 2 | 3 | on: 4 | push: 5 | branches: [main] 6 | pull_request: 7 | branches: [main] 8 | 9 | jobs: 10 | test-python: 11 | runs-on: ubuntu-latest 12 | strategy: 13 | matrix: 14 | python-version: ["3.10", "3.11", "3.12"] 15 | 16 | steps: 17 | - uses: actions/checkout@v4 18 | 19 | - name: Set up Python ${{ matrix.python-version }} 20 | uses: actions/setup-python@v5 21 | with: 22 | python-version: ${{ matrix.python-version }} 23 | 24 | - name: Set up Node.js 25 | uses: actions/setup-node@v4 26 | with: 27 | node-version: "20" 28 | cache: "npm" 29 | cache-dependency-path: frontend/package-lock.json 30 | 31 | - name: Install uv 32 | uses: astral-sh/setup-uv@v6 33 | 34 | - name: Install dependencies 35 | run: | 36 | uv sync 37 | 38 | - name: Install frontend dependencies 39 | run: | 40 | cd frontend 41 | npm ci 42 | 43 | - name: Build frontend 44 | run: | 45 | cd frontend 46 | npm run build 47 | 48 | - name: Run pre-commit 49 | run: | 50 | uv run pre-commit run --all-files 51 | 52 | - name: Run tests with pytest 53 | run: | 54 | uv run pytest --verbose --tb=short 55 | 56 | test-frontend: 57 | runs-on: ubuntu-latest 58 | defaults: 59 | run: 60 | working-directory: ./frontend 61 | 62 | steps: 63 | - uses: actions/checkout@v4 64 | 65 | - name: Set up Node.js 66 | uses: actions/setup-node@v4 67 | with: 68 | node-version: "20" 69 | cache: "npm" 70 | cache-dependency-path: frontend/package-lock.json 71 | 72 | - name: Install dependencies 73 | run: npm ci 74 | 75 | - name: Run ESLint 76 | run: npm run lint 77 | 78 | - name: Run type checking 79 | run: npm run type-check 80 | 81 | - name: Build frontend 82 | run: npm run build 83 | 84 | integration-test: 85 | runs-on: ubuntu-latest 86 | needs: [test-python, test-frontend] 87 | 88 | steps: 89 | - uses: actions/checkout@v4 90 | 91 | - name: Set up Python 3.11 92 | uses: actions/setup-python@v5 93 | with: 94 | python-version: "3.11" 95 | 96 | - name: Set up Node.js 97 | uses: actions/setup-node@v4 98 | with: 99 | node-version: "20" 100 | cache: "npm" 101 | cache-dependency-path: frontend/package-lock.json 102 | 103 | - name: Install uv 104 | uses: astral-sh/setup-uv@v6 105 | 106 | - name: Install Python dependencies 107 | run: | 108 | uv sync --all-extras 109 | 110 | - name: Install Node dependencies and build frontend 111 | run: | 112 | cd frontend 113 | npm ci 114 | npm run build 115 | 116 | - name: Run integration tests 117 | run: | 118 | uv run pytest tests/ --verbose 119 | -------------------------------------------------------------------------------- /examples/langgraph_approval_demo.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "metadata": {}, 6 | "source": [ 7 | "# LangGraph Approval Workflow Demo\n", 8 | "\n", 9 | "This notebook demonstrates the LangGraph-based AI service with approval workflows." 10 | ] 11 | }, 12 | { 13 | "cell_type": "code", 14 | "execution_count": null, 15 | "metadata": {}, 16 | "outputs": [], 17 | "source": [ 18 | "from assistant_ui_anywidget import AgentWidget\n", 19 | "\n", 20 | "# Create a widget with LangGraph and approval enabled\n", 21 | "widget = AgentWidget(\n", 22 | " require_approval=True, # Require approval for code execution\n", 23 | " provider=\"auto\",\n", 24 | ")\n", 25 | "\n", 26 | "# Display the widget\n", 27 | "widget" 28 | ] 29 | }, 30 | { 31 | "cell_type": "markdown", 32 | "metadata": {}, 33 | "source": [ 34 | "## How It Works\n", 35 | "\n", 36 | "1. **Automatic Execution** for read-only operations:\n", 37 | " - `get_variables()` - Lists all variables\n", 38 | " - `inspect_variable(name)` - Inspects a specific variable\n", 39 | " - `kernel_info()` - Shows kernel status\n", 40 | "\n", 41 | "2. **Approval Required** for code execution:\n", 42 | " - `execute_code(code)` - Requires your approval before running\n", 43 | "\n", 44 | "Try these examples:\n", 45 | "- \"Show me all variables\" (automatic)\n", 46 | "- \"What's in the variable x?\" (automatic)\n", 47 | "- \"Execute x = 42\" (requires approval)\n", 48 | "- \"Run print('Hello World')\" (requires approval)" 49 | ] 50 | }, 51 | { 52 | "cell_type": "code", 53 | "execution_count": null, 54 | "metadata": {}, 55 | "outputs": [], 56 | "source": [ 57 | "# Create some test data\n", 58 | "x = 42\n", 59 | "data = [1, 2, 3, 4, 5]\n", 60 | "message = \"Hello from Jupyter!\"" 61 | ] 62 | }, 63 | { 64 | "cell_type": "markdown", 65 | "metadata": {}, 66 | "source": [ 67 | "## Comparison with Simple Service\n", 68 | "\n", 69 | "You can also use the simple service without LangGraph:" 70 | ] 71 | }, 72 | { 73 | "cell_type": "code", 74 | "execution_count": null, 75 | "metadata": {}, 76 | "outputs": [], 77 | "source": [ 78 | "# Create a widget with simple service (no approval workflow)\n", 79 | "simple_widget = AgentWidget(\n", 80 | " require_approval=False, # No approval required\n", 81 | " provider=\"auto\",\n", 82 | ")\n", 83 | "\n", 84 | "simple_widget" 85 | ] 86 | } 87 | ], 88 | "metadata": { 89 | "kernelspec": { 90 | "display_name": "Python 3", 91 | "language": "python", 92 | "name": "python3" 93 | }, 94 | "language_info": { 95 | "codemirror_mode": { 96 | "name": "ipython", 97 | "version": 3 98 | }, 99 | "file_extension": ".py", 100 | "mimetype": "text/x-python", 101 | "name": "python", 102 | "nbconvert_exporter": "python", 103 | "pygments_lexer": "ipython3", 104 | "version": "3.11.0" 105 | } 106 | }, 107 | "nbformat": 4, 108 | "nbformat_minor": 4 109 | } 110 | -------------------------------------------------------------------------------- /frontend/scripts/screenshot.cjs: -------------------------------------------------------------------------------- 1 | const puppeteer = require("puppeteer"); 2 | const path = require("path"); 3 | const fs = require("fs"); 4 | 5 | async function takeScreenshot() { 6 | const browser = await puppeteer.launch({ 7 | headless: true, 8 | args: ["--no-sandbox", "--disable-setuid-sandbox"], 9 | }); 10 | 11 | try { 12 | const page = await browser.newPage(); 13 | 14 | // Set viewport to standard desktop size 15 | await page.setViewport({ 16 | width: 1280, 17 | height: 720, 18 | deviceScaleFactor: 2, // For high quality screenshots 19 | }); 20 | 21 | // Navigate to the demo page 22 | const url = process.env.DEMO_URL || "http://localhost:3000"; 23 | console.log(`Navigating to ${url}...`); 24 | await page.goto(url, { waitUntil: "networkidle0" }); 25 | 26 | // Wait for the widget to be visible 27 | await page.waitForSelector("#root", { visible: true }); 28 | 29 | // Optional: Add some demo messages 30 | if (process.env.ADD_DEMO_MESSAGES === "true") { 31 | // Type a message 32 | await page.type('textarea[placeholder*="Type a message"]', "Hello! Can you help me with coding?"); 33 | await page.keyboard.down("Control"); 34 | await page.keyboard.press("d"); 35 | await page.keyboard.up("Control"); 36 | 37 | // Wait for response 38 | await new Promise(resolve => setTimeout(resolve, 2000)); 39 | 40 | // Click an action button 41 | const buttons = await page.$$('button:not([type="submit"])'); 42 | if (buttons.length > 0) { 43 | await buttons[0].click(); 44 | await new Promise(resolve => setTimeout(resolve, 1500)); 45 | } 46 | } 47 | 48 | // Create screenshots directory if it doesn't exist 49 | const screenshotsDir = path.join(__dirname, "../screenshots"); 50 | if (!fs.existsSync(screenshotsDir)) { 51 | fs.mkdirSync(screenshotsDir, { recursive: true }); 52 | } 53 | 54 | // Take full page screenshot 55 | const timestamp = new Date().toISOString().replace(/[:.]/g, "-"); 56 | const fullPagePath = path.join(screenshotsDir, `widget-fullpage-${timestamp}.png`); 57 | await page.screenshot({ 58 | path: fullPagePath, 59 | fullPage: true, 60 | }); 61 | console.log(`Full page screenshot saved to: ${fullPagePath}`); 62 | 63 | // Take widget-only screenshot 64 | let widgetPath; 65 | const widgetElement = await page.$("#root > div"); 66 | if (widgetElement) { 67 | widgetPath = path.join(screenshotsDir, `widget-only-${timestamp}.png`); 68 | await widgetElement.screenshot({ 69 | path: widgetPath, 70 | }); 71 | console.log(`Widget screenshot saved to: ${widgetPath}`); 72 | } 73 | 74 | // Take a screenshot with specific interactions if requested 75 | if (process.env.INTERACTIVE_SCREENSHOT === "true") { 76 | // Focus on textarea 77 | await page.focus('textarea[placeholder*="Type a message"]'); 78 | 79 | // Take screenshot with focused input 80 | const focusedPath = path.join(screenshotsDir, `widget-focused-${timestamp}.png`); 81 | await widgetElement.screenshot({ 82 | path: focusedPath, 83 | }); 84 | console.log(`Focused widget screenshot saved to: ${focusedPath}`); 85 | } 86 | 87 | // Return the paths for programmatic use 88 | return { 89 | fullPage: fullPagePath, 90 | widget: widgetPath || null, 91 | }; 92 | } catch (error) { 93 | console.error("Error taking screenshot:", error); 94 | throw error; 95 | } finally { 96 | await browser.close(); 97 | } 98 | } 99 | 100 | // Run if called directly 101 | if (require.main === module) { 102 | takeScreenshot() 103 | .then(() => { 104 | console.log("Screenshots captured successfully!"); 105 | process.exit(0); 106 | }) 107 | .catch((error) => { 108 | console.error("Failed to capture screenshots:", error); 109 | process.exit(1); 110 | }); 111 | } 112 | 113 | module.exports = { takeScreenshot }; 114 | -------------------------------------------------------------------------------- /CLAUDE.md: -------------------------------------------------------------------------------- 1 | # Guiding Principles for Development 2 | 3 | This document outlines the core principles and practices for developing this project. Adhering to these guidelines ensures consistency, quality, and a focus on forward-looking development. 4 | 5 | ## 1. Core Philosophy 6 | 7 | - **Embrace Change, Avoid Backward Compatibility**: This project has no end-users yet. Prioritize innovation and improvement over maintaining backward compatibility. 8 | - **Simplicity is Key**: Implement the simplest possible solution. Avoid over-engineering or generalizing features prematurely. 9 | - **Focus on the Task**: Implement only the requested feature, without adding extras. 10 | - **Functional Over Classes**: Prefer a functional programming style for Python over complex class hierarchies. 11 | - **Keep it DRY**: Don't Repeat Yourself. Reuse code wherever possible. 12 | - **Be Ruthless with Code Removal**: Aggressively remove any unused code, including functions, imports, and variables. 13 | - **Prefer dataclasses**: Use `dataclasses` that can be typed over dictionaries for better type safety and clarity. 14 | - Do not wrap things in try-excepts unless it's necessary. Avoid wrapping things that should not fail. 15 | 16 | ## 2. Workflow 17 | 18 | ### Step 1: Understand the Context 19 | 20 | - **Understand Current Task**: Review the issue, PR description, or task at hand. 21 | - **Explore the Codebase**: List existing files and read the `README.md` to understand the project's structure and purpose. 22 | - **Consult Documentation**: Review the LangGraph documentation (https://langchain-ai.github.io/langgraph/llms.txt and https://langchain-ai.github.io/langgraph/llms-full.txt) to understand agent capabilities. 23 | 24 | ### Step 2: Environment & Dependencies 25 | 26 | - **Environment Setup**: Use `uv sync --all-extras` to install all dependencies and `source .venv/bin/activate` to activate the virtual environment. 27 | - **Adding Packages**: Use `uv add ` for new dependencies or `uv add --dev ` for development-only packages. 28 | 29 | ### Step 3: Development & Git 30 | 31 | - **Check for Changes**: Before starting, review the latest changes from the main branch with `git diff origin/main | cat`. Make sure to use `--no-pager`, or pipe the output to `cat`. 32 | - **Commit Frequently**: Make small, frequent commits. 33 | - **Atomic Commits**: Ensure each commit corresponds to a tested, working state. 34 | - **Targeted Adds**: **NEVER** use `git add .`. Always add files individually (`git add `) to prevent committing unrelated changes. 35 | 36 | ### Step 4: Testing & Quality 37 | 38 | - **Test Before Committing**: **NEVER** claim a task is complete without running `pytest` to ensure all tests pass. 39 | - **Run Pre-commit Hooks**: Always run `pre-commit run --all-files` before committing to enforce code style and quality. 40 | - **Handle Linter Issues**: 41 | - **False Positives**: The linter may incorrectly flag issues in `pyproject.toml`; these can be ignored. 42 | - **Test-Related Errors**: If a pre-commit fix breaks a test (e.g., by removing an unused but necessary fixture), suppress the warning with a `# noqa: ` comment. 43 | 44 | ### Step 5: Refactoring 45 | 46 | - **Be Proactive**: Continuously look for opportunities to refactor and improve the codebase for better organization and readability. 47 | - **Incremental Changes**: Refactor in small, testable steps. Run tests after each change and commit on success. 48 | 49 | ### Step 6: Viewing the Widget 50 | 51 | - **Taking Screenshots**: To view the widget without Jupyter, use `python take_screenshot.py` from the project root. 52 | - **Manual Screenshot**: From the frontend directory, run `npm run demo` then `npm run screenshot` in another terminal. 53 | - **Screenshot Location**: Screenshots are saved to `frontend/screenshots/` with timestamps. 54 | - **Use Cases**: This is helpful for visual verification, documentation, and sharing the widget appearance. 55 | 56 | ## 3. Critical "Don'ts" 57 | 58 | - **DO NOT** manually edit the CLI help messages in `README.md`. They are auto-generated. 59 | - **NEVER** use `git add .`. 60 | - **NEVER** claim a task is done without passing all `pytest` tests. 61 | -------------------------------------------------------------------------------- /tests/test_prompt_config_regression.py: -------------------------------------------------------------------------------- 1 | """Regression tests for system prompt configuration loading.""" 2 | 3 | from unittest.mock import patch, MagicMock 4 | 5 | from assistant_ui_anywidget.ai.prompt_config import SystemPromptConfig 6 | from assistant_ui_anywidget.ai.langgraph_service import LangGraphAIService 7 | from assistant_ui_anywidget.kernel_interface import KernelInterface 8 | 9 | 10 | class TestPromptConfigRegression: 11 | """Test for system prompt configuration loading issues.""" 12 | 13 | def test_pydantic_settings_yaml_loading_fixed(self) -> None: 14 | """Test that YAML loading works correctly with our fix. 15 | 16 | The issue was that pydantic-settings doesn't automatically load YAML files. 17 | We fixed it by using YamlConfigSettingsSource. 18 | """ 19 | # This should now work without errors 20 | config = SystemPromptConfig() 21 | 22 | # Verify all fields are loaded 23 | assert config.approval_note 24 | assert config.main_prompt 25 | assert config.slash_commands 26 | assert config.examples_of_proactive_behavior 27 | assert config.scientific_computing_awareness 28 | assert config.final_reminder 29 | 30 | # Test get_full_prompt method 31 | full_prompt = config.get_full_prompt(require_approval=True) 32 | assert "TOOL USAGE - BE EXTREMELY EAGER!" in full_prompt 33 | assert "You are an **EXTREMELY PROACTIVE** AI assistant" in full_prompt 34 | 35 | def test_langgraph_service_initialization_works(self) -> None: 36 | """Test that LangGraphAIService initializes correctly with fixed prompt config.""" 37 | mock_kernel = MagicMock(spec=KernelInterface) 38 | mock_kernel.is_available = True 39 | 40 | # This should work now 41 | service = LangGraphAIService(kernel=mock_kernel) 42 | assert service is not None 43 | assert service.require_approval is True 44 | 45 | def test_chat_works_with_fixed_prompt_config(self) -> None: 46 | """Test that chat operations work correctly with the fixed prompt config.""" 47 | mock_kernel = MagicMock(spec=KernelInterface) 48 | mock_kernel.is_available = True 49 | mock_kernel.get_kernel_info.return_value = { 50 | "available": True, 51 | "status": "idle", 52 | "language": "python", 53 | "execution_count": 0, 54 | "namespace_size": 0, 55 | } 56 | mock_kernel.get_namespace.return_value = {} 57 | 58 | # Create service (this should work) 59 | service = LangGraphAIService(kernel=mock_kernel) 60 | 61 | # Chat should work and include the system prompt 62 | result = service.chat("hi") 63 | 64 | # Should return a successful result 65 | assert result.success 66 | assert result.content 67 | assert result.thread_id 68 | 69 | def test_system_prompt_included_in_messages(self) -> None: 70 | """Test that the system prompt is correctly included in chat messages.""" 71 | mock_kernel = MagicMock(spec=KernelInterface) 72 | mock_kernel.is_available = True 73 | 74 | # Mock the agent to capture messages 75 | with patch( 76 | "assistant_ui_anywidget.ai.langgraph_service.create_agent_graph" 77 | ) as mock_create: 78 | mock_agent = MagicMock() 79 | captured_messages = [] 80 | 81 | def capture_invoke(payload: dict, config: dict) -> dict: 82 | captured_messages.extend(payload.get("messages", [])) 83 | return {"messages": []} 84 | 85 | mock_agent.invoke = capture_invoke 86 | mock_create.return_value = mock_agent 87 | 88 | service = LangGraphAIService(kernel=mock_kernel) 89 | service.chat("test message") 90 | 91 | # Verify system message is included 92 | assert len(captured_messages) >= 2 93 | system_msg = captured_messages[0] 94 | assert system_msg.content 95 | assert "EXTREMELY PROACTIVE" in system_msg.content 96 | assert "TOOL USAGE - BE EXTREMELY EAGER!" in system_msg.content 97 | -------------------------------------------------------------------------------- /assistant_ui_anywidget/ai/logger.py: -------------------------------------------------------------------------------- 1 | """Conversation logger for AI interactions.""" 2 | 3 | import json 4 | from datetime import datetime 5 | from pathlib import Path 6 | from typing import Any, Dict, List, Optional, TYPE_CHECKING 7 | 8 | if TYPE_CHECKING: 9 | from ..kernel_interface import KernelContext 10 | 11 | 12 | class ConversationLogger: 13 | """Logs AI conversations to timestamped files.""" 14 | 15 | def __init__(self, log_dir: Optional[str] = None): 16 | """Initialize the logger. 17 | 18 | Args: 19 | log_dir: Directory to store logs. Defaults to 'ai_conversation_logs' 20 | """ 21 | self.log_dir = Path(log_dir or "ai_conversation_logs") 22 | self.log_dir.mkdir(exist_ok=True) 23 | self.current_log_file: Optional[Path] = None 24 | self.session_start: Optional[datetime] = None 25 | 26 | def start_session(self) -> Path: 27 | """Start a new logging session with timestamp.""" 28 | self.session_start = datetime.now() 29 | timestamp = self.session_start.strftime("%Y%m%d_%H%M%S") 30 | self.current_log_file = self.log_dir / f"conversation_{timestamp}.json" 31 | 32 | # Initialize log file with metadata 33 | initial_data = { 34 | "session_start": self.session_start.isoformat(), 35 | "session_id": timestamp, 36 | "conversations": [], 37 | } 38 | 39 | with open(self.current_log_file, "w") as f: 40 | json.dump(initial_data, f, indent=2) 41 | 42 | return self.current_log_file 43 | 44 | def log_conversation( 45 | self, 46 | thread_id: str, 47 | user_message: str, 48 | ai_response: str, 49 | tool_calls: Optional[List[Dict[str, Any]]] = None, 50 | context: Optional["KernelContext"] = None, 51 | error: Optional[str] = None, 52 | ) -> None: 53 | """Log a single conversation exchange. 54 | 55 | Args: 56 | thread_id: The thread/conversation ID 57 | user_message: The user's input 58 | ai_response: The AI's response 59 | tool_calls: Any tool calls made during the conversation 60 | context: Kernel context object (will be converted to dict) 61 | error: Any error that occurred 62 | """ 63 | if not self.current_log_file: 64 | self.start_session() 65 | 66 | # Load existing data 67 | assert ( 68 | self.current_log_file is not None 69 | ) # mypy hint: guaranteed by start_session() 70 | with open(self.current_log_file, "r") as f: 71 | data = json.load(f) 72 | 73 | # Create conversation entry 74 | conversation = { 75 | "timestamp": datetime.now().isoformat(), 76 | "thread_id": thread_id, 77 | "user_message": user_message, 78 | "ai_response": ai_response, 79 | "tool_calls": tool_calls or [], 80 | "context": context.to_dict() if context else {}, 81 | "error": error, 82 | } 83 | 84 | # Append to conversations 85 | data["conversations"].append(conversation) 86 | 87 | # Update session end time 88 | data["session_end"] = datetime.now().isoformat() 89 | 90 | # Write back to file 91 | with open(self.current_log_file, "w") as f: 92 | json.dump(data, f, indent=2) 93 | 94 | def get_current_log_path(self) -> Optional[Path]: 95 | """Get the current log file path.""" 96 | return self.current_log_file 97 | 98 | def get_session_summary(self) -> Dict[str, Any]: 99 | """Get a summary of the current session.""" 100 | if not self.current_log_file or not self.current_log_file.exists(): 101 | return {"status": "No active session"} 102 | 103 | with open(self.current_log_file, "r") as f: 104 | data = json.load(f) 105 | 106 | return { 107 | "session_id": data.get("session_id"), 108 | "session_start": data.get("session_start"), 109 | "session_end": data.get("session_end"), 110 | "total_conversations": len(data.get("conversations", [])), 111 | "log_file": str(self.current_log_file), 112 | } 113 | -------------------------------------------------------------------------------- /tests/test_langgraph_approval.py: -------------------------------------------------------------------------------- 1 | """Tests for LangGraph approval workflow.""" 2 | 3 | import os 4 | from unittest.mock import MagicMock, patch 5 | 6 | 7 | from assistant_ui_anywidget.ai.langgraph_service import LangGraphAIService, ChatResult 8 | from assistant_ui_anywidget.kernel_interface import KernelInterface 9 | 10 | 11 | class TestLangGraphApproval: 12 | """Test LangGraph approval workflow functionality.""" 13 | 14 | def test_langgraph_service_creation(self) -> None: 15 | """Test that LangGraph service can be created.""" 16 | mock_kernel = MagicMock(spec=KernelInterface) 17 | mock_kernel.is_available = True 18 | 19 | service = LangGraphAIService(kernel=mock_kernel, require_approval=True) 20 | 21 | assert service is not None 22 | assert service.require_approval is True 23 | assert service.kernel is mock_kernel 24 | 25 | def test_read_only_operations_no_approval(self) -> None: 26 | """Test that read-only operations don't require approval.""" 27 | mock_kernel = MagicMock(spec=KernelInterface) 28 | mock_kernel.is_available = True 29 | mock_kernel.get_namespace.return_value = {"x": 42, "y": "hello"} 30 | mock_kernel.get_kernel_info.return_value = { 31 | "available": True, 32 | "execution_count": 5, 33 | "namespace_size": 2, 34 | } 35 | 36 | with patch.dict(os.environ, {}, clear=True): 37 | service = LangGraphAIService(kernel=mock_kernel, require_approval=True) 38 | 39 | # Test get_variables - should work without approval 40 | result = service.chat("Show me all variables") 41 | assert result.success 42 | assert not result.interrupted 43 | assert result.content # Should have some response 44 | 45 | def test_code_execution_requires_approval(self) -> None: 46 | """Test that code execution requires approval when enabled.""" 47 | mock_kernel = MagicMock(spec=KernelInterface) 48 | mock_kernel.is_available = True 49 | 50 | with patch.dict(os.environ, {}, clear=True): 51 | service = LangGraphAIService(kernel=mock_kernel, require_approval=True) 52 | 53 | # Test execute_code - should interrupt for approval 54 | result = service.chat("Execute x = 42") 55 | 56 | # Note: Without a real LLM, this might not trigger the approval flow 57 | # In real usage, the LLM would call the execute_code tool 58 | assert result is not None 59 | 60 | def test_approval_disabled(self) -> None: 61 | """Test that code execution works without approval when disabled.""" 62 | mock_kernel = MagicMock(spec=KernelInterface) 63 | mock_kernel.is_available = True 64 | 65 | with patch.dict(os.environ, {}, clear=True): 66 | service = LangGraphAIService(kernel=mock_kernel, require_approval=False) 67 | 68 | assert service.require_approval is False 69 | 70 | def test_chat_result_needs_approval(self) -> None: 71 | """Test ChatResult needs_approval property.""" 72 | # Test when needs approval 73 | result = ChatResult( 74 | content="", 75 | thread_id="123", 76 | interrupted=True, 77 | interrupt_message="Approve code execution?", 78 | ) 79 | assert result.needs_approval is True 80 | 81 | # Test when doesn't need approval 82 | result2 = ChatResult( 83 | content="Done", 84 | thread_id="123", 85 | interrupted=False, 86 | ) 87 | assert result2.needs_approval is False 88 | 89 | def test_widget_integration(self) -> None: 90 | """Test that widget can use LangGraph service.""" 91 | from assistant_ui_anywidget import AgentWidget 92 | 93 | # Create widget with LangGraph 94 | widget = AgentWidget( 95 | require_approval=True, 96 | provider="auto", 97 | ) 98 | 99 | assert widget.ai_service is not None 100 | assert isinstance(widget.ai_service, LangGraphAIService) 101 | 102 | # Create widget with LangGraph service (approval disabled) 103 | simple_widget = AgentWidget( 104 | require_approval=False, 105 | provider="auto", 106 | ) 107 | 108 | assert simple_widget.ai_service is not None 109 | assert isinstance(simple_widget.ai_service, LangGraphAIService) 110 | -------------------------------------------------------------------------------- /frontend/src/kernelApi.ts: -------------------------------------------------------------------------------- 1 | /** 2 | * Kernel API service for communicating with the Python backend 3 | */ 4 | 5 | // Simple UUID v4 generator to avoid external dependency 6 | function uuidv4(): string { 7 | return "xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx".replace(/[xy]/g, function (c) { 8 | const r = (Math.random() * 16) | 0; 9 | const v = c === "x" ? r : (r & 0x3) | 0x8; 10 | return v.toString(16); 11 | }); 12 | } 13 | import type { 14 | Request, 15 | Response, 16 | GetVariablesRequest, 17 | GetVariablesResponse, 18 | ExecuteCodeRequest, 19 | ExecuteCodeResponse, 20 | InspectVariableResponse, 21 | VariableInfo, 22 | Output, 23 | DetailedVariableInfo, 24 | } from "./types"; 25 | 26 | export class KernelAPI { 27 | private model: any; // anywidget model 28 | private pendingRequests: Map< 29 | string, 30 | { 31 | resolve: (value: any) => void; 32 | reject: (reason?: any) => void; 33 | } 34 | > = new Map(); 35 | 36 | constructor(model: any) { 37 | this.model = model; 38 | 39 | // Listen for API responses 40 | this.model.on("msg:custom", this.handleMessage.bind(this)); 41 | } 42 | 43 | private handleMessage(msg: any) { 44 | if (msg.type === "api_response" && msg.response) { 45 | const response = msg.response as Response; 46 | const pending = this.pendingRequests.get(response.request_id); 47 | 48 | if (pending) { 49 | this.pendingRequests.delete(response.request_id); 50 | 51 | if (response.success) { 52 | pending.resolve(response); 53 | } else { 54 | pending.reject(response.error); 55 | } 56 | } 57 | } 58 | } 59 | 60 | private sendRequest( 61 | request: Omit 62 | ): Promise { 63 | return new Promise((resolve, reject) => { 64 | const id = uuidv4(); 65 | const fullRequest: Request = { 66 | ...request, 67 | id, 68 | timestamp: Date.now(), 69 | version: "1.0.0", 70 | }; 71 | 72 | // Store the promise handlers 73 | this.pendingRequests.set(id, { resolve, reject }); 74 | 75 | // Send the request 76 | this.model.send({ 77 | type: "api_request", 78 | request: fullRequest, 79 | }); 80 | 81 | // Timeout after 30 seconds 82 | setTimeout(() => { 83 | if (this.pendingRequests.has(id)) { 84 | this.pendingRequests.delete(id); 85 | reject(new Error("Request timeout")); 86 | } 87 | }, 30000); 88 | }); 89 | } 90 | 91 | // Variable management 92 | async getVariables(params?: GetVariablesRequest["params"]): Promise { 93 | const response = await this.sendRequest({ 94 | type: "get_variables", 95 | params, 96 | }); 97 | 98 | return response.data.variables; 99 | } 100 | 101 | async inspectVariable(name: string, deep = false): Promise { 102 | const response = await this.sendRequest({ 103 | type: "inspect_variable", 104 | params: { name, deep }, 105 | }); 106 | 107 | return response.data.info; 108 | } 109 | 110 | // Code execution 111 | async executeCode( 112 | code: string, 113 | options?: Partial 114 | ): Promise { 115 | const response = await this.sendRequest({ 116 | type: "execute_code", 117 | params: { 118 | code, 119 | mode: "exec", 120 | capture_output: true, 121 | ...options, 122 | }, 123 | }); 124 | 125 | return response.data.outputs; 126 | } 127 | 128 | // Kernel info 129 | async getKernelInfo(): Promise { 130 | const response = await this.sendRequest({ 131 | type: "get_kernel_info", 132 | }); 133 | 134 | return response.data; 135 | } 136 | 137 | // History 138 | async getHistory(params?: any): Promise { 139 | const response = await this.sendRequest({ 140 | type: "get_history", 141 | params, 142 | }); 143 | 144 | return response.data; 145 | } 146 | 147 | // Stack trace 148 | async getStackTrace(params?: any): Promise { 149 | const response = await this.sendRequest({ 150 | type: "get_stack_trace", 151 | params, 152 | }); 153 | 154 | return response.data; 155 | } 156 | } 157 | 158 | // Helper function to create API instance 159 | export function createKernelAPI(model: any): KernelAPI { 160 | return new KernelAPI(model); 161 | } 162 | -------------------------------------------------------------------------------- /examples/demo_global_agent.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "metadata": {}, 6 | "source": [ 7 | "# Global Agent Demo - Simple Notebook Interface\n", 8 | "\n", 9 | "This notebook demonstrates the new global agent interface that prevents keyboard shortcut conflicts in Jupyter notebooks." 10 | ] 11 | }, 12 | { 13 | "cell_type": "markdown", 14 | "metadata": {}, 15 | "source": [ 16 | "## 1. Import and Get Agent (Simple Way)\n", 17 | "\n", 18 | "The new interface provides a simple way to get an AI assistant without worrying about keyboard conflicts:" 19 | ] 20 | }, 21 | { 22 | "cell_type": "code", 23 | "execution_count": null, 24 | "metadata": {}, 25 | "outputs": [], 26 | "source": [ 27 | "# The new recommended way - simple and safe!\n", 28 | "from assistant_ui_anywidget import get_agent\n", 29 | "\n", 30 | "# Get the global agent instance (creates if doesn't exist)\n", 31 | "agent = get_agent()\n", 32 | "agent" 33 | ] 34 | }, 35 | { 36 | "cell_type": "markdown", 37 | "metadata": {}, 38 | "source": [ 39 | "## 2. Key Benefits\n", 40 | "\n", 41 | "### 🔥 **Keyboard Safety**\n", 42 | "- Uses **Ctrl+D** to send messages (not Shift+Enter)\n", 43 | "- No more accidental cell execution when chatting!\n", 44 | "- Works consistently across different Jupyter environments\n", 45 | "\n", 46 | "### 🎯 **Global State**\n", 47 | "- Same agent instance across all notebook cells\n", 48 | "- Persistent chat history and configuration\n", 49 | "- No need to pass around widget instances\n", 50 | "\n", 51 | "### ⚡ **Simple API**\n", 52 | "- One line to get started: `get_agent()`\n", 53 | "- Sensible defaults (auto-approve code, auto-detect AI provider)\n", 54 | "- Thread-safe global state management" 55 | ] 56 | }, 57 | { 58 | "cell_type": "markdown", 59 | "metadata": {}, 60 | "source": [ 61 | "## 3. Create Some Test Data\n", 62 | "\n", 63 | "Let's create some variables the agent can explore:" 64 | ] 65 | }, 66 | { 67 | "cell_type": "code", 68 | "execution_count": null, 69 | "metadata": {}, 70 | "outputs": [], 71 | "source": [ 72 | "import numpy as np\n", 73 | "import pandas as pd\n", 74 | "\n", 75 | "# Create test data\n", 76 | "x = 42\n", 77 | "data = np.random.randn(100, 3)\n", 78 | "df = pd.DataFrame(\n", 79 | " {\n", 80 | " \"A\": np.random.randn(50),\n", 81 | " \"B\": np.random.randn(50),\n", 82 | " \"C\": np.random.choice([\"cat\", \"dog\", \"bird\"], 50),\n", 83 | " }\n", 84 | ")\n", 85 | "\n", 86 | "print(\"Test data created! Ask the agent to explore it.\")\n", 87 | "print(\"Try: 'Show me all my variables' or 'What's in my DataFrame?'\")" 88 | ] 89 | }, 90 | { 91 | "cell_type": "markdown", 92 | "metadata": {}, 93 | "source": [ 94 | "## 4. Custom Configuration\n", 95 | "\n", 96 | "You can still customize the agent configuration:" 97 | ] 98 | }, 99 | { 100 | "cell_type": "code", 101 | "execution_count": null, 102 | "metadata": {}, 103 | "outputs": [], 104 | "source": [ 105 | "# Create agent with custom configuration\n", 106 | "custom_agent = get_agent(\n", 107 | " require_approval=True, # Ask before executing code\n", 108 | " provider=\"openai\", # Force specific provider\n", 109 | " model=\"gpt-4\", # Force specific model\n", 110 | " temperature=0.3, # Make responses more focused\n", 111 | " reset=True, # Create fresh instance\n", 112 | ")\n", 113 | "\n", 114 | "print(\"Custom agent created with approval required!\")" 115 | ] 116 | }, 117 | { 118 | "cell_type": "markdown", 119 | "metadata": {}, 120 | "source": [ 121 | "## 💡 Usage Tips\n", 122 | "\n", 123 | "1. **Sending Messages**: Use **Ctrl+D** (or Cmd+D on Mac) to send messages\n", 124 | "2. **Global Instance**: The same agent persists across all cells - no need to recreate\n", 125 | "3. **Auto-Configuration**: The agent auto-detects available AI providers (OpenAI, Anthropic, Google)\n", 126 | "4. **Slash Commands**: Use `/vars`, `/help`, `/exec code` for direct control\n", 127 | "5. **Natural Language**: Just ask questions like \"Show me my DataFrame\" or \"Help debug this error\"\n", 128 | "\n", 129 | "**Remember**: Use **Ctrl+D** to chat, **Shift+Enter** to execute cells!" 130 | ] 131 | } 132 | ], 133 | "metadata": { 134 | "kernelspec": { 135 | "display_name": ".venv", 136 | "language": "python", 137 | "name": "python3" 138 | }, 139 | "language_info": { 140 | "codemirror_mode": { 141 | "name": "ipython", 142 | "version": 3 143 | }, 144 | "file_extension": ".py", 145 | "mimetype": "text/x-python", 146 | "name": "python", 147 | "nbconvert_exporter": "python", 148 | "pygments_lexer": "ipython3", 149 | "version": "3.11.12" 150 | } 151 | }, 152 | "nbformat": 4, 153 | "nbformat_minor": 4 154 | } 155 | -------------------------------------------------------------------------------- /tests/test_imported_modules_context.py: -------------------------------------------------------------------------------- 1 | """Test that imported modules are included in kernel context.""" 2 | 3 | import sys 4 | from unittest.mock import MagicMock 5 | 6 | from assistant_ui_anywidget.agent_widget import AgentWidget 7 | from assistant_ui_anywidget.kernel_interface import KernelInterface 8 | 9 | 10 | class TestImportedModulesContext: 11 | """Test imported modules detection and context inclusion.""" 12 | 13 | def test_get_imported_modules(self) -> None: 14 | """Test that get_imported_modules correctly identifies module types.""" 15 | # Create a mock kernel 16 | mock_kernel = KernelInterface() 17 | 18 | # Mock the shell and namespace 19 | mock_shell = MagicMock() 20 | mock_kernel.shell = mock_shell 21 | 22 | # Import types to create proper module type 23 | from types import ModuleType 24 | 25 | # Create proper module objects for testing 26 | mock_numpy = ModuleType("numpy") 27 | mock_numpy.__file__ = ( 28 | "/usr/local/lib/python3.11/site-packages/numpy/__init__.py" 29 | ) 30 | 31 | mock_custom = ModuleType("mymodule") 32 | mock_custom.__file__ = "/home/user/project/mymodule.py" 33 | 34 | # Set up namespace with modules 35 | mock_namespace = { 36 | "np": mock_numpy, 37 | "mymodule": mock_custom, 38 | "sys": sys, # Real sys module (builtin) 39 | "x": 42, # Not a module 40 | } 41 | mock_shell.user_ns = mock_namespace 42 | 43 | # Test get_imported_modules 44 | imported = mock_kernel.get_imported_modules() 45 | 46 | # Verify results 47 | assert "np" in imported 48 | assert "external" in imported["np"] 49 | assert "numpy" in imported["np"] 50 | 51 | assert "mymodule" in imported 52 | assert "user" in imported["mymodule"] 53 | 54 | assert "sys" in imported 55 | assert "builtin" in imported["sys"] 56 | 57 | # Non-module objects should not be included 58 | assert "x" not in imported 59 | 60 | def test_kernel_context_includes_modules(self) -> None: 61 | """Test that KernelContext includes imported modules.""" 62 | # Create widget 63 | widget = AgentWidget(show_help=False) 64 | 65 | # Replace kernel with a mock 66 | mock_kernel = MagicMock(spec=KernelInterface) 67 | mock_kernel.is_available = True 68 | mock_kernel.get_kernel_info.return_value = { 69 | "available": True, 70 | "status": "idle", 71 | "execution_count": 0, 72 | "namespace_size": 3, 73 | } 74 | mock_kernel.get_namespace.return_value = { 75 | "np": "module", 76 | "pd": "module", 77 | "x": 42, 78 | } 79 | mock_kernel.get_variable_info.return_value = None 80 | mock_kernel.get_notebook_state.return_value = MagicMock(cells=[]) 81 | mock_kernel.get_last_error.return_value = None 82 | mock_kernel.get_imported_modules.return_value = { 83 | "np": "numpy (external)", 84 | "pd": "pandas (external)", 85 | } 86 | 87 | # Replace the kernel 88 | widget.kernel = mock_kernel 89 | 90 | # Get context 91 | context = widget._get_kernel_context() 92 | 93 | # Verify imported modules are in context 94 | assert context.imported_modules is not None 95 | assert "np" in context.imported_modules 96 | assert "pd" in context.imported_modules 97 | assert context.imported_modules["np"] == "numpy (external)" 98 | assert context.imported_modules["pd"] == "pandas (external)" 99 | 100 | # Verify to_dict includes modules 101 | context_dict = context.to_dict() 102 | assert "imported_modules" in context_dict 103 | assert context_dict["imported_modules"] == { 104 | "np": "numpy (external)", 105 | "pd": "pandas (external)", 106 | } 107 | 108 | def test_context_message_includes_modules(self) -> None: 109 | """Test that build_context_message includes imported modules.""" 110 | from assistant_ui_anywidget.ai.langgraph_service import build_context_message 111 | from assistant_ui_anywidget.kernel_interface import KernelContext 112 | 113 | # Create context with modules 114 | context = KernelContext( 115 | kernel_info={"namespace_size": 5}, 116 | variables=[], 117 | imported_modules={ 118 | "np": "numpy (external)", 119 | "pd": "pandas (external)", 120 | "mymodule": "mymodule (user)", 121 | }, 122 | ) 123 | 124 | # Build message 125 | message = build_context_message(context) 126 | 127 | # Verify modules are included 128 | assert "IMPORTED MODULES:" in message 129 | assert "np: numpy (external)" in message 130 | assert "pd: pandas (external)" in message 131 | assert "mymodule: mymodule (user)" in message 132 | -------------------------------------------------------------------------------- /frontend/src/types.ts: -------------------------------------------------------------------------------- 1 | /** 2 | * Type definitions for the kernel interaction API 3 | */ 4 | 5 | // Base message types 6 | export interface BaseMessage { 7 | id: string; 8 | timestamp: number; 9 | type: string; 10 | version: string; 11 | } 12 | 13 | export interface Request extends BaseMessage { 14 | params?: unknown; 15 | } 16 | 17 | export interface Response extends BaseMessage { 18 | request_id: string; 19 | success: boolean; 20 | data?: unknown; 21 | error?: ErrorInfo; 22 | } 23 | 24 | export interface ErrorInfo { 25 | code: string; 26 | message: string; 27 | details?: unknown; 28 | } 29 | 30 | // Variable management types 31 | export interface VariableInfo { 32 | name: string; 33 | type: string; 34 | type_str: string; 35 | size: number | null; 36 | shape?: number[]; 37 | preview: string; 38 | is_callable: boolean; 39 | attributes?: string[]; 40 | last_modified?: number; 41 | } 42 | 43 | export interface GetVariablesRequest extends Request { 44 | type: "get_variables"; 45 | params?: { 46 | filter?: { 47 | types?: string[]; 48 | pattern?: string; 49 | exclude_private?: boolean; 50 | max_preview_size?: number; 51 | }; 52 | sort?: { 53 | by: "name" | "type" | "size" | "modified"; 54 | order: "asc" | "desc"; 55 | }; 56 | }; 57 | } 58 | 59 | export interface GetVariablesResponse extends Response { 60 | type: "get_variables"; 61 | data: { 62 | variables: VariableInfo[]; 63 | total_count: number; 64 | filtered_count: number; 65 | }; 66 | } 67 | 68 | // Code execution types 69 | export interface ExecuteCodeRequest extends Request { 70 | type: "execute_code"; 71 | params: { 72 | code: string; 73 | mode?: "exec" | "eval" | "single"; 74 | capture_output?: boolean; 75 | silent?: boolean; 76 | store_result?: boolean; 77 | timeout?: number; 78 | }; 79 | } 80 | 81 | export interface ExecuteCodeResponse extends Response { 82 | type: "execute_code"; 83 | data: { 84 | execution_count: number; 85 | outputs: Output[]; 86 | execution_time: number; 87 | variables_changed?: string[]; 88 | }; 89 | } 90 | 91 | export interface Output { 92 | type: "stream" | "display_data" | "execute_result" | "error"; 93 | name?: "stdout" | "stderr"; 94 | text?: string; 95 | data?: Record; 96 | metadata?: Record; 97 | traceback?: string[]; 98 | } 99 | 100 | // Variable inspection types 101 | export interface InspectVariableRequest extends Request { 102 | type: "inspect_variable"; 103 | params: { 104 | name: string; 105 | deep?: boolean; 106 | include_methods?: boolean; 107 | include_source?: boolean; 108 | max_depth?: number; 109 | }; 110 | } 111 | 112 | export interface InspectVariableResponse extends Response { 113 | type: "inspect_variable"; 114 | data: { 115 | name: string; 116 | info: DetailedVariableInfo; 117 | }; 118 | } 119 | 120 | export interface DetailedVariableInfo extends VariableInfo { 121 | value?: unknown; 122 | repr: string; 123 | str: string; 124 | doc?: string; 125 | source?: string; 126 | methods?: MethodInfo[]; 127 | attributes_detail?: AttributeInfo[]; 128 | memory_usage?: number; 129 | } 130 | 131 | export interface MethodInfo { 132 | name: string; 133 | type: string; 134 | } 135 | 136 | export interface AttributeInfo { 137 | name: string; 138 | type: string; 139 | value?: unknown; 140 | } 141 | 142 | // Kernel state types 143 | export interface KernelState { 144 | available: boolean; 145 | status: "idle" | "busy" | "error" | "not_connected"; 146 | execution_count: number; 147 | namespace_size: number; 148 | variables_by_type: Record; 149 | } 150 | 151 | // Message types 152 | export interface ChatMessage { 153 | role: "user" | "assistant" | "system"; 154 | content: string; 155 | timestamp?: number; 156 | metadata?: { 157 | execution_time?: number; 158 | tokens_used?: number; 159 | model?: string; 160 | }; 161 | } 162 | 163 | export interface ActionButton { 164 | text: string; 165 | color?: string; 166 | icon?: string; 167 | tooltip?: string; 168 | } 169 | 170 | // Error codes 171 | export enum ErrorCode { 172 | UNKNOWN_ERROR = "UNKNOWN_ERROR", 173 | INVALID_REQUEST = "INVALID_REQUEST", 174 | PERMISSION_DENIED = "PERMISSION_DENIED", 175 | RATE_LIMITED = "RATE_LIMITED", 176 | VARIABLE_NOT_FOUND = "VARIABLE_NOT_FOUND", 177 | VARIABLE_TOO_LARGE = "VARIABLE_TOO_LARGE", 178 | INSPECTION_FAILED = "INSPECTION_FAILED", 179 | EXECUTION_ERROR = "EXECUTION_ERROR", 180 | EXECUTION_TIMEOUT = "EXECUTION_TIMEOUT", 181 | SYNTAX_ERROR = "SYNTAX_ERROR", 182 | KERNEL_NOT_READY = "KERNEL_NOT_READY", 183 | KERNEL_DEAD = "KERNEL_DEAD", 184 | KERNEL_BUSY = "KERNEL_BUSY", 185 | } 186 | 187 | // Widget model state 188 | export interface WidgetModel { 189 | message: string; 190 | chat_history: ChatMessage[]; 191 | action_buttons: ActionButton[]; 192 | kernel_state: KernelState; 193 | code_result: ExecuteCodeResponse | null; 194 | variables_info: VariableInfo[]; 195 | debug_info: any; 196 | } 197 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["hatchling", "hatch-vcs", "hatch-jupyter-builder>=0.5.0"] 3 | build-backend = "hatchling.build" 4 | 5 | [project] 6 | name = "assistant-ui-anywidget" 7 | description = "Interactive AI assistant widget with kernel access for Jupyter notebooks using anywidget" 8 | requires-python = ">=3.10" 9 | dynamic = ["version"] 10 | maintainers = [{ name = "Bas Nijholt", email = "bas@nijho.lt" }] 11 | license = { text = "MIT" } 12 | dependencies = [ 13 | "anywidget>=0.9.0", 14 | "pydantic-settings>=2.10.1", 15 | "python-dotenv>=1.0.0", 16 | "pyyaml>=6.0", 17 | "langgraph>=0.0.20", 18 | "langchain-community>=0.3.27", 19 | "langchain>=0.1.0", 20 | "langchain-openai>=0.0.5", 21 | "langchain-anthropic>=0.1.0", 22 | # Using custom branch until PR is merged: https://github.com/langchain-ai/langchain-google/pull/1056 23 | # Fixes Gemini API issue with AIMessage + tool_calls in conversation history 24 | "langchain-google-genai @ git+https://github.com/basnijholt/langchain-google.git@fix/gemini-aimessage-tool-calls-conversation-history#subdirectory=libs/genai", 25 | ] 26 | classifiers = [ 27 | "Development Status :: 4 - Beta", 28 | "Intended Audience :: Developers", 29 | "Intended Audience :: Science/Research", 30 | "License :: OSI Approved :: MIT License", 31 | "Operating System :: OS Independent", 32 | "Programming Language :: Python", 33 | "Programming Language :: Python :: 3", 34 | "Programming Language :: Python :: 3.10", 35 | "Programming Language :: Python :: 3.11", 36 | "Programming Language :: Python :: 3.12", 37 | "Programming Language :: Python :: 3.13", 38 | "Framework :: Jupyter", 39 | "Framework :: Jupyter :: JupyterLab", 40 | "Framework :: IPython", 41 | "Topic :: Scientific/Engineering", 42 | "Topic :: Software Development :: Libraries :: Python Modules", 43 | "Topic :: Software Development :: Widget Sets", 44 | "Topic :: System :: Shells", 45 | "Typing :: Typed", 46 | ] 47 | keywords = ["jupyter", "widget", "ai", "assistant", "chat", "kernel", "interactive", "anywidget"] 48 | 49 | [project.readme] 50 | content-type = "text/markdown" 51 | file = "README.md" 52 | 53 | [project.urls] 54 | homepage = "https://github.com/basnijholt/assistant-ui-anywidget" 55 | documentation = "https://github.com/basnijholt/assistant-ui-anywidget" 56 | repository = "https://github.com/basnijholt/assistant-ui-anywidget" 57 | issues = "https://github.com/basnijholt/assistant-ui-anywidget/issues" 58 | 59 | [project.optional-dependencies] 60 | dev = ["watchfiles", "jupyterlab", "hatch-vcs", "hatch-jupyter-builder", "pytest", "pytest-cov", "pre-commit", "ruff", "mypy"] 61 | 62 | # Dependency groups (recognized by `uv`). For more details, visit: 63 | # https://peps.python.org/pep-0735/ 64 | [dependency-groups] 65 | dev = ["watchfiles", "jupyterlab", "hatch-vcs", "hatch-jupyter-builder", "pytest", "pytest-cov", "pre-commit", "ruff", "mypy"] 66 | 67 | [tool.hatch.version] 68 | source = "vcs" 69 | 70 | [tool.hatch.metadata] 71 | allow-direct-references = true 72 | 73 | [tool.hatch.build] 74 | artifacts = ["assistant_ui_anywidget/static/*"] 75 | 76 | [tool.hatch.build.hooks.jupyter-builder] 77 | build-function = "hatch_jupyter_builder.npm_builder" 78 | ensured-targets = ["assistant_ui_anywidget/static/index.js"] 79 | skip-if-exists = ["assistant_ui_anywidget/static/index.js"] 80 | dependencies = ["hatch-jupyter-builder>=0.5.0"] 81 | 82 | [tool.hatch.build.hooks.jupyter-builder.build-kwargs] 83 | npm = "npm" 84 | build_cmd = "build" 85 | path = "frontend" 86 | build_dir = "assistant_ui_anywidget/static" 87 | 88 | [tool.pytest.ini_options] 89 | minversion = "8.0" 90 | addopts = [ 91 | "-ra", 92 | "--strict-markers", 93 | "--strict-config", 94 | "--cov=assistant_ui_anywidget", 95 | "--cov-report=term-missing", 96 | "--cov-report=html", 97 | "--cov-report=xml", 98 | "--cov-fail-under=70", 99 | ] 100 | testpaths = ["tests"] 101 | python_files = ["test_*.py"] 102 | python_classes = ["Test*"] 103 | python_functions = ["test_*"] 104 | markers = [ 105 | "slow: marks tests as slow (deselect with '-m \"not slow\"')", 106 | "integration: marks tests as integration tests", 107 | "unit: marks tests as unit tests", 108 | ] 109 | 110 | [tool.coverage.run] 111 | source = ["assistant_ui_anywidget", "."] 112 | omit = [ 113 | "tests/*", 114 | "test_*", 115 | "*/__pycache__/*", 116 | "frontend/*", 117 | ".venv/*", 118 | ] 119 | 120 | [tool.coverage.report] 121 | exclude_lines = [ 122 | "pragma: no cover", 123 | "def __repr__", 124 | "if self.debug:", 125 | "if settings.DEBUG", 126 | "raise AssertionError", 127 | "raise NotImplementedError", 128 | "if 0:", 129 | "if __name__ == .__main__.:", 130 | "class .*\\bProtocol\\):", 131 | "@(abc\\.)?abstractmethod", 132 | ] 133 | 134 | [tool.mypy] 135 | python_version = "3.10" 136 | warn_return_any = true 137 | warn_unused_configs = true 138 | disallow_untyped_defs = true 139 | disallow_incomplete_defs = true 140 | check_untyped_defs = true 141 | disallow_untyped_decorators = true 142 | no_implicit_optional = true 143 | warn_redundant_casts = true 144 | warn_unused_ignores = true 145 | warn_no_return = true 146 | warn_unreachable = true 147 | strict_equality = true 148 | -------------------------------------------------------------------------------- /tests/test_global_agent_auto_detection.py: -------------------------------------------------------------------------------- 1 | """Test auto-detection functionality in global agent.""" 2 | 3 | import os 4 | from unittest.mock import MagicMock, patch 5 | 6 | from assistant_ui_anywidget import get_agent, reset_agent 7 | 8 | 9 | class TestGlobalAgentAutoDetection: 10 | """Test that get_agent() correctly handles automatic provider detection.""" 11 | 12 | def setup_method(self) -> None: 13 | """Reset the global agent before each test.""" 14 | reset_agent() 15 | 16 | def test_get_agent_auto_detection_with_openai_key(self) -> None: 17 | """Test that get_agent() auto-detects OpenAI when API key is available.""" 18 | with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}, clear=True): 19 | with patch( 20 | "assistant_ui_anywidget.ai.langgraph_service.init_chat_model" 21 | ) as mock_init: 22 | mock_init.return_value = MagicMock() 23 | 24 | # get_agent() should auto-detect OpenAI 25 | get_agent() 26 | 27 | # Should have called init_chat_model with OpenAI provider 28 | mock_init.assert_called_once() 29 | call_args = mock_init.call_args 30 | assert call_args.kwargs["model"] == "gpt-4o-mini" 31 | assert call_args.kwargs["model_provider"] == "openai" 32 | 33 | def test_get_agent_auto_detection_with_google_key(self) -> None: 34 | """Test that get_agent() auto-detects Google when API key is available.""" 35 | with patch.dict(os.environ, {"GOOGLE_API_KEY": "test-key"}, clear=True): 36 | with patch( 37 | "assistant_ui_anywidget.ai.langgraph_service.init_chat_model" 38 | ) as mock_init: 39 | mock_init.return_value = MagicMock() 40 | 41 | get_agent() 42 | 43 | # Should have called init_chat_model with Google provider 44 | mock_init.assert_called_once() 45 | call_args = mock_init.call_args 46 | assert call_args.kwargs["model"] == "gemini-2.5-flash" 47 | assert call_args.kwargs["model_provider"] == "google_genai" 48 | 49 | def test_get_agent_auto_detection_with_explicit_provider(self) -> None: 50 | """Test that explicit provider overrides auto-detection.""" 51 | with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}, clear=True): 52 | with patch( 53 | "assistant_ui_anywidget.ai.langgraph_service.init_chat_model" 54 | ) as mock_init: 55 | mock_init.return_value = MagicMock() 56 | 57 | # Override auto-detection with explicit provider 58 | get_agent(provider="openai", model="gpt-3.5-turbo") 59 | 60 | mock_init.assert_called_once() 61 | call_args = mock_init.call_args 62 | assert call_args.kwargs["model"] == "gpt-3.5-turbo" 63 | assert call_args.kwargs["model_provider"] == "openai" 64 | 65 | def test_get_agent_auto_detection_no_keys_fallback(self) -> None: 66 | """Test that get_agent() falls back to MockLLM when no API keys available.""" 67 | # Clear all API keys to ensure fallback to mock 68 | with patch.dict(os.environ, {}, clear=True): 69 | with patch("assistant_ui_anywidget.ai.langgraph_service.load_dotenv"): 70 | agent = get_agent() 71 | 72 | # Should be using MockLLM 73 | assert agent.ai_service is not None 74 | assert agent.ai_service.llm is not None 75 | assert "mock" in str(type(agent.ai_service.llm)).lower() 76 | 77 | def test_get_agent_model_inference_with_auto_provider(self) -> None: 78 | """Test that get_agent() infers provider from model name when provider is auto.""" 79 | with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}, clear=True): 80 | with patch( 81 | "assistant_ui_anywidget.ai.langgraph_service.init_chat_model" 82 | ) as mock_init: 83 | mock_init.return_value = MagicMock() 84 | 85 | # Specify model but keep provider as auto 86 | get_agent(model="gpt-3.5-turbo") 87 | 88 | mock_init.assert_called_once() 89 | call_args = mock_init.call_args 90 | assert call_args.kwargs["model"] == "gpt-3.5-turbo" 91 | assert call_args.kwargs["model_provider"] == "openai" 92 | 93 | def test_get_agent_prefers_openai_when_multiple_keys(self) -> None: 94 | """Test that get_agent() prefers OpenAI when multiple API keys are available.""" 95 | with patch.dict( 96 | os.environ, 97 | {"OPENAI_API_KEY": "test-key", "GOOGLE_API_KEY": "test-key2"}, 98 | clear=True, 99 | ): 100 | with patch( 101 | "assistant_ui_anywidget.ai.langgraph_service.init_chat_model" 102 | ) as mock_init: 103 | mock_init.return_value = MagicMock() 104 | 105 | get_agent() 106 | 107 | # Should prefer OpenAI (first in the list) 108 | mock_init.assert_called_once() 109 | call_args = mock_init.call_args 110 | assert call_args.kwargs["model"] == "gpt-4o-mini" 111 | assert call_args.kwargs["model_provider"] == "openai" 112 | -------------------------------------------------------------------------------- /assistant_ui_anywidget/module_inspector.py: -------------------------------------------------------------------------------- 1 | """Module inspection utilities for reading source code of imported modules.""" 2 | 3 | import inspect 4 | import sys 5 | from pathlib import Path 6 | from typing import Dict, List, Optional, Tuple 7 | 8 | 9 | class ModuleInspector: 10 | """Utilities for inspecting imported modules and reading their source code.""" 11 | 12 | @staticmethod 13 | def get_user_modules() -> Dict[str, str]: 14 | """Get all user-imported modules (excluding standard library and installed packages). 15 | 16 | Returns: 17 | Dict mapping module names to their file paths 18 | """ 19 | user_modules = {} 20 | 21 | for name, module in sys.modules.items(): 22 | # Skip None modules (can happen with failed imports) 23 | if module is None: 24 | continue # type: ignore[unreachable] 25 | 26 | # Skip built-in modules 27 | if not hasattr(module, "__file__") or module.__file__ is None: 28 | continue 29 | 30 | module_path = Path(module.__file__) 31 | 32 | # Skip standard library (usually in python install dir) 33 | if "site-packages" in str(module_path) or "dist-packages" in str( 34 | module_path 35 | ): 36 | continue 37 | 38 | # Skip paths that don't exist 39 | if not module_path.exists(): 40 | continue 41 | 42 | # Only include .py files (skip .so, .pyd, etc.) 43 | if module_path.suffix != ".py": 44 | continue 45 | 46 | user_modules[name] = str(module_path) 47 | 48 | return user_modules 49 | 50 | @staticmethod 51 | def get_module_source(module_name: str) -> Optional[str]: 52 | """Get the source code of a module by name. 53 | 54 | Args: 55 | module_name: Name of the module (e.g., 'my_package.my_module') 56 | 57 | Returns: 58 | Source code as string, or None if not found 59 | """ 60 | if module_name not in sys.modules: 61 | return None 62 | 63 | module = sys.modules[module_name] 64 | 65 | try: 66 | return inspect.getsource(module) 67 | except (TypeError, OSError): 68 | # Try reading the file directly 69 | if hasattr(module, "__file__") and module.__file__: 70 | try: 71 | return Path(module.__file__).read_text() 72 | except Exception: 73 | pass 74 | 75 | return None 76 | 77 | @staticmethod 78 | def get_function_source(module_name: str, function_name: str) -> Optional[str]: 79 | """Get the source code of a specific function in a module. 80 | 81 | Args: 82 | module_name: Name of the module 83 | function_name: Name of the function 84 | 85 | Returns: 86 | Source code of the function, or None if not found 87 | """ 88 | if module_name not in sys.modules: 89 | return None 90 | 91 | module = sys.modules[module_name] 92 | 93 | if not hasattr(module, function_name): 94 | return None 95 | 96 | obj = getattr(module, function_name) 97 | 98 | try: 99 | return inspect.getsource(obj) 100 | except (TypeError, OSError): 101 | return None 102 | 103 | @staticmethod 104 | def find_in_traceback(tb_lines: List[str]) -> List[Tuple[str, int, Optional[str]]]: 105 | """Extract module references from a traceback. 106 | 107 | Args: 108 | tb_lines: List of traceback lines 109 | 110 | Returns: 111 | List of (module_path, line_number, function_name) tuples 112 | """ 113 | import re 114 | 115 | references = [] 116 | pattern = r'File "([^"]+)", line (\d+), in (.+)' 117 | 118 | for line in tb_lines: 119 | match = re.match(pattern, line.strip()) 120 | if match: 121 | file_path, line_num, func_name = match.groups() 122 | references.append((file_path, int(line_num), func_name)) 123 | 124 | return references 125 | 126 | @staticmethod 127 | def read_source_around_line( 128 | file_path: str, line_number: int, context: int = 5 129 | ) -> Optional[str]: 130 | """Read source code around a specific line number. 131 | 132 | Args: 133 | file_path: Path to the source file 134 | line_number: Target line number (1-indexed) 135 | context: Number of lines to show before and after 136 | 137 | Returns: 138 | Source code snippet with line numbers 139 | """ 140 | try: 141 | path = Path(file_path) 142 | if not path.exists(): 143 | return None 144 | 145 | lines = path.read_text().splitlines() 146 | 147 | # Calculate range (1-indexed to 0-indexed) 148 | start = max(0, line_number - context - 1) 149 | end = min(len(lines), line_number + context) 150 | 151 | # Format with line numbers 152 | result = [] 153 | for i in range(start, end): 154 | line_num = i + 1 155 | marker = ">>>" if line_num == line_number else " " 156 | result.append(f"{marker} {line_num:4d} | {lines[i]}") 157 | 158 | return "\n".join(result) 159 | 160 | except Exception: 161 | return None 162 | -------------------------------------------------------------------------------- /frontend/src/VariableExplorer.test.tsx: -------------------------------------------------------------------------------- 1 | /** 2 | * Tests for VariableExplorer component 3 | */ 4 | 5 | import { render, screen, fireEvent } from "@testing-library/react"; 6 | import { describe, it, expect, vi } from "vitest"; 7 | import { VariableExplorer } from "./VariableExplorer"; 8 | import type { VariableInfo } from "./types"; 9 | 10 | describe("VariableExplorer", () => { 11 | const mockVariables: VariableInfo[] = [ 12 | { 13 | name: "df", 14 | type: "DataFrame", 15 | type_str: "pandas.DataFrame", 16 | size: 1024, 17 | shape: [100, 5], 18 | preview: "DataFrame preview", 19 | is_callable: false, 20 | attributes: [], 21 | }, 22 | { 23 | name: "x", 24 | type: "int", 25 | type_str: "builtins.int", 26 | size: 28, 27 | shape: undefined, 28 | preview: "42", 29 | is_callable: false, 30 | attributes: [], 31 | }, 32 | { 33 | name: "calculate_mean", 34 | type: "function", 35 | type_str: "function", 36 | size: null, 37 | preview: "", 38 | is_callable: true, 39 | attributes: [], 40 | }, 41 | ]; 42 | 43 | const mockOnInspect = vi.fn(); 44 | const mockOnExecute = vi.fn(); 45 | 46 | it("should render variable list", () => { 47 | render( 48 | 53 | ); 54 | 55 | expect(screen.getByText("Variable Explorer")).toBeInTheDocument(); 56 | expect(screen.getByText("df")).toBeInTheDocument(); 57 | expect(screen.getByText("x")).toBeInTheDocument(); 58 | expect(screen.getByText("calculate_mean")).toBeInTheDocument(); 59 | }); 60 | 61 | it("should show variable types and sizes", () => { 62 | render(); 63 | 64 | // Use getAllByText since "DataFrame" appears in both the filter and the table 65 | const dataFrameElements = screen.getAllByText("DataFrame"); 66 | expect(dataFrameElements.length).toBeGreaterThan(0); 67 | 68 | expect(screen.getByText("100×5")).toBeInTheDocument(); 69 | expect(screen.getByText("1.0 KB")).toBeInTheDocument(); // Updated to match actual output 70 | expect(screen.getByText("()")).toBeInTheDocument(); // Callable indicator 71 | }); 72 | 73 | it("should filter variables by search term", () => { 74 | render(); 75 | 76 | const searchInput = screen.getByPlaceholderText("🔍 Search variables..."); 77 | fireEvent.change(searchInput, { target: { value: "df" } }); 78 | 79 | expect(screen.getByText("df")).toBeInTheDocument(); 80 | expect(screen.queryByText("calculate_mean")).not.toBeInTheDocument(); 81 | }); 82 | 83 | it("should filter by type", () => { 84 | render(); 85 | 86 | const typeSelect = screen.getByRole("combobox"); 87 | fireEvent.change(typeSelect, { target: { value: "DataFrame" } }); 88 | 89 | expect(screen.getByText("df")).toBeInTheDocument(); 90 | expect(screen.queryByText("x")).not.toBeInTheDocument(); 91 | expect(screen.queryByText("calculate_mean")).not.toBeInTheDocument(); 92 | }); 93 | 94 | it("should sort variables", () => { 95 | render(); 96 | 97 | // Click on Type header to sort 98 | const typeHeader = screen.getByText("Type"); 99 | fireEvent.click(typeHeader); 100 | 101 | // Check that items are sorted by type 102 | const rows = screen.getAllByRole("row"); 103 | expect(rows.length).toBeGreaterThan(1); // Header + data rows 104 | }); 105 | 106 | it("should call onInspect when variable is clicked", () => { 107 | render(); 108 | 109 | const dfRow = screen.getByText("df").closest("tr"); 110 | if (dfRow) { 111 | fireEvent.click(dfRow); 112 | } 113 | 114 | // Verify the spy was called with the correct DataFrame variable 115 | expect(mockOnInspect).toHaveBeenCalledTimes(1); 116 | expect(mockOnInspect).toHaveBeenCalledWith( 117 | expect.objectContaining({ 118 | name: "df", 119 | type: "DataFrame", 120 | type_str: "pandas.DataFrame", 121 | }) 122 | ); 123 | }); 124 | 125 | it("should show loading state", () => { 126 | render(); 127 | 128 | expect(screen.getByText("Loading variables...")).toBeInTheDocument(); 129 | }); 130 | 131 | it("should show empty state", () => { 132 | render(); 133 | 134 | expect(screen.getByText("No variables in namespace")).toBeInTheDocument(); 135 | }); 136 | 137 | it("should show filtered empty state", () => { 138 | render(); 139 | 140 | const searchInput = screen.getByPlaceholderText("🔍 Search variables..."); 141 | fireEvent.change(searchInput, { target: { value: "nonexistent" } }); 142 | 143 | expect(screen.getByText("No variables match your filters")).toBeInTheDocument(); 144 | }); 145 | 146 | it("should display variable count in footer", () => { 147 | render(); 148 | 149 | expect(screen.getByText("3 variables")).toBeInTheDocument(); 150 | }); 151 | }); 152 | -------------------------------------------------------------------------------- /assistant_ui_anywidget/ai/mock.py: -------------------------------------------------------------------------------- 1 | """Mock LLM for testing and when no API is configured.""" 2 | 3 | from typing import Any, List, Optional, Sequence, Union 4 | from langchain_core.language_models import BaseChatModel 5 | from langchain_core.messages import AIMessage, BaseMessage 6 | from langchain_core.outputs import ChatResult, ChatGeneration 7 | from langchain_core.callbacks import CallbackManagerForLLMRun 8 | from langchain_core.tools import BaseTool 9 | from langchain_core.runnables import Runnable 10 | 11 | 12 | class MockLLM(BaseChatModel): 13 | """Mock LLM that provides helpful responses without calling an API.""" 14 | 15 | @property 16 | def _llm_type(self) -> str: 17 | """Return identifier of llm type.""" 18 | return "mock" 19 | 20 | def _generate( 21 | self, 22 | messages: List[BaseMessage], 23 | stop: Optional[List[str]] = None, 24 | run_manager: Optional[CallbackManagerForLLMRun] = None, 25 | **kwargs: Any, 26 | ) -> ChatResult: 27 | """Generate a mock response.""" 28 | # Get the last user message 29 | last_message = "" 30 | for msg in reversed(messages): 31 | if msg.type == "human": 32 | last_message = msg.content 33 | break 34 | 35 | # Generate a helpful response 36 | response = self._get_mock_response(last_message) 37 | 38 | message = AIMessage(content=response) 39 | generation = ChatGeneration(message=message) 40 | 41 | return ChatResult(generations=[generation]) 42 | 43 | def _get_mock_response(self, message: str) -> str: 44 | """Generate a mock response based on the message.""" 45 | message_lower = message.lower() 46 | 47 | if any(greeting in message_lower for greeting in ["hello", "hi", "hey"]): 48 | return ( 49 | "Hello! I'm a mock AI assistant. While I can't provide intelligent responses, " 50 | "I can help demonstrate the kernel integration features.\n\n" 51 | "To use a real AI, please:\n" 52 | "1. Install an AI provider: `pip install openai anthropic google-generativeai`\n" 53 | "2. Set your API key: `export OPENAI_API_KEY='your-key'`\n\n" 54 | "Try asking me to show variables or execute code!" 55 | ) 56 | 57 | elif "variable" in message_lower or "namespace" in message_lower: 58 | return ( 59 | "I can help you explore variables! While I'm just a mock AI, " 60 | "I have access to kernel tools. Try these:\n\n" 61 | "- Ask me to 'list all variables'\n" 62 | "- Ask me to 'inspect df' (if you have a DataFrame named df)\n" 63 | "- Ask me to 'show the type of x' (for any variable x)\n\n" 64 | "I'll use the appropriate tools to get real information from your kernel!" 65 | ) 66 | 67 | elif ( 68 | "execute" in message_lower 69 | or "run" in message_lower 70 | or "calculate" in message_lower 71 | ): 72 | return ( 73 | "I can execute code in your kernel! For example:\n\n" 74 | "- 'Calculate 2 + 2'\n" 75 | "- 'Create a variable x with value 42'\n" 76 | "- 'Import numpy as np'\n\n" 77 | "Note: As a mock AI, I'll execute simple requests but won't provide " 78 | "intelligent code generation. Install a real AI provider for that!" 79 | ) 80 | 81 | elif "help" in message_lower: 82 | return ( 83 | "I'm a mock AI assistant with kernel access. Here's what I can do:\n\n" 84 | "**Kernel Tools:**\n" 85 | "- List variables: 'Show me all variables'\n" 86 | "- Inspect data: 'What's in the variable df?'\n" 87 | "- Execute code: 'Calculate the mean of [1, 2, 3]'\n" 88 | "- Kernel info: 'What's the kernel status?'\n\n" 89 | "**To use a real AI:**\n" 90 | "1. `pip install langchain-openai` (or anthropic, google-genai)\n" 91 | "2. Set API key: `export OPENAI_API_KEY='your-key'`\n" 92 | "3. Restart the kernel\n\n" 93 | "Even as a mock, I have real access to your kernel!" 94 | ) 95 | 96 | elif any( 97 | word in message_lower for word in ["error", "debug", "problem", "issue"] 98 | ): 99 | return ( 100 | "I can help you debug! While I'm a mock AI, I can:\n\n" 101 | "- Check your variables: 'Show me all variables'\n" 102 | "- Inspect specific data: 'What's the type of my_var?'\n" 103 | "- Look at error details: 'Show me the last error'\n\n" 104 | "For intelligent debugging assistance, please set up a real AI provider." 105 | ) 106 | 107 | else: 108 | return ( 109 | f"I understood: '{message}'\n\n" 110 | "As a mock AI, I can't provide intelligent responses, but I can:\n" 111 | "- List your variables\n" 112 | "- Inspect specific data\n" 113 | "- Execute simple code\n" 114 | "- Show kernel information\n\n" 115 | "For better assistance, please configure an AI provider with an API key." 116 | ) 117 | 118 | def bind_tools( 119 | self, 120 | tools: Sequence[Union[BaseTool, type[BaseTool], dict]], 121 | **kwargs: Any, 122 | ) -> Runnable[Any, Any]: 123 | """Bind tools to the model for compatibility.""" 124 | # Return self since mock doesn't actually use tools 125 | return self 126 | -------------------------------------------------------------------------------- /frontend/src/kernelApi.test.ts: -------------------------------------------------------------------------------- 1 | /** 2 | * Tests for KernelAPI 3 | */ 4 | 5 | import { describe, it, expect, vi, beforeEach } from "vitest"; 6 | import { KernelAPI } from "./kernelApi"; 7 | 8 | describe("KernelAPI", () => { 9 | let mockModel: any; 10 | let api: KernelAPI; 11 | let messageHandler: (msg: any) => void; 12 | 13 | beforeEach(() => { 14 | // Create mock model 15 | mockModel = { 16 | send: vi.fn(), 17 | on: vi.fn((event: string, handler: (msg: any) => void) => { 18 | if (event === "msg:custom") { 19 | messageHandler = handler; 20 | } 21 | }), 22 | }; 23 | 24 | api = new KernelAPI(mockModel); 25 | }); 26 | 27 | describe("constructor", () => { 28 | it("should set up message listener", () => { 29 | expect(mockModel.on).toHaveBeenCalledWith("msg:custom", expect.any(Function)); 30 | }); 31 | }); 32 | 33 | describe("getVariables", () => { 34 | it("should send get_variables request and return variables", async () => { 35 | const mockVariables = [ 36 | { name: "x", type: "int", preview: "42" }, 37 | { name: "y", type: "str", preview: "'hello'" }, 38 | ]; 39 | 40 | // Start the request 41 | const promise = api.getVariables(); 42 | 43 | // Verify request was sent 44 | expect(mockModel.send).toHaveBeenCalledWith({ 45 | type: "api_request", 46 | request: { 47 | type: "get_variables", 48 | params: undefined, 49 | id: expect.any(String), 50 | timestamp: expect.any(Number), 51 | version: "1.0.0", 52 | }, 53 | }); 54 | 55 | // Get the request ID 56 | const sentRequest = mockModel.send.mock.calls[0][0].request; 57 | const requestId = sentRequest.id; 58 | 59 | // Simulate response 60 | messageHandler({ 61 | type: "api_response", 62 | response: { 63 | request_id: requestId, 64 | success: true, 65 | data: { 66 | variables: mockVariables, 67 | total_count: 2, 68 | filtered_count: 2, 69 | }, 70 | }, 71 | }); 72 | 73 | // Check result 74 | const result = await promise; 75 | expect(result).toEqual(mockVariables); 76 | }); 77 | 78 | it("should handle filters", async () => { 79 | const filters = { 80 | filter: { 81 | types: ["DataFrame"], 82 | exclude_private: true, 83 | }, 84 | sort: { 85 | by: "name" as const, 86 | order: "asc" as const, 87 | }, 88 | }; 89 | 90 | api.getVariables(filters); 91 | 92 | expect(mockModel.send).toHaveBeenCalledWith({ 93 | type: "api_request", 94 | request: { 95 | type: "get_variables", 96 | params: filters, 97 | id: expect.any(String), 98 | timestamp: expect.any(Number), 99 | version: "1.0.0", 100 | }, 101 | }); 102 | }); 103 | }); 104 | 105 | describe("executeCode", () => { 106 | it("should execute code and return outputs", async () => { 107 | const code = 'print("Hello")'; 108 | const mockOutputs = [{ type: "stream", name: "stdout", text: "Hello\n" }]; 109 | 110 | const promise = api.executeCode(code); 111 | 112 | // Get request ID 113 | const sentRequest = mockModel.send.mock.calls[0][0].request; 114 | const requestId = sentRequest.id; 115 | 116 | // Verify request 117 | expect(sentRequest).toMatchObject({ 118 | type: "execute_code", 119 | params: { 120 | code, 121 | mode: "exec", 122 | capture_output: true, 123 | }, 124 | }); 125 | 126 | // Send response 127 | messageHandler({ 128 | type: "api_response", 129 | response: { 130 | request_id: requestId, 131 | success: true, 132 | data: { 133 | execution_count: 1, 134 | outputs: mockOutputs, 135 | execution_time: 10, 136 | }, 137 | }, 138 | }); 139 | 140 | const result = await promise; 141 | expect(result).toEqual(mockOutputs); 142 | }); 143 | }); 144 | 145 | describe("inspectVariable", () => { 146 | it("should inspect variable and return detailed info", async () => { 147 | const mockInfo = { 148 | name: "df", 149 | type: "DataFrame", 150 | shape: [100, 5], 151 | preview: " A B C\n0 1 2 3", 152 | }; 153 | 154 | const promise = api.inspectVariable("df", true); 155 | 156 | const sentRequest = mockModel.send.mock.calls[0][0].request; 157 | expect(sentRequest).toMatchObject({ 158 | type: "inspect_variable", 159 | params: { name: "df", deep: true }, 160 | }); 161 | 162 | // Send response 163 | messageHandler({ 164 | type: "api_response", 165 | response: { 166 | request_id: sentRequest.id, 167 | success: true, 168 | data: { name: "df", info: mockInfo }, 169 | }, 170 | }); 171 | 172 | const result = await promise; 173 | expect(result).toEqual(mockInfo); 174 | }); 175 | }); 176 | 177 | describe("error handling", () => { 178 | it("should reject on error response", async () => { 179 | const promise = api.getVariables(); 180 | const sentRequest = mockModel.send.mock.calls[0][0].request; 181 | 182 | messageHandler({ 183 | type: "api_response", 184 | response: { 185 | request_id: sentRequest.id, 186 | success: false, 187 | error: { 188 | code: "KERNEL_NOT_READY", 189 | message: "Kernel is not available", 190 | }, 191 | }, 192 | }); 193 | 194 | await expect(promise).rejects.toEqual({ 195 | code: "KERNEL_NOT_READY", 196 | message: "Kernel is not available", 197 | }); 198 | }); 199 | 200 | it("should timeout after 30 seconds", async () => { 201 | vi.useFakeTimers(); 202 | const promise = api.getVariables(); 203 | 204 | // Fast-forward time 205 | vi.advanceTimersByTime(30001); 206 | 207 | await expect(promise).rejects.toThrow("Request timeout"); 208 | vi.useRealTimers(); 209 | }); 210 | }); 211 | }); 212 | -------------------------------------------------------------------------------- /tests/test_widget_basic.py: -------------------------------------------------------------------------------- 1 | """Basic pytest tests for AgentWidget functionality.""" 2 | 3 | import pathlib 4 | from collections.abc import Callable 5 | 6 | import pytest 7 | 8 | from assistant_ui_anywidget.agent_widget import AgentWidget 9 | 10 | # Type alias for the UI message factory 11 | UIMessageFactory = Callable[[str], dict[str, str]] 12 | 13 | 14 | class TestWidgetBasics: 15 | """Basic widget functionality tests.""" 16 | 17 | def test_widget_creation(self, widget: AgentWidget) -> None: 18 | """Test that widget can be created successfully.""" 19 | assert widget is not None 20 | assert hasattr(widget, "chat_history") 21 | assert hasattr(widget, "_esm") 22 | 23 | def test_initial_state(self, widget: AgentWidget) -> None: 24 | """Test widget initial state.""" 25 | assert widget.chat_history == [] 26 | assert len(widget.chat_history) == 0 27 | 28 | def test_esm_content_loaded(self, widget: AgentWidget) -> None: 29 | """Test that the ESM bundle content is loaded.""" 30 | # AnyWidget loads the file content, not the path 31 | assert isinstance(widget._esm, str), "ESM should be a string" # type: ignore[unreachable] 32 | assert len(widget._esm) > 0, "ESM content should not be empty" # type: ignore[unreachable] 33 | # Check for React content (indicating the bundle is loaded) 34 | assert "react" in widget._esm.lower(), "ESM should contain React code" 35 | 36 | def test_esm_bundle_file_exists(self) -> None: 37 | """Test that the original ESM bundle file exists.""" 38 | # The original path before AnyWidget loads it 39 | expected_path = ( 40 | pathlib.Path(__file__).parent.parent 41 | / "assistant_ui_anywidget" 42 | / "static" 43 | / "index.js" 44 | ) 45 | assert expected_path.exists(), f"ESM bundle file not found at {expected_path}" 46 | assert expected_path.is_file(), ( 47 | f"ESM bundle path is not a file: {expected_path}" 48 | ) 49 | 50 | 51 | class TestMessageAPI: 52 | """Test the message API methods.""" 53 | 54 | def test_add_message_user(self, widget: AgentWidget) -> None: 55 | """Test adding a user message.""" 56 | widget.add_message("user", "Hello") 57 | 58 | assert len(widget.chat_history) == 1 59 | message = widget.chat_history[0] 60 | assert message["role"] == "user" 61 | assert message["content"] == "Hello" 62 | 63 | def test_add_message_assistant(self, widget: AgentWidget) -> None: 64 | """Test adding an assistant message.""" 65 | widget.add_message("assistant", "Hello back!") 66 | 67 | assert len(widget.chat_history) == 1 68 | message = widget.chat_history[0] 69 | assert message["role"] == "assistant" 70 | assert message["content"] == "Hello back!" 71 | 72 | def test_add_multiple_messages(self, widget: AgentWidget) -> None: 73 | """Test adding multiple messages.""" 74 | widget.add_message("user", "First message") 75 | widget.add_message("assistant", "First response") 76 | widget.add_message("user", "Second message") 77 | 78 | expected_message_count = 3 79 | assert len(widget.chat_history) == expected_message_count 80 | assert widget.chat_history[0]["content"] == "First message" 81 | assert widget.chat_history[1]["content"] == "First response" 82 | assert widget.chat_history[2]["content"] == "Second message" 83 | 84 | def test_get_chat_history(self, widget: AgentWidget) -> None: 85 | """Test getting chat history.""" 86 | widget.add_message("user", "Test message") 87 | 88 | history = widget.get_chat_history() 89 | assert isinstance(history, list) 90 | assert len(history) == 1 91 | assert history[0]["content"] == "Test message" 92 | # Should be a copy, not the same object 93 | assert history is not widget.chat_history 94 | 95 | def test_clear_chat_history(self, widget: AgentWidget) -> None: 96 | """Test clearing chat history.""" 97 | widget.add_message("user", "Message 1") 98 | widget.add_message("assistant", "Response 1") 99 | 100 | expected_message_count = 2 101 | assert len(widget.chat_history) == expected_message_count 102 | 103 | widget.clear_chat_history() 104 | 105 | assert len(widget.chat_history) == 0 106 | assert widget.chat_history == [] 107 | 108 | 109 | class TestUIMessageHandling: 110 | """Test UI message handling.""" 111 | 112 | def test_handle_user_message( 113 | self, 114 | widget: AgentWidget, 115 | ui_message_factory: UIMessageFactory, 116 | ) -> None: 117 | """Test handling a user message from UI.""" 118 | ui_message = ui_message_factory("Hello from UI") 119 | widget._handle_message(None, ui_message) 120 | 121 | assert len(widget.chat_history) == 1 122 | assert widget.chat_history[0]["role"] == "user" 123 | assert widget.chat_history[0]["content"] == "Hello from UI" 124 | 125 | def test_handle_invalid_message_type(self, widget: AgentWidget) -> None: 126 | """Test handling invalid message type.""" 127 | invalid_message = {"type": "invalid", "text": "Should be ignored"} 128 | widget._handle_message(None, invalid_message) 129 | 130 | assert len(widget.chat_history) == 0 131 | 132 | def test_handle_empty_message( 133 | self, 134 | widget: AgentWidget, 135 | ui_message_factory: UIMessageFactory, 136 | ) -> None: 137 | """Test handling empty message.""" 138 | ui_message = ui_message_factory("") 139 | widget._handle_message(None, ui_message) 140 | 141 | assert len(widget.chat_history) == 1 142 | assert widget.chat_history[0]["content"] == "" 143 | 144 | def test_handle_message_no_auto_response( 145 | self, 146 | widget: AgentWidget, 147 | ui_message_factory: UIMessageFactory, 148 | ) -> None: 149 | """Test that handling UI messages doesn't create auto-responses.""" 150 | ui_message = ui_message_factory("Test message") 151 | widget._handle_message(None, ui_message) 152 | 153 | # Should only have the user message, no assistant response 154 | assert len(widget.chat_history) == 1 155 | assert widget.chat_history[0]["role"] == "user" 156 | 157 | 158 | if __name__ == "__main__": 159 | pytest.main([__file__, "-v"]) 160 | -------------------------------------------------------------------------------- /assistant_ui_anywidget/simple_handlers.py: -------------------------------------------------------------------------------- 1 | """Simplified message handlers for kernel interaction.""" 2 | 3 | import time 4 | from typing import Any, Dict, List, Optional 5 | 6 | from .kernel_interface import KernelInterface 7 | 8 | 9 | class SimpleHandlers: 10 | """Simple message handlers for kernel interaction.""" 11 | 12 | def __init__(self, kernel_interface: Optional[KernelInterface] = None): 13 | """Initialize handlers.""" 14 | self.kernel = kernel_interface or KernelInterface() 15 | self.execution_history: List[Dict[str, Any]] = [] 16 | 17 | def handle_message(self, message: Any) -> Dict[str, Any]: 18 | """Route message to appropriate handler.""" 19 | # Handle non-dict messages 20 | if not isinstance(message, dict): 21 | return {"success": False, "error": "Message type is required"} 22 | 23 | msg_type = message.get("type") 24 | if not msg_type: 25 | return {"success": False, "error": "Message type is required"} 26 | 27 | # Route to handler methods 28 | if msg_type == "get_variables": 29 | return self._handle_get_variables(message) 30 | elif msg_type == "inspect_variable": 31 | return self._handle_inspect_variable(message) 32 | elif msg_type == "execute_code": 33 | return self._handle_execute_code(message) 34 | elif msg_type == "get_kernel_info": 35 | return self._handle_get_kernel_info(message) 36 | elif msg_type == "get_stack_trace": 37 | return self._handle_get_stack_trace(message) 38 | elif msg_type == "get_history": 39 | return self._handle_get_history(message) 40 | else: 41 | return {"success": False, "error": f"Unknown message type: {msg_type}"} 42 | 43 | def _handle_get_variables(self, message: Dict[str, Any]) -> Dict[str, Any]: 44 | """Handle get_variables request.""" 45 | if not self.kernel.is_available: 46 | return {"success": False, "error": "Kernel is not available"} 47 | 48 | params = message.get("params", {}) 49 | filter_params = params.get("filter", {}) 50 | sort_params = params.get("sort", {}) 51 | 52 | # Get filter parameters 53 | types_filter = filter_params.get("types", []) 54 | pattern = filter_params.get("pattern") 55 | exclude_private = filter_params.get("exclude_private", True) 56 | 57 | # Get sort parameters 58 | sort_by = sort_params.get("by", "name") 59 | sort_order = sort_params.get("order", "asc") 60 | 61 | # Get all variables 62 | namespace = self.kernel.get_namespace() 63 | variables = [] 64 | 65 | for name, value in namespace.items(): 66 | # Apply filters 67 | if exclude_private and name.startswith("_"): 68 | continue 69 | 70 | if pattern: 71 | import re 72 | 73 | if not re.search(pattern, name): 74 | continue 75 | 76 | # Get variable info 77 | var_info = self.kernel.get_variable_info(name) 78 | if var_info: 79 | if types_filter and var_info.type not in types_filter: 80 | continue 81 | variables.append(var_info.to_dict()) 82 | 83 | # Sort variables 84 | if sort_by == "name": 85 | variables.sort(key=lambda x: x["name"]) 86 | elif sort_by == "type": 87 | variables.sort(key=lambda x: x["type"]) 88 | elif sort_by == "size": 89 | variables.sort(key=lambda x: x.get("size", 0) or 0) 90 | 91 | if sort_order == "desc": 92 | variables.reverse() 93 | 94 | return { 95 | "success": True, 96 | "data": { 97 | "variables": variables, 98 | "total_count": len(variables), 99 | "timestamp": time.time(), 100 | }, 101 | } 102 | 103 | def _handle_inspect_variable(self, message: Dict[str, Any]) -> Dict[str, Any]: 104 | """Handle inspect_variable request.""" 105 | if not self.kernel.is_available: 106 | return {"success": False, "error": "Kernel is not available"} 107 | 108 | params = message.get("params", {}) 109 | var_name = params.get("name") 110 | 111 | if not var_name: 112 | return {"success": False, "error": "Variable name is required"} 113 | 114 | # Get variable info 115 | deep = params.get("deep", False) 116 | var_info = self.kernel.get_variable_info(var_name, deep=deep) 117 | 118 | if not var_info: 119 | return {"success": False, "error": f"Variable '{var_name}' not found"} 120 | 121 | return {"success": True, "data": var_info.to_dict()} 122 | 123 | def _handle_execute_code(self, message: Dict[str, Any]) -> Dict[str, Any]: 124 | """Handle execute_code request.""" 125 | if not self.kernel.is_available: 126 | return {"success": False, "error": "Kernel is not available"} 127 | 128 | params = message.get("params", {}) 129 | code = params.get("code") 130 | 131 | if not code: 132 | return {"success": False, "error": "Code is required"} 133 | 134 | # Execute code 135 | silent = params.get("silent", False) 136 | store_history = params.get("store_history", True) 137 | 138 | result = self.kernel.execute_code( 139 | code, silent=silent, store_history=store_history 140 | ) 141 | 142 | # Store in execution history 143 | if store_history: 144 | self.execution_history.append( 145 | { 146 | "input": code, 147 | "timestamp": time.time(), 148 | "result": result.to_dict(), 149 | } 150 | ) 151 | 152 | if result.success: 153 | return {"success": True, "data": result.to_dict()} 154 | else: 155 | error_msg = ( 156 | result.error.get("message", "Execution failed") 157 | if result.error 158 | else "Execution failed" 159 | ) 160 | return {"success": False, "error": error_msg, "details": result.to_dict()} 161 | 162 | def _handle_get_kernel_info(self, message: Dict[str, Any]) -> Dict[str, Any]: 163 | """Handle get_kernel_info request.""" 164 | return {"success": True, "data": self.kernel.get_kernel_info()} 165 | 166 | def _handle_get_stack_trace(self, message: Dict[str, Any]) -> Dict[str, Any]: 167 | """Handle get_stack_trace request.""" 168 | if not self.kernel.is_available: 169 | return {"success": False, "error": "Kernel is not available"} 170 | 171 | params = message.get("params", {}) 172 | max_depth = params.get("max_depth", 10) 173 | 174 | # Get last error from kernel 175 | last_error = self.kernel.get_last_error() 176 | if not last_error: 177 | return { 178 | "success": True, 179 | "data": {"stack_trace": None, "message": "No recent error"}, 180 | } 181 | 182 | # Limit traceback depth 183 | traceback = last_error.get("traceback", []) 184 | if max_depth and len(traceback) > max_depth: 185 | traceback = traceback[:max_depth] 186 | 187 | return { 188 | "success": True, 189 | "data": { 190 | "stack_trace": traceback, 191 | "error_type": last_error.get("type"), 192 | "message": last_error.get("message"), 193 | }, 194 | } 195 | 196 | def _handle_get_history(self, message: Dict[str, Any]) -> Dict[str, Any]: 197 | """Handle get_history request.""" 198 | params = message.get("params", {}) 199 | limit = params.get("limit", 50) 200 | search = params.get("search") 201 | 202 | history = self.execution_history 203 | 204 | # Apply search filter 205 | if search: 206 | history = [ 207 | item for item in history if search.lower() in item["input"].lower() 208 | ] 209 | 210 | # Apply limit 211 | if limit: 212 | history = history[-limit:] 213 | 214 | return { 215 | "success": True, 216 | "data": { 217 | "items": history, 218 | "total_count": len(self.execution_history), 219 | "filtered_count": len(history), 220 | }, 221 | } 222 | -------------------------------------------------------------------------------- /frontend/demo.html: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | 6 | Assistant UI Widget - Demo 7 | 8 | 9 | 153 | 154 | 155 |

Assistant UI Widget - Demo

156 |
157 |
158 |
159 |
160 |
161 | 162 | 248 | 249 | 250 | -------------------------------------------------------------------------------- /tests/test_ai_service_regression.py: -------------------------------------------------------------------------------- 1 | """Regression tests for AI service issues.""" 2 | 3 | import os 4 | from typing import Any 5 | from unittest.mock import MagicMock, patch 6 | 7 | from assistant_ui_anywidget.ai import AIService 8 | from assistant_ui_anywidget.kernel_interface import KernelInterface 9 | 10 | 11 | class TestAIServiceRegression: 12 | """Regression tests for AI service bugs.""" 13 | 14 | def test_missing_model_argument_regression(self) -> None: 15 | """Test for bug: _init_chat_model_helper() missing 1 required positional argument: 'model'. 16 | 17 | This reproduces the error from ai_conversation_logs/conversation_20250723_115415.json 18 | where the AI service fails to initialize when no API keys are available and no model 19 | is explicitly specified. 20 | 21 | The bug occurs in the _init_llm method when: 22 | 1. Some API keys are detected (but init_chat_model still fails) 23 | 2. The code tries to call init_chat_model with model=None 24 | 3. init_chat_model requires the model parameter but gets None 25 | """ 26 | mock_kernel = MagicMock(spec=KernelInterface) 27 | mock_kernel.is_available = True 28 | 29 | # Setup: Mock init_chat_model to raise the exact error from the log 30 | def mock_init_chat_model_error(*args: Any, **kwargs: Any) -> MagicMock: 31 | # Simulate the exact error condition 32 | if kwargs.get("model") is None: 33 | raise TypeError( 34 | "_init_chat_model_helper() missing 1 required positional argument: 'model'" 35 | ) 36 | return MagicMock() 37 | 38 | # Create a scenario where we have API keys but init_chat_model fails 39 | with ( 40 | patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}, clear=True), 41 | patch("assistant_ui_anywidget.ai.langgraph_service.load_dotenv"), 42 | patch( 43 | "assistant_ui_anywidget.ai.langgraph_service.init_chat_model", 44 | side_effect=mock_init_chat_model_error, 45 | ), 46 | ): 47 | # This should NOT crash but should fall back to MockLLM when init_chat_model fails 48 | service = AIService(kernel=mock_kernel) 49 | 50 | # Verify it fell back to MockLLM after the init_chat_model error 51 | assert service.llm is not None 52 | # Should be MockLLM since init_chat_model failed 53 | assert "mock" in str(type(service.llm)).lower() 54 | 55 | def test_missing_model_argument_specific_case(self) -> None: 56 | """Test the specific case where model inference fails and model=None is passed.""" 57 | mock_kernel = MagicMock(spec=KernelInterface) 58 | mock_kernel.is_available = True 59 | 60 | # Mock a scenario where we have keys but model inference fails 61 | def mock_init_chat_model_with_none_check( 62 | *args: Any, **kwargs: Any 63 | ) -> MagicMock: 64 | model = kwargs.get("model") 65 | if model is None: 66 | raise TypeError( 67 | "_init_chat_model_helper() missing 1 required positional argument: 'model'" 68 | ) 69 | return MagicMock() 70 | 71 | with ( 72 | patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}, clear=True), 73 | patch("assistant_ui_anywidget.ai.langgraph_service.load_dotenv"), 74 | patch( 75 | "assistant_ui_anywidget.ai.langgraph_service.init_chat_model", 76 | side_effect=mock_init_chat_model_with_none_check, 77 | ), 78 | ): 79 | # Force the condition where no model is provided and provider auto-detection has issues 80 | service = AIService(kernel=mock_kernel, model=None, provider="openai") 81 | 82 | # Should fall back to MockLLM due to our fix that prevents None model/provider 83 | assert service.llm is not None 84 | assert "mock" in str(type(service.llm)).lower() 85 | 86 | def test_none_model_and_provider_prevention(self) -> None: 87 | """Test that None model and provider are prevented from being passed to init_chat_model.""" 88 | mock_kernel = MagicMock(spec=KernelInterface) 89 | mock_kernel.is_available = True 90 | 91 | # This should trigger the fix that prevents calling init_chat_model with None values 92 | with ( 93 | patch.dict(os.environ, {}, clear=True), 94 | patch("assistant_ui_anywidget.ai.langgraph_service.load_dotenv"), 95 | ): 96 | # This should never call init_chat_model because both model and provider will be None 97 | with patch( 98 | "assistant_ui_anywidget.ai.langgraph_service.init_chat_model" 99 | ) as mock_init: 100 | service = AIService(kernel=mock_kernel, model=None, provider=None) 101 | 102 | # Should not have called init_chat_model at all 103 | mock_init.assert_not_called() 104 | 105 | # Should be using MockLLM 106 | assert service.llm is not None 107 | assert "mock" in str(type(service.llm)).lower() 108 | 109 | def test_auto_detection_with_single_provider(self) -> None: 110 | """Test auto-detection works correctly with a single API key.""" 111 | mock_kernel = MagicMock(spec=KernelInterface) 112 | mock_kernel.is_available = True 113 | 114 | # Test with only OpenAI key 115 | with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}, clear=True): 116 | with patch( 117 | "assistant_ui_anywidget.ai.langgraph_service.init_chat_model" 118 | ) as mock_init: 119 | mock_init.return_value = MagicMock() 120 | 121 | AIService(kernel=mock_kernel) 122 | 123 | # Should have called init_chat_model with model and provider 124 | mock_init.assert_called_once() 125 | call_args = mock_init.call_args 126 | assert call_args.kwargs["model"] is not None 127 | assert call_args.kwargs["model_provider"] is not None 128 | 129 | def test_explicit_model_without_provider(self) -> None: 130 | """Test specifying model without provider should infer provider correctly.""" 131 | mock_kernel = MagicMock(spec=KernelInterface) 132 | mock_kernel.is_available = True 133 | 134 | with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}, clear=True): 135 | with patch( 136 | "assistant_ui_anywidget.ai.langgraph_service.init_chat_model" 137 | ) as mock_init: 138 | mock_init.return_value = MagicMock() 139 | 140 | # Should infer openai provider from gpt model name 141 | AIService(kernel=mock_kernel, model="gpt-4") 142 | 143 | mock_init.assert_called_once() 144 | call_args = mock_init.call_args 145 | assert call_args.kwargs["model"] == "gpt-4" 146 | assert call_args.kwargs["model_provider"] == "openai" 147 | 148 | def test_no_api_keys_fallback_to_mock(self) -> None: 149 | """Test that when no API keys are available, it falls back to MockLLM.""" 150 | mock_kernel = MagicMock(spec=KernelInterface) 151 | mock_kernel.is_available = True 152 | 153 | # Clear all environment variables and prevent .env loading 154 | with ( 155 | patch.dict(os.environ, {}, clear=True), 156 | patch("assistant_ui_anywidget.ai.langgraph_service.load_dotenv"), 157 | ): 158 | service = AIService(kernel=mock_kernel) 159 | 160 | # Should have fallen back to MockLLM 161 | assert service.llm is not None 162 | # MockLLM should be identifiable 163 | assert ( 164 | hasattr(service.llm, "_llm_type") 165 | or "mock" in str(type(service.llm)).lower() 166 | ) 167 | 168 | def test_chat_with_no_api_keys_should_not_crash(self) -> None: 169 | """Test that chat operations work even with no API keys (using MockLLM).""" 170 | mock_kernel = MagicMock(spec=KernelInterface) 171 | mock_kernel.is_available = True 172 | mock_kernel.get_kernel_info.return_value = { 173 | "available": True, 174 | "status": "idle", 175 | "language": "python", 176 | "execution_count": 0, 177 | "namespace_size": 0, 178 | } 179 | mock_kernel.get_namespace.return_value = {} 180 | 181 | with ( 182 | patch.dict(os.environ, {}, clear=True), 183 | patch("assistant_ui_anywidget.ai.langgraph_service.load_dotenv"), 184 | ): 185 | service = AIService(kernel=mock_kernel) 186 | 187 | # This should not crash and should return a response 188 | result = service.chat("hi") 189 | 190 | assert result is not None 191 | assert hasattr(result, "content") 192 | assert hasattr(result, "success") 193 | assert result.success is True 194 | assert result.content # Should have some response content 195 | -------------------------------------------------------------------------------- /frontend/src/VariableExplorer.tsx: -------------------------------------------------------------------------------- 1 | /** 2 | * Variable Explorer component for viewing kernel variables 3 | */ 4 | 5 | import React, { useState, useMemo } from "react"; 6 | import type { VariableInfo } from "./types"; 7 | 8 | interface VariableExplorerProps { 9 | variables: VariableInfo[]; 10 | onInspect: (variable: VariableInfo) => void; 11 | onExecute?: (code: string) => void; 12 | isLoading?: boolean; 13 | } 14 | 15 | export const VariableExplorer: React.FC = ({ 16 | variables, 17 | onInspect, 18 | onExecute: _onExecute, 19 | isLoading = false, 20 | }) => { 21 | const [searchTerm, setSearchTerm] = useState(""); 22 | const [sortBy, setSortBy] = useState<"name" | "type" | "size">("name"); 23 | const [sortOrder, setSortOrder] = useState<"asc" | "desc">("asc"); 24 | const [filterType, setFilterType] = useState("all"); 25 | 26 | // Get unique types for filter 27 | const uniqueTypes = useMemo(() => { 28 | const types = new Set(variables.map(v => v.type)); 29 | return ["all", ...Array.from(types).sort()]; 30 | }, [variables]); 31 | 32 | // Filter and sort variables 33 | const filteredVariables = useMemo(() => { 34 | let filtered = variables; 35 | 36 | // Apply search filter 37 | if (searchTerm) { 38 | const search = searchTerm.toLowerCase(); 39 | filtered = filtered.filter( 40 | v => v.name.toLowerCase().includes(search) || v.type.toLowerCase().includes(search) 41 | ); 42 | } 43 | 44 | // Apply type filter 45 | if (filterType !== "all") { 46 | filtered = filtered.filter(v => v.type === filterType); 47 | } 48 | 49 | // Sort 50 | filtered.sort((a, b) => { 51 | let aVal: string | number, bVal: string | number; 52 | 53 | switch (sortBy) { 54 | case "name": 55 | aVal = a.name; 56 | bVal = b.name; 57 | break; 58 | case "type": 59 | aVal = a.type; 60 | bVal = b.type; 61 | break; 62 | case "size": 63 | aVal = a.size || 0; 64 | bVal = b.size || 0; 65 | break; 66 | } 67 | 68 | if (sortOrder === "asc") { 69 | return aVal < bVal ? -1 : aVal > bVal ? 1 : 0; 70 | } else { 71 | return aVal > bVal ? -1 : aVal < bVal ? 1 : 0; 72 | } 73 | }); 74 | 75 | return filtered; 76 | }, [variables, searchTerm, filterType, sortBy, sortOrder]); 77 | 78 | const handleSort = (field: typeof sortBy) => { 79 | if (sortBy === field) { 80 | setSortOrder(sortOrder === "asc" ? "desc" : "asc"); 81 | } else { 82 | setSortBy(field); 83 | setSortOrder("asc"); 84 | } 85 | }; 86 | 87 | const formatSize = (bytes: number | null) => { 88 | if (bytes === null) return "-"; 89 | if (bytes < 1024) return `${bytes} B`; 90 | if (bytes < 1024 * 1024) return `${(bytes / 1024).toFixed(1)} KB`; 91 | return `${(bytes / (1024 * 1024)).toFixed(1)} MB`; 92 | }; 93 | 94 | const getVariableIcon = (type: string) => { 95 | switch (type) { 96 | case "DataFrame": 97 | return "📊"; 98 | case "ndarray": 99 | return "🔢"; 100 | case "list": 101 | case "tuple": 102 | return "📝"; 103 | case "dict": 104 | return "📚"; 105 | case "function": 106 | case "method": 107 | return "⚡"; 108 | case "str": 109 | return "📄"; 110 | case "int": 111 | case "float": 112 | return "🔢"; 113 | default: 114 | return "📦"; 115 | } 116 | }; 117 | 118 | return ( 119 |
129 | {/* Header */} 130 |
137 |

145 | Variable Explorer 146 |

147 |
148 | 149 | {/* Controls */} 150 |
157 | {/* Search */} 158 | setSearchTerm(e.target.value)} 163 | style={{ 164 | width: "100%", 165 | padding: "8px 12px", 166 | border: "1px solid #ddd", 167 | borderRadius: "6px", 168 | fontSize: "14px", 169 | marginBottom: "8px", 170 | }} 171 | /> 172 | 173 | {/* Filters */} 174 |
181 | 182 | 198 |
199 |
200 | 201 | {/* Variable list */} 202 |
208 | {isLoading ? ( 209 |
216 | Loading variables... 217 |
218 | ) : filteredVariables.length === 0 ? ( 219 |
226 | {searchTerm || filterType !== "all" 227 | ? "No variables match your filters" 228 | : "No variables in namespace"} 229 |
230 | ) : ( 231 | 237 | 238 | 245 | 258 | 271 | 284 | 285 | 286 | 287 | {filteredVariables.map(variable => ( 288 | onInspect(variable)} 291 | style={{ 292 | cursor: "pointer", 293 | backgroundColor: "#fff", 294 | borderBottom: "1px solid #f0f0f0", 295 | transition: "background-color 0.1s", 296 | }} 297 | onMouseEnter={e => { 298 | e.currentTarget.style.backgroundColor = "#f8f9fa"; 299 | }} 300 | onMouseLeave={e => { 301 | e.currentTarget.style.backgroundColor = "#fff"; 302 | }} 303 | > 304 | 325 | 345 | 355 | 356 | ))} 357 | 358 |
handleSort("name")} 247 | style={{ 248 | padding: "8px 12px", 249 | textAlign: "left", 250 | cursor: "pointer", 251 | fontSize: "13px", 252 | fontWeight: 600, 253 | color: "#666", 254 | }} 255 | > 256 | Name {sortBy === "name" && (sortOrder === "asc" ? "↑" : "↓")} 257 | handleSort("type")} 260 | style={{ 261 | padding: "8px 12px", 262 | textAlign: "left", 263 | cursor: "pointer", 264 | fontSize: "13px", 265 | fontWeight: 600, 266 | color: "#666", 267 | }} 268 | > 269 | Type {sortBy === "type" && (sortOrder === "asc" ? "↑" : "↓")} 270 | handleSort("size")} 273 | style={{ 274 | padding: "8px 12px", 275 | textAlign: "right", 276 | cursor: "pointer", 277 | fontSize: "13px", 278 | fontWeight: 600, 279 | color: "#666", 280 | }} 281 | > 282 | Size {sortBy === "size" && (sortOrder === "asc" ? "↑" : "↓")} 283 |
311 | {getVariableIcon(variable.type)} 312 | {variable.name} 313 | {variable.is_callable && ( 314 | 321 | () 322 | 323 | )} 324 | 332 | {variable.type} 333 | {variable.shape && ( 334 | 341 | {variable.shape.join("×")} 342 | 343 | )} 344 | 353 | {formatSize(variable.size)} 354 |
359 | )} 360 |
361 | 362 | {/* Footer */} 363 |
372 | {filteredVariables.length} variables 373 | {variables.length !== filteredVariables.length && ` (${variables.length} total)`} 374 |
375 |
376 | ); 377 | }; 378 | -------------------------------------------------------------------------------- /tests/test_git_native_tools.py: -------------------------------------------------------------------------------- 1 | """Tests for git-native kernel tools.""" 2 | 3 | import subprocess 4 | import tempfile 5 | from pathlib import Path 6 | from typing import Any, Generator 7 | from langchain_core.tools import BaseTool 8 | from unittest.mock import patch 9 | 10 | import pytest 11 | 12 | from assistant_ui_anywidget.kernel_tools import GitFindTool, GitGrepTool, ListFilesTool 13 | 14 | 15 | class TestGitNativeTools: 16 | """Test git-native kernel tools functionality.""" 17 | 18 | @pytest.fixture # type: ignore[misc] 19 | def git_repo(self) -> Generator[Path, None, None]: 20 | """Create a temporary git repository for testing.""" 21 | with tempfile.TemporaryDirectory() as tmpdir: 22 | repo_path = Path(tmpdir) 23 | 24 | # Initialize git repo 25 | subprocess.run(["git", "init"], cwd=repo_path, check=True) 26 | subprocess.run( 27 | ["git", "config", "user.name", "Test User"], cwd=repo_path, check=True 28 | ) 29 | subprocess.run( 30 | ["git", "config", "user.email", "test@example.com"], 31 | cwd=repo_path, 32 | check=True, 33 | ) 34 | 35 | # Create test files 36 | (repo_path / "main.py").write_text( 37 | "def main():\n print('Hello World')\n" 38 | ) 39 | (repo_path / "config.yaml").write_text("database:\n host: localhost\n") 40 | (repo_path / "test_main.py").write_text( 41 | "def test_main():\n assert True\n" 42 | ) 43 | (repo_path / "build").mkdir() 44 | (repo_path / "build" / "output.log").write_text("Build log content\n") 45 | (repo_path / "untracked.txt").write_text("This file is not tracked\n") 46 | 47 | # Add and commit tracked files 48 | subprocess.run( 49 | ["git", "add", "main.py", "config.yaml", "test_main.py"], 50 | cwd=repo_path, 51 | check=True, 52 | ) 53 | subprocess.run( 54 | ["git", "commit", "-m", "Initial commit"], cwd=repo_path, check=True 55 | ) 56 | 57 | yield repo_path 58 | 59 | def test_list_files_git_tracked_only(self, git_repo: Path) -> None: 60 | """Test listing only git-tracked files.""" 61 | tool = ListFilesTool() 62 | 63 | result = tool._run(directory=str(git_repo), git_tracked_only=True) 64 | 65 | assert "git-tracked only" in result 66 | assert "main.py" in result 67 | assert "config.yaml" in result 68 | assert "test_main.py" in result 69 | assert "untracked.txt" not in result 70 | assert "output.log" not in result 71 | 72 | def test_list_files_all_files(self, git_repo: Path) -> None: 73 | """Test listing all files including untracked.""" 74 | tool = ListFilesTool() 75 | 76 | result = tool._run(directory=str(git_repo), git_tracked_only=False) 77 | 78 | assert "all files" in result 79 | assert "main.py" in result 80 | assert "config.yaml" in result 81 | assert "test_main.py" in result 82 | assert "untracked.txt" in result 83 | 84 | def test_list_files_with_pattern(self, git_repo: Path) -> None: 85 | """Test listing files with pattern matching.""" 86 | tool = ListFilesTool() 87 | 88 | result = tool._run( 89 | directory=str(git_repo), git_tracked_only=True, pattern="*.py" 90 | ) 91 | 92 | assert "main.py" in result 93 | assert "test_main.py" in result 94 | assert "config.yaml" not in result 95 | 96 | def test_list_files_non_git_repo(self) -> None: 97 | """Test behavior when not in a git repository.""" 98 | tool = ListFilesTool() 99 | 100 | with tempfile.TemporaryDirectory() as tmpdir: 101 | result = tool._run(directory=tmpdir, git_tracked_only=True) 102 | 103 | assert "Not a git repository" in result 104 | assert "Use git_tracked_only=False" in result 105 | 106 | def test_git_grep_basic_search(self, git_repo: Path) -> None: 107 | """Test basic git grep functionality.""" 108 | tool = GitGrepTool() 109 | 110 | # Mock subprocess.run to work within the git_repo directory 111 | original_run = subprocess.run 112 | 113 | def mock_run(cmd: Any, **kwargs: Any) -> Any: 114 | if "cwd" not in kwargs: 115 | kwargs["cwd"] = str(git_repo) 116 | return original_run(cmd, **kwargs) 117 | 118 | with patch("subprocess.run", side_effect=mock_run): 119 | result = tool._run(search_term="def main") 120 | 121 | assert "Found" in result 122 | assert "main.py" in result 123 | assert "def main" in result 124 | 125 | def test_git_grep_case_sensitive(self, git_repo: Path) -> None: 126 | """Test case-sensitive git grep.""" 127 | tool = GitGrepTool() 128 | 129 | # Mock subprocess.run to work within the git_repo directory 130 | original_run = subprocess.run 131 | 132 | def mock_run(cmd: Any, **kwargs: Any) -> Any: 133 | if "cwd" not in kwargs: 134 | kwargs["cwd"] = str(git_repo) 135 | return original_run(cmd, **kwargs) 136 | 137 | with patch("subprocess.run", side_effect=mock_run): 138 | # Case-insensitive (default) 139 | result_insensitive = tool._run(search_term="HELLO") 140 | assert "Hello World" in result_insensitive 141 | 142 | # Case-sensitive 143 | result_sensitive = tool._run(search_term="HELLO", case_sensitive=True) 144 | assert "No matches found" in result_sensitive 145 | 146 | def test_git_grep_with_file_pattern(self, git_repo: Path) -> None: 147 | """Test git grep with file pattern filter.""" 148 | tool = GitGrepTool() 149 | 150 | original_run = subprocess.run 151 | 152 | def mock_run(cmd: Any, **kwargs: Any) -> Any: 153 | if "cwd" not in kwargs: 154 | kwargs["cwd"] = str(git_repo) 155 | return original_run(cmd, **kwargs) 156 | 157 | with patch("subprocess.run", side_effect=mock_run): 158 | result = tool._run(search_term="def", file_pattern="*.py") 159 | 160 | assert "def main" in result 161 | assert "def test_main" in result 162 | 163 | def test_git_grep_no_matches(self, git_repo: Path) -> None: 164 | """Test git grep when no matches found.""" 165 | tool = GitGrepTool() 166 | 167 | original_run = subprocess.run 168 | 169 | def mock_run(cmd: Any, **kwargs: Any) -> Any: 170 | if "cwd" not in kwargs: 171 | kwargs["cwd"] = str(git_repo) 172 | return original_run(cmd, **kwargs) 173 | 174 | with patch("subprocess.run", side_effect=mock_run): 175 | result = tool._run(search_term="nonexistent_text") 176 | 177 | assert "No matches found" in result 178 | 179 | def test_git_grep_non_git_repo(self) -> None: 180 | """Test git grep behavior when not in a git repository.""" 181 | tool = GitGrepTool() 182 | 183 | with tempfile.TemporaryDirectory() as tmpdir: 184 | original_run = subprocess.run 185 | 186 | def mock_run(cmd: Any, **kwargs: Any) -> Any: 187 | if "cwd" not in kwargs: 188 | kwargs["cwd"] = tmpdir 189 | return original_run(cmd, **kwargs) 190 | 191 | with patch("subprocess.run", side_effect=mock_run): 192 | result = tool._run(search_term="test") 193 | 194 | assert "Not a git repository" in result 195 | 196 | def test_git_find_basic_search(self, git_repo: Path) -> None: 197 | """Test basic git find functionality.""" 198 | tool = GitFindTool() 199 | 200 | original_run = subprocess.run 201 | 202 | def mock_run(cmd: Any, **kwargs: Any) -> Any: 203 | if "cwd" not in kwargs: 204 | kwargs["cwd"] = str(git_repo) 205 | return original_run(cmd, **kwargs) 206 | 207 | with patch("subprocess.run", side_effect=mock_run): 208 | result = tool._run(name_pattern="*.py") 209 | 210 | assert "main.py" in result 211 | assert "test_main.py" in result 212 | assert "config.yaml" not in result 213 | 214 | def test_git_find_exact_name(self, git_repo: Path) -> None: 215 | """Test finding exact filename.""" 216 | tool = GitFindTool() 217 | 218 | original_run = subprocess.run 219 | 220 | def mock_run(cmd: Any, **kwargs: Any) -> Any: 221 | if "cwd" not in kwargs: 222 | kwargs["cwd"] = str(git_repo) 223 | return original_run(cmd, **kwargs) 224 | 225 | with patch("subprocess.run", side_effect=mock_run): 226 | result = tool._run(name_pattern="config.yaml") 227 | 228 | assert "config.yaml" in result 229 | assert "main.py" not in result 230 | 231 | def test_git_find_case_sensitive(self, git_repo: Path) -> None: 232 | """Test case-sensitive git find.""" 233 | tool = GitFindTool() 234 | 235 | original_run = subprocess.run 236 | 237 | def mock_run(cmd: Any, **kwargs: Any) -> Any: 238 | if "cwd" not in kwargs: 239 | kwargs["cwd"] = str(git_repo) 240 | return original_run(cmd, **kwargs) 241 | 242 | with patch("subprocess.run", side_effect=mock_run): 243 | # Case-insensitive (default) 244 | result_insensitive = tool._run(name_pattern="MAIN.PY") 245 | assert "main.py" in result_insensitive 246 | 247 | # Case-sensitive 248 | result_sensitive = tool._run(name_pattern="MAIN.PY", case_sensitive=True) 249 | assert "No git-tracked files found" in result_sensitive 250 | 251 | def test_git_find_no_matches(self, git_repo: Path) -> None: 252 | """Test git find when no matches found.""" 253 | tool = GitFindTool() 254 | 255 | original_run = subprocess.run 256 | 257 | def mock_run(cmd: Any, **kwargs: Any) -> Any: 258 | if "cwd" not in kwargs: 259 | kwargs["cwd"] = str(git_repo) 260 | return original_run(cmd, **kwargs) 261 | 262 | with patch("subprocess.run", side_effect=mock_run): 263 | result = tool._run(name_pattern="*.nonexistent") 264 | 265 | assert "No git-tracked files found" in result 266 | 267 | def test_git_find_non_git_repo(self) -> None: 268 | """Test git find behavior when not in a git repository.""" 269 | tool = GitFindTool() 270 | 271 | with tempfile.TemporaryDirectory() as tmpdir: 272 | original_run = subprocess.run 273 | 274 | def mock_run(cmd: Any, **kwargs: Any) -> Any: 275 | if "cwd" not in kwargs: 276 | kwargs["cwd"] = tmpdir 277 | return original_run(cmd, **kwargs) 278 | 279 | with patch("subprocess.run", side_effect=mock_run): 280 | result = tool._run(name_pattern="*.py") 281 | 282 | assert "Not a git repository" in result 283 | 284 | def test_error_handling(self) -> None: 285 | """Test error handling in git tools.""" 286 | tools = [ListFilesTool(), GitGrepTool(), GitFindTool()] 287 | 288 | # Test with invalid directory 289 | for tool in tools: 290 | if hasattr(tool, "_run"): 291 | with patch( 292 | "subprocess.run", 293 | side_effect=subprocess.CalledProcessError(1, "git"), 294 | ): 295 | if isinstance(tool, ListFilesTool): 296 | result = tool._run( 297 | directory="/nonexistent", git_tracked_only=True 298 | ) 299 | elif isinstance(tool, GitGrepTool): 300 | result = tool._run(search_term="test") 301 | else: # GitFindTool 302 | result = tool._run(name_pattern="*.py") 303 | 304 | assert "Error" in result or "Not a git repository" in result 305 | 306 | def test_tools_have_correct_names(self) -> None: 307 | """Test that tools have the expected names.""" 308 | assert ListFilesTool().name == "list_files" 309 | assert GitGrepTool().name == "git_grep" 310 | assert GitFindTool().name == "git_find" 311 | 312 | def test_tools_have_descriptions(self) -> None: 313 | """Test that tools have proper descriptions.""" 314 | tools = [ListFilesTool(), GitGrepTool(), GitFindTool()] 315 | 316 | for tool in tools: 317 | assert isinstance(tool, BaseTool) 318 | assert tool.description 319 | assert len(tool.description) > 50 # Ensure substantial description 320 | assert "git" in tool.description.lower() 321 | -------------------------------------------------------------------------------- /tests/test_kernel_interface.py: -------------------------------------------------------------------------------- 1 | """Tests for kernel interface functionality.""" 2 | 3 | from unittest.mock import Mock, patch 4 | 5 | import pytest 6 | 7 | from assistant_ui_anywidget.kernel_interface import ( 8 | ExecutionResult, 9 | KernelInterface, 10 | StackFrame, 11 | VariableInfo, 12 | ) 13 | 14 | 15 | class MockIPython: 16 | """Mock IPython shell for testing.""" 17 | 18 | def __init__(self) -> None: 19 | self.user_ns = { 20 | "x": 42, 21 | "y": "hello", 22 | "df": self._create_mock_dataframe(), 23 | "arr": self._create_mock_array(), 24 | "_private": "hidden", 25 | "func": lambda x: x * 2, 26 | } 27 | self.execution_count = 10 28 | self._last_error = None 29 | 30 | def _create_mock_dataframe(self) -> Mock: 31 | """Create a mock DataFrame-like object.""" 32 | mock_df = Mock() 33 | mock_df.__class__.__name__ = "DataFrame" 34 | mock_df.__class__.__module__ = "pandas.core.frame" 35 | mock_df.__class__.__qualname__ = "DataFrame" 36 | mock_df.shape = (100, 5) 37 | mock_df.nbytes = 4000 38 | mock_df.head = Mock(return_value=" A B C\n0 1 2 3\n1 4 5 6") 39 | return mock_df 40 | 41 | def _create_mock_array(self) -> Mock: 42 | """Create a mock numpy array-like object.""" 43 | # Use a MagicMock with spec to make it non-callable 44 | mock_arr = Mock(spec=["shape", "dtype", "nbytes", "__class__", "__repr__"]) 45 | mock_arr.__class__.__name__ = "ndarray" 46 | mock_arr.__class__.__module__ = "numpy" 47 | mock_arr.__class__.__qualname__ = "ndarray" 48 | mock_arr.shape = (10, 20) 49 | mock_arr.dtype = "float64" 50 | mock_arr.nbytes = 1600 51 | # Configure repr to return expected format 52 | mock_arr.__repr__ = Mock(return_value="array(shape=(10, 20), dtype=float64)") # type: ignore[method-assign] 53 | return mock_arr 54 | 55 | def run_cell( 56 | self, code: str, silent: bool = False, store_history: bool = True 57 | ) -> Mock: 58 | """Mock run_cell method.""" 59 | result = Mock() 60 | 61 | # Simulate successful execution 62 | if code == "1 + 1": 63 | result.result = 2 64 | result.error_in_exec = None 65 | elif code == "new_var = 100": 66 | self.user_ns["new_var"] = 100 67 | result.result = None 68 | result.error_in_exec = None 69 | elif code == "raise ValueError('test error')": 70 | error = ValueError("test error") 71 | result.result = None 72 | result.error_in_exec = error 73 | self._last_error = error # type: ignore[assignment] 74 | else: 75 | result.result = None 76 | result.error_in_exec = None 77 | 78 | self.execution_count += 1 79 | return result 80 | 81 | def _get_exc_info(self) -> tuple[type, Exception, None]: 82 | """Mock for getting exception info.""" 83 | if self._last_error: 84 | return (type(self._last_error), self._last_error, None) # type: ignore[unreachable] 85 | return (None, None, None) # type: ignore[return-value] 86 | 87 | 88 | @pytest.fixture # type: ignore[misc] 89 | def mock_ipython() -> MockIPython: 90 | """Create a mock IPython instance.""" 91 | return MockIPython() 92 | 93 | 94 | @pytest.fixture # type: ignore[misc] 95 | def kernel_interface(mock_ipython: MockIPython) -> KernelInterface: 96 | """Create a KernelInterface with mock IPython.""" 97 | with patch( 98 | "assistant_ui_anywidget.kernel_interface.get_ipython", return_value=mock_ipython 99 | ): 100 | interface = KernelInterface() 101 | interface.shell = mock_ipython 102 | return interface 103 | 104 | 105 | class TestKernelInterface: 106 | """Test KernelInterface functionality.""" 107 | 108 | def test_is_available(self, kernel_interface: KernelInterface) -> None: 109 | """Test kernel availability check.""" 110 | assert kernel_interface.is_available is True 111 | 112 | # Test without IPython 113 | kernel_interface.shell = None 114 | assert kernel_interface.is_available is False 115 | 116 | def test_get_namespace(self, kernel_interface: KernelInterface) -> None: 117 | """Test getting namespace variables.""" 118 | namespace = kernel_interface.get_namespace() 119 | 120 | # Should include public variables 121 | assert "x" in namespace 122 | assert "y" in namespace 123 | assert "df" in namespace 124 | assert "arr" in namespace 125 | assert "func" in namespace 126 | 127 | # Should exclude private and IPython internals 128 | assert "_private" not in namespace 129 | assert "In" not in namespace 130 | assert "Out" not in namespace 131 | 132 | def test_get_variable_info(self, kernel_interface: KernelInterface) -> None: 133 | """Test getting variable information.""" 134 | # Test integer variable 135 | x_info = kernel_interface.get_variable_info("x") 136 | assert x_info is not None 137 | assert x_info.name == "x" 138 | assert x_info.type == "int" 139 | assert x_info.preview == "42" 140 | assert x_info.is_callable is False 141 | 142 | # Test string variable 143 | y_info = kernel_interface.get_variable_info("y") 144 | assert y_info is not None 145 | assert y_info.name == "y" 146 | assert y_info.type == "str" 147 | assert y_info.preview == "'hello'" 148 | 149 | # Test DataFrame-like variable 150 | df_info = kernel_interface.get_variable_info("df") 151 | assert df_info is not None 152 | assert df_info.name == "df" 153 | assert df_info.type == "DataFrame" 154 | assert df_info.shape == [100, 5] 155 | assert df_info.size == 4000 156 | assert "A B C" in df_info.preview 157 | 158 | # Test array-like variable 159 | arr_info = kernel_interface.get_variable_info("arr") 160 | assert arr_info is not None 161 | assert arr_info.name == "arr" 162 | assert arr_info.type == "ndarray" 163 | assert arr_info.shape == [10, 20] 164 | assert "array(shape=" in arr_info.preview 165 | assert "dtype=" in arr_info.preview 166 | 167 | # Test callable 168 | func_info = kernel_interface.get_variable_info("func") 169 | assert func_info is not None 170 | assert func_info.is_callable is True 171 | 172 | # Test non-existent variable 173 | none_info = kernel_interface.get_variable_info("nonexistent") 174 | assert none_info is None 175 | 176 | # Test deep inspection 177 | deep_info = kernel_interface.get_variable_info("x", deep=True) 178 | assert deep_info is not None 179 | assert len(deep_info.attributes) > 0 180 | 181 | def test_execute_code_success(self, kernel_interface: KernelInterface) -> None: 182 | """Test successful code execution.""" 183 | # Test expression evaluation 184 | result = kernel_interface.execute_code("1 + 1") 185 | assert result.success is True 186 | assert result.execution_count == 11 # Started at 10 187 | assert len(result.outputs) == 1 188 | assert result.outputs[0]["type"] == "execute_result" 189 | assert result.outputs[0]["data"]["text/plain"] == "2" 190 | assert result.error is None 191 | 192 | # Test variable assignment 193 | result = kernel_interface.execute_code("new_var = 100") 194 | assert result.success is True 195 | assert "new_var" in result.variables_changed 196 | assert result.outputs == [] # No output for assignment 197 | 198 | def test_execute_code_error(self, kernel_interface: KernelInterface) -> None: 199 | """Test code execution with error.""" 200 | result = kernel_interface.execute_code("raise ValueError('test error')") 201 | assert result.success is False 202 | assert result.error is not None 203 | assert result.error["type"] == "ValueError" 204 | assert result.error["message"] == "test error" 205 | assert len(result.error["traceback"]) > 0 206 | 207 | def test_get_last_error(self, kernel_interface: KernelInterface) -> None: 208 | """Test getting last error information.""" 209 | # Initially no error 210 | assert kernel_interface.get_last_error() is None 211 | 212 | # Execute code that raises error 213 | kernel_interface.execute_code("raise ValueError('test error')") 214 | 215 | # Should be able to get error info 216 | error_info = kernel_interface.get_last_error() 217 | assert error_info is not None 218 | assert error_info["type"] == "ValueError" 219 | assert error_info["message"] == "test error" 220 | 221 | def test_get_stack_trace(self, kernel_interface: KernelInterface) -> None: 222 | """Test getting stack trace.""" 223 | frames = kernel_interface.get_stack_trace(max_frames=5) 224 | assert isinstance(frames, list) 225 | assert len(frames) <= 5 226 | 227 | if frames: 228 | frame = frames[0] 229 | assert isinstance(frame, StackFrame) 230 | assert frame.is_current is True 231 | assert frame.filename is not None 232 | assert frame.line_number > 0 233 | assert frame.function_name is not None 234 | 235 | def test_get_kernel_info(self, kernel_interface: KernelInterface) -> None: 236 | """Test getting kernel information.""" 237 | info = kernel_interface.get_kernel_info() 238 | assert info["available"] is True 239 | assert info["status"] == "idle" 240 | assert info["language"] == "python" 241 | assert info["execution_count"] == 10 242 | assert info["namespace_size"] == 5 # x, y, df, arr, func (excluding _private) 243 | 244 | # Test without kernel 245 | kernel_interface.shell = None 246 | info = kernel_interface.get_kernel_info() 247 | assert info["available"] is False 248 | assert info["status"] == "not_connected" 249 | 250 | def test_preview_generation(self, kernel_interface: KernelInterface) -> None: 251 | """Test preview generation for different types.""" 252 | # Test long string preview 253 | assert kernel_interface.shell is not None 254 | kernel_interface.shell.user_ns["long_str"] = "a" * 200 255 | info = kernel_interface.get_variable_info("long_str") 256 | assert info is not None 257 | assert len(info.preview) <= 103 # 100 + "..." 258 | assert info.preview.endswith("...") 259 | 260 | # Test object without special handling 261 | assert kernel_interface.shell is not None 262 | assert kernel_interface.shell.user_ns is not None 263 | kernel_interface.shell.user_ns["obj"] = object() 264 | info = kernel_interface.get_variable_info("obj") 265 | assert info is not None 266 | assert info.preview.startswith(" None: 269 | """Test VariableInfo serialization.""" 270 | var_info = VariableInfo( 271 | name="test", 272 | type="int", 273 | type_str="builtins.int", 274 | size=28, 275 | shape=None, 276 | preview="42", 277 | is_callable=False, 278 | attributes=["bit_length", "real", "imag"], 279 | last_modified=None, 280 | ) 281 | 282 | data = var_info.to_dict() 283 | assert data["name"] == "test" 284 | assert data["type"] == "int" 285 | assert data["size"] == 28 286 | assert len(data["attributes"]) == 3 287 | 288 | def test_execution_result_to_dict(self) -> None: 289 | """Test ExecutionResult serialization.""" 290 | result = ExecutionResult( 291 | success=True, 292 | execution_count=5, 293 | outputs=[{"type": "execute_result", "data": {"text/plain": "42"}}], 294 | execution_time=0.001, 295 | variables_changed=["x"], 296 | ) 297 | 298 | data = result.to_dict() 299 | assert data["success"] is True 300 | assert data["execution_count"] == 5 301 | assert len(data["outputs"]) == 1 302 | assert data["variables_changed"] == ["x"] 303 | assert data["error"] is None 304 | 305 | 306 | class TestKernelInterfaceWithoutIPython: 307 | """Test KernelInterface when IPython is not available.""" 308 | 309 | def test_without_ipython(self) -> None: 310 | """Test behavior when IPython is not available.""" 311 | with patch( 312 | "assistant_ui_anywidget.kernel_interface.get_ipython", return_value=None 313 | ): 314 | interface = KernelInterface() 315 | 316 | assert interface.is_available is False 317 | assert interface.get_namespace() == {} 318 | assert interface.get_variable_info("x") is None 319 | 320 | result = interface.execute_code("1 + 1") 321 | assert result.success is False 322 | assert result.error is not None 323 | assert result.error["type"] == "KernelError" 324 | assert "not available" in result.error["message"] 325 | -------------------------------------------------------------------------------- /examples/demo_kernel_assistant.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "metadata": {}, 6 | "source": [ 7 | "# AI Assistant Widget with Kernel Access Demo\n", 8 | "\n", 9 | "This notebook demonstrates the capabilities of the assistant widget that can interact with your Jupyter kernel." 10 | ] 11 | }, 12 | { 13 | "cell_type": "code", 14 | "execution_count": null, 15 | "metadata": {}, 16 | "outputs": [], 17 | "source": [ 18 | "# Import the agent widget\n", 19 | "from assistant_ui_anywidget import AgentWidget\n", 20 | "import numpy as np\n", 21 | "import pandas as pd\n", 22 | "\n", 23 | "# The widget automatically loads API keys from a .env file\n", 24 | "# You can also set them manually:\n", 25 | "# os.environ['OPENAI_API_KEY'] = 'your-key-here'\n", 26 | "# os.environ['ANTHROPIC_API_KEY'] = 'your-key-here'\n", 27 | "# os.environ['GOOGLE_API_KEY'] = 'your-key-here'\n", 28 | "\n", 29 | "# Note: If only one API key is set (e.g., only Google), it will be used automatically!" 30 | ] 31 | }, 32 | { 33 | "cell_type": "markdown", 34 | "metadata": {}, 35 | "source": [ 36 | "## 1. Create Some Variables to Explore\n", 37 | "\n", 38 | "Let's create various types of variables that the assistant can inspect:" 39 | ] 40 | }, 41 | { 42 | "cell_type": "code", 43 | "execution_count": null, 44 | "metadata": {}, 45 | "outputs": [], 46 | "source": [ 47 | "# Create some sample data\n", 48 | "x = 42\n", 49 | "y = \"Hello, AI Assistant!\"\n", 50 | "numbers = [1, 2, 3, 4, 5]\n", 51 | "\n", 52 | "# Create a numpy array\n", 53 | "data_array = np.random.randn(100, 3)\n", 54 | "\n", 55 | "# Create a pandas DataFrame\n", 56 | "df = pd.DataFrame(\n", 57 | " {\n", 58 | " \"A\": np.random.randn(100),\n", 59 | " \"B\": np.random.randn(100),\n", 60 | " \"C\": np.random.choice([\"cat\", \"dog\", \"bird\"], 100),\n", 61 | " }\n", 62 | ")\n", 63 | "\n", 64 | "\n", 65 | "# Define a function\n", 66 | "def calculate_mean(values):\n", 67 | " \"\"\"Calculate the mean of a list of values.\"\"\"\n", 68 | " return sum(values) / len(values)\n", 69 | "\n", 70 | "\n", 71 | "print(\"Variables created successfully!\")" 72 | ] 73 | }, 74 | { 75 | "cell_type": "markdown", 76 | "metadata": {}, 77 | "source": [ 78 | "## 2. Create and Display the Assistant Widget" 79 | ] 80 | }, 81 | { 82 | "cell_type": "code", 83 | "execution_count": null, 84 | "metadata": {}, 85 | "outputs": [], 86 | "source": [ 87 | "# Create the agent widget with AI configuration\n", 88 | "assistant = AgentWidget(\n", 89 | " # Optional: Override automatic provider detection\n", 90 | " # provider='openai', # or 'anthropic', 'google_genai'\n", 91 | " # model='gpt-4', # or 'claude-3-opus', 'gemini-pro'\n", 92 | " require_approval=True # Auto-approve code execution for demo\n", 93 | ")\n", 94 | "\n", 95 | "# The widget automatically detects available API keys:\n", 96 | "# - If you have multiple keys, it uses them in order: OpenAI → Anthropic → Google\n", 97 | "# - If you have only one key (e.g., only Google), it uses that automatically\n", 98 | "# - If no keys are found, it falls back to a helpful mock AI\n", 99 | "\n", 100 | "# Display it\n", 101 | "assistant" 102 | ] 103 | }, 104 | { 105 | "cell_type": "markdown", 106 | "metadata": {}, 107 | "source": [ 108 | "## 3. Ways to Interact\n", 109 | "\n", 110 | "You can interact with the assistant in two ways:\n", 111 | "\n", 112 | "### Natural Language (AI-Powered)\n", 113 | "Just type naturally! The AI understands your intent and uses kernel tools automatically:\n", 114 | "- \"Show me all my variables\"\n", 115 | "- \"What's the shape of data_array?\"\n", 116 | "- \"Calculate the correlation between columns A and B in df\"\n", 117 | "- \"Create a histogram of column A\"\n", 118 | "- \"Help me understand this error\"\n", 119 | "\n", 120 | "### Direct Commands\n", 121 | "For precise control, use slash commands:\n", 122 | "\n", 123 | "**Variable Exploration**\n", 124 | "- `/vars` - List all variables in the kernel\n", 125 | "- `/inspect df` - Get detailed information about the DataFrame\n", 126 | "- `/inspect data_array` - Inspect the numpy array\n", 127 | "- `/inspect calculate_mean` - Look at the function\n", 128 | "\n", 129 | "**Code Execution**\n", 130 | "- `/exec df.describe()` - Run DataFrame describe\n", 131 | "- `/exec print(data_array.shape)` - Check array shape\n", 132 | "- `/exec result = calculate_mean(numbers)` - Execute the function\n", 133 | "- `/exec print(f\"Mean: {result}\")` - Print the result\n", 134 | "\n", 135 | "**Other Commands**\n", 136 | "- `/help` - Show all available commands\n", 137 | "- `/clear` - Clear the namespace (with confirmation)" 138 | ] 139 | }, 140 | { 141 | "cell_type": "markdown", 142 | "metadata": {}, 143 | "source": [ 144 | "## AI Integration\n", 145 | "\n", 146 | "The assistant now has AI capabilities with automatic provider detection!\n", 147 | "\n", 148 | "### Automatic Provider Detection\n", 149 | "\n", 150 | "The widget automatically:\n", 151 | "1. **Loads API keys from `.env` file** using python-dotenv\n", 152 | "2. **Detects available providers** and selects the best one\n", 153 | "3. **Uses smart defaults** - GPT-4 for OpenAI, Claude Opus for Anthropic, Gemini Pro for Google\n", 154 | "4. **Falls back gracefully** - If no API keys are found, uses a helpful mock AI\n", 155 | "\n", 156 | "### Setting Up API Keys\n", 157 | "\n", 158 | "Create a `.env` file in your project root (see `.env.example`):\n", 159 | "```bash\n", 160 | "# Any of these will work - the widget picks the first available\n", 161 | "OPENAI_API_KEY=sk-...\n", 162 | "ANTHROPIC_API_KEY=sk-ant-...\n", 163 | "GOOGLE_API_KEY=...\n", 164 | "```\n", 165 | "\n", 166 | "### Example Questions to Try:\n", 167 | "\n", 168 | "- \"What variables do I have in my namespace?\"\n", 169 | "- \"Can you show me what's in the df variable?\"\n", 170 | "- \"Calculate the mean of column A in df\"\n", 171 | "- \"Create a scatter plot of the data_array\"\n", 172 | "- \"Why am I getting an error with undefined_variable?\"\n", 173 | "\n", 174 | "### Supported AI Providers:\n", 175 | "\n", 176 | "The widget supports multiple AI providers:\n", 177 | "- **OpenAI**: GPT-4, GPT-3.5-turbo\n", 178 | "- **Anthropic**: Claude 3 Opus, Sonnet, Haiku\n", 179 | "- **Google**: Gemini Pro, Gemini Ultra\n", 180 | "\n", 181 | "If only one API key is set (e.g., only Google), it will be used automatically!" 182 | ] 183 | }, 184 | { 185 | "cell_type": "markdown", 186 | "metadata": {}, 187 | "source": [ 188 | "## 4. Programmatic Interaction\n", 189 | "\n", 190 | "You can also interact with the widget programmatically:" 191 | ] 192 | }, 193 | { 194 | "cell_type": "code", 195 | "execution_count": null, 196 | "metadata": {}, 197 | "outputs": [], 198 | "source": [ 199 | "# Add a message from Python\n", 200 | "assistant.add_message(\"assistant\", \"Let me help you explore your kernel variables!\")\n", 201 | "\n", 202 | "# Inspect a variable programmatically\n", 203 | "df_info = assistant.inspect_variable(\"df\")\n", 204 | "if df_info:\n", 205 | " print(f\"DataFrame shape: {df_info['shape']}\")\n", 206 | " print(f\"DataFrame type: {df_info['type']}\")" 207 | ] 208 | }, 209 | { 210 | "cell_type": "code", 211 | "execution_count": null, 212 | "metadata": {}, 213 | "outputs": [], 214 | "source": [ 215 | "# Execute code and show result\n", 216 | "result = assistant.execute_code(\"df['A'].mean()\", show_result=True)\n", 217 | "print(f\"\\nExecution successful: {result['success']}\")" 218 | ] 219 | }, 220 | { 221 | "cell_type": "code", 222 | "execution_count": null, 223 | "metadata": {}, 224 | "outputs": [], 225 | "source": [ 226 | "# Set action buttons for quick operations\n", 227 | "assistant.set_action_buttons(\n", 228 | " [\n", 229 | " {\"text\": \"Show Variables\", \"color\": \"#28a745\", \"icon\": \"📊\"},\n", 230 | " {\"text\": \"Run df.info()\", \"color\": \"#007bff\", \"icon\": \"▶️\"},\n", 231 | " {\"text\": \"Clear Output\", \"color\": \"#6c757d\", \"icon\": \"🧹\"},\n", 232 | " ]\n", 233 | ")" 234 | ] 235 | }, 236 | { 237 | "cell_type": "markdown", 238 | "metadata": {}, 239 | "source": [ 240 | "## 5. Create an Error for Debugging Demo" 241 | ] 242 | }, 243 | { 244 | "cell_type": "code", 245 | "execution_count": null, 246 | "metadata": {}, 247 | "outputs": [], 248 | "source": [ 249 | "# This will create an error\n", 250 | "try:\n", 251 | " undefined_variable\n", 252 | "except NameError as e:\n", 253 | " print(f\"Error caught: {e}\")\n", 254 | " # The assistant can help debug this!" 255 | ] 256 | }, 257 | { 258 | "cell_type": "markdown", 259 | "metadata": {}, 260 | "source": [ 261 | "## 6. Advanced Features\n", 262 | "\n", 263 | "### Current Kernel State" 264 | ] 265 | }, 266 | { 267 | "cell_type": "code", 268 | "execution_count": null, 269 | "metadata": {}, 270 | "outputs": [], 271 | "source": [ 272 | "# Check the kernel state\n", 273 | "print(\"Kernel State:\")\n", 274 | "print(f\"Available: {assistant.kernel_state['available']}\")\n", 275 | "print(f\"Status: {assistant.kernel_state['status']}\")\n", 276 | "print(f\"Namespace size: {assistant.kernel_state['namespace_size']}\")\n", 277 | "print(f\"Variables by type: {assistant.kernel_state['variables_by_type']}\")" 278 | ] 279 | }, 280 | { 281 | "cell_type": "markdown", 282 | "metadata": {}, 283 | "source": [ 284 | "### Variable Information" 285 | ] 286 | }, 287 | { 288 | "cell_type": "code", 289 | "execution_count": null, 290 | "metadata": {}, 291 | "outputs": [], 292 | "source": [ 293 | "# Get detailed variable information\n", 294 | "print(\"Variables in namespace:\")\n", 295 | "for var_info in assistant.variables_info[:5]: # Show first 5\n", 296 | " print(\n", 297 | " f\"- {var_info['name']}: {var_info['type']} \"\n", 298 | " f\"{'(callable)' if var_info['is_callable'] else ''}\"\n", 299 | " )" 300 | ] 301 | }, 302 | { 303 | "cell_type": "markdown", 304 | "metadata": {}, 305 | "source": [ 306 | "### Conversation Logging\n", 307 | "\n", 308 | "All conversations are automatically logged with timestamps:" 309 | ] 310 | }, 311 | { 312 | "cell_type": "code", 313 | "execution_count": null, 314 | "metadata": {}, 315 | "outputs": [], 316 | "source": [ 317 | "# Get the conversation log path\n", 318 | "log_path = assistant.get_conversation_log_path()\n", 319 | "print(f\"Conversations are being logged to: {log_path}\")\n", 320 | "\n", 321 | "# You can read the log file to see all interactions\n", 322 | "if log_path:\n", 323 | " import json\n", 324 | "\n", 325 | " with open(log_path, \"r\") as f:\n", 326 | " log_data = json.load(f)\n", 327 | "\n", 328 | " print(f\"\\nSession started: {log_data['session_start']}\")\n", 329 | " print(f\"Total conversations: {len(log_data['conversations'])}\")\n", 330 | "\n", 331 | " # Show last conversation if any\n", 332 | " if log_data[\"conversations\"]:\n", 333 | " last_conv = log_data[\"conversations\"][-1]\n", 334 | " print(\"\\nLast conversation:\")\n", 335 | " print(f\" User: {last_conv['user_message']}\")\n", 336 | " print(f\" AI: {last_conv['ai_response'][:100]}...\")\n", 337 | " print(f\" Thread ID: {last_conv['thread_id']}\")" 338 | ] 339 | }, 340 | { 341 | "cell_type": "markdown", 342 | "metadata": {}, 343 | "source": [ 344 | "## Summary\n", 345 | "\n", 346 | "The assistant widget provides:\n", 347 | "\n", 348 | "1. **AI-Powered Assistant** - Natural language understanding with multiple provider support\n", 349 | "2. **Direct Kernel Access** - Read and execute code in your notebook's kernel\n", 350 | "3. **Variable Inspection** - Deep inspection of any variable type\n", 351 | "4. **Interactive Commands** - Simple slash commands for common operations\n", 352 | "5. **Programmatic API** - Full control from Python code\n", 353 | "6. **Real-time Updates** - Kernel state synchronized with the UI\n", 354 | "\n", 355 | "The AI assistant can:\n", 356 | "- Understand your questions and provide intelligent responses\n", 357 | "- Automatically use kernel tools to inspect variables and execute code\n", 358 | "- Help debug errors and understand your data\n", 359 | "- Work even without an API key (using a helpful mock AI)\n", 360 | "\n", 361 | "This creates a powerful AI assistant that truly understands your notebook context!" 362 | ] 363 | } 364 | ], 365 | "metadata": { 366 | "kernelspec": { 367 | "display_name": ".venv", 368 | "language": "python", 369 | "name": "python3" 370 | }, 371 | "language_info": { 372 | "codemirror_mode": { 373 | "name": "ipython", 374 | "version": 3 375 | }, 376 | "file_extension": ".py", 377 | "mimetype": "text/x-python", 378 | "name": "python", 379 | "nbconvert_exporter": "python", 380 | "pygments_lexer": "ipython3", 381 | "version": "3.11.12" 382 | } 383 | }, 384 | "nbformat": 4, 385 | "nbformat_minor": 4 386 | } 387 | -------------------------------------------------------------------------------- /tests/test_message_handlers.py: -------------------------------------------------------------------------------- 1 | """Tests for message handlers.""" 2 | # mypy: disable-error-code=misc 3 | 4 | from typing import Any 5 | from unittest.mock import Mock 6 | 7 | import pytest 8 | 9 | from assistant_ui_anywidget.simple_handlers import SimpleHandlers 10 | from assistant_ui_anywidget.kernel_interface import VariableInfo, ExecutionResult 11 | 12 | 13 | class MockKernelInterface: 14 | """Mock kernel interface for testing.""" 15 | 16 | def __init__(self) -> None: 17 | self.is_available = True 18 | self.namespace = {"x": 42, "y": "hello", "data": [1, 2, 3, 4, 5]} 19 | # Mock shell with execution_count 20 | self.shell = Mock(execution_count=10) 21 | 22 | def get_namespace(self) -> dict[str, Any]: 23 | """Get mock namespace.""" 24 | return self.namespace 25 | 26 | def get_variable_info(self, name: str, deep: bool = False) -> VariableInfo | None: 27 | """Get mock variable info.""" 28 | if name not in self.namespace: 29 | return None 30 | 31 | value = self.namespace[name] 32 | return VariableInfo( 33 | name=name, 34 | type=type(value).__name__, 35 | type_str=f"{type(value).__module__}.{type(value).__name__}", 36 | size=28 if isinstance(value, int) else None, 37 | shape=None, 38 | preview=repr(value), 39 | is_callable=False, 40 | attributes=["__class__", "__str__"] if deep else [], 41 | last_modified=None, 42 | ) 43 | 44 | def execute_code( 45 | self, code: str, silent: bool = False, store_history: bool = True 46 | ) -> ExecutionResult: 47 | """Mock code execution.""" 48 | if code == "1 + 1": 49 | return ExecutionResult( 50 | success=True, 51 | execution_count=1, 52 | outputs=[ 53 | { 54 | "type": "execute_result", 55 | "data": {"text/plain": "2"}, 56 | "execution_count": 1, 57 | } 58 | ], 59 | execution_time=0.001, 60 | variables_changed=[], 61 | ) 62 | elif code == "z = 100": 63 | self.namespace["z"] = 100 64 | return ExecutionResult( 65 | success=True, 66 | execution_count=2, 67 | outputs=[], 68 | execution_time=0.001, 69 | variables_changed=["z"], 70 | ) 71 | elif code == "raise ValueError('test')": 72 | return ExecutionResult( 73 | success=False, 74 | execution_count=3, 75 | outputs=[], 76 | execution_time=0.001, 77 | variables_changed=[], 78 | error={ 79 | "type": "ValueError", 80 | "message": "test", 81 | "traceback": ["Traceback...", "ValueError: test"], 82 | }, 83 | ) 84 | else: 85 | return ExecutionResult( 86 | success=True, 87 | execution_count=4, 88 | outputs=[], 89 | execution_time=0.001, 90 | variables_changed=[], 91 | ) 92 | 93 | def get_kernel_info(self) -> dict[str, Any]: 94 | """Get mock kernel info.""" 95 | return { 96 | "available": self.is_available, 97 | "language": "python", 98 | "execution_count": 10, 99 | } 100 | 101 | def get_stack_trace( 102 | self, include_locals: bool = False, max_frames: int = 10 103 | ) -> list[Any]: 104 | """Get mock stack trace.""" 105 | return [] 106 | 107 | def get_last_error(self) -> None: 108 | """Get mock last error.""" 109 | return None 110 | 111 | 112 | @pytest.fixture 113 | def mock_kernel() -> MockKernelInterface: 114 | """Create mock kernel interface.""" 115 | return MockKernelInterface() 116 | 117 | 118 | @pytest.fixture 119 | def message_handlers(mock_kernel: Any) -> SimpleHandlers: 120 | """Create message handlers with mock kernel.""" 121 | return SimpleHandlers(mock_kernel) 122 | 123 | 124 | class TestMessageHandlers: 125 | """Test message handler functionality.""" 126 | 127 | def test_handle_invalid_message(self, message_handlers: Any) -> None: 128 | """Test handling invalid messages.""" 129 | # Test non-dict message 130 | response = message_handlers.handle_message("not a dict") 131 | assert response["success"] is False 132 | assert "type is required" in response["error"] 133 | 134 | # Test missing type 135 | response = message_handlers.handle_message({}) 136 | assert response["success"] is False 137 | assert "type is required" in response["error"] 138 | 139 | # Test unknown type 140 | response = message_handlers.handle_message( 141 | {"id": "123", "type": "unknown_type"} 142 | ) 143 | assert response["success"] is False 144 | assert "Unknown message type" in response["error"] 145 | 146 | def test_handle_get_variables(self, message_handlers: Any) -> None: 147 | """Test get_variables handler.""" 148 | # Basic request 149 | response = message_handlers.handle_message( 150 | {"id": "123", "type": "get_variables"} 151 | ) 152 | 153 | assert response["success"] is True 154 | assert len(response["data"]["variables"]) == 3 155 | assert response["data"]["total_count"] == 3 156 | 157 | # Test with filter 158 | response = message_handlers.handle_message( 159 | { 160 | "id": "124", 161 | "type": "get_variables", 162 | "params": {"filter": {"types": ["int"]}}, 163 | } 164 | ) 165 | 166 | assert response["success"] is True 167 | variables = response["data"]["variables"] 168 | assert len(variables) == 1 169 | assert variables[0]["name"] == "x" 170 | assert variables[0]["type"] == "int" 171 | 172 | # Test with pattern filter 173 | response = message_handlers.handle_message( 174 | { 175 | "id": "125", 176 | "type": "get_variables", 177 | "params": {"filter": {"pattern": "^[xy]$"}}, 178 | } 179 | ) 180 | 181 | variables = response["data"]["variables"] 182 | assert len(variables) == 2 183 | assert all(v["name"] in ["x", "y"] for v in variables) 184 | 185 | # Test sorting 186 | response = message_handlers.handle_message( 187 | { 188 | "id": "126", 189 | "type": "get_variables", 190 | "params": {"sort": {"by": "name", "order": "desc"}}, 191 | } 192 | ) 193 | 194 | variables = response["data"]["variables"] 195 | names = [v["name"] for v in variables] 196 | assert names == ["y", "x", "data"] 197 | 198 | def test_handle_inspect_variable(self, message_handlers: Any) -> None: 199 | """Test inspect_variable handler.""" 200 | # Valid variable 201 | response = message_handlers.handle_message( 202 | {"id": "200", "type": "inspect_variable", "params": {"name": "x"}} 203 | ) 204 | 205 | assert response["success"] is True 206 | info = response["data"] 207 | assert info["name"] == "x" 208 | assert info["type"] == "int" 209 | assert info["preview"] == "42" 210 | 211 | # Deep inspection 212 | response = message_handlers.handle_message( 213 | { 214 | "id": "201", 215 | "type": "inspect_variable", 216 | "params": {"name": "y", "deep": True}, 217 | } 218 | ) 219 | 220 | info = response["data"] 221 | assert len(info["attributes"]) > 0 222 | assert "name" in info 223 | assert info["name"] == "y" 224 | 225 | # Non-existent variable 226 | response = message_handlers.handle_message( 227 | {"id": "202", "type": "inspect_variable", "params": {"name": "nonexistent"}} 228 | ) 229 | 230 | assert response["success"] is False 231 | assert "not found" in response["error"] 232 | 233 | # Missing variable name 234 | response = message_handlers.handle_message( 235 | {"id": "203", "type": "inspect_variable", "params": {}} 236 | ) 237 | 238 | assert response["success"] is False 239 | assert "Variable name is required" in response["error"] 240 | 241 | def test_handle_execute_code(self, message_handlers: Any) -> None: 242 | """Test execute_code handler.""" 243 | # Successful execution 244 | response = message_handlers.handle_message( 245 | {"id": "300", "type": "execute_code", "params": {"code": "1 + 1"}} 246 | ) 247 | 248 | assert response["success"] is True 249 | data = response["data"] 250 | assert data["success"] is True 251 | assert data["execution_count"] == 1 252 | assert len(data["outputs"]) == 1 253 | assert data["outputs"][0]["data"]["text/plain"] == "2" 254 | 255 | # Variable assignment 256 | response = message_handlers.handle_message( 257 | {"id": "301", "type": "execute_code", "params": {"code": "z = 100"}} 258 | ) 259 | 260 | assert response["success"] is True 261 | data = response["data"] 262 | assert data["variables_changed"] == ["z"] 263 | 264 | # Execution error 265 | response = message_handlers.handle_message( 266 | { 267 | "id": "302", 268 | "type": "execute_code", 269 | "params": {"code": "raise ValueError('test')"}, 270 | } 271 | ) 272 | 273 | assert response["success"] is False 274 | assert "test" in response["error"] 275 | 276 | # Missing code 277 | response = message_handlers.handle_message( 278 | {"id": "303", "type": "execute_code", "params": {}} 279 | ) 280 | 281 | assert response["success"] is False 282 | assert "Code is required" in response["error"] 283 | 284 | def test_handle_get_kernel_info(self, message_handlers: Any) -> None: 285 | """Test get_kernel_info handler.""" 286 | response = message_handlers.handle_message( 287 | {"id": "400", "type": "get_kernel_info"} 288 | ) 289 | 290 | assert response["success"] is True 291 | data = response["data"] 292 | assert data["available"] is True 293 | assert data["language"] == "python" 294 | assert data["execution_count"] == 10 295 | 296 | def test_handle_get_stack_trace(self, message_handlers: Any) -> None: 297 | """Test get_stack_trace handler.""" 298 | response = message_handlers.handle_message( 299 | { 300 | "id": "500", 301 | "type": "get_stack_trace", 302 | "params": {"include_locals": True, "max_frames": 5}, 303 | } 304 | ) 305 | 306 | assert response["success"] is True 307 | data = response["data"] 308 | assert "stack_trace" in data 309 | assert data["stack_trace"] is None # Mock returns None 310 | assert "message" in data 311 | 312 | def test_handle_get_history(self, message_handlers: Any) -> None: 313 | """Test get_history handler.""" 314 | # Execute some code first 315 | message_handlers.handle_message( 316 | {"id": "600", "type": "execute_code", "params": {"code": "1 + 1"}} 317 | ) 318 | 319 | message_handlers.handle_message( 320 | {"id": "601", "type": "execute_code", "params": {"code": "z = 100"}} 321 | ) 322 | 323 | # Get history 324 | response = message_handlers.handle_message( 325 | { 326 | "id": "602", 327 | "type": "get_history", 328 | "params": {"n_items": 10, "include_output": True}, 329 | } 330 | ) 331 | 332 | assert response["success"] is True 333 | items = response["data"]["items"] 334 | assert len(items) == 2 335 | assert items[0]["input"] == "1 + 1" 336 | assert items[1]["input"] == "z = 100" 337 | 338 | # Search history 339 | response = message_handlers.handle_message( 340 | {"id": "603", "type": "get_history", "params": {"search": "z ="}} 341 | ) 342 | 343 | items = response["data"]["items"] 344 | assert len(items) == 1 345 | assert items[0]["input"] == "z = 100" 346 | 347 | def test_kernel_not_available(self, message_handlers: Any) -> None: 348 | """Test handling when kernel is not available.""" 349 | message_handlers.kernel.is_available = False 350 | 351 | response = message_handlers.handle_message( 352 | {"id": "700", "type": "get_variables"} 353 | ) 354 | 355 | assert response["success"] is False 356 | assert "not available" in response["error"] 357 | 358 | # get_kernel_info should still work 359 | response = message_handlers.handle_message( 360 | {"id": "701", "type": "get_kernel_info"} 361 | ) 362 | 363 | assert response["success"] is True 364 | assert response["data"]["available"] is False 365 | 366 | def test_response_structure(self, message_handlers: Any) -> None: 367 | """Test response message structure.""" 368 | response = message_handlers.handle_message( 369 | {"id": "800", "type": "get_kernel_info"} 370 | ) 371 | 372 | # Check required fields in simplified format 373 | assert "success" in response 374 | assert response["success"] is True 375 | 376 | # Success response should have data 377 | assert "data" in response 378 | assert "error" not in response 379 | 380 | def test_simple_response_structure(self) -> None: 381 | """Test simplified response structure.""" 382 | # Success response 383 | success_response: dict[str, Any] = {"success": True, "data": {"result": 42}} 384 | assert success_response["success"] is True 385 | assert success_response["data"]["result"] == 42 386 | 387 | # Error response 388 | error_response: dict[str, Any] = { 389 | "success": False, 390 | "error": "Something went wrong", 391 | } 392 | assert error_response["success"] is False 393 | assert "Something went wrong" in error_response["error"] 394 | --------------------------------------------------------------------------------