Metadata-Version: 2.4
Name: jupyterlab-cellsense
Version: 0.3.0
Summary: Multi-file FIM code completion for JupyterLab.
Project-URL: Homepage, https://github.com/arunkarthik11/CellSense
Project-URL: Repository, https://github.com/arunkarthik11/CellSense
Project-URL: Issues, https://github.com/arunkarthik11/CellSense/issues
Author: Arun Karthik
License: GPL 3.0
License-File: LICENSE.md
Requires-Python: >=3.10
Requires-Dist: httpx>=0.27
Requires-Dist: jupyter-server<3,>=2.0
Requires-Dist: pathspec>=0.12.1
Requires-Dist: rank-bm25>=0.2.2
Requires-Dist: watchdog>=3.0.0
Description-Content-Type: text/markdown

<p align="center">
  <img src="assets/cellsense-logo.png" alt="CellSense logo" width="140" />
</p>

<h1 align="center">CellSense</h1>

<p align="center">Ghost-text inline code completion for JupyterLab, powered by local or remote AI models.</p>

<p align="center">
  <a href="https://pypi.org/project/jupyterlab-cellsense/"><img alt="pip install jupyterlab-cellsense" src="https://img.shields.io/pypi/v/jupyterlab-cellsense?logo=pypi&label=pip%20install" /></a>
  <a href="https://huggingface.co/spaces/arun11karthik/cellsense-fim"><img alt="Hugging Face Space — CellSense-FIM" src="https://img.shields.io/badge/Hugging%20Face-CellSense--FIM-FFD21E?logo=huggingface&logoColor=black" /></a>
  <a href="https://huggingface.co/arun11karthik"><img alt="Hugging Face — model weights" src="https://img.shields.io/badge/Models-CellSense--FIM-FFD21E?logo=huggingface&logoColor=black" /></a>
  <a href="LICENSE.md"><img alt="License: GPL v3" src="https://img.shields.io/badge/License-GPLv3-blue.svg" /></a>
</p>

---

## Overview

### What is this?

CellSense is a JupyterLab 4 extension that adds FIM (Fill-In-Middle) inline completions to your notebooks and file editor. As you type, it silently assembles context from your current cell, surrounding cells, and related files across your workspace — then requests a ghost-text suggestion from your chosen model. Press Tab to accept.

![CellSense ghost-text suggestion appearing in a notebook cell](assets/screenshots/AutoCompleteSuggestion.png)

### Features

- **Ghost-text completions** — inline suggestions appear as you type, accepted with Tab
- **Multi-file retrieval** — Indexes your workspace and surfaces relevant code chunks automatically
- **Recent State Aware** - Maintains a state of recently read and editted code chunks to provide rich context to your completions
- **Import tracking** — resolves local imports and injects relevant signatures into the prompt
- **Multi-cell context** — walks outward from your cursor through surrounding cells within a configurable token budget
- **Flexible providers** — works with Ollama (local) or any OpenAI-compatible API (OpenAI, vLLM, LM Studio, Azure, etc.)
- **Multiple model families** — first-class prompt support for DeepSeek Coder, Codestral, CodeLlama, StarCoder2, and Qwen2.5-Coder
- **In-JupyterLab settings** — configure everything from a sidebar panel without touching config files

### Why use it?

Most notebook completion tools treat each cell in isolation, or at best are aware of the surrounding cells in the currently open notebook. The quality of a completion is determined by the quality of context given to the model. CellSense is built around the idea that the most accurate completions come from a model that understands the current task, the tools at your disposal, and your typical coding style. It achieves this through three design principles:

**Local awareness** — Your utility functions, domain classes, and project-specific abstractions will never appear in any model's training data. CellSense addresses this by traversing your import tree and injecting the relevant function and class signatures directly into the prompt. The model reasons about your actual code, not a generic approximation of it.

![CellSense completing code using a local utility class](assets/screenshots/LocalUtilityAware.png)

**Recent state** — When working through a problem, the files you have been reading and editing are almost always the most relevant to your current task. CellSense maintains a rolling state of recently visited and modified code chunks, so the context reflects where your attention has actually been — not just what happens to be open in the active tab.

**Relevant chunk retrieval** — Beyond recency, CellSense indexes your workspace and retrieves code chunks that are lexically related to what you are currently writing. This surfaces patterns and solutions from elsewhere in your project that a cell-local view would miss entirely.

Together, these give the model a view of your work that is closer to your own mental model — without requiring any manual context management on your part. It can run entirely on local hardware if you use Ollama, with no data ever leaving your machine.

---

## Installation

### Prerequisites

