chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,98 @@
|
||||
"""
|
||||
Compare performance before and after optimizations
|
||||
"""
|
||||
|
||||
def read_baseline():
|
||||
"""Read baseline performance metrics"""
|
||||
with open('performance_baseline.txt', 'r') as f:
|
||||
content = f.read()
|
||||
|
||||
# Extract key metrics
|
||||
metrics = {}
|
||||
lines = content.split('\n')
|
||||
for i, line in enumerate(lines):
|
||||
if 'Total Time:' in line:
|
||||
metrics['total_time'] = float(line.split(':')[1].strip().split()[0])
|
||||
elif 'Memory Used:' in line:
|
||||
metrics['memory_mb'] = float(line.split(':')[1].strip().split()[0])
|
||||
elif 'validate_coverage:' in line and i+1 < len(lines) and 'Avg Time:' in lines[i+2]:
|
||||
metrics['validate_coverage_ms'] = float(lines[i+2].split(':')[1].strip().split()[0])
|
||||
elif 'select_links:' in line and i+1 < len(lines) and 'Avg Time:' in lines[i+2]:
|
||||
metrics['select_links_ms'] = float(lines[i+2].split(':')[1].strip().split()[0])
|
||||
elif 'calculate_confidence:' in line and i+1 < len(lines) and 'Avg Time:' in lines[i+2]:
|
||||
metrics['calculate_confidence_ms'] = float(lines[i+2].split(':')[1].strip().split()[0])
|
||||
|
||||
return metrics
|
||||
|
||||
|
||||
def print_comparison(before_metrics, after_metrics):
|
||||
"""Print performance comparison"""
|
||||
print("\n" + "="*80)
|
||||
print("PERFORMANCE COMPARISON: BEFORE vs AFTER OPTIMIZATIONS")
|
||||
print("="*80)
|
||||
|
||||
# Total time
|
||||
time_improvement = (before_metrics['total_time'] - after_metrics['total_time']) / before_metrics['total_time'] * 100
|
||||
print(f"\n📊 Total Time:")
|
||||
print(f" Before: {before_metrics['total_time']:.2f} seconds")
|
||||
print(f" After: {after_metrics['total_time']:.2f} seconds")
|
||||
print(f" Improvement: {time_improvement:.1f}% faster ✅" if time_improvement > 0 else f" Slower: {-time_improvement:.1f}% ❌")
|
||||
|
||||
# Memory
|
||||
mem_improvement = (before_metrics['memory_mb'] - after_metrics['memory_mb']) / before_metrics['memory_mb'] * 100
|
||||
print(f"\n💾 Memory Usage:")
|
||||
print(f" Before: {before_metrics['memory_mb']:.2f} MB")
|
||||
print(f" After: {after_metrics['memory_mb']:.2f} MB")
|
||||
print(f" Improvement: {mem_improvement:.1f}% less memory ✅" if mem_improvement > 0 else f" More memory: {-mem_improvement:.1f}% ❌")
|
||||
|
||||
# Key operations
|
||||
print(f"\n⚡ Key Operations:")
|
||||
|
||||
# Validate coverage
|
||||
if 'validate_coverage_ms' in before_metrics and 'validate_coverage_ms' in after_metrics:
|
||||
val_improvement = (before_metrics['validate_coverage_ms'] - after_metrics['validate_coverage_ms']) / before_metrics['validate_coverage_ms'] * 100
|
||||
print(f"\n validate_coverage:")
|
||||
print(f" Before: {before_metrics['validate_coverage_ms']:.1f} ms")
|
||||
print(f" After: {after_metrics['validate_coverage_ms']:.1f} ms")
|
||||
print(f" Improvement: {val_improvement:.1f}% faster ✅" if val_improvement > 0 else f" Slower: {-val_improvement:.1f}% ❌")
|
||||
|
||||
# Select links
|
||||
if 'select_links_ms' in before_metrics and 'select_links_ms' in after_metrics:
|
||||
sel_improvement = (before_metrics['select_links_ms'] - after_metrics['select_links_ms']) / before_metrics['select_links_ms'] * 100
|
||||
print(f"\n select_links:")
|
||||
print(f" Before: {before_metrics['select_links_ms']:.1f} ms")
|
||||
print(f" After: {after_metrics['select_links_ms']:.1f} ms")
|
||||
print(f" Improvement: {sel_improvement:.1f}% faster ✅" if sel_improvement > 0 else f" Slower: {-sel_improvement:.1f}% ❌")
|
||||
|
||||
# Calculate confidence
|
||||
if 'calculate_confidence_ms' in before_metrics and 'calculate_confidence_ms' in after_metrics:
|
||||
calc_improvement = (before_metrics['calculate_confidence_ms'] - after_metrics['calculate_confidence_ms']) / before_metrics['calculate_confidence_ms'] * 100
|
||||
print(f"\n calculate_confidence:")
|
||||
print(f" Before: {before_metrics['calculate_confidence_ms']:.1f} ms")
|
||||
print(f" After: {after_metrics['calculate_confidence_ms']:.1f} ms")
|
||||
print(f" Improvement: {calc_improvement:.1f}% faster ✅" if calc_improvement > 0 else f" Slower: {-calc_improvement:.1f}% ❌")
|
||||
|
||||
print("\n" + "="*80)
|
||||
|
||||
# Overall assessment
|
||||
if time_improvement > 50:
|
||||
print("🎉 EXCELLENT OPTIMIZATION! More than 50% performance improvement!")
|
||||
elif time_improvement > 30:
|
||||
print("✅ GOOD OPTIMIZATION! Significant performance improvement!")
|
||||
elif time_improvement > 10:
|
||||
print("👍 DECENT OPTIMIZATION! Noticeable performance improvement!")
|
||||
else:
|
||||
print("🤔 MINIMAL IMPROVEMENT. Further optimization may be needed.")
|
||||
|
||||
print("="*80)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Example usage - you'll run this after implementing optimizations
|
||||
baseline = read_baseline()
|
||||
print("Baseline metrics loaded:")
|
||||
for k, v in baseline.items():
|
||||
print(f" {k}: {v}")
|
||||
|
||||
print("\n⚠️ Run the performance test again after optimizations to compare!")
|
||||
print("Then update this script with the new metrics to see the comparison.")
|
||||
@@ -0,0 +1,293 @@
|
||||
"""
|
||||
Test and demo script for Adaptive Crawler
|
||||
|
||||
This script demonstrates the progressive crawling functionality
|
||||
with various configurations and use cases.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from pathlib import Path
|
||||
import time
|
||||
from typing import Dict, List
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
from rich.progress import Progress
|
||||
from rich import print as rprint
|
||||
|
||||
# Add parent directory to path for imports
|
||||
import sys
|
||||
sys.path.append(str(Path(__file__).parent.parent))
|
||||
|
||||
from crawl4ai import (
|
||||
AsyncWebCrawler,
|
||||
AdaptiveCrawler,
|
||||
AdaptiveConfig,
|
||||
CrawlState
|
||||
)
|
||||
|
||||
|
||||
console = Console()
|
||||
|
||||
|
||||
|
||||
|
||||
def print_relevant_content(crawler: AdaptiveCrawler, top_k: int = 3):
|
||||
"""Print most relevant content found"""
|
||||
relevant = crawler.get_relevant_content(top_k=top_k)
|
||||
|
||||
if not relevant:
|
||||
console.print("[yellow]No relevant content found yet.[/yellow]")
|
||||
return
|
||||
|
||||
console.print(f"\n[bold cyan]Top {len(relevant)} Most Relevant Pages:[/bold cyan]")
|
||||
for i, doc in enumerate(relevant, 1):
|
||||
console.print(f"\n[green]{i}. {doc['url']}[/green]")
|
||||
console.print(f" Score: {doc['score']:.2f}")
|
||||
# Show snippet
|
||||
content = doc['content'] or ""
|
||||
snippet = content[:200].replace('\n', ' ') + "..." if len(content) > 200 else content
|
||||
console.print(f" [dim]{snippet}[/dim]")
|
||||
|
||||
|
||||
async def test_basic_progressive_crawl():
|
||||
"""Test basic progressive crawling functionality"""
|
||||
console.print("\n[bold yellow]Test 1: Basic Progressive Crawl[/bold yellow]")
|
||||
console.print("Testing on Python documentation with query about async/await")
|
||||
|
||||
config = AdaptiveConfig(
|
||||
confidence_threshold=0.7,
|
||||
max_pages=10,
|
||||
top_k_links=2,
|
||||
min_gain_threshold=0.1
|
||||
)
|
||||
|
||||
# Create crawler
|
||||
async with AsyncWebCrawler() as crawler:
|
||||
prog_crawler = AdaptiveCrawler(
|
||||
crawler=crawler,
|
||||
config=config
|
||||
)
|
||||
|
||||
# Start progressive crawl
|
||||
start_time = time.time()
|
||||
state = await prog_crawler.digest(
|
||||
start_url="https://docs.python.org/3/library/asyncio.html",
|
||||
query="async await context managers"
|
||||
)
|
||||
elapsed = time.time() - start_time
|
||||
|
||||
# Print results
|
||||
prog_crawler.print_stats(detailed=False)
|
||||
prog_crawler.print_stats(detailed=True)
|
||||
print_relevant_content(prog_crawler)
|
||||
|
||||
console.print(f"\n[green]Crawl completed in {elapsed:.2f} seconds[/green]")
|
||||
console.print(f"Final confidence: {prog_crawler.confidence:.2%}")
|
||||
console.print(f"URLs crawled: {list(state.crawled_urls)[:5]}...") # Show first 5
|
||||
|
||||
# Test export functionality
|
||||
export_path = "knowledge_base_export.jsonl"
|
||||
prog_crawler.export_knowledge_base(export_path)
|
||||
console.print(f"[green]Knowledge base exported to {export_path}[/green]")
|
||||
|
||||
# Clean up
|
||||
Path(export_path).unlink(missing_ok=True)
|
||||
|
||||
|
||||
async def test_with_persistence():
|
||||
"""Test state persistence and resumption"""
|
||||
console.print("\n[bold yellow]Test 2: Persistence and Resumption[/bold yellow]")
|
||||
console.print("Testing state save/load functionality")
|
||||
|
||||
state_path = "test_crawl_state.json"
|
||||
|
||||
config = AdaptiveConfig(
|
||||
confidence_threshold=0.6,
|
||||
max_pages=5,
|
||||
top_k_links=2,
|
||||
save_state=True,
|
||||
state_path=state_path
|
||||
)
|
||||
|
||||
# First crawl - partial
|
||||
async with AsyncWebCrawler() as crawler:
|
||||
prog_crawler = AdaptiveCrawler(
|
||||
crawler=crawler,
|
||||
config=config
|
||||
)
|
||||
|
||||
state1 = await prog_crawler.digest(
|
||||
start_url="https://httpbin.org",
|
||||
query="http headers response"
|
||||
)
|
||||
|
||||
console.print(f"[cyan]First crawl: {len(state1.crawled_urls)} pages[/cyan]")
|
||||
|
||||
# Resume crawl
|
||||
config.max_pages = 10 # Increase limit
|
||||
|
||||
async with AsyncWebCrawler() as crawler:
|
||||
prog_crawler = AdaptiveCrawler(
|
||||
crawler=crawler,
|
||||
config=config
|
||||
)
|
||||
|
||||
state2 = await prog_crawler.digest(
|
||||
start_url="https://httpbin.org",
|
||||
query="http headers response",
|
||||
resume_from=state_path
|
||||
)
|
||||
|
||||
console.print(f"[green]Resumed crawl: {len(state2.crawled_urls)} total pages[/green]")
|
||||
|
||||
# Clean up
|
||||
Path(state_path).unlink(missing_ok=True)
|
||||
|
||||
|
||||
async def test_different_domains():
|
||||
"""Test on different types of websites"""
|
||||
console.print("\n[bold yellow]Test 3: Different Domain Types[/bold yellow]")
|
||||
|
||||
test_cases = [
|
||||
{
|
||||
"name": "Documentation Site",
|
||||
"url": "https://docs.python.org/3/",
|
||||
"query": "decorators and context managers"
|
||||
},
|
||||
{
|
||||
"name": "API Documentation",
|
||||
"url": "https://httpbin.org",
|
||||
"query": "http authentication headers"
|
||||
}
|
||||
]
|
||||
|
||||
for test in test_cases:
|
||||
console.print(f"\n[cyan]Testing: {test['name']}[/cyan]")
|
||||
console.print(f"URL: {test['url']}")
|
||||
console.print(f"Query: {test['query']}")
|
||||
|
||||
config = AdaptiveConfig(
|
||||
confidence_threshold=0.6,
|
||||
max_pages=5,
|
||||
top_k_links=2
|
||||
)
|
||||
|
||||
async with AsyncWebCrawler() as crawler:
|
||||
prog_crawler = AdaptiveCrawler(
|
||||
crawler=crawler,
|
||||
config=config
|
||||
)
|
||||
|
||||
start_time = time.time()
|
||||
state = await prog_crawler.digest(
|
||||
start_url=test['url'],
|
||||
query=test['query']
|
||||
)
|
||||
elapsed = time.time() - start_time
|
||||
|
||||
# Summary using print_stats
|
||||
prog_crawler.print_stats(detailed=False)
|
||||
|
||||
|
||||
async def test_stopping_criteria():
|
||||
"""Test different stopping criteria"""
|
||||
console.print("\n[bold yellow]Test 4: Stopping Criteria[/bold yellow]")
|
||||
|
||||
# Test 1: High confidence threshold
|
||||
console.print("\n[cyan]4.1 High confidence threshold (0.9)[/cyan]")
|
||||
config = AdaptiveConfig(
|
||||
confidence_threshold=0.9, # Very high
|
||||
max_pages=20,
|
||||
top_k_links=3
|
||||
)
|
||||
|
||||
async with AsyncWebCrawler() as crawler:
|
||||
prog_crawler = AdaptiveCrawler(crawler=crawler, config=config)
|
||||
state = await prog_crawler.digest(
|
||||
start_url="https://docs.python.org/3/library/",
|
||||
query="python standard library"
|
||||
)
|
||||
|
||||
console.print(f"Pages needed for 90% confidence: {len(state.crawled_urls)}")
|
||||
prog_crawler.print_stats(detailed=False)
|
||||
|
||||
# Test 2: Page limit
|
||||
console.print("\n[cyan]4.2 Page limit (3 pages max)[/cyan]")
|
||||
config = AdaptiveConfig(
|
||||
confidence_threshold=0.9,
|
||||
max_pages=3, # Very low limit
|
||||
top_k_links=2
|
||||
)
|
||||
|
||||
async with AsyncWebCrawler() as crawler:
|
||||
prog_crawler = AdaptiveCrawler(crawler=crawler, config=config)
|
||||
state = await prog_crawler.digest(
|
||||
start_url="https://docs.python.org/3/library/",
|
||||
query="python standard library modules"
|
||||
)
|
||||
|
||||
console.print(f"Stopped by: {'Page limit' if len(state.crawled_urls) >= 3 else 'Other'}")
|
||||
prog_crawler.print_stats(detailed=False)
|
||||
|
||||
|
||||
async def test_crawl_patterns():
|
||||
"""Analyze crawl patterns and link selection"""
|
||||
console.print("\n[bold yellow]Test 5: Crawl Pattern Analysis[/bold yellow]")
|
||||
|
||||
config = AdaptiveConfig(
|
||||
confidence_threshold=0.7,
|
||||
max_pages=8,
|
||||
top_k_links=2,
|
||||
min_gain_threshold=0.05
|
||||
)
|
||||
|
||||
async with AsyncWebCrawler() as crawler:
|
||||
prog_crawler = AdaptiveCrawler(crawler=crawler, config=config)
|
||||
|
||||
# Track crawl progress
|
||||
console.print("\n[cyan]Crawl Progress:[/cyan]")
|
||||
|
||||
state = await prog_crawler.digest(
|
||||
start_url="https://httpbin.org",
|
||||
query="http methods post get"
|
||||
)
|
||||
|
||||
# Show crawl order
|
||||
console.print("\n[green]Crawl Order:[/green]")
|
||||
for i, url in enumerate(state.crawl_order, 1):
|
||||
console.print(f"{i}. {url}")
|
||||
|
||||
# Show new terms discovered per page
|
||||
console.print("\n[green]New Terms Discovered:[/green]")
|
||||
for i, new_terms in enumerate(state.new_terms_history, 1):
|
||||
console.print(f"Page {i}: {new_terms} new terms")
|
||||
|
||||
# Final metrics
|
||||
console.print(f"\n[yellow]Saturation reached: {state.metrics.get('saturation', 0):.2%}[/yellow]")
|
||||
|
||||
|
||||
async def main():
|
||||
"""Run all tests"""
|
||||
console.print("[bold magenta]Adaptive Crawler Test Suite[/bold magenta]")
|
||||
console.print("=" * 50)
|
||||
|
||||
try:
|
||||
# Run tests
|
||||
await test_basic_progressive_crawl()
|
||||
# await test_with_persistence()
|
||||
# await test_different_domains()
|
||||
# await test_stopping_criteria()
|
||||
# await test_crawl_patterns()
|
||||
|
||||
console.print("\n[bold green]✅ All tests completed successfully![/bold green]")
|
||||
|
||||
except Exception as e:
|
||||
console.print(f"\n[bold red]❌ Test failed with error: {e}[/bold red]")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Run the test suite
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,182 @@
|
||||
"""
|
||||
Test script for debugging confidence calculation in adaptive crawler
|
||||
Focus: Testing why confidence decreases when crawling relevant URLs
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import List, Dict
|
||||
import math
|
||||
|
||||
# Add parent directory to path for imports
|
||||
sys.path.append(str(Path(__file__).parent.parent))
|
||||
|
||||
from crawl4ai import AsyncWebCrawler
|
||||
from crawl4ai.adaptive_crawler import CrawlState, StatisticalStrategy
|
||||
from crawl4ai.models import CrawlResult
|
||||
|
||||
|
||||
class ConfidenceTestHarness:
|
||||
"""Test harness for analyzing confidence calculation"""
|
||||
|
||||
def __init__(self):
|
||||
self.strategy = StatisticalStrategy()
|
||||
self.test_urls = [
|
||||
'https://docs.python.org/3/library/asyncio.html',
|
||||
'https://docs.python.org/3/library/asyncio-runner.html',
|
||||
'https://docs.python.org/3/library/asyncio-api-index.html',
|
||||
'https://docs.python.org/3/library/contextvars.html',
|
||||
'https://docs.python.org/3/library/asyncio-stream.html'
|
||||
]
|
||||
self.query = "async await context manager"
|
||||
|
||||
async def test_confidence_progression(self):
|
||||
"""Test confidence calculation as we crawl each URL"""
|
||||
print(f"Testing confidence for query: '{self.query}'")
|
||||
print("=" * 80)
|
||||
|
||||
# Initialize state
|
||||
state = CrawlState(query=self.query)
|
||||
|
||||
# Create crawler
|
||||
async with AsyncWebCrawler() as crawler:
|
||||
for i, url in enumerate(self.test_urls, 1):
|
||||
print(f"\n{i}. Crawling: {url}")
|
||||
print("-" * 80)
|
||||
|
||||
# Crawl the URL
|
||||
result = await crawler.arun(url=url)
|
||||
|
||||
# Extract markdown content
|
||||
if hasattr(result, '_results') and result._results:
|
||||
result = result._results[0]
|
||||
|
||||
# Create a mock CrawlResult with markdown
|
||||
mock_result = type('CrawlResult', (), {
|
||||
'markdown': type('Markdown', (), {
|
||||
'raw_markdown': result.markdown.raw_markdown if hasattr(result, 'markdown') else ''
|
||||
})(),
|
||||
'url': url
|
||||
})()
|
||||
|
||||
# Update state
|
||||
state.knowledge_base.append(mock_result)
|
||||
await self.strategy.update_state(state, [mock_result])
|
||||
|
||||
# Calculate metrics
|
||||
confidence = await self.strategy.calculate_confidence(state)
|
||||
|
||||
# Get individual components
|
||||
coverage = state.metrics.get('coverage', 0)
|
||||
consistency = state.metrics.get('consistency', 0)
|
||||
saturation = state.metrics.get('saturation', 0)
|
||||
|
||||
# Analyze term frequencies
|
||||
query_terms = self.strategy._tokenize(self.query.lower())
|
||||
term_stats = {}
|
||||
for term in query_terms:
|
||||
term_stats[term] = {
|
||||
'tf': state.term_frequencies.get(term, 0),
|
||||
'df': state.document_frequencies.get(term, 0)
|
||||
}
|
||||
|
||||
# Print detailed results
|
||||
print(f"State after crawl {i}:")
|
||||
print(f" Total documents: {state.total_documents}")
|
||||
print(f" Unique terms: {len(state.term_frequencies)}")
|
||||
print(f" New terms added: {state.new_terms_history[-1] if state.new_terms_history else 0}")
|
||||
|
||||
print(f"\nQuery term statistics:")
|
||||
for term, stats in term_stats.items():
|
||||
print(f" '{term}': tf={stats['tf']}, df={stats['df']}")
|
||||
|
||||
print(f"\nMetrics:")
|
||||
print(f" Coverage: {coverage:.3f}")
|
||||
print(f" Consistency: {consistency:.3f}")
|
||||
print(f" Saturation: {saturation:.3f}")
|
||||
print(f" → Confidence: {confidence:.3f}")
|
||||
|
||||
# Show coverage calculation details
|
||||
print(f"\nCoverage calculation details:")
|
||||
self._debug_coverage_calculation(state, query_terms)
|
||||
|
||||
# Alert if confidence decreased
|
||||
if i > 1 and confidence < state.metrics.get('prev_confidence', 0):
|
||||
print(f"\n⚠️ WARNING: Confidence decreased from {state.metrics.get('prev_confidence', 0):.3f} to {confidence:.3f}")
|
||||
|
||||
state.metrics['prev_confidence'] = confidence
|
||||
|
||||
def _debug_coverage_calculation(self, state: CrawlState, query_terms: List[str]):
|
||||
"""Debug coverage calculation step by step"""
|
||||
coverage_score = 0.0
|
||||
max_possible_score = 0.0
|
||||
|
||||
for term in query_terms:
|
||||
tf = state.term_frequencies.get(term, 0)
|
||||
df = state.document_frequencies.get(term, 0)
|
||||
|
||||
if df > 0:
|
||||
idf = math.log((state.total_documents - df + 0.5) / (df + 0.5) + 1)
|
||||
doc_coverage = df / state.total_documents
|
||||
tf_boost = min(tf / df, 3.0)
|
||||
term_score = doc_coverage * idf * (1 + 0.1 * math.log1p(tf_boost))
|
||||
|
||||
print(f" '{term}': doc_cov={doc_coverage:.2f}, idf={idf:.2f}, boost={1 + 0.1 * math.log1p(tf_boost):.2f} → score={term_score:.3f}")
|
||||
coverage_score += term_score
|
||||
else:
|
||||
print(f" '{term}': not found → score=0.000")
|
||||
|
||||
max_possible_score += 1.0 * 1.0 * 1.1
|
||||
|
||||
print(f" Total: {coverage_score:.3f} / {max_possible_score:.3f} = {coverage_score/max_possible_score if max_possible_score > 0 else 0:.3f}")
|
||||
|
||||
# New coverage calculation
|
||||
print(f"\n NEW Coverage calculation (without IDF):")
|
||||
new_coverage = self._calculate_coverage_new(state, query_terms)
|
||||
print(f" → New Coverage: {new_coverage:.3f}")
|
||||
|
||||
def _calculate_coverage_new(self, state: CrawlState, query_terms: List[str]) -> float:
|
||||
"""New coverage calculation without IDF"""
|
||||
if not query_terms or state.total_documents == 0:
|
||||
return 0.0
|
||||
|
||||
term_scores = []
|
||||
max_tf = max(state.term_frequencies.values()) if state.term_frequencies else 1
|
||||
|
||||
for term in query_terms:
|
||||
tf = state.term_frequencies.get(term, 0)
|
||||
df = state.document_frequencies.get(term, 0)
|
||||
|
||||
if df > 0:
|
||||
# Document coverage: what fraction of docs contain this term
|
||||
doc_coverage = df / state.total_documents
|
||||
|
||||
# Frequency signal: normalized log frequency
|
||||
freq_signal = math.log(1 + tf) / math.log(1 + max_tf) if max_tf > 0 else 0
|
||||
|
||||
# Combined score: document coverage with frequency boost
|
||||
term_score = doc_coverage * (1 + 0.5 * freq_signal)
|
||||
|
||||
print(f" '{term}': doc_cov={doc_coverage:.2f}, freq_signal={freq_signal:.2f} → score={term_score:.3f}")
|
||||
term_scores.append(term_score)
|
||||
else:
|
||||
print(f" '{term}': not found → score=0.000")
|
||||
term_scores.append(0.0)
|
||||
|
||||
# Average across all query terms
|
||||
coverage = sum(term_scores) / len(term_scores)
|
||||
return coverage
|
||||
|
||||
|
||||
async def main():
|
||||
"""Run the confidence test"""
|
||||
tester = ConfidenceTestHarness()
|
||||
await tester.test_confidence_progression()
|
||||
|
||||
print("\n" + "=" * 80)
|
||||
print("Test complete!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,254 @@
|
||||
"""
|
||||
Performance test for Embedding Strategy optimizations
|
||||
Measures time and memory usage before and after optimizations
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
import tracemalloc
|
||||
import numpy as np
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import os
|
||||
|
||||
# Add parent directory to path for imports
|
||||
sys.path.append(str(Path(__file__).parent.parent.parent))
|
||||
|
||||
from crawl4ai import AsyncWebCrawler, AdaptiveCrawler, AdaptiveConfig
|
||||
from crawl4ai.adaptive_crawler import EmbeddingStrategy, CrawlState
|
||||
from crawl4ai.models import CrawlResult
|
||||
|
||||
|
||||
class PerformanceMetrics:
|
||||
def __init__(self):
|
||||
self.start_time = 0
|
||||
self.end_time = 0
|
||||
self.start_memory = 0
|
||||
self.peak_memory = 0
|
||||
self.operation_times = {}
|
||||
|
||||
def start(self):
|
||||
tracemalloc.start()
|
||||
self.start_time = time.perf_counter()
|
||||
self.start_memory = tracemalloc.get_traced_memory()[0]
|
||||
|
||||
def end(self):
|
||||
self.end_time = time.perf_counter()
|
||||
current, peak = tracemalloc.get_traced_memory()
|
||||
self.peak_memory = peak
|
||||
tracemalloc.stop()
|
||||
|
||||
def record_operation(self, name: str, duration: float):
|
||||
if name not in self.operation_times:
|
||||
self.operation_times[name] = []
|
||||
self.operation_times[name].append(duration)
|
||||
|
||||
@property
|
||||
def total_time(self):
|
||||
return self.end_time - self.start_time
|
||||
|
||||
@property
|
||||
def memory_used_mb(self):
|
||||
return (self.peak_memory - self.start_memory) / 1024 / 1024
|
||||
|
||||
def print_summary(self, label: str):
|
||||
print(f"\n{'='*60}")
|
||||
print(f"Performance Summary: {label}")
|
||||
print(f"{'='*60}")
|
||||
print(f"Total Time: {self.total_time:.3f} seconds")
|
||||
print(f"Memory Used: {self.memory_used_mb:.2f} MB")
|
||||
|
||||
if self.operation_times:
|
||||
print("\nOperation Breakdown:")
|
||||
for op, times in self.operation_times.items():
|
||||
avg_time = sum(times) / len(times)
|
||||
total_time = sum(times)
|
||||
print(f" {op}:")
|
||||
print(f" - Calls: {len(times)}")
|
||||
print(f" - Avg Time: {avg_time*1000:.2f} ms")
|
||||
print(f" - Total Time: {total_time:.3f} s")
|
||||
|
||||
|
||||
async def create_mock_crawl_results(n: int) -> list:
|
||||
"""Create mock crawl results for testing"""
|
||||
results = []
|
||||
for i in range(n):
|
||||
class MockMarkdown:
|
||||
def __init__(self, content):
|
||||
self.raw_markdown = content
|
||||
|
||||
class MockResult:
|
||||
def __init__(self, url, content):
|
||||
self.url = url
|
||||
self.markdown = MockMarkdown(content)
|
||||
self.success = True
|
||||
|
||||
content = f"This is test content {i} about async await coroutines event loops. " * 50
|
||||
result = MockResult(f"https://example.com/page{i}", content)
|
||||
results.append(result)
|
||||
return results
|
||||
|
||||
|
||||
async def test_embedding_performance():
|
||||
"""Test the performance of embedding strategy operations"""
|
||||
|
||||
# Configuration
|
||||
n_kb_docs = 30 # Number of documents in knowledge base
|
||||
n_queries = 10 # Number of query variations
|
||||
n_links = 50 # Number of candidate links
|
||||
n_iterations = 5 # Number of calculation iterations
|
||||
|
||||
print(f"\nTest Configuration:")
|
||||
print(f"- Knowledge Base Documents: {n_kb_docs}")
|
||||
print(f"- Query Variations: {n_queries}")
|
||||
print(f"- Candidate Links: {n_links}")
|
||||
print(f"- Iterations: {n_iterations}")
|
||||
|
||||
# Create embedding strategy
|
||||
config = AdaptiveConfig(
|
||||
strategy="embedding",
|
||||
max_pages=50,
|
||||
n_query_variations=n_queries,
|
||||
embedding_model="sentence-transformers/all-MiniLM-L6-v2" # 384 dimensions
|
||||
)
|
||||
|
||||
# Set up API key if available
|
||||
if os.getenv('OPENAI_API_KEY'):
|
||||
config.embedding_llm_config = {
|
||||
'provider': 'openai/text-embedding-3-small',
|
||||
'api_token': os.getenv('OPENAI_API_KEY'),
|
||||
'embedding_model': 'text-embedding-3-small'
|
||||
}
|
||||
else:
|
||||
config.embedding_llm_config = {
|
||||
'provider': 'openai/gpt-4o-mini',
|
||||
'api_token': 'dummy-key'
|
||||
}
|
||||
|
||||
strategy = EmbeddingStrategy(
|
||||
embedding_model=config.embedding_model,
|
||||
llm_config=config.embedding_llm_config
|
||||
)
|
||||
strategy.config = config
|
||||
|
||||
# Initialize state
|
||||
state = CrawlState()
|
||||
state.query = "async await coroutines event loops tasks"
|
||||
|
||||
# Start performance monitoring
|
||||
metrics = PerformanceMetrics()
|
||||
metrics.start()
|
||||
|
||||
# 1. Generate query embeddings
|
||||
print("\n1. Generating query embeddings...")
|
||||
start = time.perf_counter()
|
||||
query_embeddings, expanded_queries = await strategy.map_query_semantic_space(
|
||||
state.query,
|
||||
config.n_query_variations
|
||||
)
|
||||
state.query_embeddings = query_embeddings
|
||||
state.expanded_queries = expanded_queries
|
||||
metrics.record_operation("query_embedding", time.perf_counter() - start)
|
||||
print(f" Generated {len(query_embeddings)} query embeddings")
|
||||
|
||||
# 2. Build knowledge base incrementally
|
||||
print("\n2. Building knowledge base...")
|
||||
mock_results = await create_mock_crawl_results(n_kb_docs)
|
||||
|
||||
for i in range(0, n_kb_docs, 5): # Add 5 documents at a time
|
||||
batch = mock_results[i:i+5]
|
||||
start = time.perf_counter()
|
||||
await strategy.update_state(state, batch)
|
||||
metrics.record_operation("update_state", time.perf_counter() - start)
|
||||
state.knowledge_base.extend(batch)
|
||||
|
||||
print(f" Knowledge base has {len(state.kb_embeddings)} documents")
|
||||
|
||||
# 3. Test repeated confidence calculations
|
||||
print(f"\n3. Testing {n_iterations} confidence calculations...")
|
||||
for i in range(n_iterations):
|
||||
start = time.perf_counter()
|
||||
confidence = await strategy.calculate_confidence(state)
|
||||
metrics.record_operation("calculate_confidence", time.perf_counter() - start)
|
||||
print(f" Iteration {i+1}: {confidence:.3f} ({(time.perf_counter() - start)*1000:.1f} ms)")
|
||||
|
||||
# 4. Test coverage gap calculations
|
||||
print(f"\n4. Testing coverage gap calculations...")
|
||||
for i in range(n_iterations):
|
||||
start = time.perf_counter()
|
||||
gaps = strategy.find_coverage_gaps(state.kb_embeddings, state.query_embeddings)
|
||||
metrics.record_operation("find_coverage_gaps", time.perf_counter() - start)
|
||||
print(f" Iteration {i+1}: {len(gaps)} gaps ({(time.perf_counter() - start)*1000:.1f} ms)")
|
||||
|
||||
# 5. Test validation
|
||||
print(f"\n5. Testing validation coverage...")
|
||||
for i in range(n_iterations):
|
||||
start = time.perf_counter()
|
||||
val_score = await strategy.validate_coverage(state)
|
||||
metrics.record_operation("validate_coverage", time.perf_counter() - start)
|
||||
print(f" Iteration {i+1}: {val_score:.3f} ({(time.perf_counter() - start)*1000:.1f} ms)")
|
||||
|
||||
# 6. Create mock links for ranking
|
||||
from crawl4ai.models import Link
|
||||
mock_links = []
|
||||
for i in range(n_links):
|
||||
link = Link(
|
||||
href=f"https://example.com/new{i}",
|
||||
text=f"Link about async programming {i}",
|
||||
title=f"Async Guide {i}"
|
||||
)
|
||||
mock_links.append(link)
|
||||
|
||||
# 7. Test link selection
|
||||
print(f"\n6. Testing link selection with {n_links} candidates...")
|
||||
start = time.perf_counter()
|
||||
scored_links = await strategy.select_links_for_expansion(
|
||||
mock_links,
|
||||
gaps,
|
||||
state.kb_embeddings
|
||||
)
|
||||
metrics.record_operation("select_links", time.perf_counter() - start)
|
||||
print(f" Scored {len(scored_links)} links in {(time.perf_counter() - start)*1000:.1f} ms")
|
||||
|
||||
# End monitoring
|
||||
metrics.end()
|
||||
|
||||
return metrics
|
||||
|
||||
|
||||
async def main():
|
||||
"""Run performance tests before and after optimizations"""
|
||||
|
||||
print("="*80)
|
||||
print("EMBEDDING STRATEGY PERFORMANCE TEST")
|
||||
print("="*80)
|
||||
|
||||
# Test current implementation
|
||||
print("\n📊 Testing CURRENT Implementation...")
|
||||
metrics_before = await test_embedding_performance()
|
||||
metrics_before.print_summary("BEFORE Optimizations")
|
||||
|
||||
# Store key metrics for comparison
|
||||
total_time_before = metrics_before.total_time
|
||||
memory_before = metrics_before.memory_used_mb
|
||||
|
||||
# Calculate specific operation costs
|
||||
calc_conf_avg = sum(metrics_before.operation_times.get("calculate_confidence", [])) / len(metrics_before.operation_times.get("calculate_confidence", [1]))
|
||||
find_gaps_avg = sum(metrics_before.operation_times.get("find_coverage_gaps", [])) / len(metrics_before.operation_times.get("find_coverage_gaps", [1]))
|
||||
validate_avg = sum(metrics_before.operation_times.get("validate_coverage", [])) / len(metrics_before.operation_times.get("validate_coverage", [1]))
|
||||
|
||||
print(f"\n🔍 Key Bottlenecks Identified:")
|
||||
print(f" - calculate_confidence: {calc_conf_avg*1000:.1f} ms per call")
|
||||
print(f" - find_coverage_gaps: {find_gaps_avg*1000:.1f} ms per call")
|
||||
print(f" - validate_coverage: {validate_avg*1000:.1f} ms per call")
|
||||
|
||||
print("\n" + "="*80)
|
||||
print("EXPECTED IMPROVEMENTS AFTER OPTIMIZATION:")
|
||||
print("- Distance calculations: 80-90% faster (vectorization)")
|
||||
print("- Memory usage: 20-30% reduction (deduplication)")
|
||||
print("- Overall performance: 60-70% improvement")
|
||||
print("="*80)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,634 @@
|
||||
"""
|
||||
Test and demo script for Embedding-based Adaptive Crawler
|
||||
|
||||
This script demonstrates the embedding-based adaptive crawling
|
||||
with semantic space coverage and gap-driven expansion.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from pathlib import Path
|
||||
import time
|
||||
from rich.console import Console
|
||||
from rich import print as rprint
|
||||
import sys
|
||||
|
||||
# Add parent directory to path for imports
|
||||
sys.path.append(str(Path(__file__).parent.parent.parent))
|
||||
|
||||
from crawl4ai import (
|
||||
AsyncWebCrawler,
|
||||
AdaptiveCrawler,
|
||||
AdaptiveConfig,
|
||||
CrawlState
|
||||
)
|
||||
|
||||
console = Console()
|
||||
|
||||
|
||||
async def test_basic_embedding_crawl():
|
||||
"""Test basic embedding-based adaptive crawling"""
|
||||
console.print("\n[bold yellow]Test 1: Basic Embedding-based Crawl[/bold yellow]")
|
||||
console.print("Testing semantic space coverage with query expansion")
|
||||
|
||||
# Configure with embedding strategy
|
||||
config = AdaptiveConfig(
|
||||
strategy="embedding",
|
||||
confidence_threshold=0.7, # Not used for stopping in embedding strategy
|
||||
min_gain_threshold=0.01,
|
||||
max_pages=15,
|
||||
top_k_links=3,
|
||||
n_query_variations=8,
|
||||
embedding_model="sentence-transformers/all-MiniLM-L6-v2" # Fast, good quality
|
||||
)
|
||||
|
||||
# For query expansion, we need an LLM config
|
||||
llm_config = {
|
||||
'provider': 'openai/gpt-4o-mini',
|
||||
'api_token': os.getenv('OPENAI_API_KEY')
|
||||
}
|
||||
|
||||
if not llm_config['api_token']:
|
||||
console.print("[red]Warning: OPENAI_API_KEY not set. Using mock data for demo.[/red]")
|
||||
# Continue with mock for demo purposes
|
||||
|
||||
config.embedding_llm_config = llm_config
|
||||
|
||||
# Create crawler
|
||||
async with AsyncWebCrawler() as crawler:
|
||||
prog_crawler = AdaptiveCrawler(
|
||||
crawler=crawler,
|
||||
config=config
|
||||
)
|
||||
|
||||
# Start adaptive crawl
|
||||
start_time = time.time()
|
||||
console.print("\n[cyan]Starting semantic adaptive crawl...[/cyan]")
|
||||
|
||||
state = await prog_crawler.digest(
|
||||
start_url="https://docs.python.org/3/library/asyncio.html",
|
||||
query="async await coroutines event loops"
|
||||
)
|
||||
elapsed = time.time() - start_time
|
||||
|
||||
# Print results
|
||||
console.print(f"\n[green]Crawl completed in {elapsed:.2f} seconds[/green]")
|
||||
prog_crawler.print_stats(detailed=False)
|
||||
|
||||
# Show semantic coverage details
|
||||
console.print("\n[bold cyan]Semantic Coverage Details:[/bold cyan]")
|
||||
if state.expanded_queries:
|
||||
console.print(f"Query expanded to {len(state.expanded_queries)} variations")
|
||||
console.print("Sample variations:")
|
||||
for i, q in enumerate(state.expanded_queries[:3], 1):
|
||||
console.print(f" {i}. {q}")
|
||||
|
||||
if state.semantic_gaps:
|
||||
console.print(f"\nSemantic gaps identified: {len(state.semantic_gaps)}")
|
||||
|
||||
console.print(f"\nFinal confidence: {prog_crawler.confidence:.2%}")
|
||||
console.print(f"Is Sufficient: {'Yes (Validated)' if prog_crawler.is_sufficient else 'No'}")
|
||||
console.print(f"Pages needed: {len(state.crawled_urls)}")
|
||||
|
||||
|
||||
async def test_embedding_vs_statistical(use_openai=False):
|
||||
"""Compare embedding strategy with statistical strategy"""
|
||||
console.print("\n[bold yellow]Test 2: Embedding vs Statistical Strategy Comparison[/bold yellow]")
|
||||
|
||||
test_url = "https://httpbin.org"
|
||||
test_query = "http headers authentication api"
|
||||
|
||||
# Test 1: Statistical strategy
|
||||
console.print("\n[cyan]1. Statistical Strategy:[/cyan]")
|
||||
config_stat = AdaptiveConfig(
|
||||
strategy="statistical",
|
||||
confidence_threshold=0.7,
|
||||
max_pages=10
|
||||
)
|
||||
|
||||
async with AsyncWebCrawler() as crawler:
|
||||
stat_crawler = AdaptiveCrawler(crawler=crawler, config=config_stat)
|
||||
|
||||
start_time = time.time()
|
||||
state_stat = await stat_crawler.digest(start_url=test_url, query=test_query)
|
||||
stat_time = time.time() - start_time
|
||||
|
||||
stat_pages = len(state_stat.crawled_urls)
|
||||
stat_confidence = stat_crawler.confidence
|
||||
|
||||
# Test 2: Embedding strategy
|
||||
console.print("\n[cyan]2. Embedding Strategy:[/cyan]")
|
||||
config_emb = AdaptiveConfig(
|
||||
strategy="embedding",
|
||||
confidence_threshold=0.7, # Not used for stopping
|
||||
max_pages=10,
|
||||
n_query_variations=5,
|
||||
min_gain_threshold=0.01
|
||||
)
|
||||
|
||||
# Use OpenAI if available or requested
|
||||
if use_openai and os.getenv('OPENAI_API_KEY'):
|
||||
config_emb.embedding_llm_config = {
|
||||
'provider': 'openai/text-embedding-3-small',
|
||||
'api_token': os.getenv('OPENAI_API_KEY'),
|
||||
'embedding_model': 'text-embedding-3-small'
|
||||
}
|
||||
console.print("[cyan]Using OpenAI embeddings[/cyan]")
|
||||
else:
|
||||
# Default config will try sentence-transformers
|
||||
config_emb.embedding_llm_config = {
|
||||
'provider': 'openai/gpt-4o-mini',
|
||||
'api_token': os.getenv('OPENAI_API_KEY', 'dummy-key')
|
||||
}
|
||||
|
||||
async with AsyncWebCrawler() as crawler:
|
||||
emb_crawler = AdaptiveCrawler(crawler=crawler, config=config_emb)
|
||||
|
||||
start_time = time.time()
|
||||
state_emb = await emb_crawler.digest(start_url=test_url, query=test_query)
|
||||
emb_time = time.time() - start_time
|
||||
|
||||
emb_pages = len(state_emb.crawled_urls)
|
||||
emb_confidence = emb_crawler.confidence
|
||||
|
||||
# Compare results
|
||||
console.print("\n[bold green]Comparison Results:[/bold green]")
|
||||
console.print(f"Statistical: {stat_pages} pages in {stat_time:.2f}s, confidence: {stat_confidence:.2%}, sufficient: {stat_crawler.is_sufficient}")
|
||||
console.print(f"Embedding: {emb_pages} pages in {emb_time:.2f}s, confidence: {emb_confidence:.2%}, sufficient: {emb_crawler.is_sufficient}")
|
||||
|
||||
if emb_pages < stat_pages:
|
||||
efficiency = ((stat_pages - emb_pages) / stat_pages) * 100
|
||||
console.print(f"\n[green]Embedding strategy used {efficiency:.0f}% fewer pages![/green]")
|
||||
|
||||
# Show validation info for embedding
|
||||
if hasattr(state_emb, 'metrics') and 'validation_confidence' in state_emb.metrics:
|
||||
console.print(f"Embedding validation score: {state_emb.metrics['validation_confidence']:.2%}")
|
||||
|
||||
|
||||
async def test_custom_embedding_provider():
|
||||
"""Test with different embedding providers"""
|
||||
console.print("\n[bold yellow]Test 3: Custom Embedding Provider[/bold yellow]")
|
||||
|
||||
# Example with OpenAI embeddings
|
||||
config = AdaptiveConfig(
|
||||
strategy="embedding",
|
||||
confidence_threshold=0.8, # Not used for stopping
|
||||
max_pages=10,
|
||||
min_gain_threshold=0.01,
|
||||
n_query_variations=5
|
||||
)
|
||||
|
||||
# Configure to use OpenAI embeddings instead of sentence-transformers
|
||||
config.embedding_llm_config = {
|
||||
'provider': 'openai/text-embedding-3-small',
|
||||
'api_token': os.getenv('OPENAI_API_KEY'),
|
||||
'embedding_model': 'text-embedding-3-small'
|
||||
}
|
||||
|
||||
if not config.embedding_llm_config['api_token']:
|
||||
console.print("[yellow]Skipping OpenAI embedding test - no API key[/yellow]")
|
||||
return
|
||||
|
||||
async with AsyncWebCrawler() as crawler:
|
||||
prog_crawler = AdaptiveCrawler(crawler=crawler, config=config)
|
||||
|
||||
console.print("Using OpenAI embeddings for semantic analysis...")
|
||||
state = await prog_crawler.digest(
|
||||
start_url="https://httpbin.org",
|
||||
query="api endpoints json response"
|
||||
)
|
||||
|
||||
prog_crawler.print_stats(detailed=False)
|
||||
|
||||
|
||||
async def test_knowledge_export_import():
|
||||
"""Test exporting and importing semantic knowledge bases"""
|
||||
console.print("\n[bold yellow]Test 4: Semantic Knowledge Base Export/Import[/bold yellow]")
|
||||
|
||||
config = AdaptiveConfig(
|
||||
strategy="embedding",
|
||||
confidence_threshold=0.7, # Not used for stopping
|
||||
max_pages=5,
|
||||
min_gain_threshold=0.01,
|
||||
n_query_variations=4
|
||||
)
|
||||
|
||||
# First crawl
|
||||
async with AsyncWebCrawler() as crawler:
|
||||
crawler1 = AdaptiveCrawler(crawler=crawler, config=config)
|
||||
|
||||
console.print("\n[cyan]Building initial knowledge base...[/cyan]")
|
||||
state1 = await crawler1.digest(
|
||||
start_url="https://httpbin.org",
|
||||
query="http methods headers"
|
||||
)
|
||||
|
||||
# Export
|
||||
export_path = "semantic_kb.jsonl"
|
||||
crawler1.export_knowledge_base(export_path)
|
||||
console.print(f"[green]Exported {len(state1.knowledge_base)} documents with embeddings[/green]")
|
||||
|
||||
# Import and continue
|
||||
async with AsyncWebCrawler() as crawler:
|
||||
crawler2 = AdaptiveCrawler(crawler=crawler, config=config)
|
||||
|
||||
console.print("\n[cyan]Importing knowledge base...[/cyan]")
|
||||
await crawler2.import_knowledge_base(export_path)
|
||||
|
||||
# Continue with new query - should be faster
|
||||
console.print("\n[cyan]Extending with new query...[/cyan]")
|
||||
state2 = await crawler2.digest(
|
||||
start_url="https://httpbin.org",
|
||||
query="authentication oauth tokens"
|
||||
)
|
||||
|
||||
console.print(f"[green]Total knowledge base: {len(state2.knowledge_base)} documents[/green]")
|
||||
|
||||
# Cleanup
|
||||
Path(export_path).unlink(missing_ok=True)
|
||||
|
||||
|
||||
async def test_gap_visualization():
|
||||
"""Visualize semantic gaps and coverage"""
|
||||
console.print("\n[bold yellow]Test 5: Semantic Gap Analysis[/bold yellow]")
|
||||
|
||||
config = AdaptiveConfig(
|
||||
strategy="embedding",
|
||||
confidence_threshold=0.9, # Not used for stopping
|
||||
max_pages=8,
|
||||
n_query_variations=6,
|
||||
min_gain_threshold=0.01
|
||||
)
|
||||
|
||||
async with AsyncWebCrawler() as crawler:
|
||||
prog_crawler = AdaptiveCrawler(crawler=crawler, config=config)
|
||||
|
||||
# Initial crawl
|
||||
state = await prog_crawler.digest(
|
||||
start_url="https://docs.python.org/3/library/",
|
||||
query="concurrency threading multiprocessing"
|
||||
)
|
||||
|
||||
# Analyze gaps
|
||||
console.print("\n[bold cyan]Semantic Gap Analysis:[/bold cyan]")
|
||||
console.print(f"Query variations: {len(state.expanded_queries)}")
|
||||
console.print(f"Knowledge documents: {len(state.knowledge_base)}")
|
||||
console.print(f"Identified gaps: {len(state.semantic_gaps)}")
|
||||
|
||||
if state.semantic_gaps:
|
||||
console.print("\n[yellow]Gap sizes (distance from coverage):[/yellow]")
|
||||
for i, (_, distance) in enumerate(state.semantic_gaps[:5], 1):
|
||||
console.print(f" Gap {i}: {distance:.3f}")
|
||||
|
||||
# Show crawl progression
|
||||
console.print("\n[cyan]Crawl Order (gap-driven selection):[/cyan]")
|
||||
for i, url in enumerate(state.crawl_order[:5], 1):
|
||||
console.print(f" {i}. {url}")
|
||||
|
||||
|
||||
async def test_fast_convergence_with_relevant_query():
|
||||
"""Test that both strategies reach high confidence quickly with relevant queries"""
|
||||
console.print("\n[bold yellow]Test 7: Fast Convergence with Relevant Query[/bold yellow]")
|
||||
console.print("Testing that strategies reach 80%+ confidence within 2-3 batches")
|
||||
|
||||
# Test scenarios
|
||||
test_cases = [
|
||||
{
|
||||
"name": "Python Async Documentation",
|
||||
"url": "https://docs.python.org/3/library/asyncio.html",
|
||||
"query": "async await coroutines event loops tasks"
|
||||
}
|
||||
]
|
||||
|
||||
for test_case in test_cases:
|
||||
console.print(f"\n[bold cyan]Testing: {test_case['name']}[/bold cyan]")
|
||||
console.print(f"URL: {test_case['url']}")
|
||||
console.print(f"Query: {test_case['query']}")
|
||||
|
||||
# Test Embedding Strategy
|
||||
console.print("\n[yellow]Embedding Strategy:[/yellow]")
|
||||
config_emb = AdaptiveConfig(
|
||||
strategy="embedding",
|
||||
confidence_threshold=0.8,
|
||||
max_pages=9,
|
||||
top_k_links=3,
|
||||
min_gain_threshold=0.01,
|
||||
n_query_variations=5
|
||||
)
|
||||
|
||||
# Configure embeddings
|
||||
config_emb.embedding_llm_config = {
|
||||
'provider': 'openai/gpt-4o-mini',
|
||||
'api_token': os.getenv('OPENAI_API_KEY'),
|
||||
}
|
||||
|
||||
async with AsyncWebCrawler() as crawler:
|
||||
emb_crawler = AdaptiveCrawler(crawler=crawler, config=config_emb)
|
||||
|
||||
start_time = time.time()
|
||||
state = await emb_crawler.digest(
|
||||
start_url=test_case['url'],
|
||||
query=test_case['query']
|
||||
)
|
||||
|
||||
# Get batch breakdown
|
||||
total_pages = len(state.crawled_urls)
|
||||
for i in range(0, total_pages, 3):
|
||||
batch_num = (i // 3) + 1
|
||||
batch_pages = min(3, total_pages - i)
|
||||
pages_so_far = i + batch_pages
|
||||
estimated_confidence = state.metrics.get('confidence', 0) * (pages_so_far / total_pages)
|
||||
|
||||
console.print(f"Batch {batch_num}: {batch_pages} pages → Confidence: {estimated_confidence:.1%} {'✅' if estimated_confidence >= 0.8 else '❌'}")
|
||||
|
||||
final_confidence = emb_crawler.confidence
|
||||
console.print(f"[green]Final: {total_pages} pages → Confidence: {final_confidence:.1%} {'✅ (Sufficient!)' if emb_crawler.is_sufficient else '❌'}[/green]")
|
||||
|
||||
# Show learning metrics for embedding
|
||||
if 'avg_min_distance' in state.metrics:
|
||||
console.print(f"[dim]Avg gap distance: {state.metrics['avg_min_distance']:.3f}[/dim]")
|
||||
if 'validation_confidence' in state.metrics:
|
||||
console.print(f"[dim]Validation score: {state.metrics['validation_confidence']:.1%}[/dim]")
|
||||
|
||||
# Test Statistical Strategy
|
||||
console.print("\n[yellow]Statistical Strategy:[/yellow]")
|
||||
config_stat = AdaptiveConfig(
|
||||
strategy="statistical",
|
||||
confidence_threshold=0.8,
|
||||
max_pages=9,
|
||||
top_k_links=3,
|
||||
min_gain_threshold=0.01
|
||||
)
|
||||
|
||||
async with AsyncWebCrawler() as crawler:
|
||||
stat_crawler = AdaptiveCrawler(crawler=crawler, config=config_stat)
|
||||
|
||||
# Track batch progress
|
||||
batch_results = []
|
||||
current_pages = 0
|
||||
|
||||
# Custom batch tracking
|
||||
start_time = time.time()
|
||||
state = await stat_crawler.digest(
|
||||
start_url=test_case['url'],
|
||||
query=test_case['query']
|
||||
)
|
||||
|
||||
# Get batch breakdown (every 3 pages)
|
||||
total_pages = len(state.crawled_urls)
|
||||
for i in range(0, total_pages, 3):
|
||||
batch_num = (i // 3) + 1
|
||||
batch_pages = min(3, total_pages - i)
|
||||
# Estimate confidence at this point (simplified)
|
||||
pages_so_far = i + batch_pages
|
||||
estimated_confidence = state.metrics.get('confidence', 0) * (pages_so_far / total_pages)
|
||||
|
||||
console.print(f"Batch {batch_num}: {batch_pages} pages → Confidence: {estimated_confidence:.1%} {'✅' if estimated_confidence >= 0.8 else '❌'}")
|
||||
|
||||
final_confidence = stat_crawler.confidence
|
||||
console.print(f"[green]Final: {total_pages} pages → Confidence: {final_confidence:.1%} {'✅ (Sufficient!)' if stat_crawler.is_sufficient else '❌'}[/green]")
|
||||
|
||||
|
||||
|
||||
|
||||
async def test_irrelevant_query_behavior():
|
||||
"""Test how embedding strategy handles completely irrelevant queries"""
|
||||
console.print("\n[bold yellow]Test 8: Irrelevant Query Behavior[/bold yellow]")
|
||||
console.print("Testing embedding strategy with a query that has no semantic relevance to the content")
|
||||
|
||||
# Test with irrelevant query on Python async documentation
|
||||
test_case = {
|
||||
"name": "Irrelevant Query on Python Docs",
|
||||
"url": "https://docs.python.org/3/library/asyncio.html",
|
||||
"query": "how to cook fried rice with vegetables"
|
||||
}
|
||||
|
||||
console.print(f"\n[bold cyan]Testing: {test_case['name']}[/bold cyan]")
|
||||
console.print(f"URL: {test_case['url']} (Python async documentation)")
|
||||
console.print(f"Query: '{test_case['query']}' (completely irrelevant)")
|
||||
console.print("\n[dim]Expected behavior: Low confidence, high distances, no convergence[/dim]")
|
||||
|
||||
# Configure embedding strategy
|
||||
config_emb = AdaptiveConfig(
|
||||
strategy="embedding",
|
||||
confidence_threshold=0.8,
|
||||
max_pages=9,
|
||||
top_k_links=3,
|
||||
min_gain_threshold=0.01,
|
||||
n_query_variations=5,
|
||||
embedding_min_relative_improvement=0.05, # Lower threshold to see more iterations
|
||||
embedding_min_confidence_threshold=0.1 # Will stop if confidence < 10%
|
||||
)
|
||||
|
||||
# Configure embeddings using the correct format
|
||||
config_emb.embedding_llm_config = {
|
||||
'provider': 'openai/gpt-4o-mini',
|
||||
'api_token': os.getenv('OPENAI_API_KEY'),
|
||||
}
|
||||
|
||||
async with AsyncWebCrawler() as crawler:
|
||||
emb_crawler = AdaptiveCrawler(crawler=crawler, config=config_emb)
|
||||
|
||||
start_time = time.time()
|
||||
state = await emb_crawler.digest(
|
||||
start_url=test_case['url'],
|
||||
query=test_case['query']
|
||||
)
|
||||
elapsed = time.time() - start_time
|
||||
|
||||
# Analyze results
|
||||
console.print(f"\n[bold]Results after {elapsed:.1f} seconds:[/bold]")
|
||||
|
||||
# Basic metrics
|
||||
total_pages = len(state.crawled_urls)
|
||||
final_confidence = emb_crawler.confidence
|
||||
|
||||
console.print(f"\nPages crawled: {total_pages}")
|
||||
console.print(f"Final confidence: {final_confidence:.1%} {'✅' if emb_crawler.is_sufficient else '❌'}")
|
||||
|
||||
# Distance metrics
|
||||
if 'avg_min_distance' in state.metrics:
|
||||
console.print(f"\n[yellow]Distance Metrics:[/yellow]")
|
||||
console.print(f" Average minimum distance: {state.metrics['avg_min_distance']:.3f}")
|
||||
console.print(f" Close neighbors (<0.3): {state.metrics.get('avg_close_neighbors', 0):.1f}")
|
||||
console.print(f" Very close neighbors (<0.2): {state.metrics.get('avg_very_close_neighbors', 0):.1f}")
|
||||
|
||||
# Interpret distances
|
||||
avg_dist = state.metrics['avg_min_distance']
|
||||
if avg_dist > 0.8:
|
||||
console.print(f" [red]→ Very poor match (distance > 0.8)[/red]")
|
||||
elif avg_dist > 0.6:
|
||||
console.print(f" [yellow]→ Poor match (distance > 0.6)[/yellow]")
|
||||
elif avg_dist > 0.4:
|
||||
console.print(f" [blue]→ Moderate match (distance > 0.4)[/blue]")
|
||||
else:
|
||||
console.print(f" [green]→ Good match (distance < 0.4)[/green]")
|
||||
|
||||
# Show sample expanded queries
|
||||
if state.expanded_queries:
|
||||
console.print(f"\n[yellow]Sample Query Variations Generated:[/yellow]")
|
||||
for i, q in enumerate(state.expanded_queries[:3], 1):
|
||||
console.print(f" {i}. {q}")
|
||||
|
||||
# Show crawl progression
|
||||
console.print(f"\n[yellow]Crawl Progression:[/yellow]")
|
||||
for i, url in enumerate(state.crawl_order[:5], 1):
|
||||
console.print(f" {i}. {url}")
|
||||
if len(state.crawl_order) > 5:
|
||||
console.print(f" ... and {len(state.crawl_order) - 5} more")
|
||||
|
||||
# Validation score
|
||||
if 'validation_confidence' in state.metrics:
|
||||
console.print(f"\n[yellow]Validation:[/yellow]")
|
||||
console.print(f" Validation score: {state.metrics['validation_confidence']:.1%}")
|
||||
|
||||
# Why it stopped
|
||||
if 'stopped_reason' in state.metrics:
|
||||
console.print(f"\n[yellow]Stopping Reason:[/yellow] {state.metrics['stopped_reason']}")
|
||||
if state.metrics.get('is_irrelevant', False):
|
||||
console.print("[red]→ Query and content are completely unrelated![/red]")
|
||||
elif total_pages >= config_emb.max_pages:
|
||||
console.print(f"\n[yellow]Stopping Reason:[/yellow] Reached max pages limit ({config_emb.max_pages})")
|
||||
|
||||
# Summary
|
||||
console.print(f"\n[bold]Summary:[/bold]")
|
||||
if final_confidence < 0.2:
|
||||
console.print("[red]✗ As expected: Query is completely irrelevant to content[/red]")
|
||||
console.print("[green]✓ The embedding strategy correctly identified no semantic match[/green]")
|
||||
else:
|
||||
console.print(f"[yellow]⚠ Unexpected: Got {final_confidence:.1%} confidence for irrelevant query[/yellow]")
|
||||
console.print("[yellow] This may indicate the query variations are too broad[/yellow]")
|
||||
|
||||
|
||||
async def test_high_dimensional_handling():
|
||||
"""Test handling of high-dimensional embedding spaces"""
|
||||
console.print("\n[bold yellow]Test 6: High-Dimensional Embedding Space Handling[/bold yellow]")
|
||||
console.print("Testing how the system handles 384+ dimensional embeddings")
|
||||
|
||||
config = AdaptiveConfig(
|
||||
strategy="embedding",
|
||||
confidence_threshold=0.8, # Not used for stopping
|
||||
max_pages=5,
|
||||
n_query_variations=8, # Will create 9 points total
|
||||
min_gain_threshold=0.01,
|
||||
embedding_model="sentence-transformers/all-MiniLM-L6-v2" # 384 dimensions
|
||||
)
|
||||
|
||||
# Use OpenAI if available, otherwise mock
|
||||
if os.getenv('OPENAI_API_KEY'):
|
||||
config.embedding_llm_config = {
|
||||
'provider': 'openai/text-embedding-3-small',
|
||||
'api_token': os.getenv('OPENAI_API_KEY'),
|
||||
'embedding_model': 'text-embedding-3-small'
|
||||
}
|
||||
else:
|
||||
config.embedding_llm_config = {
|
||||
'provider': 'openai/gpt-4o-mini',
|
||||
'api_token': 'mock-key'
|
||||
}
|
||||
|
||||
async with AsyncWebCrawler() as crawler:
|
||||
prog_crawler = AdaptiveCrawler(crawler=crawler, config=config)
|
||||
|
||||
console.print("\n[cyan]Testing with high-dimensional embeddings (384D)...[/cyan]")
|
||||
|
||||
try:
|
||||
state = await prog_crawler.digest(
|
||||
start_url="https://httpbin.org",
|
||||
query="api endpoints json"
|
||||
)
|
||||
|
||||
console.print(f"[green]✓ Successfully handled {len(state.expanded_queries)} queries in 384D space[/green]")
|
||||
console.print(f"Coverage shape type: {type(state.coverage_shape)}")
|
||||
|
||||
if isinstance(state.coverage_shape, dict):
|
||||
console.print(f"Coverage model: centroid + radius")
|
||||
console.print(f" - Center shape: {state.coverage_shape['center'].shape if 'center' in state.coverage_shape else 'N/A'}")
|
||||
console.print(f" - Radius: {state.coverage_shape.get('radius', 'N/A'):.3f}")
|
||||
|
||||
except Exception as e:
|
||||
console.print(f"[red]Error: {e}[/red]")
|
||||
console.print("[yellow]This demonstrates why alpha shapes don't work in high dimensions[/yellow]")
|
||||
|
||||
|
||||
async def main():
|
||||
"""Run all embedding strategy tests"""
|
||||
console.print("[bold magenta]Embedding-based Adaptive Crawler Test Suite[/bold magenta]")
|
||||
console.print("=" * 60)
|
||||
|
||||
try:
|
||||
# Check if we have required dependencies
|
||||
has_sentence_transformers = True
|
||||
has_numpy = True
|
||||
|
||||
try:
|
||||
import numpy
|
||||
console.print("[green]✓ NumPy installed[/green]")
|
||||
except ImportError:
|
||||
has_numpy = False
|
||||
console.print("[red]Missing numpy[/red]")
|
||||
|
||||
# Try to import sentence_transformers but catch numpy compatibility errors
|
||||
try:
|
||||
import sentence_transformers
|
||||
console.print("[green]✓ Sentence-transformers installed[/green]")
|
||||
except (ImportError, RuntimeError, ValueError) as e:
|
||||
has_sentence_transformers = False
|
||||
console.print(f"[yellow]Warning: sentence-transformers not available[/yellow]")
|
||||
console.print("[yellow]Tests will use OpenAI embeddings if available or mock data[/yellow]")
|
||||
|
||||
# Run tests based on available dependencies
|
||||
if has_numpy:
|
||||
# Check if we should use OpenAI for embeddings
|
||||
use_openai = not has_sentence_transformers and os.getenv('OPENAI_API_KEY')
|
||||
|
||||
if not has_sentence_transformers and not os.getenv('OPENAI_API_KEY'):
|
||||
console.print("\n[red]Neither sentence-transformers nor OpenAI API key available[/red]")
|
||||
console.print("[yellow]Please set OPENAI_API_KEY or fix sentence-transformers installation[/yellow]")
|
||||
return
|
||||
|
||||
# Run all tests
|
||||
# await test_basic_embedding_crawl()
|
||||
# await test_embedding_vs_statistical(use_openai=use_openai)
|
||||
|
||||
# Run the fast convergence test - this is the most important one
|
||||
# await test_fast_convergence_with_relevant_query()
|
||||
|
||||
# Test with irrelevant query
|
||||
await test_irrelevant_query_behavior()
|
||||
|
||||
# Only run OpenAI-specific test if we have API key
|
||||
# if os.getenv('OPENAI_API_KEY'):
|
||||
# await test_custom_embedding_provider()
|
||||
|
||||
# # Skip tests that require sentence-transformers when it's not available
|
||||
# if has_sentence_transformers:
|
||||
# await test_knowledge_export_import()
|
||||
# await test_gap_visualization()
|
||||
# else:
|
||||
# console.print("\n[yellow]Skipping tests that require sentence-transformers due to numpy compatibility issues[/yellow]")
|
||||
|
||||
# This test should work with mock data
|
||||
# await test_high_dimensional_handling()
|
||||
else:
|
||||
console.print("\n[red]Cannot run tests without NumPy[/red]")
|
||||
return
|
||||
|
||||
console.print("\n[bold green]✅ All tests completed![/bold green]")
|
||||
|
||||
except Exception as e:
|
||||
console.print(f"\n[bold red]❌ Test failed: {e}[/bold red]")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,154 @@
|
||||
import asyncio
|
||||
import os
|
||||
from crawl4ai import AsyncWebCrawler, AdaptiveCrawler, AdaptiveConfig, LLMConfig
|
||||
|
||||
|
||||
async def test_configuration(name: str, config: AdaptiveConfig, url: str, query: str):
|
||||
"""Test a specific configuration"""
|
||||
print(f"\n{'='*60}")
|
||||
print(f"Configuration: {name}")
|
||||
print(f"{'='*60}")
|
||||
|
||||
async with AsyncWebCrawler(verbose=False) as crawler:
|
||||
adaptive = AdaptiveCrawler(crawler, config)
|
||||
result = await adaptive.digest(start_url=url, query=query)
|
||||
|
||||
print("\n" + "="*50)
|
||||
print("CRAWL STATISTICS")
|
||||
print("="*50)
|
||||
adaptive.print_stats(detailed=False)
|
||||
|
||||
# Get the most relevant content found
|
||||
print("\n" + "="*50)
|
||||
print("MOST RELEVANT PAGES")
|
||||
print("="*50)
|
||||
|
||||
relevant_pages = adaptive.get_relevant_content(top_k=5)
|
||||
for i, page in enumerate(relevant_pages, 1):
|
||||
print(f"\n{i}. {page['url']}")
|
||||
print(f" Relevance Score: {page['score']:.2%}")
|
||||
|
||||
# Show a snippet of the content
|
||||
content = page['content'] or ""
|
||||
if content:
|
||||
snippet = content[:200].replace('\n', ' ')
|
||||
if len(content) > 200:
|
||||
snippet += "..."
|
||||
print(f" Preview: {snippet}")
|
||||
|
||||
print(f"\n{'='*50}")
|
||||
print(f"Pages crawled: {len(result.crawled_urls)}")
|
||||
print(f"Final confidence: {adaptive.confidence:.1%}")
|
||||
print(f"Stopped reason: {result.metrics.get('stopped_reason', 'max_pages')}")
|
||||
|
||||
if result.metrics.get('is_irrelevant', False):
|
||||
print("⚠️ Query detected as irrelevant!")
|
||||
|
||||
return result
|
||||
|
||||
|
||||
async def llm_embedding():
|
||||
"""Demonstrate various embedding configurations"""
|
||||
|
||||
print("EMBEDDING STRATEGY CONFIGURATION EXAMPLES")
|
||||
print("=" * 60)
|
||||
|
||||
# Base URL and query for testing
|
||||
test_url = "https://docs.python.org/3/library/asyncio.html"
|
||||
|
||||
openai_llm_config = LLMConfig(
|
||||
provider='openai/text-embedding-3-small',
|
||||
api_token=os.getenv('OPENAI_API_KEY'),
|
||||
temperature=0.7,
|
||||
max_tokens=2000
|
||||
)
|
||||
config_openai = AdaptiveConfig(
|
||||
strategy="embedding",
|
||||
max_pages=10,
|
||||
|
||||
# Use OpenAI embeddings
|
||||
embedding_llm_config=openai_llm_config,
|
||||
# embedding_llm_config={
|
||||
# 'provider': 'openai/text-embedding-3-small',
|
||||
# 'api_token': os.getenv('OPENAI_API_KEY')
|
||||
# },
|
||||
|
||||
# OpenAI embeddings are high quality, can be stricter
|
||||
embedding_k_exp=4.0,
|
||||
n_query_variations=12
|
||||
)
|
||||
|
||||
await test_configuration(
|
||||
"OpenAI Embeddings",
|
||||
config_openai,
|
||||
test_url,
|
||||
# "event-driven architecture patterns"
|
||||
"async await context managers coroutines"
|
||||
)
|
||||
return
|
||||
|
||||
|
||||
|
||||
async def basic_adaptive_crawling():
|
||||
"""Basic adaptive crawling example"""
|
||||
|
||||
# Initialize the crawler
|
||||
async with AsyncWebCrawler(verbose=True) as crawler:
|
||||
# Create an adaptive crawler with default settings (statistical strategy)
|
||||
adaptive = AdaptiveCrawler(crawler)
|
||||
|
||||
# Note: You can also use embedding strategy for semantic understanding:
|
||||
# from crawl4ai import AdaptiveConfig
|
||||
# config = AdaptiveConfig(strategy="embedding")
|
||||
# adaptive = AdaptiveCrawler(crawler, config)
|
||||
|
||||
# Start adaptive crawling
|
||||
print("Starting adaptive crawl for Python async programming information...")
|
||||
result = await adaptive.digest(
|
||||
start_url="https://docs.python.org/3/library/asyncio.html",
|
||||
query="async await context managers coroutines"
|
||||
)
|
||||
|
||||
# Display crawl statistics
|
||||
print("\n" + "="*50)
|
||||
print("CRAWL STATISTICS")
|
||||
print("="*50)
|
||||
adaptive.print_stats(detailed=False)
|
||||
|
||||
# Get the most relevant content found
|
||||
print("\n" + "="*50)
|
||||
print("MOST RELEVANT PAGES")
|
||||
print("="*50)
|
||||
|
||||
relevant_pages = adaptive.get_relevant_content(top_k=5)
|
||||
for i, page in enumerate(relevant_pages, 1):
|
||||
print(f"\n{i}. {page['url']}")
|
||||
print(f" Relevance Score: {page['score']:.2%}")
|
||||
|
||||
# Show a snippet of the content
|
||||
content = page['content'] or ""
|
||||
if content:
|
||||
snippet = content[:200].replace('\n', ' ')
|
||||
if len(content) > 200:
|
||||
snippet += "..."
|
||||
print(f" Preview: {snippet}")
|
||||
|
||||
# Show final confidence
|
||||
print(f"\n{'='*50}")
|
||||
print(f"Final Confidence: {adaptive.confidence:.2%}")
|
||||
print(f"Total Pages Crawled: {len(result.crawled_urls)}")
|
||||
print(f"Knowledge Base Size: {len(adaptive.state.knowledge_base)} documents")
|
||||
|
||||
|
||||
if adaptive.confidence >= 0.8:
|
||||
print("✓ High confidence - can answer detailed questions about async Python")
|
||||
elif adaptive.confidence >= 0.6:
|
||||
print("~ Moderate confidence - can answer basic questions")
|
||||
else:
|
||||
print("✗ Low confidence - need more information")
|
||||
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(llm_embedding())
|
||||
# asyncio.run(basic_adaptive_crawling())
|
||||
@@ -0,0 +1,284 @@
|
||||
"""
|
||||
E2E tests for separate embedding and query LLM configs (issue #1682).
|
||||
|
||||
Tests that AdaptiveConfig.query_llm_config flows correctly through
|
||||
AdaptiveCrawler → EmbeddingStrategy → map_query_semantic_space,
|
||||
and that the right config is used for embeddings vs query expansion.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch, MagicMock, AsyncMock
|
||||
import numpy as np
|
||||
|
||||
sys.path.append(str(Path(__file__).parent.parent.parent))
|
||||
|
||||
from crawl4ai import AdaptiveConfig, LLMConfig
|
||||
from crawl4ai.adaptive_crawler import EmbeddingStrategy, AdaptiveCrawler
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 1: Config plumbing — AdaptiveConfig → AdaptiveCrawler → EmbeddingStrategy
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_config_plumbing():
|
||||
"""query_llm_config flows from AdaptiveConfig through _create_strategy."""
|
||||
config = AdaptiveConfig(
|
||||
strategy="embedding",
|
||||
embedding_llm_config=LLMConfig(provider="openai/text-embedding-3-small", api_token="emb-key"),
|
||||
query_llm_config=LLMConfig(provider="openai/gpt-4o-mini", api_token="query-key"),
|
||||
)
|
||||
|
||||
# Simulate what AdaptiveCrawler.__init__ does
|
||||
with patch("crawl4ai.adaptive_crawler.AsyncWebCrawler"):
|
||||
crawler_mock = MagicMock()
|
||||
adaptive = AdaptiveCrawler(crawler=crawler_mock, config=config)
|
||||
|
||||
strategy = adaptive.strategy
|
||||
assert isinstance(strategy, EmbeddingStrategy)
|
||||
|
||||
# Strategy should have both configs
|
||||
assert strategy.query_llm_config is not None
|
||||
query_dict = strategy._get_query_llm_config_dict()
|
||||
assert query_dict["provider"] == "openai/gpt-4o-mini"
|
||||
assert query_dict["api_token"] == "query-key"
|
||||
|
||||
emb_dict = strategy._get_embedding_llm_config_dict()
|
||||
assert emb_dict["provider"] == "openai/text-embedding-3-small"
|
||||
assert emb_dict["api_token"] == "emb-key"
|
||||
|
||||
print("PASS: test_config_plumbing")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 2: Backward compat — no query_llm_config falls back to llm_config
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_backward_compat_fallback():
|
||||
"""When query_llm_config is not set, falls back to llm_config (legacy)."""
|
||||
strategy = EmbeddingStrategy(
|
||||
embedding_model="sentence-transformers/all-MiniLM-L6-v2",
|
||||
llm_config={"provider": "openai/gpt-4o-mini", "api_token": "shared-key"},
|
||||
query_llm_config=None,
|
||||
)
|
||||
# No AdaptiveConfig attached → should fall back to llm_config
|
||||
result = strategy._get_query_llm_config_dict()
|
||||
assert result["provider"] == "openai/gpt-4o-mini"
|
||||
assert result["api_token"] == "shared-key"
|
||||
print("PASS: test_backward_compat_fallback")
|
||||
|
||||
|
||||
def test_backward_compat_no_config():
|
||||
"""When nothing is set, returns None (caller uses hardcoded defaults)."""
|
||||
strategy = EmbeddingStrategy()
|
||||
result = strategy._get_query_llm_config_dict()
|
||||
assert result is None
|
||||
print("PASS: test_backward_compat_no_config")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 3: Fallback priority chain
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_fallback_priority():
|
||||
"""Explicit query_llm_config beats AdaptiveConfig beats llm_config."""
|
||||
config = AdaptiveConfig(
|
||||
strategy="embedding",
|
||||
query_llm_config={"provider": "config-level", "api_token": "cfg"},
|
||||
)
|
||||
strategy = EmbeddingStrategy(
|
||||
llm_config={"provider": "legacy-level", "api_token": "leg"},
|
||||
query_llm_config={"provider": "strategy-level", "api_token": "strat"},
|
||||
)
|
||||
strategy.config = config
|
||||
|
||||
# Strategy-level should win
|
||||
result = strategy._get_query_llm_config_dict()
|
||||
assert result["provider"] == "strategy-level"
|
||||
|
||||
# Remove strategy-level → config-level should win
|
||||
strategy.query_llm_config = None
|
||||
result = strategy._get_query_llm_config_dict()
|
||||
assert result["provider"] == "config-level"
|
||||
|
||||
# Remove config-level → legacy llm_config should win
|
||||
config.query_llm_config = None
|
||||
result = strategy._get_query_llm_config_dict()
|
||||
assert result["provider"] == "legacy-level"
|
||||
|
||||
# Remove everything → None
|
||||
strategy.llm_config = None
|
||||
result = strategy._get_query_llm_config_dict()
|
||||
assert result is None
|
||||
|
||||
print("PASS: test_fallback_priority")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 4: E2E — map_query_semantic_space uses query config, not embedding config
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
async def test_map_query_uses_query_config():
|
||||
"""map_query_semantic_space should call perform_completion_with_backoff
|
||||
with the query LLM config (chat model), NOT the embedding config."""
|
||||
|
||||
config = AdaptiveConfig(
|
||||
strategy="embedding",
|
||||
embedding_llm_config=LLMConfig(
|
||||
provider="openai/text-embedding-3-small",
|
||||
api_token="emb-key",
|
||||
base_url="https://emb.example.com",
|
||||
),
|
||||
query_llm_config=LLMConfig(
|
||||
provider="openai/gpt-4o-mini",
|
||||
api_token="query-key",
|
||||
base_url="https://query.example.com",
|
||||
),
|
||||
)
|
||||
|
||||
strategy = EmbeddingStrategy(
|
||||
embedding_model="sentence-transformers/all-MiniLM-L6-v2",
|
||||
llm_config=config.embedding_llm_config,
|
||||
query_llm_config=config.query_llm_config,
|
||||
)
|
||||
strategy.config = config
|
||||
|
||||
# Mock perform_completion_with_backoff to capture its arguments
|
||||
mock_response = MagicMock()
|
||||
mock_response.choices = [MagicMock()]
|
||||
mock_response.choices[0].message.content = json.dumps({
|
||||
"queries": [f"variation {i}" for i in range(13)]
|
||||
})
|
||||
|
||||
captured_kwargs = {}
|
||||
|
||||
def mock_completion(**kwargs):
|
||||
# Also accept positional-style
|
||||
captured_kwargs.update(kwargs)
|
||||
return mock_response
|
||||
|
||||
# Also mock _get_embeddings to avoid real embedding calls
|
||||
fake_embeddings = np.random.rand(11, 384).astype(np.float32)
|
||||
|
||||
with patch("crawl4ai.utils.perform_completion_with_backoff", side_effect=mock_completion):
|
||||
with patch.object(strategy, "_get_embeddings", new_callable=AsyncMock, return_value=fake_embeddings):
|
||||
await strategy.map_query_semantic_space("test query", n_synthetic=10)
|
||||
|
||||
# Verify the query config was used, NOT the embedding config
|
||||
assert captured_kwargs["provider"] == "openai/gpt-4o-mini", \
|
||||
f"Expected query model, got {captured_kwargs['provider']}"
|
||||
assert captured_kwargs["api_token"] == "query-key", \
|
||||
f"Expected query-key, got {captured_kwargs['api_token']}"
|
||||
assert captured_kwargs["base_url"] == "https://query.example.com", \
|
||||
f"Expected query base_url, got {captured_kwargs['base_url']}"
|
||||
|
||||
# Verify backoff params are passed (bug fix)
|
||||
assert "base_delay" in captured_kwargs
|
||||
assert "max_attempts" in captured_kwargs
|
||||
assert "exponential_factor" in captured_kwargs
|
||||
|
||||
print("PASS: test_map_query_uses_query_config")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 5: E2E — legacy single-config still works for query expansion
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
async def test_legacy_single_config_for_query():
|
||||
"""When only embedding_llm_config is set (old usage), query expansion
|
||||
falls back to it via llm_config → still works."""
|
||||
|
||||
single_config = LLMConfig(
|
||||
provider="openai/gpt-4o-mini",
|
||||
api_token="single-key",
|
||||
)
|
||||
|
||||
config = AdaptiveConfig(
|
||||
strategy="embedding",
|
||||
embedding_llm_config=single_config,
|
||||
# No query_llm_config — legacy usage
|
||||
)
|
||||
|
||||
strategy = EmbeddingStrategy(
|
||||
embedding_model="sentence-transformers/all-MiniLM-L6-v2",
|
||||
llm_config=config.embedding_llm_config, # This is how _create_strategy passes it
|
||||
# No query_llm_config
|
||||
)
|
||||
strategy.config = config
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.choices = [MagicMock()]
|
||||
mock_response.choices[0].message.content = json.dumps({
|
||||
"queries": [f"variation {i}" for i in range(13)]
|
||||
})
|
||||
|
||||
captured_kwargs = {}
|
||||
|
||||
def mock_completion(**kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
return mock_response
|
||||
|
||||
fake_embeddings = np.random.rand(11, 384).astype(np.float32)
|
||||
|
||||
with patch("crawl4ai.utils.perform_completion_with_backoff", side_effect=mock_completion):
|
||||
with patch.object(strategy, "_get_embeddings", new_callable=AsyncMock, return_value=fake_embeddings):
|
||||
await strategy.map_query_semantic_space("test query", n_synthetic=10)
|
||||
|
||||
# Should fall back to llm_config (the single shared config)
|
||||
assert captured_kwargs["provider"] == "openai/gpt-4o-mini"
|
||||
assert captured_kwargs["api_token"] == "single-key"
|
||||
|
||||
print("PASS: test_legacy_single_config_for_query")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 6: LLMConfig.to_dict() includes backoff params (bug fix verification)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_to_dict_includes_backoff():
|
||||
"""_embedding_llm_config_dict now uses to_dict() which includes backoff params."""
|
||||
config = AdaptiveConfig(
|
||||
embedding_llm_config=LLMConfig(
|
||||
provider="openai/text-embedding-3-small",
|
||||
api_token="test",
|
||||
backoff_base_delay=5,
|
||||
backoff_max_attempts=10,
|
||||
backoff_exponential_factor=3,
|
||||
),
|
||||
)
|
||||
d = config._embedding_llm_config_dict
|
||||
assert d["backoff_base_delay"] == 5
|
||||
assert d["backoff_max_attempts"] == 10
|
||||
assert d["backoff_exponential_factor"] == 3
|
||||
print("PASS: test_to_dict_includes_backoff")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
async def main():
|
||||
print("=" * 60)
|
||||
print("E2E Tests: Separate Embedding & Query LLM Configs (#1682)")
|
||||
print("=" * 60)
|
||||
|
||||
# Sync tests
|
||||
test_config_plumbing()
|
||||
test_backward_compat_fallback()
|
||||
test_backward_compat_no_config()
|
||||
test_fallback_priority()
|
||||
test_to_dict_includes_backoff()
|
||||
|
||||
# Async tests
|
||||
await test_map_query_uses_query_config()
|
||||
await test_legacy_single_config_for_query()
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
print("ALL TESTS PASSED")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
Reference in New Issue
Block a user