import requests
import json

# Test query with only_need_context to see what's being retrieved
base_url = 'http://localhost:3015'

print('🔍 Testing query with only_need_context to see retrieved content...')

# Login first
login_data = {
    'username': 'jleu3482',
    'password': 'jleu1212'
}

login_response = requests.post(f'{base_url}/login', data=login_data)
if login_response.status_code == 200:
    token = login_response.json().get('access_token')
    headers = {'Authorization': f'Bearer {token}'}
    print('✅ Login successful')
    
    # Test query with only_need_context to see what's retrieved
    print('🔍 Testing query with only_need_context=True...')
    query_data = {
        'query': 'railway maintenance', 
        'top_k': 10,
        'only_need_context': True
    }
    query_response = requests.post(f'{base_url}/query', json=query_data, headers=headers)
    
    if query_response.status_code == 200:
        context_results = query_response.json()
        print(f'✅ Context retrieved: {json.dumps(context_results, indent=2)}')
        
        # Check if we have chunks, entities, or relationships
        chunks = context_results.get('chunks', [])
        entities = context_results.get('entities', [])
        relationships = context_results.get('relationships', [])
        
        print(f'📄 Chunks found: {len(chunks)}')
        print(f'🏷️ Entities found: {len(entities)}')
        print(f'🔗 Relationships found: {len(relationships)}')
        
        if chunks:
            print('📝 First chunk content:')
            for i, chunk in enumerate(chunks[:3]):
                content = chunk.get("content", "")
                print(f'  Chunk {i+1}: {content[:200]}...')
                
    else:
        print(f'❌ Query failed: {query_response.status_code} - {query_response.text}')
        
else:
    print(f'❌ Login failed: {login_response.status_code} - {login_response.text}')