test_debug_output.py

#!/usr/bin/env python3
"""
Test script to verify:
1. BAAI model loading works without timeout interference
2. Debug output is properly saved
"""

import sys
from pathlib import Path

# Add project paths
project_root = Path('.').resolve()
src_path = project_root / 'layered-context-graph' / 'src'
sys.path.insert(0, str(src_path))

# Import required modules
from partitioning.unified_partition_manager import UnifiedPartitionManager, PartitionConfig
from models.baai_model import BAAIModel
from models.qwq_model import QwQModel
import os

print("=== Testing Debug Output ===\n")

# Initialize models
print("1. Loading BAAI model...")
baai_model = BAAIModel(model_path='bge-en-icl')
print("✓ BAAI model loaded successfully (no timeout!)")

print("\n2. Initializing QwQ model with Ollama...")
os.environ['OLLAMA_HOST'] = 'http://172.17.0.1:11434'
qwq_model = QwQModel(qwq_model_path='QwQ_LCoT_7B_Instruct')

# Configure for Ollama
try:
    import ollama
    qwq_model.ollama_client = ollama.Client(host=os.environ['OLLAMA_HOST'])
    qwq_model.heavy_model_loaded = True
    qwq_model.use_heavy_model = True
    qwq_model.set_fixed_context_size(31000)
    print("✓ QwQ model configured for Ollama")
except Exception as e:
    print(f"✗ Failed to configure Ollama: {e}")
    sys.exit(1)

# Test document
test_text = """
# Test Document

This is a test document with code:

```python
def hello():
    print("Hello, World!")
```

And some text after the code.
"""

print("\n3. Testing debug output with small document...")

# Create config with debug enabled
config = PartitionConfig(
    segmentation_method='linked_list',
    segmentation_params={
        'create_summaries': True,
        'min_segment_size': 100,
        'max_segment_size': 500,
        'merge_code_text': True,
        'debug_output': True  # Enable debug output
    }
)

# Process document
manager = UnifiedPartitionManager(config)

try:
    print("   Processing document...")
    root_id = manager.process_as_linked_list(
        test_text,
        baai_model=baai_model,
        llm_model=qwq_model,
        debug_mode=True
    )
    print(f"   ✓ Document processed successfully!")
    print(f"   Root ID: {root_id}")
    print(f"   Segments: {len(manager.segments)}")
    
    # Check if debug files were created
    debug_dir = Path("../segmentation_debug_outputs")
    if debug_dir.exists():
        sessions = list(debug_dir.glob("session_*"))
        if sessions:
            latest_session = max(sessions, key=lambda p: p.stat().st_mtime)
            files = list(latest_session.glob("*.txt"))
            print(f"\n   ✓ Debug output saved to: {latest_session}")
            print(f"   Files created: {len(files)}")
            for f in files[:5]:  # Show first 5 files
                print(f"     - {f.name}")
        else:
            print("   ⚠ No debug session folders found")
    else:
        print("   ⚠ Debug directory not created")
        
except Exception as e:
    print(f"   ✗ Error: {e}")
    import traceback
    traceback.print_exc()

print("\n=== Test Complete ===")