chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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