Context Branching SDK

A Python SDK for checkpoint-based conversation branching in LLM applications. Solve context pollution by creating isolated exploration branches without losing your place in the main conversation.

Overview

Problem: When exploring alternative approaches with LLMs, you pollute the conversation context, making it impossible to return to your original line of reasoning without starting fresh.

Solution: Context Branching SDK provides Git-like branching for LLM conversations:

• Checkpoint conversation state at decision points
• Branch to explore alternatives in isolation
• Switch between branches freely
• Inject insights from explorations back to main
• Export and share reproducible conversation states

Quick Start

Installation

pip install context-branching

# Or install from source
git clone https://github.com/glanzz/context-branching
cd context-branching
pip install -e .

Basic Usage

from context_branching import ContextBranchingSDK, Message

# Initialize SDK
sdk = ContextBranchingSDK(storage_backend="file", storage_path="./data")
workspace = sdk.create_workspace("my_session")

# Have a conversation
workspace.add_message(Message(role="user", content="Help me optimize this code"))
workspace.add_message(Message(role="assistant", content="Here are some options..."))

# Create checkpoint at decision point
checkpoint_id = workspace.create_checkpoint("optimization-decision")

# Branch to explore alternative
workspace.create_branch(checkpoint_id, "explore-rust")
workspace.switch_branch("explore-rust")
workspace.add_message(Message(role="user", content="What about Rust?"))

# Get context for LLM
context = workspace.get_current_context()  # Returns List[Message]
# ... send to your LLM ...

# Return to main and inject insights
workspace.switch_branch("main")
workspace.inject_messages("explore-rust", [0])  # Cherry-pick specific messages

Features

Core Capabilities

• Checkpoint Management: Create immutable snapshots of conversation state
• Branch Creation: Explore multiple paths without pollution
• Context Isolation: Each branch maintains its own message history
• Selective Injection: Cherry-pick insights to bring back to main
• Export/Import: Share checkpoints with team members
• Storage Agnostic: File, memory, or custom backends

Built-in Storage Backends

# File storage (default)
sdk = ContextBranchingSDK(storage_backend="file", storage_path="./my_data")

# In-memory storage (for testing)
sdk = ContextBranchingSDK(storage_backend="memory")

# Custom storage (implement StorageBackend interface)
from my_storage import PostgreSQLStorage
sdk = ContextBranchingSDK(custom_backend=PostgreSQLStorage())

Use Cases

1. Code Refactoring Decisions

# Explore Rust rewrite vs Python optimization in parallel
checkpoint = workspace.create_checkpoint("refactoring-decision")

# Branch 1: Explore Rust
workspace.create_branch(checkpoint, "rust-option")
workspace.switch_branch("rust-option")
# ... explore Rust approach ...

# Branch 2: Explore Python optimization
workspace.switch_branch("main")
workspace.create_branch(checkpoint, "python-option")
workspace.switch_branch("python-option")
# ... explore Python approach ...

# Return to main with insights from both
workspace.switch_branch("main")
workspace.inject_messages("rust-option", [1, 3])
workspace.inject_messages("python-option", [1, 2])

2. Debugging Investigations

Explore multiple hypotheses in parallel without polluting the main debugging conversation.

Example Applications

CLI Chat Application

python app/simple_chat.py --session my_session

# Commands available:
# /checkpoint <name>       - Create checkpoint
# /branch <name>          - Create and switch to branch
# /switch <name>          - Switch to existing branch
# /inject <branch> <idx>  - Inject messages from branch
# /status                 - Show workspace status

Web Chat Application

python app/web_chat.py --port 5000

# Visit http://localhost:5000 for web interface

Running Tests

# Install development dependencies
pip install -e ".[dev]"

# Run all tests
pytest tests/ -v

# Run with coverage
pytest tests/ --cov=sdk --cov-report=html

# Run specific test file
pytest tests/test_workspace.py -v

API Reference

Core Classes

ContextBranchingSDK

Main SDK interface.

sdk = ContextBranchingSDK(
    storage_backend="file",  # "file", "memory", or "custom"
    storage_path="./data",   # Path for file storage
    custom_backend=None      # Custom StorageBackend implementation
)

Workspace

Manages branches, checkpoints, and conversation state.

Methods:

• add_message(message: Message) - Add message to current branch
• create_checkpoint(name: str) -> str - Create checkpoint, returns ID
• create_branch(from_checkpoint: str, name: str) -> Branch - Create new branch
• switch_branch(name: str) - Switch to different branch
• get_current_context() -> List[Message] - Get full context for LLM
• inject_messages(from_branch: str, indices: List[int]) - Inject messages
• list_branches() -> List[Dict] - List all branches
• list_checkpoints() -> List[Dict] - List all checkpoints
• export_checkpoint(checkpoint_id: str, filepath: str) - Export checkpoint
• import_checkpoint(filepath: str) -> str - Import checkpoint
• get_statistics() -> Dict - Get workspace statistics

Message

Represents a single conversation message.

message = Message(
    role="user",           # "user", "assistant", or "system"
    content="Hello",       # Message content
    timestamp=datetime.now(),  # Optional, auto-generated
    metadata={}            # Optional metadata
)

Custom Storage Backend

Implement StorageBackend interface for custom storage:

from context_branching import StorageBackend

class MyCustomStorage(StorageBackend):
    def save_workspace(self, session_id: str, data: dict) -> None:
        # Implement save logic
        pass

    def load_workspace(self, session_id: str) -> dict:
        # Implement load logic
        pass

    def list_workspaces(self) -> List[str]:
        # Implement list logic
        pass

    def delete_workspace(self, session_id: str) -> None:
        # Implement delete logic
        pass

# Use custom backend
sdk = ContextBranchingSDK(custom_backend=MyCustomStorage())

Evaluations

The evaluation design and setup is present in results/quantitative_context_pollution. The folder contains the 30 conversation scenerios along with evaluated data results.

Development

Project Structure

context-branching-sdk/
├── sdk/                      # Core SDK library
│   ├── __init__.py
│   ├── sdk.py               # Main SDK class
│   ├── workspace.py         # Workspace & Message
│   ├── checkpoint.py        # Checkpoint implementation
│   ├── branch.py            # Branch implementation
│   └── storage.py           # Storage backends
├── app/                     # Example applications
│   ├── simple_chat.py       # CLI chat
│   └── web_chat.py          # Web chat
├── tests/                   # Unit and integration tests
├── docs/                    # Documentation
├── README.md
├── setup.py
└── requirements.txt

Contributing

1. Fork the repository
2. Create a feature branch
3. Make your changes
4. Run tests: pytest tests/
5. Submit a pull request

License

MIT License - see LICENSE file for details

Contact

For questions, issues, or contributions:

• GitHub Issues: https://github.com/glanzz/context-branching/issues
• Email: chickmagalurnanjun.b[at]northeastern[dot]edu

Status

Version: 0.1.0 (Alpha) Status: Research Prototype Python: 3.9+

---

Star ⭐ this repository if you find it useful!