feat: High-leverage production improvements

Major improvements based on whitepaper priority:

1.  Neo4j Connection Reliability
   - Added retry logic with exponential backoff
   - Connection pooling (50 connections, 1hr max lifetime)
   - Automatic connection recovery on failures
   - Comprehensive error handling for auth/network issues
   - Wrap all DB operations with _execute_with_retry()

2.  INHERITS & Scope Resolution
   - Added full inheritance chain traversal (ancestors/descendants)
   - New resolve_symbol_scope() for smart symbol resolution
   - Scope search order: local -> direct imports -> transitive
   - get_inheritance_chain() for complete class hierarchies
   - Reduces hallucinations by providing proper context

3.  Jamba Encoder Integration (Frozen)
   - Reasoner.generate_refactor_plan() uses real AI21 Jamba
   - use_deep_context=True activates Phase 3 encoder
   - Full integration tests with AI21 Cloud API
   - 4/4 tests passing with real Jamba compression
   - Produces validated ContextBlock outputs

4.  Tree-Sitter AST-Based Masking
   - src/diffusion/masking.py (415 lines)
   - Deterministic masking anchored to AST node spans
   - Replaces heuristic token masking
   - Supports Python, JavaScript, TypeScript
   - 7 masking strategies (function_body, identifiers, types, etc.)
   - Preserves syntactic validity during diffusion
   - 15/15 tests passing

Test Results:
- test_jamba_integration.py: 4/4 PASS
- test_ast_masking.py: 15/15 PASS
- Real AI21 API compression verified

Ready for Phase 5 (outlines/CFG validation + Qwen fallback)

Author: vivek5200

.env.example

@@ -6,3 +6,35 @@ NEO4J_PASSWORD=ouroboros123
 # Provenance Metadata
 MODEL_NAME=ouroboros-librarian
 MODEL_VERSION=1.0.0
+
+# ============================================
+# Phase 2: LLM API Keys (The Reasoner)
+# ============================================
+
+# Claude (Anthropic)
+ANTHROPIC_API_KEY=your_anthropic_api_key_here
+
+# Gemini (Google)
+GOOGLE_API_KEY=your_google_api_key_here
+
+# OpenAI
+OPENAI_API_KEY=your_openai_api_key_here
+
+# Jamba (AI21)
+AI21_API_KEY=your_ai21_api_key_here
+
+# ============================================
+# Phase 3: Context Encoder Configuration
+# ============================================
+
+# AI21 Jamba Configuration
+# Get your API key: https://studio.ai21.com/account/api-key
+AI21_API_KEY=your_ai21_api_key_here
+
+# Context Encoder Mode
+# - "cloud": Use AI21 Cloud API (recommended, reliable)
+# - "local": Use LM Studio local inference (free, requires setup)
+JAMBA_MODE=cloud
+
+# LM Studio Configuration (only needed if JAMBA_MODE=local)
+LMSTUDIO_BASE_URL=http://localhost:1234/v1

docs/AI21_SETUP.md

