Skip to content

IBM/chuk-ai-session-manager

Repository files navigation

CHUK AI Session Manager

A powerful session management system for AI applications

Python 3.11+ License: Apache 2.0

Automatic conversation tracking, token usage monitoring, tool call logging, infinite context support with automatic summarization, and hierarchical session relationships. Perfect for AI applications that need reliable session management.

πŸš€ Quick Start

Installation Options

# Basic installation (memory storage only)
pip install chuk-ai-session-manager

# With Redis support for persistent storage
pip install chuk-ai-session-manager[redis]

# With enhanced token counting
pip install chuk-ai-session-manager[tiktoken]

# Full installation with all optional features
pip install chuk-ai-session-manager[all]

# Development installation
pip install chuk-ai-session-manager[dev]

Quick Example

from chuk_ai_session_manager import track_conversation
from chuk_ai_session_manager.config import DEFAULT_TOKEN_MODEL

# Track any conversation automatically
session_id = await track_conversation(
    user_message="What's the weather like?",
    ai_response="I don't have access to real-time weather data.",
    model=DEFAULT_TOKEN_MODEL,
    provider="openai"
)

print(f"Conversation tracked in session: {session_id}")

That's it! Zero configuration required.

⚑ Major Features

🧠 AI Virtual Memory

OS-style memory management for AI context windows. Pages, working sets, faults, eviction, and compression β€” giving conversations the illusion of infinite memory.

from chuk_ai_session_manager import SessionManager
from chuk_ai_session_manager.config import DEFAULT_TOKEN_MODEL
from chuk_ai_session_manager.memory import (
    MemoryManager, CompressorRegistry, ImportanceWeightedLRU,
    PageType, VMMode, WorkingSetConfig,
)

# Zero-config: just enable VM on SessionManager
sm = SessionManager(enable_vm=True, vm_mode=VMMode.STRICT)

# Or fully customize eviction and compression
vm = MemoryManager(
    session_id="my_session",
    config=WorkingSetConfig(max_l0_tokens=32_000),
    eviction_policy=ImportanceWeightedLRU(),       # Swappable strategy
    compressor_registry=CompressorRegistry.default(), # Per-modality compression
)

# Create pages, add to working set
page = vm.create_page("Decision: Use JWT for auth", page_type=PageType.CLAIM)
await vm.add_to_working_set(page)

# Build context for LLM call
ctx = vm.build_context(system_prompt="You are helpful.")
# ctx["developer_message"] has VM:RULES + VM:MANIFEST_JSON + VM:CONTEXT

Eviction policies: ImportanceWeightedLRU (default), LRUEvictionPolicy, ModalityAwareLRU β€” or implement the EvictionPolicy protocol for custom strategies.

Compression: Pages compress through FULL β†’ REDUCED β†’ ABSTRACT β†’ REFERENCE before eviction, saving tokens without losing context. Text, image, and passthrough compressors included; plug in custom summarize_fn for LLM-based compression.

See AI Virtual Memory docs for full documentation.

🎯 Zero-Configuration Tracking

from chuk_ai_session_manager import SessionManager
from chuk_ai_session_manager.config import DEFAULT_TOKEN_MODEL

# Just start using it
sm = SessionManager()
await sm.user_says("Hello!")
await sm.ai_responds("Hi there!", model=DEFAULT_TOKEN_MODEL)

# Get stats instantly
stats = await sm.get_stats()
print(f"Tokens: {stats['total_tokens']}, Cost: ${stats['estimated_cost']:.4f}")

πŸ”„ Infinite Context

# Automatically handles conversations longer than token limits
sm = SessionManager(infinite_context=True, token_threshold=4000)
await sm.user_says("Tell me about the history of computing...")
await sm.ai_responds("Computing history begins with...", model=DEFAULT_TOKEN_MODEL)
# Session will auto-segment when limits are reached

βš™οΈ Storage Backends

Installation Storage Use Case Performance
pip install chuk-ai-session-manager Memory Development, testing 1.8M ops/sec
pip install chuk-ai-session-manager[redis] Redis Persistent, distributed 20K ops/sec

πŸ›‘οΈ Conversation Guards and Tool State

Runtime guardrails that prevent runaway tool loops, track value bindings, and enforce grounded tool calls.

from chuk_ai_session_manager.guards import get_tool_state, ToolStateManager

# Get the singleton tool state manager
tool_state = get_tool_state()

# Track tool calls and bind results as $v1, $v2, ...
binding = tool_state.bind_value("sqrt", {"x": 16}, 4.0)
# LLM can now reference $v1 in subsequent calls

# Check for runaway tool loops
status = tool_state.check_runaway()

# Detect ungrounded calls (missing $vN references)
check = tool_state.check_ungrounded_call("normal_cdf", {"mean": 0, "std": 1, "x": 1.5})

# Reset state for a new prompt
tool_state.reset_for_new_prompt()

Guard components:

  • ToolStateManager - Coordinator for all guards, bindings, and cache
  • BindingManager - $vN reference system for tracking tool results
  • ResultCache - Tool result caching for deduplication
  • UngroundedGuard - Detects calls with missing computed-value references
  • Runtime guards (budget, runaway, per-tool limits) from chuk-tool-processor

🧩 Procedural Memory

Learn from tool call history to improve future tool use.

from chuk_ai_session_manager import ToolMemoryManager, ProceduralContextFormatter

# Record tool outcomes
memory = ToolMemoryManager()
await memory.record("calculator", {"op": "add", "a": 5, "b": 3}, result=8, success=True)

