diff --git a/.env.azure.example b/.env.azure.example new file mode 100644 index 0000000..4e452e2 --- /dev/null +++ b/.env.azure.example @@ -0,0 +1,31 @@ +# Azure OpenAI Configuration for Jupyter AI Agents +# +# Copy this file to .env and fill in your values: +# cp .env.azure.example .env +# +# Then load it: +# source .env +# +# Or use direnv: +# direnv allow + +# Azure OpenAI API Key (required) +export AZURE_OPENAI_API_KEY="your-api-key-here" + +# Azure OpenAI Endpoint - Base URL only (required) +# Format: https://your-resource-name.openai.azure.com +# Do NOT include /openai/deployments/... or query parameters +export AZURE_OPENAI_ENDPOINT="https://your-resource-name.openai.azure.com" + +# Azure OpenAI API Version (optional, defaults to latest) +# Use the version that matches your deployment +export AZURE_OPENAI_API_VERSION="2024-08-01-preview" + +# Example deployment names you might use with --model-name: +# - gpt-4o-mini +# - gpt-4o +# - gpt-35-turbo +# - Your custom deployment names + +# Test your configuration: +# python -c "import os; print('API Key:', 'SET' if os.getenv('AZURE_OPENAI_API_KEY') else 'NOT SET'); print('Endpoint:', os.getenv('AZURE_OPENAI_ENDPOINT')); print('Version:', os.getenv('AZURE_OPENAI_API_VERSION', 'default'))" diff --git a/.yarnrc.yml b/.yarnrc.yml index d713341..a23e3b3 100644 --- a/.yarnrc.yml +++ b/.yarnrc.yml @@ -1,4 +1,4 @@ -# Copyright (c) Datalayer, Inc. https://datalayer.io +# Copyright (c) Datalayer, Inc. https://datalayer.ai # Distributed under the terms of the Datalayer License. enableImmutableInstalls: false diff --git a/Makefile b/Makefile index 2b4db98..1d128f9 100755 --- a/Makefile +++ b/Makefile @@ -42,22 +42,35 @@ example-fastapi: ## example-fastapi server @exec echo python -m uvicorn jupyter_ai_agents.examples.fastapi.main:main --reload --port 4400 -jupyter-ai-agents-prompt: - jupyter-ai-agents prompt \ +repl: + jupyter-ai-agents repl \ --url http://localhost:8888 \ --token MY_TOKEN \ - --model-provider azure-openai \ - --model-name gpt-4o-mini \ + --model azure-openai:gpt-4o-mini + +# Note: For Azure OpenAI, ensure these environment variables are set: +# - AZURE_OPENAI_API_KEY +# - AZURE_OPENAI_ENDPOINT (base URL, e.g., https://your-resource.openai.azure.com) +# - AZURE_OPENAI_API_VERSION (optional, defaults to latest) +# Adjust --max-requests based on your Azure tier (CLI default: 4; lower if you hit rate limits) +prompt: + jupyter-ai-agents prompt \ + --verbose \ + --mcp-servers http://localhost:8888/mcp \ + --model anthropic:claude-sonnet-4-20250514 \ --path notebook.ipynb \ + --max-requests 20 \ + --max-tool-calls 10 \ --input "Create a matplotlib example" -jupyter-ai-agents-explain-error: +explain-error: jupyter-ai-agents explain-error \ - --url http://localhost:8888 \ - --token MY_TOKEN \ - --model-provider azure-openai \ - --model-name gpt-4o-mini \ - --path notebook.ipynb + --verbose \ + --mcp-servers http://localhost:8888/mcp \ + --model anthropic:claude-sonnet-4-20250514 \ + --path notebook.ipynb \ + --max-requests 20 \ + --max-tool-calls 10 publish-pypi: # publish the pypi package git clean -fdx && \ diff --git a/README.md b/README.md index d8e8d9b..4886e6b 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ ~ BSD 3-Clause License --> -[![Datalayer](https://assets.datalayer.tech/datalayer-25.svg)](https://datalayer.io) +[![Datalayer](https://assets.datalayer.tech/datalayer-25.svg)](https://datalayer.ai) [![Become a Sponsor](https://img.shields.io/static/v1?label=Become%20a%20Sponsor&message=%E2%9D%A4&logo=GitHub&style=flat&color=1ABC9C)](https://github.com/sponsors/datalayer) @@ -75,7 +75,6 @@ You can also use Jupyter AI Agents through the command line interface for automa ![Jupyter AI Agents CLI](https://assets.datalayer.tech/jupyter-ai-agent/ai-agent-prompt-demo-terminal.gif) - ### Basic Installation To install Jupyter AI Agents, run the following command: @@ -110,7 +109,7 @@ pip install datalayer_pycrdt==0.12.17 ``` ### Examples -We put here a quick example for a Out-Kernel Stateless Agent via CLI helping your JupyterLab session. +Jupyter AI Agents provides CLI commands to help your JupyterLab session using **Pydantic AI agents** with **Model Context Protocol (MCP)** for tool integration. Start JupyterLab, setting a `port` and a `token` to be reused by the agent, and create a notebook `notebook.ipynb`. @@ -121,49 +120,186 @@ jupyter lab --port 8888 --IdentityProvider.token MY_TOKEN Jupyter AI Agents supports multiple AI model providers (more information can be found on [this documentation page](https://jupyter-ai-agents.datalayer.tech/docs/models)). -The following takes you through an example with the Azure OpenAI provider. Read the [Azure Documentation](https://learn.microsoft.com/en-us/azure/ai-services/openai) to get the needed credentials and make sure you define them in the following `.env` file. +### API Keys Configuration + +Set the appropriate API key for your chosen provider: + +**OpenAI:** +```bash +export OPENAI_API_KEY='your-api-key-here' +``` + +**Anthropic:** +```bash +export ANTHROPIC_API_KEY='your-api-key-here' +``` + +**Azure OpenAI:** +```bash +export AZURE_OPENAI_API_KEY='your-api-key-here' +export AZURE_OPENAI_ENDPOINT='https://your-resource.openai.azure.com' +export AZURE_OPENAI_API_VERSION='2024-08-01-preview' # optional +``` + +**Important for Azure OpenAI:** +- The `AZURE_OPENAI_ENDPOINT` should be just the base URL (e.g., `https://your-resource.openai.azure.com`) +- Do NOT include `/openai/deployments/...` or query parameters in the endpoint +- The deployment name is specified via the `--model-name` parameter +- See `.env.azure.example` for a complete configuration template +**Other providers:** ```bash -cat << EOF >>.env -OPENAI_API_VERSION="..." -AZURE_OPENAI_ENDPOINT="..." -AZURE_OPENAI_API_KEY="..." -EOF +export GOOGLE_API_KEY='your-api-key-here' # For Google/Gemini +export COHERE_API_KEY='your-api-key-here' # For Cohere +export GROQ_API_KEY='your-api-key-here' # For Groq +export MISTRAL_API_KEY='your-api-key-here' # For Mistral +# AWS credentials for Bedrock +export AWS_ACCESS_KEY_ID='your-key' +export AWS_SECRET_ACCESS_KEY='your-secret' +export AWS_REGION='us-east-1' ``` -**Prompt Agent** +### Model Specification + +You can specify the model in two ways: + +1. **Using `--model` with full string** (recommended): + ```bash + --model "openai:gpt-4o" + --model "anthropic:claude-sonnet-4-0" + --model "azure-openai:deployment-name" + ``` + +2. **Using `--model-provider` and `--model-name`**: + ```bash + --model-provider openai --model-name gpt-4o + --model-provider anthropic --model-name claude-sonnet-4-0 + ``` + +Supported providers: `openai`, `anthropic`, `azure-openai`, `github-copilot`, `google`, `bedrock`, `groq`, `mistral`, `cohere` + +### Prompt Agent -To use the Jupyter AI Agents, an easy way is to launch a CLI (update the Azure deployment name based on your setup). +Create and execute code based on user instructions: ```bash -# Prompt agent example. -# make jupyter-ai-agents-prompt +# Using full model string (recommended) jupyter-ai-agents prompt \ --url http://localhost:8888 \ --token MY_TOKEN \ - --model-provider azure-openai \ - --model-name gpt-4o-mini \ + --model "anthropic:claude-sonnet-4-0" \ --path notebook.ipynb \ --input "Create a matplotlib example" + +# Using provider and model name +jupyter-ai-agents prompt \ + --url http://localhost:8888 \ + --token MY_TOKEN \ + --model-provider anthropic \ + --model-name claude-sonnet-4-0 \ + --path notebook.ipynb \ + --input "Create a pandas dataframe with sample data and plot it" ``` -![Jupyter AI Agents](https://assets.datalayer.tech/jupyter-ai-agent/ai-agent-prompt-demo-terminal.gif) +![Jupyter AI Agents - Prompt](https://assets.datalayer.tech/jupyter-ai-agent/ai-agent-prompt-demo-terminal.gif) + +### Explain Error Agent -**Explain Error Agent** +Analyze and fix notebook errors: ```bash -# Explain Error agent example. -# make jupyter-ai-agents-explain-error jupyter-ai-agents explain-error \ --url http://localhost:8888 \ --token MY_TOKEN \ - --model-provider azure-openai \ - --model-name gpt-4o-mini \ - --path notebook.ipynb + --model "anthropic:claude-sonnet-4-0" \ + --path notebook.ipynb \ + --current-cell-index 5 +``` + +![Jupyter AI Agents - Explain Error](https://assets.datalayer.tech/jupyter-ai-agent/ai-agent-explainerror-demo-terminal.gif) + +### REPL Mode (Interactive) + +For an interactive experience with direct access to all Jupyter MCP tools, use the REPL mode: + +```bash +jupyter-ai-agents repl \ + --url http://localhost:8888 \ + --token MY_TOKEN \ + --model "anthropic:claude-sonnet-4-0" ``` -![Jupyter AI Agents](https://assets.datalayer.tech/jupyter-ai-agent/ai-agent-explainerror-demo-terminal.gif) +In REPL mode, you can directly ask the AI to: +- List notebooks in directories +- Read and analyze notebook contents +- Execute code in cells +- Insert new cells +- Modify existing cells +- Install Python packages + +Example REPL interactions: + +``` +> List all notebooks in the current directory +> Create a new notebook called analysis.ipynb +> In analysis.ipynb, create a cell that imports pandas and loads data.csv +> Execute the cell and show me the first 5 rows +> Add a matplotlib plot showing the distribution of the 'age' column +``` + +The REPL provides special commands: +- `/exit`: Exit the session +- `/markdown`: Show last response in markdown format +- `/multiline`: Toggle multiline input mode (use Ctrl+D to submit) +- `/cp`: Copy last response to clipboard + +You can also use a custom system prompt: + +```bash +jupyter-ai-agents repl \ + --url http://localhost:8888 \ + --token MY_TOKEN \ + --model "anthropic:claude-sonnet-4-0" \ + --system-prompt "You are a data science expert specializing in pandas and matplotlib." +``` + +### Prompt Agent + +Create and execute code based on user instructions: + +```bash +# Using full model string (recommended) +jupyter-ai-agents prompt \ + --url http://localhost:8888 \ + --token MY_TOKEN \ + --model "anthropic:claude-sonnet-4-0" \ + --path notebook.ipynb \ + --input "Create a matplotlib example" + +# Using provider and model name +jupyter-ai-agents prompt \ + --url http://localhost:8888 \ + --token MY_TOKEN \ + --model-provider anthropic \ + --model-name claude-sonnet-4-0 \ + --path notebook.ipynb \ + --input "Create a pandas dataframe with sample data and plot it" +``` + +### Explain Error Agent + +Analyze and fix notebook errors: + +```bash +jupyter-ai-agents explain-error \ + --url http://localhost:8888 \ + --token MY_TOKEN \ + --model "anthropic:claude-sonnet-4-0" \ + --path notebook.ipynb \ + --current-cell-index 5 +``` +## Uninstall ### About the Technology diff --git a/dev/README.md b/dev/README.md index 29c14e3..c16409d 100644 --- a/dev/README.md +++ b/dev/README.md @@ -4,6 +4,6 @@ ~ BSD 3-Clause License --> -[![Datalayer](https://assets.datalayer.tech/datalayer-25.svg)](https://datalayer.io) +[![Datalayer](https://assets.datalayer.tech/datalayer-25.svg)](https://datalayer.ai) [![Become a Sponsor](https://img.shields.io/static/v1?label=Become%20a%20Sponsor&message=%E2%9D%A4&logo=GitHub&style=flat&color=1ABC9C)](https://github.com/sponsors/datalayer) diff --git a/dev/content/README.md b/dev/content/README.md index 29c14e3..c16409d 100644 --- a/dev/content/README.md +++ b/dev/content/README.md @@ -4,6 +4,6 @@ ~ BSD 3-Clause License --> -[![Datalayer](https://assets.datalayer.tech/datalayer-25.svg)](https://datalayer.io) +[![Datalayer](https://assets.datalayer.tech/datalayer-25.svg)](https://datalayer.ai) [![Become a Sponsor](https://img.shields.io/static/v1?label=Become%20a%20Sponsor&message=%E2%9D%A4&logo=GitHub&style=flat&color=1ABC9C)](https://github.com/sponsors/datalayer) diff --git a/dev/content/notebook.ipynb b/dev/content/notebook.ipynb index f30aee2..e8334c7 100644 --- a/dev/content/notebook.ipynb +++ b/dev/content/notebook.ipynb @@ -3,15 +3,7 @@ { "cell_type": "code", "execution_count": null, - "id": "a582b625-e712-4581-b611-3def5f40bd59", - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "84ca6269-06aa-42c3-815c-292980f031c0", + "id": "6e785ae0-34d9-462e-aa9a-22d7b40e15d8", "metadata": {}, "outputs": [], "source": [] diff --git a/docs/.yarnrc.yml b/docs/.yarnrc.yml index 99f1f38..682f6af 100644 --- a/docs/.yarnrc.yml +++ b/docs/.yarnrc.yml @@ -1,4 +1,4 @@ -# Copyright (c) Datalayer, Inc. https://datalayer.io +# Copyright (c) Datalayer, Inc. https://datalayer.ai # Distributed under the terms of the MIT License. enableImmutableInstalls: false diff --git a/docs/Makefile b/docs/Makefile index 9467751..c24bf84 100644 --- a/docs/Makefile +++ b/docs/Makefile @@ -1,4 +1,4 @@ -# Copyright (c) Datalayer, Inc. https://datalayer.io +# Copyright (c) Datalayer, Inc. https://datalayer.ai # Distributed under the terms of the MIT License. SHELL=/bin/bash diff --git a/docs/PYDANTIC_AI_AGENTS.md b/docs/PYDANTIC_AI_AGENTS.md new file mode 100644 index 0000000..2660c5c --- /dev/null +++ b/docs/PYDANTIC_AI_AGENTS.md @@ -0,0 +1,301 @@ +# Pydantic AI Agents with MCP Integration + +This document describes the new pydantic-ai based agents that use the Model Context Protocol (MCP) for Jupyter integration. + +## Overview + +The pydantic-ai agents represent a modern approach to AI-powered Jupyter notebook manipulation: + +- **Technology Stack**: Built with [pydantic-ai](https://ai.pydantic.dev/), a Python agent framework +- **Communication Protocol**: Uses [Model Context Protocol (MCP)](https://modelcontextprotocol.io/) for tool integration +- **Architecture**: Direct server-to-server communication (no RTC/nbmodel client needed) +- **CLI Support**: Includes pydantic-ai's built-in interactive CLI + +## Architecture Comparison + +### Legacy LangChain Architecture +``` +CLI Tool → NbModel Client → RTC (WebSocket) → Jupyter Server → Notebook + → Kernel Client → Kernel → Execution +``` + +### New Pydantic AI Architecture +``` +CLI Tool → MCP Client → HTTP/SSE → jupyter-mcp-server → Jupyter Server → Notebook + → Kernel → Execution +``` + +## Key Benefits + +1. **Simplified Communication**: Direct HTTP communication instead of WebSocket RTC +2. **Standard Protocol**: MCP is becoming the standard for tool integration with AI +3. **Better Tool Discovery**: MCP provides automatic tool discovery and schema validation +4. **Modern Framework**: pydantic-ai offers better type safety and async support +5. **Interactive Mode**: Built-in CLI with conversation history and markdown rendering + +## Available Agents + +### 1. Prompt Agent + +Creates and executes code based on natural language instructions. + +**Features**: +- Breaks down complex requests into steps +- Installs required packages automatically +- Adds explanatory markdown cells +- Executes code to verify it works +- Maintains cell indexing and execution flow + +**Usage**: +```bash +jupyter-ai-agents-pydantic prompt \ + --url http://localhost:8888 \ + --token YOUR_TOKEN \ + --model-provider openai \ + --model-name gpt-4o \ + --path notebook.ipynb \ + --input "Create a pandas dataframe and plot it with matplotlib" +``` + +**Options**: +- `--url`: Jupyter server URL +- `--token`: Authentication token +- `--path`: Notebook path +- `--input`: User instruction +- `--model-provider`: AI provider (openai, anthropic, azure-openai, etc.) +- `--model-name`: Model identifier +- `--current-cell-index`: Where user made the request (default: -1) +- `--full-context`: Include full notebook content in context +- `--verbose`: Enable debug logging + +### 2. Explain Error Agent + +Analyzes errors in notebooks and provides fixes. + +**Features**: +- Parses error tracebacks +- Explains the root cause +- Inserts corrected code cells +- Executes fixes to verify they work +- Adds explanatory comments + +**Usage**: +```bash +jupyter-ai-agents-pydantic explain-error \ + --url http://localhost:8888 \ + --token YOUR_TOKEN \ + --model-provider openai \ + --model-name gpt-4o \ + --path notebook.ipynb \ + --current-cell-index 5 +``` + +**Options**: +- `--current-cell-index`: Cell with error (-1 to find first error) +- Other options same as prompt agent + +### 3. Interactive Mode + +Conversational interface for notebook manipulation. + +**Usage**: +```bash +jupyter-ai-agents-pydantic interactive \ + --url http://localhost:8888 \ + --token YOUR_TOKEN \ + --model-provider openai \ + --model-name gpt-4o \ + --path notebook.ipynb +``` + +**Special Commands**: +- `/exit`: Exit the session +- `/markdown`: Show last response in markdown +- `/multiline`: Toggle multiline input mode (use Ctrl+D to submit) +- `/cp`: Copy last response to clipboard + +## Model Providers + +### OpenAI +```bash +export OPENAI_API_KEY='your-key' +jupyter-ai-agents-pydantic prompt \ + --model-provider openai \ + --model-name gpt-4o \ + ... +``` + +### Anthropic +```bash +export ANTHROPIC_API_KEY='your-key' +jupyter-ai-agents-pydantic prompt \ + --model-provider anthropic \ + --model-name claude-sonnet-4-0 \ + ... +``` + +### Azure OpenAI +```bash +export AZURE_OPENAI_API_KEY='your-key' +export AZURE_OPENAI_ENDPOINT='https://your-endpoint.openai.azure.com/' +jupyter-ai-agents-pydantic prompt \ + --model-provider azure-openai \ + --model-name gpt-4o \ + ... +``` + +### GitHub Copilot +```bash +export GITHUB_TOKEN='your-github-token' +jupyter-ai-agents-pydantic prompt \ + --model-provider github-copilot \ + --model-name gpt-4o \ + ... +``` + +## MCP Tools Available + +The agents have access to the following tools through jupyter-mcp-server: + +### Notebook Tools +- `notebook_insert_cell`: Insert a new cell at a specific index +- `notebook_insert_code_cell`: Insert and optionally execute a code cell +- `notebook_insert_markdown_cell`: Insert a markdown cell +- `notebook_get_cell`: Read a specific cell's content +- `notebook_get_cells`: List all cells +- `notebook_update_cell`: Modify an existing cell +- `notebook_delete_cell`: Remove a cell +- `notebook_run_cell`: Execute a specific cell +- `notebook_run_all_cells`: Execute all cells in sequence + +### Kernel Tools +- `kernel_execute`: Execute code in the notebook kernel +- `kernel_interrupt`: Interrupt running code +- `kernel_restart`: Restart the kernel +- `kernel_get_info`: Get kernel information + +### File System Tools (if enabled) +- `fs_read_file`: Read file contents +- `fs_write_file`: Write to a file +- `fs_list_directory`: List directory contents + +## Configuration + +### Environment Variables + +The agents respect standard pydantic-ai environment variables: + +- `OPENAI_API_KEY`: OpenAI API key +- `ANTHROPIC_API_KEY`: Anthropic API key +- `AZURE_OPENAI_API_KEY`: Azure OpenAI key +- `AZURE_OPENAI_ENDPOINT`: Azure endpoint +- `GITHUB_TOKEN`: GitHub token for Copilot +- `PYDANTIC_AI_LOG_LEVEL`: Logging level (DEBUG, INFO, WARNING, ERROR) + +### jupyter-mcp-server + +Make sure jupyter-mcp-server is installed and running: + +```bash +pip install jupyter-mcp-server +``` + +The MCP server is automatically started with JupyterLab when jupyter_ai_agents extension is loaded. + +## Programmatic Usage + +You can also use the agents programmatically: + +```python +import asyncio +from jupyter_ai_agents.agents.cli.prompt_agent import ( + create_prompt_agent, + run_prompt_agent, +) +from jupyter_ai_agents.tools import create_mcp_server + +async def main(): + # Create MCP connection + mcp_server = create_mcp_server( + "http://localhost:8888", + token="YOUR_TOKEN" + ) + + # Create agent + agent = create_prompt_agent( + model="openai:gpt-4o", + mcp_server=mcp_server, + ) + + # Run agent + result = await run_prompt_agent( + agent, + "Create a simple plot", + ) + + print(result) + +asyncio.run(main()) +``` + +## Troubleshooting + +### Connection Issues + +If the agent can't connect to Jupyter: + +1. Verify JupyterLab is running: `jupyter lab list` +2. Check the URL and token are correct +3. Ensure jupyter-mcp-server is installed +4. Check if the MCP endpoint is accessible: `curl http://localhost:8888/mcp` + +### API Key Issues + +If you get authentication errors: + +1. Verify your API key is set: `echo $OPENAI_API_KEY` +2. Check key permissions and quotas +3. Try a different model provider + +### Tool Execution Issues + +If tools aren't working: + +1. Enable verbose logging: `--verbose` +2. Check jupyter-mcp-server logs +3. Verify notebook path is correct +4. Ensure notebook is not locked by another process + +## Migration from LangChain Agents + +If you're migrating from the legacy LangChain agents: + +| LangChain | Pydantic AI | +|-----------|-------------| +| `jupyter-ai-agents prompt` | `jupyter-ai-agents-pydantic prompt` | +| `jupyter-ai-agents explain-error` | `jupyter-ai-agents-pydantic explain-error` | +| Uses NbModel Client + Kernel Client | Uses MCP Client | +| RTC WebSocket communication | HTTP/SSE communication | +| LangChain provider setup | Pydantic AI model strings | + +**Key Differences**: +1. No need to pass `--agent` parameter (built into command) +2. Model provider format: `--model-provider openai --model-name gpt-4o` instead of LangChain specific configs +3. No need for separate kernel client setup +4. Simpler architecture with MCP + +## Contributing + +To add new agent capabilities: + +1. Define tools in jupyter-mcp-server +2. Agents automatically discover new tools via MCP +3. Update agent instructions to use new tools +4. Test with `--verbose` flag + +## References + +- [Pydantic AI Documentation](https://ai.pydantic.dev/) +- [MCP Protocol](https://modelcontextprotocol.io/) +- [jupyter-mcp-server](https://github.com/datalayer/jupyter-mcp-server) +- [Pydantic AI CLI](https://ai.pydantic.dev/cli/) diff --git a/docs/PYDANTIC_IMPLEMENTATION.md b/docs/PYDANTIC_IMPLEMENTATION.md new file mode 100644 index 0000000..be21feb --- /dev/null +++ b/docs/PYDANTIC_IMPLEMENTATION.md @@ -0,0 +1,239 @@ +# Pydantic AI Agents Implementation Summary + +## Overview + +Successfully implemented new pydantic-ai based agents for Jupyter AI Agents that use the Model Context Protocol (MCP) for communication, replacing the legacy LangChain + NbModel approach. + +## What Was Created + +### 1. **Agent Implementations** + +Created two new pydantic-ai agents in `jupyter_ai_agents/agents/pydantic/cli/`: + +#### **Prompt Agent** (`prompt_agent.py`) +- Creates and executes code based on natural language instructions +- Automatically installs required packages +- Adds explanatory markdown cells +- Executes code to verify it works +- Features: + - `create_prompt_agent()`: Factory function to create agent + - `run_prompt_agent()`: Execute agent with user input + - `PromptAgentDeps`: Type-safe dependencies class + - System prompt optimized for code generation + +#### **Explain Error Agent** (`explain_error_agent.py`) +- Analyzes notebook errors and provides fixes +- Explains root causes clearly +- Inserts corrected code cells +- Verifies fixes by execution +- Features: + - `create_explain_error_agent()`: Factory function + - `run_explain_error_agent()`: Execute agent for error analysis + - `ExplainErrorAgentDeps`: Type-safe dependencies + - System prompt optimized for debugging + +### 2. **CLI Application** + +Created new CLI in `jupyter_ai_agents/cli/pydantic_app.py` with three commands: + +#### **`prompt` Command** +```bash +jupyter-ai-agents-pydantic prompt \ + --url http://localhost:8888 \ + --token MY_TOKEN \ + --model-provider openai \ + --model-name gpt-4o \ + --path notebook.ipynb \ + --input "Create a matplotlib example" +``` + +#### **`explain-error` Command** +```bash +jupyter-ai-agents-pydantic explain-error \ + --url http://localhost:8888 \ + --token MY_TOKEN \ + --model-provider openai \ + --model-name gpt-4o \ + --path notebook.ipynb +``` + +#### **`interactive` Command** +```bash +jupyter-ai-agents-pydantic interactive \ + --url http://localhost:8888 \ + --token MY_TOKEN \ + --model-provider openai \ + --model-name gpt-4o \ + --path notebook.ipynb +``` +Uses pydantic-ai's built-in CLI for conversational experience. + +### 3. **Model Provider Support** + +Implemented intelligent model string conversion supporting: +- OpenAI (`openai:gpt-4o`) +- Anthropic (`anthropic:claude-sonnet-4-0`) +- Azure OpenAI (`openai:gpt-4o` with Azure credentials) +- GitHub Copilot (`openai:gpt-4o` with GitHub token) +- Bedrock and others + +### 4. **Documentation** + +Created comprehensive documentation: + +#### **README.md Updates** +- Added "CLI with Pydantic AI and MCP (Recommended)" section +- Documented all three commands with examples +- Explained benefits over legacy approach +- Kept legacy LangChain section for backward compatibility + +#### **New Documentation File** (`docs/PYDANTIC_AI_AGENTS.md`) +- Architecture comparison diagrams +- Detailed feature descriptions +- All model provider configurations +- Available MCP tools list +- Programmatic usage examples +- Troubleshooting guide +- Migration guide from LangChain + +### 5. **Configuration** + +Updated `pyproject.toml`: +```toml +[project.scripts] +jupyter-ai-agents-pydantic = "jupyter_ai_agents.cli.pydantic_app:main" +``` + +## Architecture Improvements + +### Old (LangChain): +``` +CLI → NbModel Client (RTC) → WebSocket → Jupyter → Notebook + → Kernel Client → Kernel +``` + +### New (Pydantic AI): +``` +CLI → MCP Client → HTTP/SSE → jupyter-mcp-server → Jupyter → Notebook + → Kernel +``` + +## Key Benefits + +1. **Simplified Architecture**: No RTC/WebSocket complexity +2. **Standard Protocol**: MCP is becoming industry standard +3. **Better Tool Discovery**: Automatic via MCP +4. **Type Safety**: Pydantic models throughout +5. **Modern Framework**: Async-first with pydantic-ai +6. **Interactive Mode**: Built-in CLI with conversation history +7. **Easy Testing**: Simpler to mock and test + +## Files Created/Modified + +### Created: +- `jupyter_ai_agents/agents/pydantic/cli/__init__.py` +- `jupyter_ai_agents/agents/pydantic/cli/prompt_agent.py` +- `jupyter_ai_agents/agents/pydantic/cli/explain_error_agent.py` +- `jupyter_ai_agents/cli/pydantic_app.py` +- `docs/PYDANTIC_AI_AGENTS.md` + +### Modified: +- `README.md` - Added pydantic-ai CLI documentation +- `pyproject.toml` - Added new CLI entry point +- `jupyter_ai_agents/handlers/chat.py` - Fixed pydantic-ai API (from earlier) + +## Usage Examples + +### Prompt Agent +```bash +# Simple example +jupyter-ai-agents-pydantic prompt \ + --url http://localhost:8888 \ + --token MY_TOKEN \ + --input "Create a pandas dataframe with sample data and visualize it" + +# With full context +jupyter-ai-agents-pydantic prompt \ + --url http://localhost:8888 \ + --token MY_TOKEN \ + --full-context \ + --input "Continue the analysis from above" +``` + +### Explain Error Agent +```bash +# Find and fix first error +jupyter-ai-agents-pydantic explain-error \ + --url http://localhost:8888 \ + --token MY_TOKEN \ + --path notebook.ipynb + +# Fix specific cell error +jupyter-ai-agents-pydantic explain-error \ + --url http://localhost:8888 \ + --token MY_TOKEN \ + --current-cell-index 5 +``` + +### Interactive Mode +```bash +jupyter-ai-agents-pydantic interactive \ + --url http://localhost:8888 \ + --token MY_TOKEN +``` + +## Testing + +Verified: +- ✅ Module imports work correctly +- ✅ CLI commands are registered +- ✅ Help messages display properly +- ✅ Model string conversion works for all providers +- ✅ Type annotations are correct + +## Next Steps + +To make the agents fully functional, you need to: + +1. **Implement Notebook Content Fetching**: Currently placeholder + - Add MCP tool or direct API call to fetch notebook content + - Extract cells and error information + - Pass to agents + +2. **Test with Real Jupyter Server**: + - Start JupyterLab with jupyter-mcp-server + - Test prompt agent with actual notebook + - Test explain-error agent with real errors + - Test interactive mode + +3. **Add More Features**: + - Streaming responses during execution + - Progress indicators + - Better error handling + - Support for more MCP tools + +4. **Documentation**: + - Add video demos + - Create example notebooks + - Add to main documentation site + +## Migration Guide + +For users of the old LangChain CLI: + +| Old Command | New Command | +|-------------|-------------| +| `jupyter-ai-agents prompt` | `jupyter-ai-agents-pydantic prompt` | +| `jupyter-ai-agents explain-error` | `jupyter-ai-agents-pydantic explain-error` | +| N/A | `jupyter-ai-agents-pydantic interactive` | + +**Breaking Changes**: None - old commands still work + +**Environment Variables**: Same API key variables work with both + +## References + +- [Pydantic AI Documentation](https://ai.pydantic.dev/) +- [MCP Protocol](https://modelcontextprotocol.io/) +- [Pydantic AI CLI](https://ai.pydantic.dev/cli/) +- [jupyter-mcp-server](https://github.com/datalayer/jupyter-mcp-server) diff --git a/docs/README.md b/docs/README.md index d47c6e7..274121a 100644 --- a/docs/README.md +++ b/docs/README.md @@ -1,10 +1,10 @@ -[![Datalayer](https://assets.datalayer.tech/datalayer-25.svg)](https://datalayer.io) +[![Datalayer](https://assets.datalayer.tech/datalayer-25.svg)](https://datalayer.ai) [![Become a Sponsor](https://img.shields.io/static/v1?label=Become%20a%20Sponsor&message=%E2%9D%A4&logo=GitHub&style=flat&color=1ABC9C)](https://github.com/sponsors/datalayer) # Jupyter AI Agents Docs -> Source code for the [Jupyter AI Agents Documentation](https://datalayer.io), built with [Docusaurus](https://docusaurus.io). +> Source code for the [Jupyter AI Agents Documentation](https://datalayer.ai), built with [Docusaurus](https://docusaurus.io). ```bash # Install the dependencies. diff --git a/examples/README.md b/examples/README.md new file mode 100644 index 0000000..d53066a --- /dev/null +++ b/examples/README.md @@ -0,0 +1,11 @@ + + +[![Datalayer](https://assets.datalayer.tech/datalayer-25.svg)](https://datalayer.ai) + +[![Become a Sponsor](https://img.shields.io/static/v1?label=Become%20a%20Sponsor&message=%E2%9D%A4&logo=GitHub&style=flat&color=1ABC9C)](https://github.com/sponsors/datalayer) + +# ✨ Jupyter AI Agents diff --git a/examples/jupyter-repl/Makefile b/examples/jupyter-repl/Makefile new file mode 100644 index 0000000..70739f3 --- /dev/null +++ b/examples/jupyter-repl/Makefile @@ -0,0 +1,35 @@ +.PHONY: help jupyterlab repl clean + +# Default target +# anthropic:claude-sonnet-4-20250514 +# azure-openai:gpt-4o-mini +help: + @echo "Jupyter MCP Example - Makefile" + @echo "" + @echo "Available targets:" + @echo " jupyterlab - Start JupyterLab server on http://localhost:8888" + @echo " repl - Start jupyter-ai-agents REPL connected to Jupyter MCP server" + @echo " clean - Clean up temporary files" + @echo "" + @echo "Usage:" + @echo " Terminal 1: make jupyterlab # Start JupyterLab with MCP server" + @echo " Terminal 2: make repl # Connect REPL to Jupyter MCP server" + @echo "" + @echo "Note: The REPL connects to jupyter-mcp-server at http://localhost:8888/mcp" + @echo "" + +jupyterlab: ## jupyterlab + jupyter lab \ + --port 8888 \ + --ServerApp.root_dir ./../../dev/content \ + --IdentityProvider.token= + +# Start REPL connected to MCP Jupyter +repl: + @echo "🤖 Starting jupyter-ai-agents REPL connected to Jupyter Server" + @echo "Make sure the Jupyter server is running first (make jupyterlab in another terminal)" + @echo "" + jupyter-ai-agents repl \ + --verbose \ + --mcp-servers http://localhost:8888/mcp \ + --model anthropic:claude-sonnet-4-20250514 diff --git a/examples/jupyter-repl/README.md b/examples/jupyter-repl/README.md new file mode 100644 index 0000000..7e7a655 --- /dev/null +++ b/examples/jupyter-repl/README.md @@ -0,0 +1,11 @@ + + +[![Datalayer](https://assets.datalayer.tech/datalayer-25.svg)](https://datalayer.ai) + +[![Become a Sponsor](https://img.shields.io/static/v1?label=Become%20a%20Sponsor&message=%E2%9D%A4&logo=GitHub&style=flat&color=1ABC9C)](https://github.com/sponsors/datalayer) + +# Jupyter REPL Example diff --git a/examples/mcp-repl/.gitignore b/examples/mcp-repl/.gitignore new file mode 100644 index 0000000..c323a1f --- /dev/null +++ b/examples/mcp-repl/.gitignore @@ -0,0 +1,27 @@ +# MCP Server Composer - Generated Files +.composer.pid +composer.log +nohup.out + +# Python +__pycache__/ +*.py[cod] +*$py.class +*.so +.Python +*.egg-info/ +dist/ +build/ + +# Logs +*.log + +# IDE +.vscode/ +.idea/ +*.swp +*.swo + +# OS +.DS_Store +Thumbs.db diff --git a/examples/mcp-repl/Makefile b/examples/mcp-repl/Makefile new file mode 100644 index 0000000..deb51a5 --- /dev/null +++ b/examples/mcp-repl/Makefile @@ -0,0 +1,104 @@ +.PHONY: help install mcp1 mcp2 repl1 repl2 clean + +# Default target +# anthropic:claude-sonnet-4-20250514 +# azure-openai:gpt-4o-mini +help: + @echo "Simple MCP Example - Makefile" + @echo "" + @echo "Available targets:" + @echo " install - Install dependencies (mcp, uvicorn, starlette, sse-starlette, jupyter-ai-agents)" + @echo " mcp1 - Start MCP server 1 (Calculator) on http://localhost:8001" + @echo " mcp2 - Start MCP server 2 (Echo) on http://localhost:8002" + @echo " repl1 - Start jupyter-ai-agents REPL connected to MCP server 1" + @echo " repl2 - Start jupyter-ai-agents REPL connected to MCP server 2" + @echo " clean - Clean up temporary files" + @echo "" + @echo "Usage:" + @echo " Terminal 1: make mcp1 # Start calculator server on port 8001" + @echo " Terminal 2: make repl1 # Connect REPL to calculator server" + @echo "" + @echo " Or:" + @echo " Terminal 1: make mcp2 # Start echo server on port 8002" + @echo " Terminal 2: make repl2 # Connect REPL to echo server" + @echo "" + @echo "Note: These MCP servers use SSE transport (HTTP/SSE)" + @echo "" + +# Install dependencies +install: + @echo "Installing dependencies..." + pip install mcp uvicorn starlette sse-starlette + pip install -e ../.. + @echo "" + @echo "✓ Installation complete!" + @echo "" + @echo "Available MCP servers:" + @echo " • mcp1.py - Calculator tools (add, subtract, multiply, divide) on port 8001" + @echo " • mcp2.py - Echo tools (ping, echo, reverse, uppercase, etc.) on port 8002" + @echo "" + @echo "These servers use SSE transport (HTTP/SSE) compatible with jupyter-ai-agents" + @echo "" + +# Start MCP server 1 (Calculator) +mcp1: + @echo "🚀 Starting MCP Server 1 (Calculator) on http://localhost:8001" + @echo "" + @echo "Available tools:" + @echo " • add(a, b) - Add two numbers" + @echo " • subtract(a, b) - Subtract two numbers" + @echo " • multiply(a, b) - Multiply two numbers" + @echo " • divide(a, b) - Divide two numbers" + @echo "" + @echo "MCP Endpoint (Streamable HTTP):" + @echo " • http://localhost:8001/mcp" + @echo "" + python mcp1.py + +# Start MCP server 2 (Echo) +mcp2: + @echo "🚀 Starting MCP Server 2 (Echo) on http://localhost:8002" + @echo "" + @echo "Available tools:" + @echo " • ping() - Test server connectivity" + @echo " • echo(message) - Echo back the message" + @echo " • reverse(text) - Reverse the text" + @echo " • uppercase(text) - Convert to uppercase" + @echo " • lowercase(text) - Convert to lowercase" + @echo " • count_words(text) - Count words in text" + @echo "" + @echo "MCP Endpoint (Streamable HTTP):" + @echo " • http://localhost:8002/mcp" + @echo "" + python mcp2.py + +# Start REPL connected to MCP server 1 (Calculator) +repl1: + @echo "🤖 Starting jupyter-ai-agents REPL connected to Calculator Server" + @echo " MCP Endpoint: http://localhost:8001/mcp" + @echo "" + @echo "Make sure the calculator server is running first (make mcp1 in another terminal)" + @echo "" + jupyter-ai-agents repl \ + --verbose \ + --mcp-servers http://localhost:8001/mcp \ + --model azure-openai:gpt-4o-mini + +# Start REPL connected to MCP server 2 (Echo) +repl2: + @echo "🤖 Starting jupyter-ai-agents REPL connected to Echo Server" + @echo " MCP Endpoint: http://localhost:8002/mcp" + @echo "" + @echo "Make sure the echo server is running first (make mcp2 in another terminal)" + @echo "" + jupyter-ai-agents repl \ + --verbose \ + --mcp-servers http://localhost:8002/mcp \ + --model azure-openai:gpt-4o-mini + +# Clean up +clean: + @echo "Cleaning up..." + @rm -f *.log nohup.out + @rm -rf __pycache__ + @echo "✓ Cleanup complete" diff --git a/examples/mcp-repl/README.md b/examples/mcp-repl/README.md new file mode 100644 index 0000000..f79f4eb --- /dev/null +++ b/examples/mcp-repl/README.md @@ -0,0 +1,114 @@ + + +[![Datalayer](https://assets.datalayer.tech/datalayer-25.svg)](https://datalayer.ai) + +[![Become a Sponsor](https://img.shields.io/static/v1?label=Become%20a%20Sponsor&message=%E2%9D%A4&logo=GitHub&style=flat&color=1ABC9C)](https://github.com/sponsors/datalayer) + +# Simple MCP Server Examples + +This directory contains two simple MCP servers that demonstrate how to create and use MCP servers with jupyter-ai-agents. + +## 🎯 Overview + +This example demonstrates standalone MCP servers using **FastMCP** with **Streamable HTTP transport** (MCP specification 2025-06-18): + +1. **Calculator Server** (`mcp1.py`) - Math operations on port 8001 +2. **Echo Server** (`mcp2.py`) - String operations on port 8002 + +Both servers use **FastMCP** (from the standard MCP SDK) which automatically handles: +- HTTP/SSE transport complexity +- GET (SSE streaming) and POST (JSON-RPC messages) at `/mcp` endpoint +- Tool registration and execution +- Input/output schema validation + +This is the same pattern used by `jupyter-mcp-server`, providing a clean, high-level abstraction for MCP server development. + +## 📋 MCP Servers + +### Calculator Server (mcp1.py) - Port 8001 +- `add(a, b)` - Add two numbers +- `subtract(a, b)` - Subtract two numbers +- `multiply(a, b)` - Multiply two numbers +- `divide(a, b)` - Divide two numbers + +### Echo Server (mcp2.py) - Port 8002 +- `ping()` - Test connectivity +- `echo(message)` - Echo back message +- `reverse(text)` - Reverse string +- `uppercase(text)` - Convert to uppercase +- `lowercase(text)` - Convert to lowercase +- `count_words(text)` - Count words + +## 🚀 Quick Start + +**1. Install:** +```bash +make install +``` + +**2. Terminal 1 - Start server:** +```bash +make mcp1 # or make mcp2 +``` + +**3. Terminal 2 - Connect REPL:** +```bash +make repl1 # or make repl2 +``` + +## 🛠️ Commands + +| Command | Description | +|---------|-------------| +| `make install` | Install dependencies | +| `make mcp1` | Start Calculator (8001) | +| `make mcp2` | Start Echo (8002) | +| `make repl1` | Connect to Calculator | +| `make repl2` | Connect to Echo | +| `make clean` | Clean temp files | + +## 📝 Manual Usage + +If you want to run commands manually: + +**Start servers:** +```bash +python mcp1.py # Calculator on port 8001 +python mcp2.py # Echo on port 8002 +``` + +**Connect REPL (using the `/mcp` endpoint):** +```bash +# Single server +jupyter-ai-agents repl --mcp-servers http://localhost:8001/mcp + +# Multiple servers +jupyter-ai-agents repl --mcp-servers "http://localhost:8001/mcp,http://localhost:8002/mcp" + +# Connect to Jupyter MCP Server (jupyter-mcp-server) +jupyter-ai-agents repl --mcp-servers http://localhost:8888/mcp +``` + +**Important:** All MCP servers use Streamable HTTP transport with the `/mcp` endpoint. The `--mcp-servers` parameter accepts comma-separated URLs for connecting to any MCP servers (standalone servers like calculator/echo, or jupyter-mcp-server). + +## �💬 REPL Examples + +``` +> Add 5 and 7 +> What is 15 divided by 3? +> Reverse the text "jupyter" +> Count words in "hello world" +``` + +## 📚 Resources + +- [MCP Specification](https://modelcontextprotocol.io/) +- [Pydantic AI](https://ai.pydantic.dev/) +- [Jupyter AI Agents](https://github.com/datalayer/jupyter-ai-agents) + +--- +**Made with ❤️ by [Datalayer](https://datalayer.ai)** diff --git a/examples/mcp-repl/mcp1.py b/examples/mcp-repl/mcp1.py new file mode 100755 index 0000000..78a69f2 --- /dev/null +++ b/examples/mcp-repl/mcp1.py @@ -0,0 +1,70 @@ +#!/usr/bin/env python3 +""" +MCP Server 1 - Calculator Tools + +Simple MCP server providing calculator operations using FastMCP +with Streamable HTTP transport (MCP specification 2025-06-18). +""" + +from typing import Annotated +from pydantic import Field +from mcp.server import FastMCP + + +# Create FastMCP server +mcp = FastMCP(name="Calculator Server") + + +@mcp.tool() +async def add( + a: Annotated[float, Field(description="First number")], + b: Annotated[float, Field(description="Second number")], +) -> float: + """Add two numbers together""" + return a + b + + +@mcp.tool() +async def subtract( + a: Annotated[float, Field(description="First number")], + b: Annotated[float, Field(description="Second number")], +) -> float: + """Subtract b from a""" + return a - b + + +@mcp.tool() +async def multiply( + a: Annotated[float, Field(description="First number")], + b: Annotated[float, Field(description="Second number")], +) -> float: + """Multiply two numbers""" + return a * b + + +@mcp.tool() +async def divide( + a: Annotated[float, Field(description="Numerator")], + b: Annotated[float, Field(description="Denominator")], +) -> float: + """Divide a by b""" + if b == 0: + raise ValueError("Cannot divide by zero") + return a / b + + +if __name__ == "__main__": + import uvicorn + + print("🚀 Starting Calculator MCP Server on http://localhost:8001") + print(" MCP endpoint: http://localhost:8001/mcp (Streamable HTTP)") + print(" Tools: add, subtract, multiply, divide") + print("") + + # Get the Starlette app with Streamable HTTP transport + app = mcp.streamable_http_app() + + # Run with uvicorn + uvicorn.run(app, host="0.0.0.0", port=8001) + + diff --git a/examples/mcp-repl/mcp2.py b/examples/mcp-repl/mcp2.py new file mode 100755 index 0000000..8289904 --- /dev/null +++ b/examples/mcp-repl/mcp2.py @@ -0,0 +1,234 @@ +#!/usr/bin/env python3 +""" +MCP Server 2 - Echo Tools + +Simple MCP server providing string manipulation tools using FastMCP +with Streamable HTTP transport (MCP specification 2025-06-18). +""" + +from typing import Annotated +from pydantic import Field +from mcp.server import FastMCP + + +# Create FastMCP server +mcp = FastMCP(name="Echo Server") + + +@mcp.tool() +async def ping() -> str: + """Simple ping tool to test connectivity""" + return "pong" + + +@mcp.tool() +async def echo( + message: Annotated[str, Field(description="Message to echo back")] +) -> str: + """Echo back a message""" + return message + + +@mcp.tool() +async def reverse( + text: Annotated[str, Field(description="Text to reverse")] +) -> str: + """Reverse a string""" + return text[::-1] + + +@mcp.tool() +async def uppercase( + text: Annotated[str, Field(description="Text to convert to uppercase")] +) -> str: + """Convert text to uppercase""" + return text.upper() + + +@mcp.tool() +async def lowercase( + text: Annotated[str, Field(description="Text to convert to lowercase")] +) -> str: + """Convert text to lowercase""" + return text.lower() + + +@mcp.tool() +async def count_words( + text: Annotated[str, Field(description="Text to count words in")] +) -> int: + """Count the number of words in text""" + return len(text.split()) + + +if __name__ == "__main__": + import uvicorn + + print("🚀 Starting Echo MCP Server on http://localhost:8002") + print(" MCP endpoint: http://localhost:8002/mcp (Streamable HTTP)") + print(" Tools: ping, echo, reverse, uppercase, lowercase, count_words") + print("") + + # Get the Starlette app with Streamable HTTP transport + app = mcp.streamable_http_app() + + # Run with uvicorn + uvicorn.run(app, host="0.0.0.0", port=8002) + +from mcp.server import Server +from mcp.server.sse import SseServerTransport +from mcp.types import Tool, TextContent +from starlette.applications import Starlette +from starlette.routing import Route +from starlette.responses import Response +import uvicorn + + +# Create MCP server +server = Server("echo-server") + + +@server.list_tools() +async def list_tools() -> list[Tool]: + """List available string manipulation tools.""" + return [ + Tool( + name="ping", + description="Simple ping that returns 'pong'", + inputSchema={ + "type": "object", + "properties": {}, + }, + ), + Tool( + name="echo", + description="Echo back the provided message", + inputSchema={ + "type": "object", + "properties": { + "message": {"type": "string", "description": "Message to echo back"}, + }, + "required": ["message"], + }, + ), + Tool( + name="reverse", + description="Reverse a string", + inputSchema={ + "type": "object", + "properties": { + "text": {"type": "string", "description": "Text to reverse"}, + }, + "required": ["text"], + }, + ), + Tool( + name="uppercase", + description="Convert text to uppercase", + inputSchema={ + "type": "object", + "properties": { + "text": {"type": "string", "description": "Text to convert"}, + }, + "required": ["text"], + }, + ), + Tool( + name="lowercase", + description="Convert text to lowercase", + inputSchema={ + "type": "object", + "properties": { + "text": {"type": "string", "description": "Text to convert"}, + }, + "required": ["text"], + }, + ), + Tool( + name="count_words", + description="Count the number of words in text", + inputSchema={ + "type": "object", + "properties": { + "text": {"type": "string", "description": "Text to analyze"}, + }, + "required": ["text"], + }, + ), + ] + + +@server.call_tool() +async def call_tool(name: str, arguments: dict) -> list[TextContent]: + """Handle tool calls.""" + + if name == "ping": + return [TextContent(type="text", text="pong")] + elif name == "echo": + message = arguments.get("message", "") + return [TextContent(type="text", text=message)] + elif name == "reverse": + text = arguments.get("text", "") + return [TextContent(type="text", text=text[::-1])] + elif name == "uppercase": + text = arguments.get("text", "") + return [TextContent(type="text", text=text.upper())] + elif name == "lowercase": + text = arguments.get("text", "") + return [TextContent(type="text", text=text.lower())] + elif name == "count_words": + text = arguments.get("text", "") + word_count = len(text.split()) + return [TextContent(type="text", text=str(word_count))] + else: + return [TextContent(type="text", text=f"Unknown tool: {name}")] + + +# Create SSE transport for the single MCP endpoint +sse = SseServerTransport("/mcp") + + +async def handle_mcp_endpoint(request): + """ + Handle the MCP endpoint according to Streamable HTTP transport specification. + + This single endpoint handles both: + - POST requests: Client sends JSON-RPC messages + - GET requests: Client opens SSE stream for server messages + """ + if request.method == "GET": + # GET request: Open SSE stream for server-to-client messages + async with sse.connect_sse( + request.scope, request.receive, request._send + ) as streams: + await server.run( + streams[0], streams[1], server.create_initialization_options() + ) + elif request.method == "POST": + # POST request: Client sends JSON-RPC message + # handle_post_message handles the response internally via ASGI callbacks + await sse.handle_post_message( + request.scope, request.receive, request._send + ) + else: + # Method not allowed + return Response(status_code=405) + + +# Create Starlette app with single MCP endpoint +app = Starlette( + routes=[ + Route("/mcp", endpoint=handle_mcp_endpoint, methods=["GET", "POST"]), + ] +) + + +if __name__ == "__main__": + print("🚀 Starting Echo MCP Server on http://localhost:8002") + print(" MCP endpoint: http://localhost:8002/mcp (Streamable HTTP)") + print(" Supports: GET (SSE stream) and POST (JSON-RPC messages)") + print("") + print(" Tools: ping, echo, reverse, uppercase, lowercase, count_words") + uvicorn.run(app, host="0.0.0.0", port=8002) + + diff --git a/jupyter_ai_agents/agents/pydantic/__init__.py b/jupyter_ai_agents/agents/__init__.py similarity index 100% rename from jupyter_ai_agents/agents/pydantic/__init__.py rename to jupyter_ai_agents/agents/__init__.py diff --git a/jupyter_ai_agents/agents/langchain/__init__.py b/jupyter_ai_agents/agents/chat/__init__.py similarity index 100% rename from jupyter_ai_agents/agents/langchain/__init__.py rename to jupyter_ai_agents/agents/chat/__init__.py diff --git a/jupyter_ai_agents/agents/pydantic/chat/agent.py b/jupyter_ai_agents/agents/chat/agent.py similarity index 71% rename from jupyter_ai_agents/agents/pydantic/chat/agent.py rename to jupyter_ai_agents/agents/chat/agent.py index 2cb55ad..17fda53 100644 --- a/jupyter_ai_agents/agents/pydantic/chat/agent.py +++ b/jupyter_ai_agents/agents/chat/agent.py @@ -2,35 +2,63 @@ # # BSD 3-Clause License -"""Pydantic AI agent for JupyterLab chat.""" +"""AI agent for JupyterLab chat.""" from typing import Any from pydantic_ai import Agent from pydantic_ai.mcp import MCPServerStreamableHTTP +from jupyter_ai_agents.utils.model import create_model_with_provider + def create_chat_agent( - model: str = "anthropic:claude-sonnet-4-5", + model: str | None = None, + model_provider: str = "anthropic", + model_name: str = "claude-sonnet-4-5", + timeout: float = 60.0, mcp_server: MCPServerStreamableHTTP | None = None, ) -> Agent: """ Create the main chat agent for JupyterLab. Args: - model: The model identifier to use (default: Claude Sonnet 4-5) + model: Optional full model string (e.g., "openai:gpt-4o", "azure-openai:gpt-4o-mini"). + If not provided, uses model_provider and model_name. + model_provider: Model provider name (default: "anthropic") + model_name: Model/deployment name (default: "claude-sonnet-4-5") + timeout: HTTP timeout in seconds for API requests (default: 60.0) mcp_server: Optional MCP server connection for Jupyter tools Returns: Configured Pydantic AI agent + + Note: + For Azure OpenAI, requires these environment variables: + - AZURE_OPENAI_API_KEY + - AZURE_OPENAI_ENDPOINT (base URL only, e.g., https://your-resource.openai.azure.com) + - AZURE_OPENAI_API_VERSION (optional, defaults to latest) """ + # Determine model to use + if model: + # User provided full model string + if model.startswith('azure-openai:'): + # Special handling for Azure OpenAI format + deployment_name = model.split(':', 1)[1] + model_obj = create_model_with_provider('azure-openai', deployment_name, timeout) + else: + model_obj = model + else: + # Create model object with provider-specific configuration + model_obj = create_model_with_provider(model_provider, model_name, timeout) + # Create toolsets list toolsets = [] if mcp_server: toolsets.append(mcp_server) agent = Agent( - model, + model_obj, toolsets=toolsets, instructions="""You are a helpful AI assistant integrated into JupyterLab. diff --git a/jupyter_ai_agents/agents/pydantic/chat/config.py b/jupyter_ai_agents/agents/chat/config.py similarity index 97% rename from jupyter_ai_agents/agents/pydantic/chat/config.py rename to jupyter_ai_agents/agents/chat/config.py index 994cf97..3305069 100644 --- a/jupyter_ai_agents/agents/pydantic/chat/config.py +++ b/jupyter_ai_agents/agents/chat/config.py @@ -8,7 +8,7 @@ from pathlib import Path from typing import List, Optional -from jupyter_ai_agents.agents.pydantic.models import MCPServer +from jupyter_ai_agents.agents.models import MCPServer class ChatConfig: diff --git a/jupyter_ai_agents/agents/pydantic/chat/models.py b/jupyter_ai_agents/agents/chat/models.py similarity index 100% rename from jupyter_ai_agents/agents/pydantic/chat/models.py rename to jupyter_ai_agents/agents/chat/models.py diff --git a/jupyter_ai_agents/agents/explain_error/__init__.py b/jupyter_ai_agents/agents/explain_error/__init__.py new file mode 100644 index 0000000..4aff33a --- /dev/null +++ b/jupyter_ai_agents/agents/explain_error/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) 2024-2025 Datalayer, Inc. +# +# BSD 3-Clause License + +"""Pydantic AI agents to explain error.""" + +from jupyter_ai_agents.agents.explain_error.explain_error_agent import create_explain_error_agent, run_explain_error_agent + +__all__ = [ + "create_explain_error_agent", + "run_explain_error_agent", +] diff --git a/jupyter_ai_agents/agents/explain_error/explain_error_agent.py b/jupyter_ai_agents/agents/explain_error/explain_error_agent.py new file mode 100644 index 0000000..713fa29 --- /dev/null +++ b/jupyter_ai_agents/agents/explain_error/explain_error_agent.py @@ -0,0 +1,233 @@ +# Copyright (c) 2024-2025 Datalayer, Inc. +# +# BSD 3-Clause License + +"""Pydantic AI Explain Error Agent - analyzes and fixes notebook errors.""" + +import logging +from typing import Any + +from pydantic_ai import Agent, RunContext +from pydantic_ai.mcp import MCPServerStreamableHTTP + +from jupyter_ai_agents.tools import create_mcp_server + +logger = logging.getLogger(__name__) + + +SYSTEM_PROMPT = """You are a powerful coding assistant specialized in debugging Jupyter notebooks. +Your goal is to help users understand coding errors and provide corrections. + +When you receive notebook content and an error: +1. Analyze the error traceback carefully +2. Identify the root cause of the problem +3. Explain the error in clear, concise terms +4. Provide a corrected code cell +5. Add comments to explain what was wrong and how you fixed it + +Important guidelines: +- Use the available Jupyter MCP tools to insert corrected code cells +- Execute the corrected code to verify it works +- Ensure updates to cell indexing when new cells are inserted +- Maintain the logical flow of execution by adjusting cell index as needed +- Be concise but thorough in your explanations + +Available tools through MCP: +- notebook tools for inserting/modifying cells +- kernel tools for executing code +- file system tools if needed + +Your response should: +1. Briefly explain what caused the error +2. Insert a corrected code cell at the appropriate location +3. Execute it to verify the fix works + +IMPORTANT: When you have completed the fix successfully, provide a final text response summarizing what you did WITHOUT making any more tool calls. This signals completion and allows the program to exit properly.""" + + +class ExplainErrorAgentDeps: + """Dependencies for the explain error agent.""" + + def __init__( + self, + notebook_content: str = "", + error_info: dict[str, Any] | None = None, + error_cell_index: int = -1, + ): + """Initialize dependencies. + + Args: + notebook_content: Content of notebook cells leading up to error + error_info: Information about the error (traceback, cell content, etc.) + error_cell_index: Index of the cell where error occurred + """ + self.notebook_content = notebook_content + self.error_info = error_info or {} + self.error_cell_index = error_cell_index + + +def create_explain_error_agent( + model: str, + mcp_server: MCPServerStreamableHTTP, + notebook_content: str = "", + error_info: dict[str, Any] | None = None, + error_cell_index: int = -1, + max_tool_calls: int = 10, +) -> Agent[ExplainErrorAgentDeps, str]: + """ + Create the Explain Error Agent for analyzing and fixing errors. + + Args: + model: Model identifier (e.g., 'anthropic:claude-sonnet-4-0', 'openai:gpt-4o') + mcp_server: MCP server connection to jupyter-mcp-server + notebook_content: Content of notebook cells leading up to the error + error_info: Information about the error + error_cell_index: Index of the cell where error occurred + max_tool_calls: Maximum number of tool calls to make (default: 10) + + Returns: + Configured Pydantic AI agent + """ + # Enhance system prompt with notebook and error context + system_prompt = SYSTEM_PROMPT + + if notebook_content: + system_prompt += f"\n\nNotebook content (cells leading to error):\n{notebook_content}" + + if error_cell_index != -1: + system_prompt += f"\n\nError occurred at cell index: {error_cell_index}" + + # Add reminder to be efficient and complete properly + system_prompt += ( + "\n\nBe efficient: Fix the error in as few steps as possible. " + "When the fix is complete, provide a final summary WITHOUT tool calls to signal completion." + ) + + # Create agent with MCP toolset + agent = Agent( + model, + toolsets=[mcp_server], + deps_type=ExplainErrorAgentDeps, + system_prompt=system_prompt, + ) + + logger.info(f"Created explain error agent with model {model} (max_tool_calls={max_tool_calls})") + + return agent + + +async def run_explain_error_agent( + agent: Agent[ExplainErrorAgentDeps, str], + error_description: str, + notebook_content: str = "", + error_info: dict[str, Any] | None = None, + error_cell_index: int = -1, + notebook_path: str = "", + max_tool_calls: int = 10, + max_requests: int = 3, +) -> str: + """ + Run the explain error agent to analyze and fix an error. + + Args: + agent: The configured explain error agent + error_description: Description of the error (traceback, message, etc.) + notebook_content: Content of notebook cells + error_info: Additional error information + error_cell_index: Index where error occurred + notebook_path: Path to the notebook file + max_tool_calls: Maximum number of tool calls to prevent excessive API usage + max_requests: Maximum number of API requests (default: 3) + + Returns: + Agent's response with explanation and fix + """ + import asyncio + import os + from pydantic_ai import UsageLimitExceeded, UsageLimits + + deps = ExplainErrorAgentDeps( + notebook_content=notebook_content, + error_info=error_info, + error_cell_index=error_cell_index, + ) + + logger.info(f"Running explain error agent for error: {error_description[:50]}... (max_tool_calls={max_tool_calls}, max_requests={max_requests})") + + # Prepend notebook connection instruction if path is provided + if notebook_path: + notebook_name = os.path.splitext(os.path.basename(notebook_path))[0] + + # Prepend instruction to connect to the notebook first + enhanced_description = ( + f"First, use the use_notebook tool to connect to the notebook at path '{notebook_path}' " + f"with notebook_name '{notebook_name}' and mode 'connect'. " + f"Then, analyze and fix this error: {error_description}" + ) + logger.info(f"Enhanced input to connect to notebook: {notebook_path}") + else: + enhanced_description = error_description + + try: + # Create usage limits to prevent excessive API calls + # Use strict limits: fewer requests to avoid rate limiting + usage_limits = UsageLimits( + tool_calls_limit=max_tool_calls, + request_limit=max_requests, # Strict limit to avoid rate limiting + ) + + # Add timeout to prevent hanging on retries + result = await asyncio.wait_for( + agent.run(enhanced_description, deps=deps, usage_limits=usage_limits), + timeout=120.0 # 2 minute timeout + ) + logger.info("Explain error agent completed successfully") + return result.response + except asyncio.TimeoutError: + logger.error("Explain error agent timed out after 120 seconds") + return "Error: Operation timed out. The agent may have hit rate limits or is taking too long." + except UsageLimitExceeded as e: + logger.error(f"Explain error agent hit usage limits: {e}") + return ( + "Error: Reached the configured usage limits.\n" + f"Increase --max-requests (currently {max_requests}) or --max-tool-calls (currently {max_tool_calls}) " + "if your model provider allows more usage." + ) + except UsageLimitExceeded as e: + logger.error(f"Explain error agent hit usage limits: {e}") + return ( + "Error: Reached the configured usage limits.\n" + f"Increase --max-requests (currently {max_requests}) or --max-tool-calls (currently {max_tool_calls}) " + "if your model provider allows more usage." + ) + except Exception as e: + logger.error(f"Error running explain error agent: {e}", exc_info=True) + return f"Error: {str(e)}" + + +def create_explain_error_agent_sync( + base_url: str, + token: str, + model: str, + notebook_content: str = "", + error_info: dict[str, Any] | None = None, + error_cell_index: int = -1, +) -> Agent[ExplainErrorAgentDeps, str]: + """ + Create explain error agent with MCP server connection (synchronous wrapper). + + Args: + base_url: Jupyter server base URL + token: Authentication token + model: Model identifier + notebook_content: Notebook content + error_info: Error information + error_cell_index: Error cell index + + Returns: + Configured agent + """ + mcp_server = create_mcp_server(base_url, token) + return create_explain_error_agent( + model, mcp_server, notebook_content, error_info, error_cell_index + ) diff --git a/jupyter_ai_agents/agents/langchain/base/__init__.py b/jupyter_ai_agents/agents/langchain/base/__init__.py deleted file mode 100644 index 3d32561..0000000 --- a/jupyter_ai_agents/agents/langchain/base/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -# Copyright (c) 2024-2025 Datalayer, Inc. -# -# BSD 3-Clause License - -# Copyright (c) 2023-2024 Datalayer, Inc. -# -# BSD 3-Clause License diff --git a/jupyter_ai_agents/agents/langchain/base/agent_base.py b/jupyter_ai_agents/agents/langchain/base/agent_base.py deleted file mode 100644 index 723ae66..0000000 --- a/jupyter_ai_agents/agents/langchain/base/agent_base.py +++ /dev/null @@ -1,666 +0,0 @@ -# Copyright (c) 2024-2025 Datalayer, Inc. -# -# BSD 3-Clause License - -from __future__ import annotations - -import asyncio -import os - -from datetime import datetime, timezone -from enum import IntEnum -from logging import Logger -from typing import Any, Literal, TypedDict, cast - -from pycrdt import Awareness, Map - -from jupyter_nbmodel_client.client import REQUEST_TIMEOUT, NbModelClient -from jupyter_nbmodel_client._version import VERSION - -""" - -This module provides a base class AI agent to interact with collaborative Jupyter notebook. - -The following json schema describes the data model used in cells and notebook metadata to communicate between user clients and an Jupyter AI Agent. - -```json -{ - "datalayer": { - "type": "object", - "properties": { - "ai": { - "type": "object", - "properties": { - "prompts": { - "type": "array", - "items": { - "type": "object", - "properties": { - "id": { - "title": "Prompt unique identifier", - "type": "string" - }, - "prompt": { - "title": "User prompt", - "type": "string" - }, - "username": { - "title": "Unique identifier of the user making the prompt.", - "type": "string" - }, - "timestamp": { - "title": "Number of milliseconds elapsed since the epoch; i.e. January 1st, 1970 at midnight UTC.", - "type": "integer" - } - }, - "required": ["id", "prompt"] - } - }, - "messages": { - "type": "array", - "items": { - "type": "object", - "properties": { - "parent_id": { - "title": "Prompt unique identifier", - "type": "string" - }, - "message": { - "title": "AI reply", - "type": "string" - }, - "type": { - "title": "Type message", - "enum": [0, 1, 2] - }, - "timestamp": { - "title": "Number of milliseconds elapsed since the epoch; i.e. January 1st, 1970 at midnight UTC.", - "type": "integer" - } - }, - "required": ["id", "prompt"] - } - } - } - } - } - } -} -``` -""" - -class AIMessageType(IntEnum): - """Type of AI agent message.""" - ERROR = -1 - """Error message.""" - ACKNOWLEDGE = 0 - """Prompt is being processed.""" - REPLY = 1 - """AI reply.""" - - -class PeerChanges(TypedDict): - added: list[int] - """List of peer ids added.""" - removed: list[int] - """List of peer ids removed.""" - updated: list[int] - """List of peer ids updated.""" - - -class PeerEvent(TypedDict): - # Only trigger on change to avoid high callback pressure - # type: Literal["change", "update"] - # """Event type; "change" if the peer state changes, "update" if the peer state is updated even if unchanged.""" - changes: PeerChanges - """Peer changes.""" - origin: Literal["local"] | str - """Event origin; "local" if emitted by itself, the peer id otherwise.""" - -""" -def _debug_print_changes(part: str, changes: Any) -> None: - print(f"{part}") - def print_change(changes): - if isinstance(changes, MapEvent): - print(f"{type(changes.target)} {changes.target} {changes.keys} {changes.path}") - elif isinstance(changes, ArrayEvent): - print(f"{type(changes.target)} {changes.target} {changes.delta} {changes.path}") - else: - print(changes) - if isinstance(changes, list): - for c in changes: - print_change(c) - else: - print_change(changes) -""" - -def timestamp() -> int: - """Return the current timestamp in milliseconds since epoch.""" - return int(datetime.now(timezone.utc).timestamp() * 1000.0) - - -class NbModelBaseAgent(NbModelClient): - """Base class to react to user prompt and notebook changes based on CRDT changes. - - Notes: - - Agents are expected to extend this base class and override either - - method:`async _on_user_prompt(self, cell_id: str, prompt: str, username: str | None = None) -> str | None`: - Callback on user prompt, it may return an AI reply and must raise an error in case of failure - - method:`async _on_cell_source_changes(self, cell_id: str, new_source: str, old_source: str, username: str | None = None) -> None`: - Callback on cell source changes, it must raise an error in case of failure - - Agents can sent transient messages to users through the method:`async notify(self, message: str, cell_id: str = "", message_type: AIMessageType = AIMessageType.ACKNOWLEDGE) -> None` - - Args: - ws_url: Endpoint to connect to the collaborative Jupyter notebook. - path: [optional] Notebook path relative to the server root directory; default None - username: [optional] Client user name; default to environment variable USER - timeout: [optional] Request timeout in seconds; default to environment variable REQUEST_TIMEOUT - log: [optional] Custom logger; default local logger - ws_max_body_size: [optional] Maximum WebSocket body size in bytes; default 16MB - close_timeout: [optional] Timeout for propagating final changes on close; default to timeout value - - Examples: - - When connection to a Jupyter notebook server, you can leverage the get_jupyter_notebook_websocket_url - helper: - - >>> from jupyter_nbmodel_client import NbModelClient, get_jupyter_notebook_websocket_url - >>> client = NbModelClient( - >>> get_jupyter_notebook_websocket_url( - >>> "http://localhost:8888", - >>> "path/to/notebook.ipynb", - >>> "your-server-token" - >>> ) - >>> ) - """ - - user_agent: str = f"Datalayer-BaseNbAgent/{VERSION}" - """User agent used to identify the nbmodel client type in the awareness state.""" - - def __init__( - self, - websocket_url: str, - path: str | None = None, - username: str = os.environ.get("USER", "username"), - timeout: float = REQUEST_TIMEOUT, - log: Logger | None = None, - ws_max_body_size: int | None = None, - close_timeout: float | None = None, - ) -> None: - super().__init__(websocket_url, path, username, timeout, log, ws_max_body_size, close_timeout) - self._doc_events: asyncio.Queue[dict] = asyncio.Queue() - self._peer_events: asyncio.Queue[PeerEvent] = asyncio.Queue() - - async def run(self) -> None: - self._doc.observe(self._on_notebook_changes) - awareness_callback = cast(Awareness, self._doc.awareness).observe(self._on_peer_changes) - doc_events_worker: asyncio.Task | None = None - peer_events_worker: asyncio.Task | None = None - try: - doc_events_worker = asyncio.create_task(self._process_doc_events()) - peer_events_worker = asyncio.create_task(self._process_peer_events()) - await super().run() - finally: - self._log.info("Stop the agent.") - cast(Awareness, self._doc.awareness).unobserve(awareness_callback) - if doc_events_worker and not doc_events_worker.done(): - if not self._doc_events.empty(): - await self._doc_events.join() - if doc_events_worker.cancel(): - await asyncio.wait([doc_events_worker]) - else: - try: - while True: - self._doc_events.get_nowait() - except asyncio.QueueEmpty: - ... - - if peer_events_worker and not peer_events_worker.done(): - if not self._peer_events.empty(): - await self._peer_events.join() - if peer_events_worker.cancel(): - await asyncio.wait([peer_events_worker]) - else: - try: - while True: - self._peer_events.get_nowait() - except asyncio.QueueEmpty: - ... - - async def __handle_cell_source_changes( - self, - cell_id: str, - new_source: str, - old_source: str, - username: str | None = None, - ) -> None: - self._log.info("Process user [%s] cell [%s] source changes.", username, cell_id) - - # # Acknowledge through awareness - # await self.notify( - # "Analyzing source changes…", - # cell_id=cell_id, - # ) - try: - await self._on_cell_source_changes(cell_id, new_source, old_source, username) - except asyncio.CancelledError: - raise - except BaseException as e: - error_message = f"Error while analyzing cell source: {e!s}" - self._log.error(error_message) - # await self.notify(error_message, cell_id=cell_id, message_type=AIMessageType.ERROR) - else: - self._log.info("AI processed successfully cell [%s] source changes.", cell_id) - # await self.notify( - # "Source changes analyzed.", - # cell_id=cell_id, - # ) - - async def __handle_user_prompt( - self, - cell_id: str, - prompt_id: str, - prompt: str, - username: str | None = None, - timestamp: int | None = None, - ) -> None: - self._log.info("Received user [%s] prompt [%s].", username, prompt_id) - self._log.debug( - "Prompt: timestamp [%s] / cell_id [%s] / prompt [%s]", - timestamp, - cell_id, - prompt[:20], - ) - - # Acknowledge - await self.save_ai_message( - AIMessageType.ACKNOWLEDGE, - "Requesting AI…", - cell_id=cell_id, - parent_id=prompt_id, - ) - try: - reply = await self._on_user_prompt(cell_id, prompt_id, prompt, username, timestamp) - except asyncio.CancelledError: - await self.save_ai_message( - AIMessageType.ERROR, - "Prompt request cancelled.", - cell_id=cell_id, - parent_id=prompt_id, - ) - raise - except BaseException as e: - error_message = "Error while processing user prompt" - self._log.error(error_message + " [%s].", prompt_id, exc_info=e) - await self.save_ai_message( - AIMessageType.ERROR, - error_message + f": {e!s}", - cell_id=cell_id, - parent_id=prompt_id, - ) - else: - self._log.info("AI replied successfully to prompt [%s]: [%s]", prompt_id, reply) - if reply is not None: - await self.save_ai_message( - AIMessageType.REPLY, reply, cell_id=cell_id, parent_id=prompt_id - ) - else: - await self.save_ai_message( - AIMessageType.ACKNOWLEDGE, - "AI has successfully processed the prompt.", - cell_id=cell_id, - parent_id=prompt_id, - ) - - - async def _process_doc_events(self) -> None: - self._log.debug("Starting listening on document [%s] changes…", self.path) - while True: - try: - event = await self._doc_events.get() - event_type = event.pop("type") - if event_type == "user": - await self.__handle_user_prompt(**event) - if event_type == "source": - await self.__handle_cell_source_changes(**event) - except asyncio.CancelledError: - raise - except BaseException as e: - self._log.error("Error while processing document events: %s", exc_info=e) - else: - # Sleep to get a chance to propagate changes through the websocket - await asyncio.sleep(0) - - - def _on_notebook_changes( - self, - part: Literal["state"] | Literal["meta"] | Literal["cells"] | str, - all_changes: Any, - ) -> None: - # _debug_print_changes(part, all_changes) - - if part == "cells": - for changes in all_changes: - transaction_origin = changes.transaction.origin - if transaction_origin == self._changes_origin: - continue - else: - self._log.debug( - "Document changes from origin [%s] != agent origin [%s].", - transaction_origin, - self._changes_origin, - ) - path_length = len(changes.path) - if path_length == 0: - # Change is on the cell list - for delta in changes.delta: - if "insert" in delta: - # New cells got added - for cell in delta["insert"]: - if "metadata" in cell: - new_metadata = cell["metadata"] - datalayer_ia = new_metadata.get("datalayer", {}).get("ai", {}) - prompts = datalayer_ia.get("prompts", []) - prompt_ids = {prompt["id"] for prompt in prompts} - new_prompts = prompt_ids.difference( - message["parent_id"] - for message in datalayer_ia.get("messages", []) - ) - if new_prompts: - for prompt in filter( - lambda p: p.get("id") in new_prompts, - prompts, - ): - self._doc_events.put_nowait( - { - "type": "user", - "cell_id": cell["id"], - "prompt_id": prompt["id"], - "prompt": prompt["prompt"], - "username": prompt.get("user"), - "timestamp": prompt.get("timestamp"), - } - ) - if "source" in cell: - self._doc_events.put_nowait( - { - "type": "source", - "cell_id": cell["id"], - "new_source": cell["source"].to_py(), - "old_source": "", - } - ) - elif path_length == 1: - # Change is on one cell - for key, change in changes.keys.items(): - if key == "source": - if change["action"] == "add": - self._doc_events.put_nowait( - { - "type": "source", - "cell_id": changes.target["id"], - "new_source": change["newValue"], - "old_source": change.get("oldValue", ""), - } - ) - elif change["action"] == "update": - self._doc_events.put_nowait( - { - "type": "source", - "cell_id": changes.target["id"], - "new_source": change["newValue"], - "old_source": change["oldValue"], - } - ) - elif change["action"] == "delete": - self._doc_events.put_nowait( - { - "type": "source", - "cell_id": changes.target["id"], - "new_source": change.get("newValue", ""), - "old_source": change["oldValue"], - } - ) - elif key == "metadata": - new_metadata = change.get("newValue", {}) - datalayer_ia = new_metadata.get("datalayer", {}).get("ai", {}) - prompts = datalayer_ia.get("prompts", []) - prompt_ids = {prompt["id"] for prompt in prompts} - new_prompts = prompt_ids.difference( - message["parent_id"] for message in datalayer_ia.get("messages", []) - ) - if new_prompts and change["action"] in {"add", "update"}: - for prompt in filter(lambda p: p.get("id") in new_prompts, prompts): - self._doc_events.put_nowait( - { - "type": "user", - "cell_id": changes.target["id"], - "prompt_id": prompt["id"], - "prompt": prompt["prompt"], - "username": prompt.get("user"), - "timestamp": prompt.get("timestamp"), - } - ) - # elif change["action"] == "delete": - # ... - # elif key == "outputs": - # # TODO - # ... - elif ( - path_length == 2 - and isinstance(changes.path[0], int) - and changes.path[1] == "metadata" - ): - # Change in cell metadata - for key, change in changes.keys.items(): - if key == "datalayer": - new_metadata = change.get("newValue", {}) - datalayer_ia = new_metadata.get("ai", {}) - prompts = datalayer_ia.get("prompts") - prompt_ids = {prompt["id"] for prompt in prompts} - new_prompts = prompt_ids.difference( - message["parent_id"] for message in datalayer_ia.get("messages", []) - ) - if new_prompts and change["action"] in {"add", "update"}: - for prompt in filter(lambda p: p.get("id") in new_prompts, prompts): - self._doc_events.put_nowait( - { - "type": "user", - "cell_id": self._doc.ycells[changes.path[0]]["id"], - "prompt_id": prompt["id"], - "prompt": prompt["prompt"], - "username": prompt.get("user"), - "timestamp": prompt.get("timestamp"), - } - ) - # elif change["action"] == "delete": - # ... - # elif part == "meta": - # # FIXME handle notebook metadata - - def _reset_y_model(self) -> None: - try: - self._doc.unobserve() - except AttributeError: - pass - finally: - super()._reset_y_model() - - async def _on_user_prompt( - self, - cell_id: str, - prompt_id: str, - prompt: str, - username: str | None = None, - timestamp: int | None = None, - ) -> str | None: - """Callback on user prompt. - - Args: - cell_id: Cell ID on which an user prompt is set; empty if the user prompt is at the notebook level. - prompt_id: Prompt unique ID - prompt: User prompt - username: User name - timestamp: Prompt creation timestamp - - Returns: - Optional agent reply to display to the user. - """ - username = username or self._username - self._log.debug("New AI prompt sets by user [%s] in [%s]: [%s].", username, cell_id, prompt) - - async def _on_cell_source_changes( - self, - cell_id: str, - new_source: str, - old_source: str, - username: str | None = None, - ) -> None: - username = username or self._username - self._log.debug("New cell source sets by user [%s] in [%s].", username, cell_id) - - # async def _on_cell_outputs_changes(self, *args) -> None: - # print(args) - - async def _process_peer_events(self) -> None: - while True: - try: - event = await self._peer_events.get() - await self._on_peer_event(event) - except asyncio.CancelledError: - raise - except BaseException as e: - self._log.error("Error while processing peer events: %s", exc_info=e) - else: - # Sleep to get a chance to propagate changes through the websocket - await asyncio.sleep(0) - - def _on_peer_changes(self, event_type: str, changes: tuple[dict, Any]) -> None: - if event_type != "udpate": - self._peer_events.put_nowait( - cast(PeerEvent, {"changes": changes[0], "origin": changes[1]}) - ) - - async def _on_peer_event(self, event: PeerEvent) -> None: - """Callback on peer awareness events.""" - self._log.debug( - "New event from peer [%s]: %s", event["origin"], event["changes"] - ) - - def get_cell(self, cell_id: str) -> Map | None: - """Find the cell with the given ID. - - If the cell cannot be found it will return ``None``. - - Args: - cell_id: str - Returns: - Cell or None - """ - for cell in self._doc.ycells: - if cell["id"] == cell_id: - return cast(Map, cell) - - return None - - def get_cell_index(self, cell_id: str) -> int: - """Find the cell with the given ID. - - If the cell cannot be found it will return ``-1``. - - Args: - cell_id: str - Returns: - Cell index or -1 - """ - for index, cell in enumerate(self._doc.ycells): - if cell["id"] == cell_id: - return index - - return -1 - - async def save_ai_message( - self, - message_type: AIMessageType, - message: str, - cell_id: str = "", - parent_id: str | None = None, - ) -> None: - """Update the document. - - If a message with the same ``parent_id`` already exists, it will be - overwritten. - - Args: - message_type: Type of message to insert in the document - message: Message to insert - cell_id: Cell targeted by the update; if empty, the notebook is the target - parent_id: Parent message id - """ - message_dict = { - "parent_id": parent_id, - "message": message, - "type": int(message_type), - "timestamp": timestamp(), - } - - def set_message(metadata: dict, message: dict): - dla = metadata.get("datalayer") or {"ai": {"prompts": [], "messages": []}} - dlai = dla.get("ai", {"prompts": [], "messages": []}) - dlmsg = dlai.get("messages", []) - - messages = list( - filter( - lambda m: not m.get("parent_id") or m["parent_id"] != parent_id, - dlmsg, - ) - ) - messages.append(message) - dlai["messages"] = messages - dla["ai"] = dlai - if "datalayer" in metadata: - del metadata["datalayer"] # FIXME upstream - update of key is not possible 😱 - metadata["datalayer"] = dla.copy() - - if cell_id: - cell = self.get_cell(cell_id) - if not cell: - raise ValueError(f"Cell [{cell_id}] not found.") - with self._doc._ydoc.transaction(origin=self._changes_origin): - if "metadata" not in cell: - cell["metadata"] = Map({"datalayer": {"ai": {"prompts": [], "messages": []}}}) - set_message(cell["metadata"], message_dict) - self._log.debug("Add ai message in cell [%s] metadata: [%s].", cell_id, message_dict) - - else: - notebook_metadata = self._doc._ymeta["metadata"] - with self._doc._ydoc.transaction(origin=self._changes_origin): - set_message(notebook_metadata, message_dict) - self._log.debug("Add ai message in notebook metadata: [%s].", cell_id, message_dict) - - # Sleep to get a chance to propagate the changes through the websocket - await asyncio.sleep(0) - - async def notify( - self, - message: str, - cell_id: str = "", - message_type: AIMessageType = AIMessageType.ACKNOWLEDGE, - ) -> None: - """Send a transient message to users. - - Args: - message: Notification message - cell_id: Cell targeted by the notification; if empty the notebook is the target - """ - self.set_local_state_field( - "notification", - { - "message": message, - "message_type": message_type, - "timestamp": timestamp(), - "cell_id": cell_id, - }, - ) - # Sleep to get a chance to propagate the changes through the websocket - await asyncio.sleep(0) diff --git a/jupyter_ai_agents/agents/langchain/base/agent_base_runtime.py b/jupyter_ai_agents/agents/langchain/base/agent_base_runtime.py deleted file mode 100644 index 7967af9..0000000 --- a/jupyter_ai_agents/agents/langchain/base/agent_base_runtime.py +++ /dev/null @@ -1,55 +0,0 @@ -# Copyright (c) 2024-2025 Datalayer, Inc. -# -# BSD 3-Clause License - -# Copyright (c) 2023-2024 Datalayer, Inc. -# -# Datalayer License - -from __future__ import annotations - -import logging -import os -from logging import Logger - -from jupyter_kernel_client import KernelClient -from jupyter_ai_agents.agents.langchain.base.agent_base import NbModelBaseAgent -from jupyter_nbmodel_client.constants import REQUEST_TIMEOUT - - -logger = logging.getLogger(__name__) - - -class NbModelBaseRuntimeAgent(NbModelBaseAgent): - """A base nbmodel agent connected to a runtime client.""" - - def __init__( - self, - websocket_url: str, - path: str | None = None, - runtime_client: KernelClient | None = None, - username: str = os.environ.get("USER", "username"), - timeout: float = REQUEST_TIMEOUT, - log: Logger | None = None, - ) -> None: - super().__init__(websocket_url, path, username, timeout, log) - self._runtime_client: KernelClient | None = runtime_client - - - @property - def runtime_client(self) -> KernelClient | None: - """Runtime client""" - return self._runtime_client - - - @runtime_client.setter - def runtime_client(self, client: KernelClient) -> None: - if self._runtime_client: - self._runtime_client.stop() - self._runtime_client = client - - - async def stop(self) -> None: - await super().stop() - if self._runtime_client: - self._runtime_client.stop() diff --git a/jupyter_ai_agents/agents/langchain/explain_error_agent.py b/jupyter_ai_agents/agents/langchain/explain_error_agent.py deleted file mode 100644 index db3e181..0000000 --- a/jupyter_ai_agents/agents/langchain/explain_error_agent.py +++ /dev/null @@ -1,91 +0,0 @@ -# Copyright (c) 2024-2025 Datalayer, Inc. -# -# BSD 3-Clause License - -import logging - -from langchain.agents import AgentExecutor, tool - -from jupyter_kernel_client import KernelClient -from jupyter_nbmodel_client import NbModelClient - -from jupyter_ai_agents.agents.langchain.providers import create_langchain_agent -from jupyter_ai_agents.utils import ( - retrieve_cells_content_error, - retrieve_cells_content_until_first_error, - insert_execute_code_cell_tool, -) - - -logger = logging.getLogger(__name__) - - -SYSTEM_PROMPT = """You are a powerful coding assistant. -Your goal is to help the user understand the coding error in a notebook and provide a correction. -You will receive the notebook content and error and you will need to insert code cells with the correction and comments to explain the error in a concise way. -Ensure updates to cell indexing when new cells are inserted. Maintain the logical flow of execution by adjusting cell index as needed. -""" - - -def _create_agent( - notebook: NbModelClient, - kernel: KernelClient, - model_provider: str, - model_name: str, - current_cell_index: int, -) -> AgentExecutor: - """Explain and correct an error in a notebook based on the prior cells.""" - - @tool - def insert_execute_code_cell(cell_index: int, cell_content: str) -> str: - """Add a Python code cell to the notebook at the given index with a content and execute it.""" - insert_execute_code_cell_tool(notebook, kernel, cell_content, cell_index) - return "Code cell added and executed." - - tools = [insert_execute_code_cell] - - if current_cell_index != -1: - cells_content_until_error, error = retrieve_cells_content_error( - notebook, current_cell_index - ) - - system_prompt_final = f""" - {SYSTEM_PROMPT} - - Notebook content: {cells_content_until_error} - """ - agent_input = f"Error: {str(error)}" - - else: - cells_content_until_first_error, first_error = retrieve_cells_content_until_first_error( - notebook - ) - - system_prompt_final = f""" - {SYSTEM_PROMPT} - - Notebook content: {cells_content_until_first_error} - """ - agent_input = f"Error: {str(first_error)}" - - logger.debug("Prompt with content: %s", system_prompt_final) - logger.debug("Input: %s", agent_input) - - ai_agent = create_langchain_agent(model_provider, model_name, system_prompt_final, tools) - - return (ai_agent, agent_input) - - -async def explain_error( - notebook: NbModelClient, - kernel: KernelClient, - model_provider: str, - model_name: str, - current_cell_index: int, -) -> list: - """Explain and correct an error in a notebook based on the prior cells.""" - (agent, agent_input) = _create_agent(notebook, kernel, model_provider, model_name, current_cell_index) - replies = [] - async for reply in agent.astream({"input": agent_input}): - replies.append(reply) - return replies diff --git a/jupyter_ai_agents/agents/langchain/manager/__init__.py b/jupyter_ai_agents/agents/langchain/manager/__init__.py deleted file mode 100644 index 3d32561..0000000 --- a/jupyter_ai_agents/agents/langchain/manager/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -# Copyright (c) 2024-2025 Datalayer, Inc. -# -# BSD 3-Clause License - -# Copyright (c) 2023-2024 Datalayer, Inc. -# -# BSD 3-Clause License diff --git a/jupyter_ai_agents/agents/langchain/manager/agents_manager.py b/jupyter_ai_agents/agents/langchain/manager/agents_manager.py deleted file mode 100644 index 2323bfe..0000000 --- a/jupyter_ai_agents/agents/langchain/manager/agents_manager.py +++ /dev/null @@ -1,122 +0,0 @@ -# Copyright (c) 2024-2025 Datalayer, Inc. -# -# BSD 3-Clause License - -from __future__ import annotations - -import asyncio -import logging -from collections import Counter - -from jupyter_ai_agents.agents.langchain.base.agent_base_runtime import NbModelBaseRuntimeAgent - - -logger = logging.getLogger(__name__) - - -SPACER_AGENT = "DatalayerSpacer" - - -"""Delay in seconds before stopping an agent.""" -DELAY_FOR_STOPPING_AGENT = 20 * 60 - - -def _stop_agent(agent: NbModelBaseRuntimeAgent, room: str) -> None: - try: - if agent.runtime_client is not None: - agent.runtime_client.stop() - agent.stop() - except BaseException as e: - logger.error("Failed to stop AI Agent for room [%s].", room, exc_info=e) - - -class AIAgentsManager: - """AI Agents manager.""" - - def __init__(self) -> None: - self._agents: dict[str, NbModelBaseRuntimeAgent] = {} - self._background_tasks: list[asyncio.Task] = [] - self._agents_to_stop: set[str] = set() - self._to_stop_counter: Counter[str] = Counter() - # A usefull task will be set when the first agent is added. - self._stop_task: asyncio.Task = asyncio.create_task(asyncio.sleep(0)) - - def __contains__(self, key: str) -> bool: - return key in self._agents - - def __getitem__(self, key: str) -> NbModelBaseRuntimeAgent: - return self._agents[key] - - async def _stop_lonely_agents(self) -> None: - """Periodically check if an agent as connected peer. - - If it the only peer is the spacer server, kill the agent after some delay. - """ - while True: - await asyncio.sleep(DELAY_FOR_STOPPING_AGENT * 0.25) - for key, agent in self._agents.items(): - peers = agent.get_connected_peers() - if len(peers) == 1: - peer_state = agent.get_peer_state(peers[0]) or {} - if (peer_state.get("user", {}).get("agent", "").startswith(SPACER_AGENT)): - self._agents_to_stop.add(key) - self._to_stop_counter.update([key]) - to_stop = [] - for key, count in self._to_stop_counter.most_common(): - if count < 4: - break - self._agents_to_stop.remove(key) - to_stop.append(key) - await asyncio.shield( - asyncio.gather( - *( - _stop_agent(self._agents.pop(room), room) - for room in to_stop - if room in self._agents - ) - ) - ) - for key in to_stop: - self._to_stop_counter.pop(key, None) - logger.info("AI Agent for room [%s] stopped.", key) - - async def stop_all(self) -> None: - """Stop all background tasks and reset the state.""" - if self._stop_task.cancel(): - await asyncio.wait([self._stop_task]) - all_tasks = asyncio.gather(*self._background_tasks) - if all_tasks.cancel(): - await asyncio.wait([all_tasks]) - await asyncio.shield( - asyncio.gather(*(_stop_agent(agent, room) for room, agent in self._agents.items())) - ) - self._agents.clear() - self._agents_to_stop.clear() - self._to_stop_counter.clear() - - def get_user_agents(self, user: str) -> list[str]: - return [k for k, a in self._agents.items() if a._username == user] - - def register_ai_agent(self, key: str, agent: NbModelBaseRuntimeAgent): - self._agents[key] = agent - - async def track_agent(self, key: str, agent: NbModelBaseRuntimeAgent) -> None: - """Add an agent and start it.""" - if self._stop_task.done(): - self._stop_task = asyncio.create_task(self._stop_lonely_agents()) - self.register_ai_agent(key, agent) - start = asyncio.create_task(await agent.start()) - self._background_tasks.append(start) - start.add_done_callback(lambda task: self._background_tasks.remove(task)) - if agent.runtime_client is not None: - agent.runtime_client.start() - - def forget_agent(self, key: str) -> None: - if key in self: - logger.info("Removing AI Agent in room [%s].", key) - agent = self._agents.pop(key) - if key in self._agents_to_stop: - self._agents_to_stop.remove(key) - if key in self._to_stop_counter: - self._to_stop_counter.pop(key) - _stop_agent(agent, key) diff --git a/jupyter_ai_agents/agents/langchain/models.py b/jupyter_ai_agents/agents/langchain/models.py deleted file mode 100644 index 162a0c7..0000000 --- a/jupyter_ai_agents/agents/langchain/models.py +++ /dev/null @@ -1,21 +0,0 @@ -# Copyright (c) 2024-2025 Datalayer, Inc. -# -# BSD 3-Clause License - -from __future__ import annotations - -from typing import Optional - -from pydantic import BaseModel - - -class RuntimeModel(BaseModel): - ingress: Optional[str] = None - token: Optional[str] = None - kernel_id: Optional[str] = None - jupyter_pod_name: Optional[str] = None - - -class NbModelAgentRequestModel(BaseModel): - room_id: Optional[str] = None - runtime: Optional[RuntimeModel] = None diff --git a/jupyter_ai_agents/agents/langchain/prompt_agent.py b/jupyter_ai_agents/agents/langchain/prompt_agent.py deleted file mode 100644 index d9eeb16..0000000 --- a/jupyter_ai_agents/agents/langchain/prompt_agent.py +++ /dev/null @@ -1,144 +0,0 @@ -# Copyright (c) 2024-2025 Datalayer, Inc. -# -# BSD 3-Clause License - -from langchain.agents import AgentExecutor, tool - -from jupyter_kernel_client import KernelClient -from jupyter_nbmodel_client import NbModelClient - -from jupyter_ai_agents.agents.langchain.base.agent_base_runtime import NbModelBaseRuntimeAgent -from jupyter_ai_agents.agents.langchain.providers import create_langchain_agent -from jupyter_ai_agents.utils import ( - insert_execute_code_cell_tool, - insert_markdown_cell_tool, - retrieve_cells_content, -) - - -SYSTEM_PROMPT = """You are a powerful coding assistant. -Create and execute code in a notebook based on user instructions. -Add markdown cells to explain the code and structure the notebook clearly. -Assume that no packages are installed in the notebook, so install them using !pip install. -Ensure updates to cell indexing when new cells are inserted. Maintain the logical flow of execution by adjusting cell index as needed. -""" - - -def _create_agent( - notebook: NbModelClient, - kernel: KernelClient, - model_provider: str, - model_name: str, - full_context: bool, - current_cell_index: int, -) -> AgentExecutor: - """From a given instruction, code and markdown cells are added to a notebook.""" - - @tool - def insert_execute_code_cell(cell_index: int, cell_content: str) -> str: - """Add a Python code cell to the notebook at the given index with a content and execute it.""" - insert_execute_code_cell_tool(notebook, kernel, cell_content, cell_index) - return "Code cell added and executed." - - @tool - def insert_markdown_cell(cell_index: int, cell_content: str) -> str: - """Add a Markdown cell to the notebook at the given index with a content.""" - insert_markdown_cell_tool(notebook, cell_content, cell_index) - return "Markdown cell added." - - tools = [ - insert_execute_code_cell, - insert_markdown_cell, - ] - - if full_context: - system_prompt_enriched = f""" - {SYSTEM_PROMPT} - - Notebook content: {retrieve_cells_content(notebook)} - """ - else: - system_prompt_enriched = SYSTEM_PROMPT - - if current_cell_index != -1: - system_prompt_final = f""" - {system_prompt_enriched} - - Cell index on which the user instruction was given: {current_cell_index} - """ - else: - system_prompt_final = system_prompt_enriched - - return create_langchain_agent(model_provider, model_name, system_prompt_final, tools) - - -async def prompt( - notebook: NbModelClient, - kernel: KernelClient, - input: str, - model_provider: str, - model_name: str, - full_context: bool, - current_cell_index: int, -) -> list: - agent = _create_agent( - notebook, kernel, model_provider, model_name, full_context, current_cell_index - ) - replies = [] - async for reply in agent.astream({"input": input}): - replies.append(reply) - return replies - - -class PromptAgent(NbModelBaseRuntimeAgent): - """AI Agent replying to user prompt.""" - - model_provider = "azure-openai" - model_name = "gpt-40-mini" - full_context = False - - async def _on_user_prompt( - self, - cell_id: str, - prompt_id: str, - prompt: str, - username: str | None = None, - timestamp: int | None = None, - **kwargs, - ) -> str | None: - """Callback on user prompt. - - Args: - cell_id: Cell ID on which an user prompt is set; empty if the user prompt is at the notebook level. - prompt_id: Prompt unique ID - prompt: User prompt - username: User name - timestamp: Prompt creation timestamp - - Returns: - Optional agent reply to display to the user. - """ - document_client = self - runtime_client = self.runtime_client - current_cell_index = self.get_cell_index(cell_id) - agent_executor = _create_agent( - document_client, - runtime_client, - self.model_provider, - self.model_name, - self.full_context, - current_cell_index, - ) - output = None - try: - await self.notify("Thinking…", cell_id=cell_id) - async for reply in agent_executor.astream({"input": prompt}): - output = reply.get("output", "") - if not output: - output = reply["messages"][-1].content - self._log.debug( - "Got a reply for prompt [%s]: [%s].", prompt_id, (output or "")[:30] - ) - finally: - await self.notify("Done", cell_id=cell_id) - return output diff --git a/jupyter_ai_agents/agents/langchain/providers/__init__.py b/jupyter_ai_agents/agents/langchain/providers/__init__.py deleted file mode 100644 index 22300b5..0000000 --- a/jupyter_ai_agents/agents/langchain/providers/__init__.py +++ /dev/null @@ -1,41 +0,0 @@ -# Copyright (c) 2024-2025 Datalayer, Inc. -# -# BSD 3-Clause License - -from langchain.agents import AgentExecutor - -from jupyter_ai_agents.agents.langchain.providers.anthropic import create_anthropic_langchain_agent -from jupyter_ai_agents.agents.langchain.providers.azure_openai import create_azure_openai_langchain_agent -from jupyter_ai_agents.agents.langchain.providers.bedrock import create_bedrock_langchain_agent -from jupyter_ai_agents.agents.langchain.providers.github_copilot import create_github_copilot_langchain_agent -from jupyter_ai_agents.agents.langchain.providers.openai import create_openai_langchain_agent - - -def create_langchain_agent( - model_provider: str, model_name: str, system_prompt_final: str, tools: list -) -> AgentExecutor: - """Create an AI Agent based on the model provider.""" - if model_provider == "azure-openai": - langchain_agent = create_azure_openai_langchain_agent(model_name, system_prompt_final, tools) - elif model_provider == "github-copilot": - langchain_agent = create_github_copilot_langchain_agent(model_name, system_prompt_final, tools) - elif model_provider == "openai": - langchain_agent = create_openai_langchain_agent(model_name, system_prompt_final, tools) - elif model_provider == "anthropic": - langchain_agent = create_anthropic_langchain_agent(model_name, system_prompt_final, tools) - elif model_provider == "bedrock": - langchain_agent = create_bedrock_langchain_agent(model_name, system_prompt_final, tools) - else: - raise ValueError(f"Model provider {model_provider} is not supported.") - return langchain_agent - - -__all__ = [ - "create_bedrock_langchain_agent", - "create_langchain_agent", - 'create_anthropic_langchain_agent', - 'create_azure_openai_langchain_agent', - 'create_github_copilot_langchain_agent', - 'create_langchain_agent', - 'create_openai_langchain_agent', -] diff --git a/jupyter_ai_agents/agents/langchain/providers/anthropic.py b/jupyter_ai_agents/agents/langchain/providers/anthropic.py deleted file mode 100644 index cc131f8..0000000 --- a/jupyter_ai_agents/agents/langchain/providers/anthropic.py +++ /dev/null @@ -1,62 +0,0 @@ -# Copyright (c) 2024-2025 Datalayer, Inc. -# -# BSD 3-Clause License - -from dotenv import load_dotenv -from typing import List - -from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder -from langchain_core.tools import BaseTool -from langchain.agents.agent import AgentExecutor -from langchain.agents.tool_calling_agent.base import create_tool_calling_agent -from langchain_anthropic import ChatAnthropic - - -def create_anthropic_langchain_agent(model_name: str, system_prompt: str, tools: List[BaseTool]) -> AgentExecutor: - """Create an agent from a set of tools using Anthropic's Claude API. - - Args: - model_name: The name of the Claude model to use (e.g., "claude-3-haiku-20240307") - system_prompt: The system prompt to use for the agent - tools: A list of tools for the agent to use - - Returns: - An agent executor that can use tools via Claude - """ - - load_dotenv() - - # Create the Anthropic LLM - llm = ChatAnthropic(model=model_name) - - # Create a prompt template for the agent with enhanced instructions - enhanced_system_prompt = f""" -{system_prompt} - -When you use tools, please include the results in your response to the user. -Be sure to always provide a text response, even if it's just to acknowledge the tool's output. -After using a tool, explain what the result means in a clear and helpful way. -""" - - # Create prompt template - prompt = ChatPromptTemplate.from_messages( - [ - ("system", enhanced_system_prompt), - ("user", "{input}"), - MessagesPlaceholder(variable_name="agent_scratchpad"), - ]) - - # Create a tool-calling agent - agent = create_tool_calling_agent(llm, tools, prompt) - - # Create an agent executor with output parsing - agent_executor = AgentExecutor( - name="AnthropicToolAgent", - agent=agent, - tools=tools, - verbose=True, - handle_parsing_errors=True, - return_intermediate_steps=True # Include intermediate steps in the output - ) - - return agent_executor diff --git a/jupyter_ai_agents/agents/langchain/providers/azure_openai.py b/jupyter_ai_agents/agents/langchain/providers/azure_openai.py deleted file mode 100644 index 39430ab..0000000 --- a/jupyter_ai_agents/agents/langchain/providers/azure_openai.py +++ /dev/null @@ -1,38 +0,0 @@ -# Copyright (c) 2024-2025 Datalayer, Inc. -# -# BSD 3-Clause License - -from dotenv import load_dotenv - -from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder -from langchain.agents import AgentExecutor, create_openai_tools_agent -from langchain_openai import AzureChatOpenAI - - -def create_azure_openai_langchain_agent(model_name: str, system_prompt: str, tools: list) -> AgentExecutor: - """Create an agent from a set of tools and an Azure deployment""" - - load_dotenv() - - llm = AzureChatOpenAI(azure_deployment=model_name) - - prompt = ChatPromptTemplate.from_messages( - [ - ("system", system_prompt), - ("user", "{input}"), - MessagesPlaceholder(variable_name="agent_scratchpad"), - ] - ) - - # Create the agent using LangChain's built-in tool handling - agent = create_openai_tools_agent(llm, tools, prompt) - - agent_executor = AgentExecutor( - name="NotebookPromptAgent", - agent=agent, - tools=tools, - verbose=True, - handle_parsing_errors=True, - ) - - return agent_executor diff --git a/jupyter_ai_agents/agents/langchain/providers/bedrock.py b/jupyter_ai_agents/agents/langchain/providers/bedrock.py deleted file mode 100644 index 4428b61..0000000 --- a/jupyter_ai_agents/agents/langchain/providers/bedrock.py +++ /dev/null @@ -1,62 +0,0 @@ -# Copyright (c) 2024-2025 Datalayer, Inc. -# -# BSD 3-Clause License - -from dotenv import load_dotenv -from typing import List - -from langchain_core.tools import BaseTool -from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder -from langchain.agents.tool_calling_agent.base import create_tool_calling_agent -from langchain.agents.agent import AgentExecutor -from langchain_aws import ChatBedrockConverse - - -def create_bedrock_langchain_agent(model_name: str, system_prompt: str, tools: List[BaseTool]) -> AgentExecutor: - """Create an agent from a set of tools using Anthropic's Claude API. - - Args: - model_name: The name of the Claude model to use (e.g., "claude-3-haiku-20240307") - system_prompt: The system prompt to use for the agent - tools: A list of tools for the agent to use - - Returns: - An agent executor that can use tools via Claude - """ - - load_dotenv() - - # Create the Anthropic LLM - llm = ChatBedrockConverse(model_id=model_name) - - # Create a prompt template for the agent with enhanced instructions - enhanced_system_prompt = f""" -{system_prompt} - -When you use tools, please include the results in your response to the user. -Be sure to always provide a text response, even if it's just to acknowledge the tool's output. -After using a tool, explain what the result means in a clear and helpful way. -""" - - # Create prompt template - prompt = ChatPromptTemplate.from_messages( - [ - ("system", enhanced_system_prompt), - ("user", "{input}"), - MessagesPlaceholder(variable_name="agent_scratchpad"), - ]) - - # Create a tool-calling agent - agent = create_tool_calling_agent(llm, tools, prompt) - - # Create an agent executor with output parsing - agent_executor = AgentExecutor( - name="AnthropicToolAgent", - agent=agent, - tools=tools, - verbose=True, - handle_parsing_errors=True, - return_intermediate_steps=True # Include intermediate steps in the output - ) - - return agent_executor diff --git a/jupyter_ai_agents/agents/langchain/providers/github_copilot.py b/jupyter_ai_agents/agents/langchain/providers/github_copilot.py deleted file mode 100644 index 36b7b66..0000000 --- a/jupyter_ai_agents/agents/langchain/providers/github_copilot.py +++ /dev/null @@ -1,38 +0,0 @@ -# Copyright (c) 2024-2025 Datalayer, Inc. -# -# BSD 3-Clause License - -from dotenv import load_dotenv - -from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder -from langchain.agents import AgentExecutor, create_openai_tools_agent -from langchain_github_copilot import ChatGitHubCopilot - - -def create_github_copilot_langchain_agent(model_name: str, system_prompt: str, tools: list) -> AgentExecutor: - """Create an agent from a set of tools and a Github Copilot model""" - - load_dotenv() - - llm = ChatGitHubCopilot(model_name=model_name) - - prompt = ChatPromptTemplate.from_messages( - [ - ("system", system_prompt), - ("user", "{input}"), - MessagesPlaceholder(variable_name="agent_scratchpad"), - ] - ) - - # Create the agent using LangChain's built-in tool handling - agent = create_openai_tools_agent(llm, tools, prompt) - - agent_executor = AgentExecutor( - name="NotebookPromptAgent", - agent=agent, - tools=tools, - verbose=True, - handle_parsing_errors=True, - ) - - return agent_executor diff --git a/jupyter_ai_agents/agents/langchain/providers/openai.py b/jupyter_ai_agents/agents/langchain/providers/openai.py deleted file mode 100644 index 7a389c9..0000000 --- a/jupyter_ai_agents/agents/langchain/providers/openai.py +++ /dev/null @@ -1,31 +0,0 @@ -# Copyright (c) 2024-2025 Datalayer, Inc. -# -# BSD 3-Clause License - -from dotenv import load_dotenv - -from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder -from langchain.agents import AgentExecutor, create_openai_tools_agent -from langchain_openai import ChatOpenAI - - -def create_openai_langchain_agent(model_name: str, system_prompt: str, tools: list) -> AgentExecutor: - """Create an agent from a set of tools using OpenAI's API.""" - - load_dotenv() - - llm = ChatOpenAI(model_name=model_name) - - prompt = ChatPromptTemplate.from_messages( - [ - ("system", system_prompt), - ("user", "{input}"), - MessagesPlaceholder(variable_name="agent_scratchpad"), - ]) - - # Create the agent using LangChain's built-in tool handling - agent = create_openai_tools_agent(llm, tools, prompt) - - agent_executor = AgentExecutor(name="NotebookPromptAgent", agent=agent, tools=tools, verbose=True, handle_parsing_errors=True) - - return agent_executor diff --git a/jupyter_ai_agents/agents/pydantic/mcp.py b/jupyter_ai_agents/agents/mcp.py similarity index 98% rename from jupyter_ai_agents/agents/pydantic/mcp.py rename to jupyter_ai_agents/agents/mcp.py index 26d9d86..51505c6 100644 --- a/jupyter_ai_agents/agents/pydantic/mcp.py +++ b/jupyter_ai_agents/agents/mcp.py @@ -7,7 +7,7 @@ import httpx from typing import Any, Dict, List -from jupyter_ai_agents.agents.pydantic.models import MCPServer +from jupyter_ai_agents.agents.models import MCPServer class MCPClient: diff --git a/jupyter_ai_agents/agents/pydantic/models.py b/jupyter_ai_agents/agents/models.py similarity index 100% rename from jupyter_ai_agents/agents/pydantic/models.py rename to jupyter_ai_agents/agents/models.py diff --git a/jupyter_ai_agents/agents/prompt/__init__.py b/jupyter_ai_agents/agents/prompt/__init__.py new file mode 100644 index 0000000..0e2b187 --- /dev/null +++ b/jupyter_ai_agents/agents/prompt/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) 2024-2025 Datalayer, Inc. +# +# BSD 3-Clause License + +"""Pydantic AI agents for prompt handling.""" + +from jupyter_ai_agents.agents.prompt.prompt_agent import create_prompt_agent, run_prompt_agent + +__all__ = [ + "create_prompt_agent", + "run_prompt_agent", +] diff --git a/jupyter_ai_agents/agents/prompt/prompt_agent.py b/jupyter_ai_agents/agents/prompt/prompt_agent.py new file mode 100644 index 0000000..1350550 --- /dev/null +++ b/jupyter_ai_agents/agents/prompt/prompt_agent.py @@ -0,0 +1,208 @@ +# Copyright (c) 2024-2025 Datalayer, Inc. +# +# BSD 3-Clause License + +"""Pydantic AI Prompt Agent - creates and executes code based on user instructions.""" + +import logging +from typing import Any + +from pydantic_ai import Agent +from pydantic_ai.mcp import MCPServerStreamableHTTP + +from jupyter_ai_agents.tools import create_mcp_server + +logger = logging.getLogger(__name__) + + +SYSTEM_PROMPT = """You are a powerful coding assistant for Jupyter notebooks. +Create and execute code in a notebook based on user instructions. +Add markdown cells to explain the code and structure the notebook clearly. + +Important guidelines: +- Assume that no packages are installed in the notebook, so install them using code cells with !pip install +- Ensure updates to cell indexing when new cells are inserted +- Maintain the logical flow of execution by adjusting cell index as needed +- Use the available Jupyter MCP tools to interact with the notebook +- Always execute code cells after inserting them to verify they work + +Available tools through MCP: +- notebook tools for inserting/modifying cells +- kernel tools for executing code +- file system tools for reading/writing files + +When the user asks you to create something, break it down into steps: +1. Install any required packages +2. Import necessary libraries +3. Write the main code +4. Add markdown explanations + +Execute each code cell as you create it to ensure it works properly. + +IMPORTANT: When you have completed the task successfully, provide a final text response summarizing what you did WITHOUT making any more tool calls. This signals completion and allows the program to exit properly.""" + + +class PromptAgentDeps: + """Dependencies for the prompt agent.""" + + def __init__(self, notebook_context: dict[str, Any] | None = None): + """Initialize dependencies. + + Args: + notebook_context: Optional context about the notebook (path, cells, etc.) + """ + self.notebook_context = notebook_context or {} + self.current_cell_index = notebook_context.get('current_cell_index', -1) if notebook_context else -1 + self.full_context = notebook_context.get('full_context', False) if notebook_context else False + self.notebook_content = notebook_context.get('notebook_content', '') if notebook_context else '' + + +def create_prompt_agent( + model: str, + mcp_server: MCPServerStreamableHTTP, + notebook_context: dict[str, Any] | None = None, + max_tool_calls: int = 10, +) -> Agent[PromptAgentDeps, str]: + """ + Create the Prompt Agent for handling user instructions. + + Args: + model: Model identifier (e.g., 'anthropic:claude-sonnet-4-0', 'openai:gpt-4o') + mcp_server: MCP server connection to jupyter-mcp-server + notebook_context: Optional context about the notebook + max_tool_calls: Maximum number of tool calls to make (default: 10) + + Returns: + Configured Pydantic AI agent + """ + # Enhance system prompt with notebook context if available + system_prompt = SYSTEM_PROMPT + + if notebook_context: + if notebook_context.get('full_context') and notebook_context.get('notebook_content'): + system_prompt += f"\n\nCurrent notebook content:\n{notebook_context['notebook_content']}" + + if notebook_context.get('current_cell_index', -1) != -1: + system_prompt += f"\n\nUser instruction was given at cell index: {notebook_context['current_cell_index']}" + + # Add reminder to be efficient and complete properly + system_prompt += ( + "\n\nBe efficient: Complete the task in as few steps as possible. " + "Don't over-verify or re-check unnecessarily. " + "When the task is done, provide a final summary WITHOUT tool calls to signal completion." + ) + + # Create agent with MCP toolset + agent = Agent( + model, + toolsets=[mcp_server], + deps_type=PromptAgentDeps, + system_prompt=system_prompt, + ) + + logger.info(f"Created prompt agent with model {model} (max_tool_calls={max_tool_calls})") + + return agent + + +async def run_prompt_agent( + agent: Agent[PromptAgentDeps, str], + user_input: str, + notebook_context: dict[str, Any] | None = None, + max_tool_calls: int = 10, + max_requests: int = 2, +) -> str: + """ + Run the prompt agent with user input. + + Args: + agent: The configured prompt agent + user_input: User's instruction/prompt + notebook_context: Optional notebook context (should include 'notebook_path') + max_tool_calls: Maximum number of tool calls to prevent excessive API usage + max_requests: Maximum number of API requests (default: 4, lower if needed for strict rate limits) + + Returns: + Agent's response + """ + import asyncio + import os + from pydantic_ai import UsageLimitExceeded, UsageLimits + + deps = PromptAgentDeps(notebook_context) + + logger.info(f"Running prompt agent with input: {user_input[:50]}... (max_tool_calls={max_tool_calls}, max_requests={max_requests})") + + # Prepend notebook connection instruction if path is provided + if notebook_context and notebook_context.get('notebook_path'): + notebook_path = notebook_context['notebook_path'] + notebook_name = os.path.splitext(os.path.basename(notebook_path))[0] + + # Prepend instruction to connect to the notebook first + enhanced_input = ( + f"First, use the use_notebook tool to connect to the notebook at path '{notebook_path}' " + f"with notebook_name '{notebook_name}' and mode 'connect'. " + f"Then, {user_input}" + ) + logger.info(f"Enhanced input to connect to notebook: {notebook_path}") + else: + enhanced_input = user_input + + # Warn about low limits + if max_requests <= 2: + logger.warning( + f"Using very conservative request limit ({max_requests}). " + "This may result in incomplete responses. " + "Increase --max-requests if your Azure tier allows." + ) + + try: + # Create usage limits to prevent excessive API calls + # Use strict limits to avoid rate limiting + usage_limits = UsageLimits( + tool_calls_limit=max_tool_calls, + request_limit=max_requests, # Strict limit to avoid rate limiting + ) + + # Add timeout to prevent hanging on retries + result = await asyncio.wait_for( + agent.run(enhanced_input, deps=deps, usage_limits=usage_limits), + timeout=120.0 # 2 minute timeout + ) + logger.info("Prompt agent completed successfully") + return result.response + except asyncio.TimeoutError: + logger.error("Prompt agent timed out after 120 seconds") + return "Error: Operation timed out. The agent may have hit rate limits or is taking too long." + except UsageLimitExceeded as e: + logger.error(f"Prompt agent hit usage limits: {e}") + return ( + "Error: Reached the configured usage limits.\n" + f"Increase --max-requests (currently {max_requests}) or --max-tool-calls (currently {max_tool_calls}) " + "if your model provider allows more usage." + ) + except Exception as e: + logger.error(f"Error running prompt agent: {e}", exc_info=True) + return f"Error: {str(e)}" + + +def create_prompt_agent_sync( + base_url: str, + token: str, + model: str, + notebook_context: dict[str, Any] | None = None, +) -> Agent[PromptAgentDeps, str]: + """ + Create prompt agent with MCP server connection (synchronous wrapper). + + Args: + base_url: Jupyter server base URL + token: Authentication token + model: Model identifier + notebook_context: Optional notebook context + + Returns: + Configured agent + """ + mcp_server = create_mcp_server(base_url, token) + return create_prompt_agent(model, mcp_server, notebook_context) diff --git a/jupyter_ai_agents/agents/pydantic/chat/__init__.py b/jupyter_ai_agents/agents/pydantic/chat/__init__.py deleted file mode 100644 index 3d32561..0000000 --- a/jupyter_ai_agents/agents/pydantic/chat/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -# Copyright (c) 2024-2025 Datalayer, Inc. -# -# BSD 3-Clause License - -# Copyright (c) 2023-2024 Datalayer, Inc. -# -# BSD 3-Clause License diff --git a/jupyter_ai_agents/cli/app.py b/jupyter_ai_agents/cli/app.py index adae195..134feae 100644 --- a/jupyter_ai_agents/cli/app.py +++ b/jupyter_ai_agents/cli/app.py @@ -2,118 +2,596 @@ # # BSD 3-Clause License +"""Jupyter AI Agents CLI - Pydantic AI agents with MCP integration.""" + from __future__ import annotations import typer import asyncio -import os +import logging +import httpx + +# Pydantic AI agents +from jupyter_ai_agents.agents.prompt.prompt_agent import ( + create_prompt_agent, + run_prompt_agent, +) +from jupyter_ai_agents.agents.explain_error.explain_error_agent import ( + create_explain_error_agent, + run_explain_error_agent, +) +from jupyter_ai_agents.agents.prompt.prompt_agent import ( + create_prompt_agent, + run_prompt_agent, +) +from jupyter_ai_agents.tools import create_mcp_server + +# Import tools for REPL +# Import model utilities +from jupyter_ai_agents.utils.model import create_model_with_provider + +# Configure logging +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' +) +logger = logging.getLogger(__name__) -from jupyter_kernel_client import KernelClient -from jupyter_nbmodel_client import NbModelClient, get_jupyter_notebook_websocket_url -from jupyter_ai_agents.agents.langchain.prompt_agent import prompt as prompt_agent -from jupyter_ai_agents.agents.langchain.explain_error_agent import explain_error as explain_error_agent +logging.getLogger("httpx").setLevel(logging.ERROR) +logging.getLogger("anthropic").setLevel(logging.ERROR) +logging.getLogger("openai").setLevel(logging.ERROR) +def enable_verbose_logging(): + """Enable verbose logging for debugging API calls and retries.""" + # logging.getLogger().setLevel(logging.DEBUG) + # Enable detailed HTTP logging to see retry reasons + logging.getLogger("httpx").setLevel(logging.DEBUG) + logging.getLogger("anthropic").setLevel(logging.DEBUG) + logging.getLogger("openai").setLevel(logging.DEBUG) + logger.debug("Verbose logging enabled - will show detailed HTTP requests, responses, and retry reasons") -app = typer.Typer(help="The Jupyter AI Agents application.") +app = typer.Typer(help="Jupyter AI Agents - AI-powered notebook manipulation with Pydantic AI and MCP.") @app.command() def prompt( - url: str = typer.Option("http://localhost:8888", help="URL to the Jupyter Server."), - token: str = typer.Option("", help="Jupyter Server token."), + mcp_servers: str = typer.Option( + "http://localhost:8888/mcp", + help="Comma-separated list of MCP server URLs (e.g., 'http://localhost:8888/mcp' for jupyter-mcp-server)." + ), path: str = typer.Option("", help="Jupyter Notebook path."), - agent: str = typer.Option("prompt", help="Agent name."), - input: str = typer.Option("", help="Input."), - model_provider: str = typer.Option("github-copilot", help="Model provider can be 'azure-openai', 'github-copilot', or 'openai'."), - openai_api_version: str = typer.Option(os.environ.get("OPENAI_API_VERSION", None), help="OpenAI API version."), - azure_openai_version: str = typer.Option(os.environ.get("AZURE_OPENAI_VERSION", None), help="Azure OpenAI version."), - azure_openai_api_key: str = typer.Option(os.environ.get("AZURE_OPENAI_API_KEY", None), help="Azure OpenAI key."), - openai_api_key: str = typer.Option(os.environ.get("OPENAI_API_KEY", None), help="OpenAI API key."), - anthropic_api_key: str = typer.Option(os.environ.get("ANTHROPIC_API_KEY", None), help="Anthropic API key."), - github_token: str = typer.Option(os.environ.get("GITHUB_TOKEN", None), help="Github token."), - model_name: str = typer.Option("gpt-4o", help="Model name (deployment name for Azure/OpenAI/Copilot)."), + input: str = typer.Option("", help="User instruction/prompt."), + model: str = typer.Option( + None, + help="Full model string (e.g., 'openai:gpt-4o', 'anthropic:claude-sonnet-4-0', 'azure-openai:gpt-4o-mini'). If not provided, uses --model-provider and --model-name." + ), + model_provider: str = typer.Option( + "openai", + help="Model provider: 'openai', 'anthropic', 'azure-openai', 'github-copilot', 'google', 'bedrock', 'groq', 'mistral', 'cohere'." + ), + model_name: str = typer.Option("gpt-4o", help="Model name or deployment name."), + timeout: float = typer.Option(60.0, help="HTTP timeout in seconds for API requests (default: 60.0)."), current_cell_index: int = typer.Option(-1, help="Index of the cell where the prompt is asked."), - full_context: bool = typer.Option(False, help="Flag to provide the full notebook context to the AI model."), + full_context: bool = typer.Option(False, help="Flag to provide full notebook context to the AI model."), + max_tool_calls: int = typer.Option(10, help="Maximum number of tool calls per agent run (prevents excessive API usage)."), + max_requests: int = typer.Option(4, help="Maximum number of API requests per run (defaults to 4; lower for strict rate limits)."), + verbose: bool = typer.Option(False, help="Enable verbose logging."), ): - """From a given instruction, code and markdown cells are added to a notebook.""" + """ + Execute user instructions in a Jupyter notebook using AI. + + The agent will create code and markdown cells based on your instructions, + execute the code, and verify it works properly. + + You can specify the model in two ways: + 1. Using --model with full string: --model "openai:gpt-4o" + 2. Using --model-provider and --model-name: --model-provider openai --model-name gpt-4o + + For Azure OpenAI, you can use either: + - --model "azure-openai:gpt-4o-mini" + - --model-provider azure-openai --model-name gpt-4o-mini + + Required environment variables for Azure OpenAI: + - AZURE_OPENAI_API_KEY + - AZURE_OPENAI_ENDPOINT (base URL, e.g., https://your-resource.openai.azure.com) + - AZURE_OPENAI_API_VERSION (optional) + + Examples: + # Using full model string + jupyter-ai-agents prompt \\ + --mcp-servers http://localhost:8888/mcp \\ + --path notebook.ipynb \\ + --model "openai:gpt-4o" \\ + --input "Create a matplotlib example" + + # Using provider and name + jupyter-ai-agents prompt \\ + --mcp-servers http://localhost:8888/mcp \\ + --path notebook.ipynb \\ + --model-provider anthropic \\ + --model-name claude-sonnet-4-0 \\ + --input "Create a matplotlib example" + """ + if verbose: + enable_verbose_logging() + async def _run(): - kernel = KernelClient(server_url=url, token=token) - kernel.start() - notebook = NbModelClient(get_jupyter_notebook_websocket_url(server_url=url, token=token, path=path)) - await notebook.start() try: - reply = await prompt_agent(notebook, kernel, input, model_provider, model_name, full_context, current_cell_index) - # Print only the final summary or message - if isinstance(reply, list) and reply: - last = reply[-1] - if isinstance(last, dict) and 'output' in last: - typer.echo(last['output']) - elif isinstance(last, dict) and 'messages' in last and last['messages']: - # LangChain style: last message content - msg = last['messages'][-1] - if hasattr(msg, 'content'): - typer.echo(msg.content) - elif isinstance(msg, dict) and 'content' in msg: - typer.echo(msg['content']) - else: - typer.echo(msg) + # Create MCP server connection(s) + from jupyter_ai_agents.tools import MCPServerStreamableHTTP + + server_urls = [s.strip() for s in mcp_servers.split(',')] + logger.info(f"Connecting to {len(server_urls)} MCP server(s)") + + toolsets = [] + for server_url in server_urls: + logger.info(f" - {server_url}") + mcp_client = MCPServerStreamableHTTP(server_url) + toolsets.append(mcp_client) + + # Use first MCP server for backward compatibility with create_prompt_agent + mcp_server = toolsets[0] if toolsets else None + + # Determine model - handle azure-openai:deployment format or use provider+name + if model: + # Check if model string is in azure-openai:deployment format + if model.startswith('azure-openai:'): + from pydantic_ai.models.openai import OpenAIChatModel + deployment_name = model.split(':', 1)[1] + model_obj = OpenAIChatModel(deployment_name, provider='azure') + logger.info(f"Using Azure OpenAI deployment: {deployment_name}") + elif model.startswith('anthropic:'): + # Parse anthropic:model-name format and use create_model_with_provider + model_name_part = model.split(':', 1)[1] + model_obj = create_model_with_provider('anthropic', model_name_part, timeout) + logger.info(f"Using Anthropic model: {model_name_part} (timeout: {timeout}s)") else: - typer.echo(last) + # User provided full model string + model_obj = model + logger.info(f"Using explicit model: {model_obj}") else: - typer.echo(reply) - finally: - await notebook.stop() - kernel.stop() + # Create model object with provider-specific configuration + model_obj = create_model_with_provider(model_provider, model_name, timeout) + if isinstance(model_obj, str): + logger.info(f"Using model: {model_obj} (from {model_provider} + {model_name})") + else: + logger.info(f"Using {model_provider} model: {model_name} (timeout: {timeout}s)") + + # Prepare notebook context + notebook_context = { + 'notebook_path': path, + 'current_cell_index': current_cell_index, + 'full_context': full_context, + } + + # Create and run agent + logger.info("Creating prompt agent...") + agent = create_prompt_agent(model_obj, mcp_server, notebook_context, max_tool_calls=max_tool_calls) + + logger.info("Running prompt agent...") + result = await run_prompt_agent(agent, input, notebook_context, max_tool_calls=max_tool_calls, max_requests=max_requests) + + # Print result + typer.echo("\n" + "="*60) + typer.echo("AI Agent Response:") + typer.echo("="*60) + typer.echo(result) + typer.echo("="*60 + "\n") + + except Exception as e: + logger.error(f"Error running prompt agent: {e}", exc_info=True) + typer.echo(f"Error: {str(e)}", err=True) + raise typer.Exit(code=1) + asyncio.run(_run()) @app.command() def explain_error( - url: str = typer.Option("http://localhost:8888", help="URL to the Jupyter Server."), - token: str = typer.Option("", help="Jupyter Server token."), + mcp_servers: str = typer.Option( + "http://localhost:8888/mcp", + help="Comma-separated list of MCP server URLs (e.g., 'http://localhost:8888/mcp' for jupyter-mcp-server)." + ), path: str = typer.Option("", help="Jupyter Notebook path."), - agent: str = typer.Option("prompt", help="Agent name."), - input: str = typer.Option("", help="Input."), - model_provider: str = typer.Option("github-copilot", help="Model provider can be 'azure-openai', 'github-copilot', or 'openai'."), - openai_api_version: str = typer.Option(os.environ.get("OPENAI_API_VERSION", None), help="OpenAI API version."), - azure_openai_version: str = typer.Option(os.environ.get("AZURE_OPENAI_VERSION", None), help="Azure OpenAI version."), - azure_openai_api_key: str = typer.Option(os.environ.get("AZURE_OPENAI_API_KEY", None), help="Azure OpenAI key."), - openai_api_key: str = typer.Option(os.environ.get("OPENAI_API_KEY", None), help="OpenAI API key."), - anthropic_api_key: str = typer.Option(os.environ.get("ANTHROPIC_API_KEY", None), help="Anthropic API key."), - github_token: str = typer.Option(os.environ.get("GITHUB_TOKEN", None), help="Github token."), - model_name: str = typer.Option("gpt-4o", help="Model name (deployment name for Azure/OpenAI/Copilot)."), - current_cell_index: int = typer.Option(-1, help="Index of the cell where the prompt is asked."), + model: str = typer.Option( + None, + help="Full model string (e.g., 'openai:gpt-4o', 'anthropic:claude-sonnet-4-0', 'azure-openai:gpt-4o-mini'). If not provided, uses --model-provider and --model-name." + ), + model_provider: str = typer.Option( + "openai", + help="Model provider: 'openai', 'anthropic', 'azure-openai', 'github-copilot', 'google', 'bedrock', 'groq', 'mistral', 'cohere'." + ), + model_name: str = typer.Option("gpt-4o", help="Model name or deployment name."), + timeout: float = typer.Option(60.0, help="HTTP timeout in seconds for API requests (default: 60.0)."), + current_cell_index: int = typer.Option(-1, help="Index of the cell with the error (-1 for first error)."), + max_tool_calls: int = typer.Option(10, help="Maximum number of tool calls per agent run (prevents excessive API usage)."), + max_requests: int = typer.Option(3, help="Maximum number of API requests per run (defaults to 3 for error fixing)."), + verbose: bool = typer.Option(False, help="Enable verbose logging."), ): - """Explain and correct an error in a notebook based on the prior cells.""" + """ + Explain and fix errors in a Jupyter notebook using AI. + + The agent will analyze the error, explain what went wrong, and insert + corrected code cells to fix the issue. + + You can specify the model in two ways: + 1. Using --model with full string: --model "openai:gpt-4o" + 2. Using --model-provider and --model-name: --model-provider openai --model-name gpt-4o + + For Azure OpenAI, you can use either: + - --model "azure-openai:gpt-4o-mini" + - --model-provider azure-openai --model-name gpt-4o-mini + + Required environment variables for Azure OpenAI: + - AZURE_OPENAI_API_KEY + - AZURE_OPENAI_ENDPOINT (base URL, e.g., https://your-resource.openai.azure.com) + - AZURE_OPENAI_API_VERSION (optional) + + Examples: + # Using full model string + jupyter-ai-agents explain-error \\ + --mcp-servers http://localhost:8888/mcp \\ + --path notebook.ipynb \\ + --model "anthropic:claude-sonnet-4-0" \\ + --current-cell-index 5 + + # Using provider and name + jupyter-ai-agents explain-error \\ + --mcp-servers http://localhost:8888/mcp \\ + --path notebook.ipynb \\ + --model-provider openai \\ + --model-name gpt-4o + """ + if verbose: + enable_verbose_logging() + async def _run(): - kernel = KernelClient(server_url=url, token=token) - kernel.start() - notebook = NbModelClient(get_jupyter_notebook_websocket_url(server_url=url, token=token, path=path)) - await notebook.start() try: - reply = await explain_error_agent(notebook, kernel, model_provider, model_name, current_cell_index) - # Print only the final summary or message - if isinstance(reply, list) and reply: - last = reply[-1] - if isinstance(last, dict) and 'output' in last: - typer.echo(last['output']) - elif isinstance(last, dict) and 'messages' in last and last['messages']: - msg = last['messages'][-1] - if hasattr(msg, 'content'): - typer.echo(msg.content) - elif isinstance(msg, dict) and 'content' in msg: - typer.echo(msg['content']) - else: - typer.echo(msg) + # Create MCP server connection(s) + from jupyter_ai_agents.tools import MCPServerStreamableHTTP + + server_urls = [s.strip() for s in mcp_servers.split(',')] + logger.info(f"Connecting to {len(server_urls)} MCP server(s)") + + toolsets = [] + for server_url in server_urls: + logger.info(f" - {server_url}") + mcp_client = MCPServerStreamableHTTP(server_url) + toolsets.append(mcp_client) + + # Use first MCP server for backward compatibility with create_explain_error_agent + mcp_server = toolsets[0] if toolsets else None + + # Determine model - handle azure-openai:deployment format or use provider+name + if model: + # Check if model string is in azure-openai:deployment format + if model.startswith('azure-openai:'): + from pydantic_ai.models.openai import OpenAIChatModel + deployment_name = model.split(':', 1)[1] + model_obj = OpenAIChatModel(deployment_name, provider='azure') + logger.info(f"Using Azure OpenAI deployment: {deployment_name}") + elif model.startswith('anthropic:'): + # Parse anthropic:model-name format and use create_model_with_provider + model_name_part = model.split(':', 1)[1] + model_obj = create_model_with_provider('anthropic', model_name_part, timeout) + logger.info(f"Using Anthropic model: {model_name_part} (timeout: {timeout}s)") else: - typer.echo(last) + model_obj = model + logger.info(f"Using explicit model: {model_obj}") else: - typer.echo(reply) - finally: - await notebook.stop() - kernel.stop() + model_obj = create_model_with_provider(model_provider, model_name, timeout) + if isinstance(model_obj, str): + logger.info(f"Using model: {model_obj} (from {model_provider} + {model_name})") + else: + logger.info(f"Using {model_provider} model: {model_name} (timeout: {timeout}s)") + + # In a real implementation, we would: + # 1. Fetch notebook content from server + # 2. Extract error information + # 3. Pass to agent + + # For now, create a placeholder + # TODO: Implement notebook content fetching via MCP or direct API + notebook_content = "# Notebook content would be fetched here" + error_description = "Error: Please implement notebook error fetching" + + logger.info("Creating explain error agent...") + agent = create_explain_error_agent( + model_obj, + mcp_server, + notebook_content=notebook_content, + error_cell_index=current_cell_index, + max_tool_calls=max_tool_calls, + ) + + logger.info("Running explain error agent...") + result = await run_explain_error_agent( + agent, + error_description, + notebook_content=notebook_content, + error_cell_index=current_cell_index, + notebook_path=path, + max_tool_calls=max_tool_calls, + max_requests=max_requests, + ) + + # Print result + typer.echo("\n" + "="*60) + typer.echo("AI Agent Error Analysis:") + typer.echo("="*60) + typer.echo(result) + typer.echo("="*60 + "\n") + + except Exception as e: + logger.error(f"Error running explain error agent: {e}", exc_info=True) + typer.echo(f"Error: {str(e)}", err=True) + raise typer.Exit(code=1) + asyncio.run(_run()) +@app.command() +def repl( + mcp_servers: str = typer.Option( + ..., + help="Comma-separated list of MCP server URLs (e.g., 'http://localhost:8001/mcp,http://localhost:8002/mcp' for standalone servers, or 'http://localhost:8888/mcp' for jupyter-mcp-server)." + ), + model: str = typer.Option( + None, + help="Full model string (e.g., 'openai:gpt-4o', 'anthropic:claude-sonnet-4-0', 'azure-openai:gpt-4o-mini'). If not provided, uses --model-provider and --model-name." + ), + model_provider: str = typer.Option( + "openai", + help="Model provider: 'openai', 'anthropic', 'azure-openai', 'github-copilot', 'google', 'bedrock', 'groq', 'mistral', 'cohere'." + ), + model_name: str = typer.Option("gpt-4o", help="Model name or deployment name."), + timeout: float = typer.Option(60.0, help="HTTP timeout in seconds for API requests (default: 60.0)."), + system_prompt: str = typer.Option( + None, + help="Custom system prompt. If not provided, uses a default prompt based on the MCP servers being used." + ), + verbose: bool = typer.Option(False, help="Enable verbose logging."), +): + """ + Start an interactive REPL with access to MCP tools. + + This command launches pydantic-ai's built-in CLI with MCP server tools. + You can connect to any MCP servers implementing the Streamable HTTP transport: + - Jupyter MCP server (jupyter-mcp-server): http://localhost:8888/mcp + - Standalone MCP servers: http://localhost:8001/mcp, http://localhost:8002/mcp, etc. + + You can specify the model in two ways: + 1. Using --model with full string: --model "openai:gpt-4o" + 2. Using --model-provider and --model-name: --model-provider openai --model-name gpt-4o + + For Azure OpenAI, you can use either: + - --model "azure-openai:gpt-4o-mini" + - --model-provider azure-openai --model-name gpt-4o-mini + + Required environment variables for Azure OpenAI: + - AZURE_OPENAI_API_KEY + - AZURE_OPENAI_ENDPOINT (base URL, e.g., https://your-resource.openai.azure.com) + - AZURE_OPENAI_API_VERSION (optional) + + Special commands in the REPL: + - /exit: Exit the session + - /markdown: Show last response in markdown format + - /multiline: Toggle multiline input mode (use Ctrl+D to submit) + - /cp: Copy last response to clipboard + + Examples: + # Connect to Jupyter MCP Server with OpenAI + jupyter-ai-agents repl \\ + --mcp-servers "http://localhost:8888/mcp" \\ + --model "openai:gpt-4o" + + # Connect to standalone MCP servers with Anthropic + # Note: Use correct Anthropic model names like claude-sonnet-4-20250514 + jupyter-ai-agents repl \\ + --mcp-servers "http://localhost:8001/mcp,http://localhost:8002/mcp" \\ + --model-provider anthropic \\ + --model-name claude-sonnet-4-20250514 + + # Set custom timeout (useful for slow connections) + jupyter-ai-agents repl \\ + --mcp-servers "http://localhost:8888/mcp" \\ + --timeout 120.0 + + # Then in the REPL, you can ask: + > List all notebooks in the current directory (with jupyter-mcp-server) + > Add 5 and 7 (with calculator server) + > Reverse the text "hello world" (with echo server) + """ + if verbose: + enable_verbose_logging() + + # Separate function to initialize and list tools + async def list_tools_async(): + """List all tools available from the MCP server(s)""" + try: + from jupyter_ai_agents.tools import MCPServerStreamableHTTP + + server_urls = [s.strip() for s in mcp_servers.split(',')] + + for server_url in server_urls: + try: + mcp_client = MCPServerStreamableHTTP(server_url) + + # Use context manager to connect and list tools + async with mcp_client: + tools = await mcp_client.list_tools() + + if not tools or len(tools) == 0: + typer.echo("\n No tools available") + continue + + typer.echo(f"\n Available Tools ({len(tools)}):") + for tool in tools: + name = tool.name + description = tool.description or "" + schema = tool.inputSchema + + # Build parameter list + params = [] + if schema and "properties" in schema: + for param_name, param_info in schema["properties"].items(): + param_type = param_info.get("type", "any") + params.append(f"{param_name}: {param_type}") + + param_str = f"({', '.join(params)})" if params else "()" + desc_first_line = description.split('\n')[0] if description else "No description" + typer.echo(f" • {name}{param_str} - {desc_first_line}") + except Exception as e: + logger.warning(f"Could not connect to {server_url}: {e}") + typer.echo(f"\n ⚠️ Could not list tools from {server_url}") + + except Exception as e: + logger.warning(f"Could not list tools: {e}") + typer.echo(f"\n ⚠️ Could not list tools: {e}") + + + try: + from pydantic_ai import Agent + + # Determine model - handle azure-openai:deployment format or use provider+name + model_display_name = None # Track the display name for welcome message + + if model: + # Check if model string is in azure-openai:deployment format + if model.startswith('azure-openai:'): + from pydantic_ai.models.openai import OpenAIChatModel + from pydantic_ai.providers import infer_provider + from pydantic_ai.providers.openai import OpenAIProvider + from openai.lib.azure import AsyncAzureOpenAI + + deployment_name = model.split(':', 1)[1] + http_timeout = httpx.Timeout(timeout, connect=30.0) + + # Infer Azure provider first to get proper configuration (API key, API version, etc.) + azure_provider_base = infer_provider('azure') + + # Extract base URL - remove /openai suffix since AsyncAzureOpenAI adds it + base_url = str(azure_provider_base.client.base_url) + # base_url is like: https://xxx.openai.azure.com/openai/ + # AsyncAzureOpenAI expects: https://xxx.openai.azure.com (it adds /openai automatically) + azure_endpoint = base_url.rstrip('/').rsplit('/openai', 1)[0] + + logger.info(f"Azure OpenAI endpoint: {azure_endpoint}") + logger.info(f"Azure OpenAI API version: {azure_provider_base.client.default_query}") + logger.info(f"Azure OpenAI timeout: {http_timeout}") + + # Create Azure OpenAI client with custom timeout + azure_client = AsyncAzureOpenAI( + azure_endpoint=azure_endpoint, + azure_deployment=deployment_name, + api_version=azure_provider_base.client.default_query.get('api-version'), + api_key=azure_provider_base.client.api_key, + timeout=http_timeout, + ) + + # Create provider with the configured client + from pydantic_ai.providers.openai import OpenAIProvider + azure_provider = OpenAIProvider(openai_client=azure_client) + + model_obj = OpenAIChatModel( + deployment_name, + provider=azure_provider + ) + model_display_name = model # azure-openai:deployment-name + logger.info(f"Using Azure OpenAI deployment: {deployment_name}") + elif model.startswith('anthropic:'): + # Parse anthropic:model-name format and use create_model_with_provider + model_name_part = model.split(':', 1)[1] + model_obj = create_model_with_provider('anthropic', model_name_part, timeout) + model_display_name = model + logger.info(f"Using Anthropic model: {model_name_part} (timeout: {timeout}s)") + else: + model_obj = model + model_display_name = model + logger.info(f"Using explicit model: {model_obj}") + else: + model_obj = create_model_with_provider(model_provider, model_name, timeout) + if isinstance(model_obj, str): + model_display_name = model_obj + logger.info(f"Using model: {model_obj} (from {model_provider} + {model_name})") + else: + model_display_name = f"{model_provider}:{model_name}" + logger.info(f"Using {model_provider} model: {model_name} (timeout: {timeout}s)") + + # Create MCP server connection(s) + from jupyter_ai_agents.tools import MCPServerStreamableHTTP + + server_urls = [s.strip() for s in mcp_servers.split(',')] + logger.info(f"Connecting to {len(server_urls)} MCP server(s)") + + toolsets = [] + for server_url in server_urls: + logger.info(f" - {server_url}") + mcp_client = MCPServerStreamableHTTP(server_url) + toolsets.append(mcp_client) + + # Display welcome message + typer.echo("="*70) + typer.echo("🪐 ✨ Jupyter AI Agents - Interactive REPL") + typer.echo("="*70) + typer.echo(f"Model: {model_display_name}") + + typer.echo(f"MCP Servers: {len(server_urls)} connected") + for server_url in server_urls: + typer.echo(f" - {server_url}") + + # List tools inline in welcome message + asyncio.run(list_tools_async()) + + # Create default system prompt if not provided + if system_prompt is None: + instructions = """You are a helpful AI assistant with access to various MCP tools. + +Use the available tools to help the user accomplish their tasks. +Be proactive in suggesting what you can do with the available tools. +""" + else: + instructions = system_prompt + + # Create agent with MCP toolset(s) + logger.info("Creating agent with MCP tools...") + agent = Agent( + model_obj, + toolsets=toolsets, + system_prompt=instructions, + ) + + typer.echo("="*70) + typer.echo("\nSpecial commands:") + typer.echo(" /exit - Exit the session") + typer.echo(" /markdown - Show last response in markdown") + typer.echo(" /multiline - Toggle multiline mode (Ctrl+D to submit)") + typer.echo(" /cp - Copy last response to clipboard") + typer.echo("="*70 + "\n") + + # Launch the CLI interface (separate asyncio.run call with context manager) + async def _run_cli() -> None: + assert agent is not None + async with agent: + await agent.to_cli(prog_name='jupyter-ai-agents') + + asyncio.run(_run_cli()) + + except KeyboardInterrupt: + typer.echo("\n\n🛑 Agent stopped by user") + except asyncio.CancelledError: + # Handle cancellation from Ctrl+C during SDK retries + logger.info("REPL session cancelled") + typer.echo("\n\n🛑 Session cancelled") + except BaseExceptionGroup as exc: + typer.echo("\n❌ Encountered errors while running the CLI:") + for idx, sub_exc in enumerate(exc.exceptions, start=1): + typer.echo(f" [{idx}] {type(sub_exc).__name__}: {sub_exc}") + raise typer.Exit(code=1) + except Exception as e: + logger.error(f"Error in REPL: {e}", exc_info=True) + typer.echo(f"\n❌ Error: {str(e)}", err=True) + raise typer.Exit(code=1) + def main(): app() diff --git a/jupyter_ai_agents/cli/consoleapp.py b/jupyter_ai_agents/cli/console_app.py similarity index 100% rename from jupyter_ai_agents/cli/consoleapp.py rename to jupyter_ai_agents/cli/console_app.py diff --git a/jupyter_ai_agents/extension.py b/jupyter_ai_agents/extension.py index 1c12a52..bea4029 100644 --- a/jupyter_ai_agents/extension.py +++ b/jupyter_ai_agents/extension.py @@ -18,16 +18,15 @@ from jupyter_ai_agents.handlers.index import IndexHandler from jupyter_ai_agents.handlers.config import ConfigHandler -from jupyter_ai_agents.handlers.nbmodel import AINbModelAgentsHandler, AINbModelAgentsInstanceHandler from jupyter_ai_agents.handlers.chat import ChatHandler from jupyter_ai_agents.handlers.configure import ConfigureHandler from jupyter_ai_agents.handlers.mcp import ( MCPServersHandler, MCPServerHandler, ) -from jupyter_ai_agents.agents.pydantic.mcp import MCPToolManager -from jupyter_ai_agents.agents.pydantic.chat.config import ChatConfig -from jupyter_ai_agents.agents.pydantic.chat.agent import create_chat_agent +from jupyter_ai_agents.agents.mcp import MCPToolManager +from jupyter_ai_agents.agents.chat.config import ChatConfig +from jupyter_ai_agents.agents.chat.agent import create_chat_agent from jupyter_ai_agents.tools import create_mcp_server from jupyter_ai_agents.__version__ import __version__ @@ -164,8 +163,6 @@ def initialize_handlers(self): (url_path_join(self.name), IndexHandler), (url_path_join(self.name, "config"), ConfigHandler), (url_path_join(self.name, "configure"), ConfigureHandler), - (url_path_join(self.name, "agents"), AINbModelAgentsHandler), - (url_path_join(self.name, r"agents/(.+)$"), AINbModelAgentsInstanceHandler), (url_path_join("api", "chat"), ChatHandler), (url_path_join("api", "mcp/servers"), MCPServersHandler), (url_path_join("api", r"mcp/servers/([^/]+)"), MCPServerHandler), diff --git a/jupyter_ai_agents/handlers/__init__.py b/jupyter_ai_agents/handlers/__init__.py index 52a519d..97c15cc 100644 --- a/jupyter_ai_agents/handlers/__init__.py +++ b/jupyter_ai_agents/handlers/__init__.py @@ -1,7 +1,3 @@ # Copyright (c) 2024-2025 Datalayer, Inc. # # BSD 3-Clause License - -# Copyright (c) 2023-2024 Datalayer, Inc. -# -# Datalayer License diff --git a/jupyter_ai_agents/handlers/configure.py b/jupyter_ai_agents/handlers/configure.py index 0a5ea27..9fd31d6 100644 --- a/jupyter_ai_agents/handlers/configure.py +++ b/jupyter_ai_agents/handlers/configure.py @@ -10,7 +10,7 @@ from jupyter_server.base.handlers import APIHandler -from jupyter_ai_agents.agents.pydantic.models import ( +from jupyter_ai_agents.agents.models import ( FrontendConfig, AIModel, BuiltinTool, @@ -83,7 +83,7 @@ async def get(self): models = [ AIModel( id="anthropic:claude-sonnet-4-5", - name="Claude Sonnet 4.0", + name="Claude Sonnet 4.5", builtin_tools=tool_ids # Associate all available tools ) ] diff --git a/jupyter_ai_agents/handlers/nbmodel.py b/jupyter_ai_agents/handlers/nbmodel.py deleted file mode 100644 index dd18835..0000000 --- a/jupyter_ai_agents/handlers/nbmodel.py +++ /dev/null @@ -1,138 +0,0 @@ -# Copyright (c) 2024-2025 Datalayer, Inc. -# -# BSD 3-Clause License - -from __future__ import annotations - -import json -import logging - -from concurrent import futures -from concurrent.futures import as_completed -from anyio import create_task_group, sleep -from anyio.from_thread import start_blocking_portal - -from jupyter_server.utils import url_path_join -from jupyter_server.base.handlers import APIHandler -from jupyter_kernel_client import KernelClient - -from jupyter_ai_agents.agents.langchain.prompt_agent import PromptAgent -from jupyter_ai_agents.agents.langchain.models import NbModelAgentRequestModel -from jupyter_ai_agents.agents.langchain.manager.agents_manager import AIAgentsManager -from jupyter_ai_agents.utils import http_to_ws -from jupyter_ai_agents import __version__ - - -logger = logging.getLogger(__name__) - - -EXECUTOR = futures.ThreadPoolExecutor(8) - - -AI_AGENTS_MANAGER: AIAgentsManager | None = None - - -def prompt_ai_nbmodel_agent(room_id, jupyter_ingress, jupyter_token, kernel_id): - async def long_running_prompt(): - global AI_AGENTS_MANAGER - room_ws_url = http_to_ws(url_path_join(jupyter_ingress, "/api/collaboration/room", room_id)) - logger.info("AI Agent will connect to room [%s]…", room_ws_url) - has_runtime = jupyter_ingress and jupyter_token and kernel_id - prompt_agent = PromptAgent( - websocket_url=room_ws_url, - runtime_client=KernelClient( - server_url=jupyter_ingress, - token=jupyter_token, - kernel_id=kernel_id, - ) if has_runtime else None, - log=logger, - ) - logger.info("Starting AI Agent for room [%s]…", room_id) - async def prompt_task() -> None: - logger.info('Starting Prompt Agent.') - await prompt_agent.start() - if prompt_agent.runtime_client is not None: - prompt_agent.runtime_client.start() - logger.info('Prompt Agent is started.') - async with create_task_group() as tg: - tg.start_soon(prompt_task) - AI_AGENTS_MANAGER.register_ai_agent(room_id, prompt_agent) - # Sleep forever to keep the ai agent alive. - # TODO Replace with AI_AGENTS_MANAGER - while True: - await sleep(10) -# await AI_AGENTS_MANAGER.track_agent(room_id, prompt_agent) - return 'Prompt task is finished.' - with start_blocking_portal() as portal: - futures = [portal.start_task_soon(long_running_prompt)] - for future in as_completed(futures): - logger.info("Future is completed with result [%s]", future.result()) - - -class AINbModelAgentsInstanceHandler(APIHandler): - -# @web.authenticated - async def get(self, matched_part=None, *args, **kwargs): - global AI_AGENTS_MANAGER - if AI_AGENTS_MANAGER is None: - AI_AGENTS_MANAGER = AIAgentsManager() - self.write({ - "success": True, - "matched_part": matched_part, - }) - -# @web.authenticated - async def post(self, matched_part=None, *args, **kwargs): - global AI_AGENTS_MANAGER - if AI_AGENTS_MANAGER is None: - AI_AGENTS_MANAGER = AIAgentsManager() - body_data = json.loads(self.request.body) - logger.info("Body data", body_data) - self.write({ - "success": True, - "matched_part": matched_part, - }) - - -class AINbModelAgentsHandler(APIHandler): - -# @web.authenticated - async def get(self, *args, **kwargs): - global AI_AGENTS_MANAGER - if AI_AGENTS_MANAGER is None: - AI_AGENTS_MANAGER = AIAgentsManager() - self.write({ - "success": True, - }) - -# @web.authenticated - async def post(self, *args, **kwargs): - """Endpoint creating an AI Agent for a given room.""" - global AI_AGENTS_MANAGER - if AI_AGENTS_MANAGER is None: - AI_AGENTS_MANAGER = AIAgentsManager() - request_body = json.loads(self.request.body) - agent_request = NbModelAgentRequestModel(**request_body) - self.log.info("AI Agents create handler requested with [%s]", agent_request.model_dump()) - room_id = agent_request.room_id - if room_id in AI_AGENTS_MANAGER: - self.log.info("AI Agent for room [%s] already exists.", room_id) - # TODO check the ai agent. - return { - "success": True, - "message": "AI Agent already exists", - } - else: - self.log.info("Creating AI Agent for room [%s]…", room_id) - runtime = agent_request.runtime - jupyter_ingress = runtime.ingress - jupyter_token = runtime.token - kernel_id = runtime.kernel_id - # Start AI Agent in a ThreadPoolExecutor. - EXECUTOR.submit(prompt_ai_nbmodel_agent, room_id, jupyter_ingress, jupyter_token, kernel_id) - res = json.dumps({ - "success": True, - "message": f"AI Agent is started for room '{room_id}'.", - }) - logger.info("AI Agent create request exiting with reponse [%s]", res) - self.finish(res) diff --git a/jupyter_ai_agents/static/README.md b/jupyter_ai_agents/static/README.md index b7f1d9e..d0a4033 100644 --- a/jupyter_ai_agents/static/README.md +++ b/jupyter_ai_agents/static/README.md @@ -4,4 +4,4 @@ ~ BSD 3-Clause License --> -[![Datalayer](https://assets.datalayer.tech/datalayer-25.svg)](https://datalayer.io) +[![Datalayer](https://assets.datalayer.tech/datalayer-25.svg)](https://datalayer.ai) diff --git a/jupyter_ai_agents/tests/test_agent.py b/jupyter_ai_agents/tests/test_agent.py index 7442be5..6f3cd61 100644 --- a/jupyter_ai_agents/tests/test_agent.py +++ b/jupyter_ai_agents/tests/test_agent.py @@ -6,134 +6,3 @@ import uuid from unittest.mock import AsyncMock -from jupyter_nbmodel_client import NbModelClient -from jupyter_ai_agents.agents.langchain.base.agent_base import NbModelBaseAgent - - -async def test_default_content(ws_server): - room = uuid.uuid4().hex - async with NbModelBaseAgent(f"{ws_server}/{room}") as agent: - await asyncio.sleep(0) - default_content = agent.as_dict() - - assert default_content == {"cells": [], "metadata": {}, "nbformat": 4, "nbformat_minor": 5} - - -async def test_set_user_prompt(ws_server): - room = uuid.uuid4().hex - room_url = f"{ws_server}/{room}" - async with NbModelClient(room_url) as client: - async with NbModelBaseAgent(room_url) as agent: - agent._on_user_prompt = AsyncMock(return_value="hello") - idx = client.add_code_cell("print('hello')") - client.set_cell_metadata( - idx, - "datalayer", - {"ai": {"prompts": [{"id": "12345", "prompt": "Once upon a time"}]}}, - ) - - await asyncio.sleep(0.5) - - content = agent.as_dict() - assert content == { - "cells": [ - { - "cell_type": "code", - "execution_count": None, - "metadata": { - "datalayer": { - "ai": { - "messages": [ - { - "message": "hello", - "parent_id": "12345", - "timestamp": content["cells"][0]["metadata"][ - "datalayer" - ]["ai"]["messages"][0]["timestamp"], - "type": 1, - }, - ], - "prompts": [{"id": "12345", "prompt": "Once upon a time"}], - } - } - }, - "outputs": [], - "source": "print('hello')", - "id": client[idx]["id"], - } - ], - "metadata": {}, - "nbformat": 4, - "nbformat_minor": 5, - } - - assert agent._on_user_prompt.called - args, kwargs = agent._on_user_prompt.call_args - assert args == ( - client[idx]["id"], - "12345", - "Once upon a time", - None, - None, - ) - - -async def test_set_cell_with_user_prompt(ws_server): - room = uuid.uuid4().hex - room_url = f"{ws_server}/{room}" - async with NbModelClient(room_url) as client: - async with NbModelBaseAgent(room_url) as agent: - agent._on_user_prompt = AsyncMock() - client.add_code_cell( - "print('hello')", - metadata={ - "datalayer": { - "ai": {"prompts": [{"id": "12345", "prompt": "Once upon a time"}]} - } - }, - ) - - await asyncio.sleep(0.5) - - content = agent.as_dict() - assert content == { - "cells": [ - { - "cell_type": "code", - "execution_count": None, - "metadata": { - "datalayer": { - "ai": { - "messages": [ - { - "message": None, - "parent_id": "12345", - "timestamp": content["cells"][0]["metadata"][ - "datalayer" - ]["ai"]["messages"][0]["timestamp"], - "type": 1, - }, - ], - "prompts": [{"id": "12345", "prompt": "Once upon a time"}], - } - } - }, - "outputs": [], - "source": "print('hello')", - "id": client[0]["id"], - } - ], - "metadata": {}, - "nbformat": 4, - "nbformat_minor": 5, - } - - assert agent._on_user_prompt.called - args, kwargs = agent._on_user_prompt.call_args - assert args == ( - client[0]["id"], - "12345", - "Once upon a time", - None, - None, - ) diff --git a/jupyter_ai_agents/utils/__init__.py b/jupyter_ai_agents/utils/__init__.py new file mode 100644 index 0000000..f302d54 --- /dev/null +++ b/jupyter_ai_agents/utils/__init__.py @@ -0,0 +1,12 @@ +# Copyright (c) 2024-2025 Datalayer, Inc. +# +# BSD 3-Clause License + +"""Utility functions for Jupyter AI Agents.""" + +from .model import create_model_with_provider, get_model_string + +__all__ = [ + 'create_model_with_provider', + 'get_model_string', +] diff --git a/jupyter_ai_agents/utils/model.py b/jupyter_ai_agents/utils/model.py new file mode 100644 index 0000000..bdb237b --- /dev/null +++ b/jupyter_ai_agents/utils/model.py @@ -0,0 +1,159 @@ +# Copyright (c) 2024-2025 Datalayer, Inc. +# +# BSD 3-Clause License + +"""Model creation utilities for Jupyter AI Agents.""" + +import logging +import httpx + +logger = logging.getLogger(__name__) + + +def get_model_string(model_provider: str, model_name: str) -> str: + """ + Convert model provider and name to pydantic-ai model string format. + + Args: + model_provider: Provider name (azure-openai, openai, anthropic, github-copilot, etc.) + model_name: Model/deployment name + + Returns: + Model string in format 'provider:model' + For Azure OpenAI, returns the model name and sets provider via create_model_with_provider() + + Note: + For Azure OpenAI, the returned string is just the model name. + The Azure provider configuration is handled separately via OpenAIModel(provider='azure'). + Required env vars for Azure: + - AZURE_OPENAI_API_KEY + - AZURE_OPENAI_ENDPOINT (base URL only, e.g., https://your-resource.openai.azure.com) + - AZURE_OPENAI_API_VERSION (optional, defaults to latest) + """ + # For Azure OpenAI, we return just the model name + # The provider will be set to 'azure' when creating the OpenAIModel + if model_provider.lower() == 'azure-openai': + return model_name + + # Map provider names to pydantic-ai format for other providers + provider_map = { + 'openai': 'openai', + 'anthropic': 'anthropic', + 'github-copilot': 'openai', # GitHub Copilot uses OpenAI models + 'bedrock': 'bedrock', + 'google': 'google', + 'gemini': 'google', + 'groq': 'groq', + 'mistral': 'mistral', + 'cohere': 'cohere', + } + + provider = provider_map.get(model_provider.lower(), model_provider) + return f"{provider}:{model_name}" + + +def create_model_with_provider( + model_provider: str, + model_name: str, + timeout: float = 60.0, +): + """ + Create a pydantic-ai model object with the appropriate provider configuration. + + This is necessary for providers like Azure OpenAI that need special initialization + and timeout configuration. + + Args: + model_provider: Provider name (e.g., 'azure-openai', 'openai', 'anthropic') + model_name: Model/deployment name + timeout: HTTP timeout in seconds (default: 60.0) + + Returns: + Model object or string for pydantic-ai Agent + + Note: + For Azure OpenAI, requires these environment variables: + - AZURE_OPENAI_API_KEY + - AZURE_OPENAI_ENDPOINT (base URL only, e.g., https://your-resource.openai.azure.com) + - AZURE_OPENAI_API_VERSION (optional, defaults to latest) + """ + # Create httpx timeout configuration with generous connect timeout + # connect timeout is separate from read/write timeout + import httpx + http_timeout = httpx.Timeout(timeout, connect=30.0) + + logger.info(f"Creating model with timeout: {timeout}s (read/write), connect: 30.0s") + + if model_provider == 'azure-openai' or model_provider == 'azure': + from pydantic_ai.models.openai import OpenAIChatModel + from pydantic_ai.providers import infer_provider + from openai import AsyncAzureOpenAI + from pydantic_ai.providers.openai import OpenAIProvider + + # Infer Azure provider to get configuration + azure_provider = infer_provider('azure') + + # Extract base URL - remove /openai suffix since AsyncAzureOpenAI adds it + base_url = str(azure_provider.client.base_url) + # base_url is like: https://xxx.openai.azure.com/openai/ + # AsyncAzureOpenAI expects: https://xxx.openai.azure.com (it adds /openai automatically) + azure_endpoint = base_url.rstrip('/').rsplit('/openai', 1)[0] + + # Create AsyncAzureOpenAI client with custom timeout + azure_client = AsyncAzureOpenAI( + azure_endpoint=azure_endpoint, + azure_deployment=model_name, + api_version=azure_provider.client.default_query.get('api-version'), + api_key=azure_provider.client.api_key, + timeout=http_timeout, + ) + + # Wrap in OpenAIProvider + azure_provider_with_timeout = OpenAIProvider(openai_client=azure_client) + + return OpenAIChatModel(model_name, provider=azure_provider_with_timeout) + elif model_provider.lower() == 'anthropic': + from pydantic_ai.models.anthropic import AnthropicModel + from pydantic_ai.providers.anthropic import AnthropicProvider + from anthropic import AsyncAnthropic + + # Create Anthropic client with custom timeout and longer connect timeout + # Note: Many corporate networks block Anthropic API, use Azure/OpenAI if connection fails + anthropic_client = AsyncAnthropic( + timeout=httpx.Timeout(timeout, connect=60.0), # Longer connect timeout for slow/restricted networks + max_retries=2 + ) + + # Wrap in AnthropicProvider + anthropic_provider = AnthropicProvider(anthropic_client=anthropic_client) + + return AnthropicModel( + model_name, + provider=anthropic_provider + ) + elif model_provider.lower() in ['openai', 'github-copilot']: + from pydantic_ai.models.openai import OpenAIChatModel + from pydantic_ai.providers import infer_provider + from pydantic_ai.providers.openai import OpenAIProvider + + # For OpenAI, create OpenAIChatModel with custom http_client via provider + # First infer the OpenAI provider to get base_url, then pass custom http_client + http_client = httpx.AsyncClient(timeout=http_timeout, follow_redirects=True) + + # Infer OpenAI provider first to get proper configuration + openai_provider_base = infer_provider('openai') + + # Create new provider with same base_url but custom http_client + openai_provider = OpenAIProvider( + base_url=str(openai_provider_base.client.base_url), + http_client=http_client + ) + + return OpenAIChatModel( + model_name, + provider=openai_provider + ) + else: + # For other providers, use the standard string format + # Note: String format doesn't allow custom timeout configuration + return get_model_string(model_provider, model_name) diff --git a/package.json b/package.json index dca5b71..277a075 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "@datalayer/jupyter-ai-agents", - "version": "0.17.0", + "version": "0.18.0", "description": "Jupyter AI Agents.", "keywords": [ "ai", @@ -18,7 +18,7 @@ "license": "BSD-3-Clause", "author": { "name": "Datalayer", - "email": "info@datalayer.io" + "email": "info@datalayer.ai" }, "files": [ "lib/**/*.{d.ts,eot,gif,html,jpg,js,js.map,json,png,svg,woff2,ttf}", diff --git a/pyproject.toml b/pyproject.toml index cb2eae0..9cd0cde 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,19 +31,12 @@ dependencies = [ "datalayer_core", "fastapi", "httpx>=0.25.0", - "jupyter_kernel_client", "jupyter_mcp_tools", "jupyter_mcp_server", "jupyter-collaboration==4.0.2", "datalayer_pycrdt==0.12.17", - "jupyter_nbmodel_client", "jupyter_server>=2.10,<3", "jupyterlab>=4.0.0,<5", - "langchain", - "langchain-anthropic", - "langchain-aws", - "langchain-github-copilot", - "langchain-openai", "mcp", "pydantic-ai-slim[anthropic,openai,cli,mcp]>=1.9.0", "pydantic-settings", @@ -63,7 +56,7 @@ typing = ["mypy>=0.990"] [project.scripts] jaa = "jupyter_ai_agents.cli.app:main" jupyter-ai-agents = "jupyter_ai_agents.cli.app:main" -jupyter-ai-agents-console = "jupyter_ai_agents.cli.consoleapp:main" +jupyter-ai-agents-console = "jupyter_ai_agents.cli.console_app:main" [project.license] file = "LICENSE"