@@ -0,0 +1,226 @@
+# AI21 Jamba Cloud Setup Guide
+
+## Why AI21 Cloud?
+
+✅ **Recommended for Production**
+- Reliable 99.9% uptime
+- No local GPU required
+- Automatic scaling
+- 256k token context window
+- Official support from AI21
+
+## Step 1: Get Your API Key
+
+1. Go to [AI21 Studio](https://studio.ai21.com/)
+2. Sign up / Log in
+3. Navigate to **Account → API Key**
+4. Click **"Create New API Key"**
+5. Copy your API key (starts with `AI21_...`)
+
+**Free Tier:**
+- $10 free credits on signup
+- ~200,000 tokens free
+- Good for testing Phase 3
+
+**Pricing (after free tier):**
+- Jamba-1.5-Mini: $0.20 per 1M input tokens
+- Jamba-1.5-Large: $2.00 per 1M input tokens
+
+## Step 2: Configure Ouroboros
+
+### Option A: Environment Variable (Recommended)
+
+```bash
+# Windows PowerShell
+$env:AI21_API_KEY="your_api_key_here"
+
+# Linux/Mac
+export AI21_API_KEY="your_api_key_here"
+```
+
+### Option B: .env File
+
+```bash
+# Copy the example file
+cp .env.example .env
+
+# Edit .env and add your key
+AI21_API_KEY=your_actual_api_key_here
+JAMBA_MODE=cloud
+```
+
+## Step 3: Test the Connection
+
+```python
+from src.context_encoder import ContextEncoder, ContextEncoderConfig
+from src.context_encoder.config import EncoderProvider, JambaConfig
+
+# Configure for AI21 Cloud
+jamba_config = JambaConfig(
+    use_cloud=True,
+    cloud_api_key="your_api_key_here"  # Or load from env
+)
+
+config = ContextEncoderConfig(
+    provider=EncoderProvider.JAMBA_CLOUD,
+    jamba=jamba_config
+)
+
+encoder = ContextEncoder(config)
+
+# Test compression
+sample_code = """
+export class AuthService {
+    async login(email: string, password: string) {
+        // Authentication logic
+    }
+}
+"""
+
+compressed = encoder.compress(
+    codebase_context=sample_code,
+    target_files=["auth.ts"]
+)
+
+print(f"✅ Compression successful!")
+print(f"Input: {compressed.tokens_in} tokens")
+print(f"Output: {compressed.tokens_out} tokens")
+print(f"Ratio: {compressed.compression_ratio:.1f}x")
+print(f"\nSummary:\n{compressed.summary}")
+```
+
+## Step 4: Use with Reasoner (Phase 2 Integration)
+
+```python
+from src.reasoner import Reasoner, ReasonerConfig
+from src.reasoner.config import LLMProvider
+
+# Initialize Reasoner
+config = ReasonerConfig(provider=LLMProvider.GEMINI)
+reasoner = Reasoner(config)
+
+# Generate refactor plan with deep context
+plan = reasoner.generate_refactor_plan(
+    task_description="Refactor authentication system",
+    target_file="src/auth/login.ts",
+    use_deep_context=True  # 🔥 Uses AI21 Jamba for 256k context
+)
+
+print(f"Plan ID: {plan.plan_id}")
+print(f"Impact: {plan.estimated_impact}")
+```
+
+## Troubleshooting
+
+### Error: "AI21_API_KEY environment variable not set"
+
+**Solution:** Set the environment variable or pass it explicitly:
+
+```python
+from src.context_encoder.config import JambaConfig
+
+jamba_config = JambaConfig(
+    use_cloud=True,
+    cloud_api_key="your_key_here"
+)
+```
+
+### Error: "Failed to initialize Jamba client"
+
+**Possible causes:**
+1. Invalid API key
+2. No internet connection
+3. AI21 API is down (check [status.ai21.com](https://status.ai21.com))
+
+**Solution:** Verify your API key at [studio.ai21.com](https://studio.ai21.com)
+
+### Error: Rate limit exceeded
+
+**Solution:** You've used your free credits. Options:
+1. Add payment method to AI21 account
+2. Switch to local mode (free): `JAMBA_MODE=local`
+3. Use mock provider for testing: `EncoderProvider.MOCK`
+
+### Slow response times
+
+**Normal behavior:**
+- First request: 10-30 seconds (cold start)
+- Subsequent requests: 3-10 seconds
+- Large context (100k+ tokens): 15-45 seconds
+
+**If consistently slow:**
+- Check your internet connection
+- Try a smaller context first
+- Consider using `max_output_tokens` to limit summary length
+
+## Switching Between Cloud and Local
+
+### Use Cloud (Recommended)
+
+```python
+config = ContextEncoderConfig(
+    provider=EncoderProvider.JAMBA_CLOUD,
+    jamba=JambaConfig(use_cloud=True)
+)
+```
+
+### Use Local (Free, Requires LM Studio)
+
+```python
+config = ContextEncoderConfig(
+    provider=EncoderProvider.JAMBA_LOCAL,
+    jamba=JambaConfig(
+        use_cloud=False,
+        local_base_url="http://localhost:1234/v1"
+    )
+)
+```
+
+See [LMSTUDIO_SETUP.md](./LMSTUDIO_SETUP.md) for local setup instructions.
+
+## Cost Estimation
+
+**Jamba-1.5-Mini Pricing:**
+
+| Context Size | Input Tokens | Output Tokens | Cost per Request |
+|--------------|--------------|---------------|------------------|
+| Small (10k)  | 10,000       | 2,000         | $0.002           |
+| Medium (50k) | 50,000       | 4,000         | $0.010           |
+| Large (100k) | 100,000      | 4,000         | $0.020           |
+| Massive (256k)| 256,000     | 4,000         | $0.051           |
+
+**Free tier gives you:**
+- ~5,000 requests (small context)
+- ~1,000 requests (medium context)
+- ~500 requests (large context)
+- ~200 requests (massive context)
+
+## Best Practices
+
+1. **Start with mock provider** for development:
+   ```python
+   config = ContextEncoderConfig(provider=EncoderProvider.MOCK)
+   ```
+
+2. **Use cloud for production** (reliable, scalable)
+
+3. **Use local for experimentation** (free, private)
+
+4. **Monitor your usage** at [studio.ai21.com](https://studio.ai21.com)
+
+5. **Cache compressed contexts** to avoid redundant API calls
+
+## Security
+
+⚠️ **Never commit your API key to Git!**
+
+- ✅ Use environment variables
+- ✅ Use `.env` file (in `.gitignore`)
+- ❌ Don't hardcode keys in source code
+- ❌ Don't share keys in screenshots/logs
+
+## Support
+
+- **AI21 Documentation:** https://docs.ai21.com/
+- **AI21 Discord:** https://discord.gg/ai21labs
+- **GitHub Issues:** https://github.com/vivek5200/ouroboros/issues

docs/index.md

@@ -21,15 +21,17 @@ Ouroboros is an autonomous agent designed to self-correct and refactor codebases
 ### Key Features
 
 * **⚡ Automated Refactoring**: Intelligent code analysis.
-* **🧠 LLM Integration**: Powered by LM Studio / Local LLMs.
+* **🧠 LLM Integration**: 6 providers + AI21 Jamba (256k context).
 * **🔄 Self-Healing**: Detects errors and proposes fixes iteratively.
+* **🚀 Deep Context**: AI21 Cloud integration for massive codebases.
 
 ### Documentation Status
 
 | Module | Status | Description |
 |:-------|:-------|:------------|
 | **Phase 1** | ✅ Complete | Code parsing and graph database construction. |
 | **Phase 2** | ✅ Complete | LLM-powered refactor plan generation with 6 providers. |
+| **Phase 3** | ✅ Complete | Jamba context encoder with 256k token window (AI21 Cloud). |
 
 ---
 

src/context_encoder/config.py

@@ -12,42 +12,66 @@
 
 class EncoderProvider(Enum):
     """Supported context encoder providers."""
-    JAMBA = "jamba"  # AI21 Jamba-1.5-Mini (Hybrid Mamba-Transformer)
-    MOCK = "mock"    # Mock encoder for testing
+    JAMBA_CLOUD = "jamba_cloud"  # AI21 Cloud API (recommended for production)
+    JAMBA_LOCAL = "jamba_local"  # LM Studio local inference
+    MOCK = "mock"                # Mock encoder for testing
 
 
 @dataclass
 class JambaConfig:
     """Configuration for AI21 Jamba-1.5-Mini model."""
 
-    model_name: str = "ai21/jamba-1.5-mini"
+    # Model Configuration
+    model_name: str = "jamba-mini"  # AI21 model name (jamba-mini or jamba-large)
     context_window: int = 256_000  # 256k tokens
     max_output_tokens: int = 4096  # For context summary
     temperature: float = 0.3  # Low temperature for deterministic summaries
 
-    # API Configuration
-    base_url: str = "http://localhost:1234/v1"  # LM Studio default
-    api_key: Optional[str] = None
-    timeout: int = 300  # 5 minutes for large context
+    # API Configuration (Cloud vs Local)
+    use_cloud: bool = True  # True: AI21 Cloud, False: LM Studio local
+    
+    # AI21 Cloud Configuration
+    cloud_api_url: str = "https://api.ai21.com/studio/v1"
+    cloud_api_key: Optional[str] = None  # Set via AI21_API_KEY env var
+    
+    # LM Studio Local Configuration
+    local_base_url: str = "http://localhost:1234/v1"  # LM Studio default
+    local_api_key: Optional[str] = None  # Not needed for local
 
-    # Performance
+    # Common Configuration
+    timeout: int = 300  # 5 minutes for large context
     batch_size: int = 1
     use_streaming: bool = False
 
     def __post_init__(self):
-        """Validate configuration."""
+        """Validate configuration and load from environment."""
         if self.max_output_tokens > self.context_window:
             raise ValueError(
                 f"max_output_tokens ({self.max_output_tokens}) cannot exceed "
                 f"context_window ({self.context_window})"
             )
+        
+        # Load API key from environment if not set
+        if self.use_cloud and not self.cloud_api_key:
+            import os
+            self.cloud_api_key = os.getenv("AI21_API_KEY")
+    
+    @property
+    def base_url(self) -> str:
+        """Get the active base URL based on cloud/local mode."""
+        return self.cloud_api_url if self.use_cloud else self.local_base_url
+    
+    @property
+    def api_key(self) -> Optional[str]:
+        """Get the active API key based on cloud/local mode."""
+        return self.cloud_api_key if self.use_cloud else self.local_api_key
 
 
 @dataclass
 class ContextEncoderConfig:
     """Main configuration for the Context Encoder."""
 
-    provider: EncoderProvider = EncoderProvider.JAMBA
+    provider: EncoderProvider = EncoderProvider.JAMBA_CLOUD
     jamba: JambaConfig = field(default_factory=JambaConfig)
 
     # Context Compression Settings