# Format learned patterns for the model's context
formatter = ProceduralContextFormatter()
context = formatter.format(memory.get_patterns())

πŸ› οΈ Tool Integration

# Automatic tool call tracking
await sm.tool_used(
    tool_name="calculator",
    arguments={"operation": "add", "a": 5, "b": 3},
    result={"result": 8}
)

πŸ’‘ Common Use Cases

Web App Conversation Tracking

from chuk_ai_session_manager import track_conversation
from chuk_ai_session_manager.config import DEFAULT_TOKEN_MODEL

# In your chat endpoint
session_id = await track_conversation(
    user_message=request.message,
    ai_response=ai_response,
    model=DEFAULT_TOKEN_MODEL,
    provider="openai",
    session_id=request.session_id  # Continue existing conversation
)

LLM Wrapper with Automatic Tracking

from chuk_ai_session_manager import track_llm_call
from chuk_ai_session_manager.config import DEFAULT_TOKEN_MODEL
import openai

async def my_openai_call(prompt):
    response = await openai.chat.completions.create(
        model=DEFAULT_TOKEN_MODEL,
        messages=[{"role": "user", "content": prompt}]
    )
    return response.choices[0].message.content

# Automatically tracked
response, session_id = await track_llm_call(
    user_input="Explain machine learning",
    llm_function=my_openai_call,
    model=DEFAULT_TOKEN_MODEL,
    provider="openai"
)

Long Conversations with Auto-Segmentation

from chuk_ai_session_manager import track_infinite_conversation
from chuk_ai_session_manager.config import DEFAULT_TOKEN_MODEL # gpt-4o, gpt-4o-mini, etc.

# Start a conversation
session_id = await track_infinite_conversation(
    user_message="Tell me about the history of computing",
    ai_response="Computing history begins with ancient calculating devices...",
    model=DEFAULT_TOKEN_MODEL,
    token_threshold=4000  # Auto-segment after 4000 tokens
)

# Continue the conversation - will auto-segment if needed
session_id = await track_infinite_conversation(
    user_message="What about quantum computers?",
    ai_response="Quantum computing represents a fundamental shift...",
    session_id=session_id,
    model=DEFAULT_TOKEN_MODEL
)

πŸ”§ Configuration

Storage Configuration

# Memory provider (default) - fast, no persistence
export SESSION_PROVIDER=memory

# Redis provider - persistent, distributed (requires redis extra)
export SESSION_PROVIDER=redis
export SESSION_REDIS_URL=redis://localhost:6379/0

Model Configuration

DEFAULT_TOKEN_MODEL resolves to the default LLM model defined in:

  • config.py
  • or the CHUK_DEFAULT_MODEL environment variable

By default (as shipped), it is set to:

  • gpt-4o-mini

You can override it explicitly:

await sm.ai_responds(
    "Hello!",
    model="gpt-4o"
)

The selected model must be compatible with the specified provider.

Installation Matrix

Command Memory Redis Token Counting Use Case
pip install chuk-ai-session-manager βœ… ❌ Basic Development
pip install chuk-ai-session-manager[redis] βœ… βœ… Basic Persistent
pip install chuk-ai-session-manager[tiktoken] βœ… ❌ Enhanced Better accuracy
pip install chuk-ai-session-manager[all] βœ… βœ… Enhanced Full features

πŸ“Š Monitoring & Analytics

# Get comprehensive session analytics
stats = await sm.get_stats(include_all_segments=True)

print(f"""
πŸš€ Session Analytics Dashboard
============================
Session ID: {stats['session_id']}
Total Messages: {stats['total_messages']}
User Messages: {stats['user_messages']}
AI Messages: {stats['ai_messages']}
Tool Calls: {stats['tool_calls']}
Total Tokens: {stats['total_tokens']}
Total Cost: ${stats['estimated_cost']:.6f}
Session Segments: {stats.get('session_segments', 1)}
""")

πŸ—οΈ Why CHUK AI Session Manager?

  • Zero Configuration: Start tracking conversations in 3 lines of code
  • Infinite Context: Never worry about token limits again
  • Universal: Works with any LLM provider (OpenAI, Anthropic, etc.)
  • Robust: Built-in persistence, monitoring, and error handling
  • Token Aware: Automatic cost tracking across all providers
  • Tool Friendly: Seamless tool call logging and retry mechanisms
  • Guardrails: Runtime guards prevent runaway tool loops and ungrounded calls
  • Procedural Memory: Learn from tool call history to improve future use

πŸ›‘οΈ Error Handling

from chuk_ai_session_manager import (
    SessionManagerError,
    SessionNotFound,
    TokenLimitExceeded
)

try:
    session_id = await track_conversation("Hello", "Hi there")
except SessionNotFound as e:
    print(f"Session not found: {e}")
except TokenLimitExceeded as e:
    print(f"Token limit exceeded: {e}")
except SessionManagerError as e:
    print(f"General session error: {e}")

πŸ”„ Dependencies

  • Required: chuk-sessions (session storage), pydantic (data models), chuk-tool-processor (tool integration)
  • Optional: redis (Redis storage), tiktoken (accurate token counting)

πŸ“„ License

Apache 2.0 - build amazing AI applications with confidence!


Ready to build better AI applications?

pip install chuk-ai-session-manager

Start tracking conversations in 30 seconds!

About

No description, website, or topics provided.

Resources

License

Security policy

Stars

Watchers

Forks

Packages

 
 
 

Contributors