- Python 3.10 or higher
- JupyterLab 4.0 or higher
- For local inference: [Ollama](https://ollama.com) installed and running
- For cloud inference: an API key for an OpenAI-compatible provider

### Install

#### Option 1 — pip (recommended)

```bash
pip install jupyterlab-cellsense
```

#### Option 2 — Clone the repo and run the install script

Clone the repository, then run the script for your platform from inside the project directory:

**macOS / Linux**
```bash
git clone https://github.com/arunkarthik11/CellSense.git
cd CellSense
./install.sh
```

**Windows (Command Prompt)**
```bat
git clone https://github.com/arunkarthik11/CellSense.git
cd CellSense
install.bat
```

**Windows (PowerShell)**
```powershell
git clone https://github.com/arunkarthik11/CellSense.git
cd CellSense
.\install.ps1
```

> Note: On first run, PowerShell may require you to allow script execution:
> `Set-ExecutionPolicy -Scope CurrentUser RemoteSigned`

#### Option 3 — GitHub release

Download the `.whl` file from the [latest release](https://github.com/arunkarthik11/CellSense/releases/latest) and install it directly:

```bash
pip install jupyterlab_cellsense-<version>-py3-none-any.whl
```

#### Option 4 — pip with git

Install the latest commit directly from GitHub without cloning:

```bash
pip install git+https://github.com/arunkarthik11/CellSense.git
```

### Verify installation

```bash
jupyter labextension list
```

You should see `jupyterlab-cellsense` in the output. Then start JupyterLab:

```bash
jupyter lab
```

---

## CellSense-FIM models

CellSense works with any FIM-capable model, but we also trained a family of in-house models — **CellSense-FIM** — that understand the long-context, repo-, import-, and task-aware prompt structure CellSense assembles *natively*. They are fine-tuned from Qwen2.5-Coder on real Jupyter notebook repositories, share one prompt format across every size, and come in four sizes so you can pick the one that fits your hardware. All sizes support a 32K context window.

> **These models are a recommendation, not a requirement.** CellSense remains fully model-agnostic — you can keep using any Ollama or OpenAI-compatible model you like (DeepSeek Coder, Codestral, CodeLlama, StarCoder2, Qwen2.5-Coder, etc.). The CellSense-FIM family is simply tuned to get the most out of the context CellSense builds.

The full interactive showcase lives in the [CellSense-FIM Hugging Face Space](https://huggingface.co/spaces/arun11karthik/cellsense-fim), and the weights are on the [Hugging Face hub](https://huggingface.co/arun11karthik).

### Model family

| Model | Params | Context | Weights |
|-------|--------|---------|---------|
| CellSense-FIM 0.5B | 0.5B | 32K | [`arun11karthik/cellsense-fim-0.5b`](https://huggingface.co/arun11karthik/cellsense-fim-0.5b) |
| CellSense-FIM 1.5B | 1.5B | 32K | [`arun11karthik/cellsense-fim-1.5b`](https://huggingface.co/arun11karthik/cellsense-fim-1.5b) |
| CellSense-FIM 3B | 3B | 32K | [`arun11karthik/cellsense-fim-3b`](https://huggingface.co/arun11karthik/cellsense-fim-3b) |
| CellSense-FIM 7B | 7B | 32K | [`arun11karthik/cellsense-fim-7b`](https://huggingface.co/arun11karthik/cellsense-fim-7b) |

### Evaluation

Held-out `test` split of the CellSense FIM dataset, scored against each Qwen2.5-Coder base model and a same-size, later-generation general model (Qwen3). Higher is better on all four completion-quality metrics.

![CellSense-FIM 7B vs. base models across Edit Similarity, BLEU, CodeBLEU, and Token Accuracy](assets/model_metrics_7b.png)

BLEU climbs sharply once the model understands the notebook-native FIM format — CellSense-FIM stays roughly 10× above the base models at every size:

![BLEU score across model size from 0.5B to 7B](assets/bleu_by_size.png)

### Bits per byte — lower is better

How tightly each model predicts held-out notebook code. CellSense-FIM is the most confident at every size.

| Size | Qwen2.5-Coder | Qwen3 | CellSense-FIM |
|------|---------------|-------|---------------|
| 0.5B | 0.305 | 1.082 | **0.254** |
| 1.5B | 0.273 | 1.097 | **0.197** |
| 3B | 0.250 | 1.052 | **0.182** |
| 7B | 0.234 | 1.041 | **0.169** |

### Running a CellSense-FIM model

Serve it with vLLM (OpenAI-compatible) or run it locally with Ollama, then point CellSense at it.

**Serve with vLLM**
```bash
vllm serve arun11karthik/cellsense-fim-7b \
  --served-model-name cellsense-fim-7b \
  --max-model-len 32768 --port 8000
```
Then in **CellSense → Basic Settings**: Provider `OpenAI Compatible`, Base URL `http://localhost:8000/v1`, Model `cellsense-fim-7b`, Model Family `cellsense`.

**Run locally with Ollama**

Make sure the Ollama server is running, then pull the model:
```bash
ollama serve
ollama pull hf.co/arun11karthik/cellsense-fim-7b-GGUF:Q5_K_M
```

The Hugging Face model tag is long, so copy it to a simpler local name with `ollama cp`:
```bash
ollama cp hf.co/arun11karthik/cellsense-fim-7b-GGUF:Q5_K_M cellsense-fim-7b
```
You can then just use `cellsense-fim-7b` as the model tag in Ollama.

Then in **CellSense → Basic Settings**: Provider `Ollama`, Base URL `http://localhost:11434`, Model `cellsense-fim-7b` (or the full `hf.co/...` tag), Model Family `cellsense`.

Swap `7b` for `0.5b` / `1.5b` / `3b`. Ollama quantization tags: `Q5_K_M`, `Q8_0`, `BF16`.

---

## Quick Start

1. Start Ollama and pull a FIM model:
   ```bash
   ollama pull qwen2.5-coder:14b
   ```
2. Launch JupyterLab and open a notebook.
3. Open the CellSense panel from the left sidebar.
4. Confirm the provider is set to **Ollama** and the model matches what you pulled.
5. Click **Save & Apply**.
6. Start typing in a cell — a ghost-text suggestion will appear after a short pause. Press **Tab** to accept.

![CellSense settings panel configured with Ollama and qwen2.5-coder](assets/screenshots/BasicSettingsConfiguration.png)

---

## Configuration

Open the CellSense settings panel from the left sidebar, or via the Command Palette (`Ctrl+Shift+C` / `Cmd+Shift+C`) → **CellSense Settings**.

![CellSense advanced settings panel showing completion and retrieval options](assets/screenshots/AdvancedSettings.png)

### Available Options

#### Provider

| Option | Values | Default | Description |
|--------|--------|---------|-------------|
| `provider` | `ollama`, `openai_compatible` | `ollama` | Which backend to use for completions |

#### Ollama

| Option | Default | Description |
|--------|---------|-------------|
| `base_url` | `http://localhost:11434` | Address of your Ollama server |
| `model_family` | `deepseek-coder` | FIM prompt format — must match your model. Options: `deepseek-coder`, `codestral`, `codellama`, `starcoder2`, `qwen2.5-coder`, `other` |
| `model` | `deepseek-coder:6.7b` | Full Ollama model tag |

#### OpenAI Compatible

| Option | Default | Description |
|--------|---------|-------------|
| `base_url` | `https://api.openai.com/v1` | API endpoint — change for vLLM, LM Studio, Codestral, etc. |
| `api_key` | *Provided API KEY* | API key — The API key provided by the Model Provider |
| `model` | `gpt-4o-mini` | Model identifier |

#### Completion

| Option | Default | Description |
|--------|---------|-------------|
| `debounce_ms` | `500` | Milliseconds of idle time before a completion request is sent |
| `max_context_tokens` | `1500` | Token budget for cell-walking — cells are included outward from your cursor until this limit is reached |
| `max_tokens` | `256` | Maximum tokens the model can generate per completion |
| `one_line_completions` | `false` | Restrict suggestions to a single line. This is generally faster and more accurate, but generates fewer tokens. |

![One-line completion mode enabled in the advanced settings panel](assets/screenshots/OneLineCompletions.png)

#### Retrieval (RAG)

| Option | Default | Description |
|--------|---------|-------------|
| `enabled` | `true` | Enable BM25 workspace retrieval to inject relevant code chunks into the prompt |
| `top_k` | `3` | Number of code chunks retrieved from the workspace index (1–20) |
| `token_budget` | `2000` | Maximum tokens allocated to retrieved context |

**Warning**: If you are using an external model provider who trains on your data, be aware that enabling RAG could retrieve code chunks from your repository or workspace root to enrich the context. This means that context from other files could also be sent. If this is not something you are comfortable with, either switch to a provider who allows you to disable training or disable RAG for this feature. However, this would still traverse the import tree and retrieve the relevant signatures.

---

## Usage

### Basic Workflow

1. Open a notebook and start typing code in a cell.
2. Pause briefly — CellSense will assemble context from your notebook and workspace, then request a completion.
3. A ghost-text suggestion appears inline.
4. Press **Tab** to accept the full suggestion.
5. To dismiss without accepting, keep typing or press **Escape**.

![CellSense inline completion in a Python file](assets/screenshots/PythonFileCompatible.png)

### Commands / UI Elements

| Element | How to access |
|---------|---------------|
| Settings sidebar panel | Left sidebar icon |
| CellSense Settings (full dialog) | Command Palette → "CellSense Settings" |
| Accept completion | Tab |
| Dismiss completion | Escape or continue typing |

The settings panel is divided into three sections:

- **Basic Settings** — provider selection, model configuration
- **Advanced: Completion** — debounce, token budgets, one-line mode
- **Advanced: Retrieval** — enable/disable RAG, top-K, token budget

Click **Save & Apply** after making changes. Settings take effect immediately without restarting JupyterLab.

---

## Additional examples

### Using Ollama with Codestral

1. Pull the model: `ollama pull codestral`
2. In the settings panel, set **Provider** to `Ollama`, **Model Family** to `codestral`, and **Model** to `codestral`.
3. Save & Apply.

### Using an OpenAI-compatible local server (LM Studio)

1. Start your LM Studio server.
2. Set **Provider** to `OpenAI Compatible`, **Base URL** to `http://localhost:1234/v1`, leave **API Key** blank, and set **Model** to the model name shown in LM Studio.
3. Save & Apply.

### Disabling retrieval for faster completions

In the settings panel, expand **Advanced: Retrieval** and toggle **Enable Retrieval** off. This skips the BM25 workspace index and sends only the notebook cell context.

---

## Troubleshooting

**No completions appear**
- Check the health endpoint in the browser: `http://localhost:8888/cellsense/health`
- If using Ollama, confirm it is running: `ollama list` and check that the right port is provided
- Verify the model name in settings exactly matches the tag in Ollama.

**Completions are slow**
- Reduce `max_tokens` (try 64–128).
- Reduce `max_context_tokens` to limit the amount of notebook context sent.
- Enable One Line completions
- Disable retrieval if your workspace is large.

**Tab key doesn't accept**
- Ensure the cell is in edit mode (click inside the cell first).
- Check that JupyterLab's inline completer is enabled: **Settings → Inline Completer → Enable**.

**`jupyter labextension list` doesn't show CellSense**
- Run `pip show jupyterlab-cellsense` to confirm the package is installed in the active environment.
- Ensure you're running `jupyter lab` from the same environment where you ran `pip install`.

**PowerShell script is blocked**
- Run: `Set-ExecutionPolicy -Scope CurrentUser RemoteSigned` and try again.

---

## Compatibility

| Component | Version |
|-----------|---------|
| JupyterLab | 4.0+ |
| Python | 3.10+ |
| Jupyter Server | 2.x |
| Ollama | Any recent version |

**Supported model families**

| Family | Example models |
|--------|---------------|
| DeepSeek Coder | `deepseek-coder:6.7b`, `deepseek-coder:33b` |
| Codestral | `codestral` |
| CodeLlama | `codellama:7b`, `codellama:13b` |
| StarCoder2 | `starcoder2:3b`, `starcoder2:7b` |
| Qwen2.5-Coder | `qwen2.5-coder:7b` |
| CellSense-FIM | `cellsense-fim-0.5b`, `cellsense-fim-1.5b`, `cellsense-fim-3b`, `cellsense-fim-7b` |
| Generic FIM | Any FIM model not listed above |
| OpenAI compatible | `gpt-4o-mini`, vLLM-served models, LM Studio, etc. |

---

## Contributing

Contributions are welcome. Please open an issue before submitting a pull request for significant changes.

1. Fork the repository
2. Create a feature branch: `git checkout -b feature/my-feature`
3. Commit your changes and open a pull request against `master`

For bug reports and feature requests, use the [issue tracker](https://github.com/arunkarthik11/CellSense/issues).

---

## Inspirations

**[Fill-In-Middle (FIM)](https://arxiv.org/abs/2207.14255)** — Bavarian et al., OpenAI, 2022. The foundational technique CellSense uses to structure every prompt: the model is given a prefix and a suffix and asked to generate the text that belongs in between, rather than simply continuing from the end of a document.

**[StarCoder](https://arxiv.org/abs/2305.06161)** — Li et al., 2023. StarCoder popularised open, openly licensed FIM-capable code models and established the `<fim_prefix>` / `<fim_suffix>` / `<fim_middle>` token conventions that CellSense's prompt templates are built around.

**[JupyterLab](https://github.com/jupyterlab/jupyterlab)** — CellSense is built directly on JupyterLab 4's `IInlineCompletionProvider` API for ghost-text suggestions and the `jupyter_server` extension system for server-side request handling. None of this required patching JupyterLab internals — both are stable, public extension points.

### Related projects

If you are looking for a **chat-based** AI coding assistant for JupyterLab rather than inline completions, [NotebookIntelligence](https://github.com/notebook-intelligence/notebook-intelligence) is an awesome tool. CellSense focuses exclusively on rich context based ghost-text completion; the two tools are complementary.

---

## License

CellSense is licensed under the [GNU General Public License v3.0](LICENSE.md).
