Merge pull request #67 from j-sperling/feat/query-type-source-tiering

Add query-type-aware source tiering and scoring
This commit is contained in:
Matt Van Horn
2026-03-14 07:30:06 -07:00
committed by GitHub
11 changed files with 819 additions and 113 deletions
+28 -20
View File
@@ -160,6 +160,7 @@ from lib import (
websearch, websearch,
xai_x, xai_x,
youtube_yt, youtube_yt,
query_type as qt,
) )
@@ -631,8 +632,10 @@ def _search_web(
topic, from_date, to_date, config["PARALLEL_API_KEY"], depth=depth, topic, from_date, to_date, config["PARALLEL_API_KEY"], depth=depth,
) )
elif backend == "brave": elif backend == "brave":
use_llm_ctx = os.environ.get("BRAVE_LLM_CONTEXT", "").strip() == "1"
raw_results = brave_search.search_web( raw_results = brave_search.search_web(
topic, from_date, to_date, config["BRAVE_API_KEY"], depth=depth, topic, from_date, to_date, config["BRAVE_API_KEY"],
depth=depth, use_llm_context=use_llm_ctx,
) )
elif backend == "openrouter": elif backend == "openrouter":
raw_results = openrouter_search.search_web( raw_results = openrouter_search.search_web(
@@ -1692,14 +1695,19 @@ def main():
else: else:
mode = sources mode = sources
# Detect query type for source tiering and scoring adjustments
query_type = qt.detect_query_type(args.topic)
# Apply --search flag: restrict sources to the specified subset # Apply --search flag: restrict sources to the specified subset
search_do_hackernews = True # Source defaults are query-type-aware (Truth Social always opt-in,
search_do_bluesky = has_bluesky # Bluesky only for query types where it adds signal)
search_do_truthsocial = has_truthsocial search_do_hackernews = qt.is_source_enabled("hn", query_type) if not args.search else True
search_do_polymarket = True search_do_bluesky = has_bluesky and qt.is_source_enabled("bluesky", query_type)
search_run_youtube = has_ytdlp search_do_truthsocial = False # Always opt-in (requires --search truthsocial)
search_run_tiktok = has_tiktok search_do_polymarket = qt.is_source_enabled("polymarket", query_type)
search_run_instagram = has_instagram search_run_youtube = has_ytdlp and qt.is_source_enabled("youtube", query_type)
search_run_tiktok = has_tiktok and qt.is_source_enabled("tiktok", query_type)
search_run_instagram = has_instagram and qt.is_source_enabled("instagram", query_type)
search_run_xiaohongshu = has_xiaohongshu search_run_xiaohongshu = has_xiaohongshu
if args.search: if args.search:
search_sources = parse_search_flag(args.search) search_sources = parse_search_flag(args.search)
@@ -1795,19 +1803,19 @@ def main():
scored_bsky = score.score_bluesky_items(filtered_bsky) if filtered_bsky else [] scored_bsky = score.score_bluesky_items(filtered_bsky) if filtered_bsky else []
scored_ts = score.score_truthsocial_items(filtered_ts) if filtered_ts else [] scored_ts = score.score_truthsocial_items(filtered_ts) if filtered_ts else []
scored_pm = score.score_polymarket_items(filtered_pm) if filtered_pm else [] scored_pm = score.score_polymarket_items(filtered_pm) if filtered_pm else []
scored_web = score.score_websearch_items(filtered_web) if filtered_web else [] scored_web = score.score_websearch_items(filtered_web, query_type=query_type) if filtered_web else []
# Sort items # Sort items (query-type-aware tiebreaker ordering)
sorted_reddit = score.sort_items(scored_reddit) sorted_reddit = score.sort_items(scored_reddit, query_type=query_type)
sorted_x = score.sort_items(scored_x) sorted_x = score.sort_items(scored_x, query_type=query_type)
sorted_youtube = score.sort_items(scored_youtube) if scored_youtube else [] sorted_youtube = score.sort_items(scored_youtube, query_type=query_type) if scored_youtube else []
sorted_tiktok = score.sort_items(scored_tiktok) if scored_tiktok else [] sorted_tiktok = score.sort_items(scored_tiktok, query_type=query_type) if scored_tiktok else []
sorted_ig = score.sort_items(scored_ig) if scored_ig else [] sorted_ig = score.sort_items(scored_ig, query_type=query_type) if scored_ig else []
sorted_hn = score.sort_items(scored_hn) if scored_hn else [] sorted_hn = score.sort_items(scored_hn, query_type=query_type) if scored_hn else []
sorted_bsky = score.sort_items(scored_bsky) if scored_bsky else [] sorted_bsky = score.sort_items(scored_bsky, query_type=query_type) if scored_bsky else []
sorted_ts = score.sort_items(scored_ts) if scored_ts else [] sorted_ts = score.sort_items(scored_ts, query_type=query_type) if scored_ts else []
sorted_pm = score.sort_items(scored_pm) if scored_pm else [] sorted_pm = score.sort_items(scored_pm, query_type=query_type) if scored_pm else []
sorted_web = score.sort_items(scored_web) if scored_web else [] sorted_web = score.sort_items(scored_web, query_type=query_type) if scored_web else []
# Dedupe items # Dedupe items
deduped_reddit = dedupe.dedupe_reddit(sorted_reddit) deduped_reddit = dedupe.dedupe_reddit(sorted_reddit)
+119 -3
View File
@@ -1,7 +1,12 @@
"""Brave Search web search for last30days skill. """Brave Search web search for last30days skill.
Uses the Brave Search API as a fallback web search backend. Uses the Brave Search API as a web search backend.
Simple, cheap (free tier: 2,000 queries/month), widely available. Requires a paid Brave Search subscription.
Two modes:
- Standard: /res/v1/web/search — returns URLs + snippets (default)
- LLM Context: /res/v1/llm/context — returns pre-extracted text chunks
optimized for LLM consumption. Enable with BRAVE_LLM_CONTEXT=1 env var.
API docs: https://api-dashboard.search.brave.com/app/documentation/web-search/get-started API docs: https://api-dashboard.search.brave.com/app/documentation/web-search/get-started
""" """
@@ -16,6 +21,7 @@ from urllib.parse import urlencode, urlparse
from . import http from . import http
ENDPOINT = "https://api.search.brave.com/res/v1/web/search" ENDPOINT = "https://api.search.brave.com/res/v1/web/search"
LLM_CONTEXT_ENDPOINT = "https://api.search.brave.com/res/v1/llm/context"
# Freshness codes: pd=24h, pw=7d, pm=31d # Freshness codes: pd=24h, pw=7d, pm=31d
FRESHNESS_MAP = {1: "pd", 7: "pw", 31: "pm"} FRESHNESS_MAP = {1: "pd", 7: "pw", 31: "pm"}
@@ -33,6 +39,7 @@ def search_web(
to_date: str, to_date: str,
api_key: str, api_key: str,
depth: str = "default", depth: str = "default",
use_llm_context: bool = False,
) -> List[Dict[str, Any]]: ) -> List[Dict[str, Any]]:
"""Search the web via Brave Search API. """Search the web via Brave Search API.
@@ -42,6 +49,7 @@ def search_web(
to_date: End date (YYYY-MM-DD) to_date: End date (YYYY-MM-DD)
api_key: Brave Search API key api_key: Brave Search API key
depth: 'quick', 'default', or 'deep' depth: 'quick', 'default', or 'deep'
use_llm_context: Use LLM Context endpoint for pre-extracted content
Returns: Returns:
List of result dicts with keys: url, title, snippet, source_domain, date, relevance List of result dicts with keys: url, title, snippet, source_domain, date, relevance
@@ -49,6 +57,9 @@ def search_web(
Raises: Raises:
http.HTTPError: On API errors http.HTTPError: On API errors
""" """
if use_llm_context:
return _search_llm_context(topic, from_date, to_date, api_key, depth)
count = {"quick": 8, "default": 15, "deep": 25}.get(depth, 15) count = {"quick": 8, "default": 15, "deep": 25}.get(depth, 15)
# Calculate days for freshness filter # Calculate days for freshness filter
@@ -81,6 +92,48 @@ def search_web(
return _normalize_results(response, from_date, to_date) return _normalize_results(response, from_date, to_date)
def _search_llm_context(
topic: str,
from_date: str,
to_date: str,
api_key: str,
depth: str = "default",
) -> List[Dict[str, Any]]:
"""Search via Brave LLM Context endpoint for pre-extracted web content.
Returns results in the same schema as search_web() for downstream compatibility.
Snippets contain actual page content instead of short descriptions.
"""
count = {"quick": 5, "default": 20, "deep": 50}.get(depth, 20)
max_tokens = {"quick": 2048, "default": 8192, "deep": 16384}.get(depth, 8192)
days = _days_between(from_date, to_date)
freshness = _brave_freshness(days)
params = {
"q": topic,
"count": count,
"maximum_number_of_tokens": max_tokens,
"context_threshold_mode": "balanced",
}
if freshness:
params["freshness"] = freshness
url = f"{LLM_CONTEXT_ENDPOINT}?{urlencode(params)}"
sys.stderr.write(f"[Web] Searching Brave LLM Context for: {topic}\n")
sys.stderr.flush()
response = http.request(
"GET",
url,
headers={"X-Subscription-Token": api_key},
timeout=30,
)
return _normalize_llm_context(response)
def _days_between(from_date: str, to_date: str) -> int: def _days_between(from_date: str, to_date: str) -> int:
"""Calculate days between two YYYY-MM-DD dates.""" """Calculate days between two YYYY-MM-DD dates."""
try: try:
@@ -138,7 +191,7 @@ def _normalize_results(
continue continue
if domain.startswith("www."): if domain.startswith("www."):
domain = domain[4:] domain = domain[4:]
except Exception: except (ValueError, TypeError):
domain = "" domain = ""
title = _clean_html(str(result.get("title", "")).strip()) title = _clean_html(str(result.get("title", "")).strip())
@@ -169,6 +222,69 @@ def _normalize_results(
return items return items
def _normalize_llm_context(response: Dict[str, Any]) -> List[Dict[str, Any]]:
"""Convert Brave LLM Context response to websearch item schema.
LLM Context returns grounding.generic[] with url, title, snippets[].
Sources metadata provides hostname and age for each URL.
"""
items = []
grounding = response.get("grounding", {})
sources = response.get("sources", {})
for i, result in enumerate(grounding.get("generic", [])):
if not isinstance(result, dict):
continue
url = result.get("url", "")
if not url:
continue
# Skip excluded domains
try:
domain = urlparse(url).netloc.lower()
if domain in EXCLUDED_DOMAINS:
continue
if domain.startswith("www."):
domain = domain[4:]
except (ValueError, TypeError):
domain = ""
title = str(result.get("title", "")).strip()
snippets = result.get("snippets", [])
snippet = "\n".join(str(s).strip() for s in snippets if s)
if not title and not snippet:
continue
# Parse date from sources metadata
source_meta = sources.get(url, {})
age_list = source_meta.get("age") or []
date = None
for age_str in age_list:
date = _parse_brave_date(age_str, None)
if date:
break
date_confidence = "med" if date else "low"
items.append({
"id": f"W{i+1}",
"title": title[:200],
"url": url,
"source_domain": source_meta.get("hostname", domain),
"snippet": snippet[:1500], # LLM Context returns richer content
"date": date,
"date_confidence": date_confidence,
"relevance": 0.7, # LLM Context pre-filters for relevance
"why_relevant": "",
})
sys.stderr.write(f"[Web] Brave LLM Context: {len(items)} results\n")
sys.stderr.flush()
return items
def _clean_html(text: str) -> str: def _clean_html(text: str) -> str:
"""Remove HTML tags and decode entities.""" """Remove HTML tags and decode entities."""
text = re.sub(r"<[^>]*>", "", text) text = re.sub(r"<[^>]*>", "", text)
+53 -21
View File
@@ -1,19 +1,34 @@
"""Model auto-selection for last30days skill.""" """Model auto-selection for last30days skill.
Model selection philosophy: this tool uses LLM APIs exclusively for
search tool invocation + structured JSON extraction. This is not
reasoning-heavy or creative work — mini models handle it equally well
at ~3-5x lower cost. We prefer the newest-generation mini model, falling
back to mainline only when mini isn't available.
xAI non-reasoning variant preferred: same pricing as reasoning, but
faster (skips thinking phase, saves reasoning token output costs).
"""
import re import re
import sys
from typing import Dict, List, Optional, Tuple from typing import Dict, List, Optional, Tuple
from . import cache, http, env from . import cache, http, env
# OpenAI API # OpenAI API
OPENAI_MODELS_URL = "https://api.openai.com/v1/models" OPENAI_MODELS_URL = "https://api.openai.com/v1/models"
OPENAI_FALLBACK_MODELS = ["gpt-5.2", "gpt-5.1", "gpt-5", "gpt-4.1", "gpt-4o"] # Ordered by cost-efficiency for web_search + JSON extraction tasks.
# Mini models first: same structured extraction quality at ~3x lower cost.
OPENAI_FALLBACK_MODELS = ["gpt-5-mini", "gpt-4.1-mini", "gpt-4.1", "gpt-4o"]
CODEX_FALLBACK_MODELS = ["gpt-5.1-codex-mini", "gpt-5.2"] CODEX_FALLBACK_MODELS = ["gpt-5.1-codex-mini", "gpt-5.2"]
# xAI API - Agent Tools API requires grok-4 family # xAI API - Agent Tools API requires grok-4 family
# Non-reasoning: same price, faster, no unnecessary thinking tokens.
# Both variants support function calling and structured outputs.
XAI_MODELS_URL = "https://api.x.ai/v1/models" XAI_MODELS_URL = "https://api.x.ai/v1/models"
XAI_ALIASES = { XAI_ALIASES = {
"latest": "grok-4-1-fast-non-reasoning", # Explicit: bare grok-4-1-fast aliases to reasoning variant "latest": "grok-4-1-fast-non-reasoning",
"stable": "grok-4-1-fast-non-reasoning", "stable": "grok-4-1-fast-non-reasoning",
} }
@@ -32,30 +47,45 @@ def parse_version(model_id: str) -> Optional[Tuple[int, ...]]:
return None return None
def is_mainline_openai_model(model_id: str) -> bool: def is_search_capable_model(model_id: str) -> bool:
"""Check if model is a mainline GPT model (not mini/nano/chat/codex/pro).""" """Check if model supports Responses API web_search with domain filtering.
Includes mini variants (same structured extraction quality, lower cost).
Excludes: nano (no web_search), gpt-4o-mini (no domain filtering),
chat/codex/pro/preview/turbo/search (specialized variants).
"""
model_lower = model_id.lower() model_lower = model_id.lower()
# Must be gpt-4o, gpt-4.1+, or gpt-5+ series (mainline, not mini/nano/etc) # gpt-4o-mini does NOT support web_search with filters — exclude it
if not re.match(r'^gpt-(?:4o|4\.1|5)(\.\d+)*$', model_lower): if model_lower.startswith("gpt-4o-mini"):
return False return False
# Exclude variants # Must be gpt-4o, gpt-4.1[-mini], or gpt-5[-mini] series
excludes = ['mini', 'nano', 'chat', 'codex', 'pro', 'preview', 'turbo'] if not re.match(r'^gpt-(?:4o|4\.1|5)(\.\d+)*(-mini)?$', model_lower):
for exc in excludes: return False
# Exclude unsupported variants
for exc in ['nano', 'chat', 'codex', 'pro', 'preview', 'turbo', 'search']:
if exc in model_lower: if exc in model_lower:
return False return False
return True return True
# Backward compat alias
is_mainline_openai_model = is_search_capable_model
def select_openai_model( def select_openai_model(
api_key: str, api_key: str,
policy: str = "auto", policy: str = "auto",
pin: Optional[str] = None, pin: Optional[str] = None,
mock_models: Optional[List[Dict]] = None, mock_models: Optional[List[Dict]] = None,
) -> str: ) -> str:
"""Select the best OpenAI model based on policy. """Select the most cost-efficient OpenAI model for web_search + JSON extraction.
Prefers mini models within the newest generation available, since the task
is structured extraction (not reasoning or creative work).
Args: Args:
api_key: OpenAI API key api_key: OpenAI API key
@@ -82,27 +112,29 @@ def select_openai_model(
headers = {"Authorization": f"Bearer {api_key}"} headers = {"Authorization": f"Bearer {api_key}"}
response = http.get(OPENAI_MODELS_URL, headers=headers) response = http.get(OPENAI_MODELS_URL, headers=headers)
models = response.get("data", []) models = response.get("data", [])
except http.HTTPError: except http.HTTPError as e:
# Fall back to known models sys.stderr.write(f"[Models] Failed to fetch OpenAI models: {e}")
if hasattr(e, 'status_code') and e.status_code in (401, 403):
sys.stderr.write(" — API key may be invalid or lack permissions")
sys.stderr.write(f", using fallback {OPENAI_FALLBACK_MODELS[0]}\n")
return OPENAI_FALLBACK_MODELS[0] return OPENAI_FALLBACK_MODELS[0]
# Filter to mainline models candidates = [m for m in models if is_search_capable_model(m.get("id", ""))]
candidates = [m for m in models if is_mainline_openai_model(m.get("id", ""))]
if not candidates: if not candidates:
# No gpt-5 models found, use fallback
return OPENAI_FALLBACK_MODELS[0] return OPENAI_FALLBACK_MODELS[0]
# Sort by version (descending), then by created timestamp # Sort: newest generation first, prefer mini within same generation
def sort_key(m): def sort_key(m):
version = parse_version(m.get("id", "")) or (0,) model_id = m.get("id", "")
created = m.get("created", 0) version = parse_version(model_id) or (0,)
return (version, created) major = version[0] if version else 0
is_mini = 1 if "mini" in model_id.lower() else 0
return (major, is_mini, version)
candidates.sort(key=sort_key, reverse=True) candidates.sort(key=sort_key, reverse=True)
selected = candidates[0]["id"] selected = candidates[0]["id"]
# Cache the selection
cache.set_cached_model("openai", selected) cache.set_cached_model("openai", selected)
return selected return selected
+4 -3
View File
@@ -7,9 +7,10 @@ from typing import Any, Dict, List, Optional
from . import http, env from . import http, env
# Fallback models when the selected model isn't accessible (e.g., org not verified for GPT-5) # Fallback models when the selected model isn't accessible (e.g., org not verified).
# Note: gpt-4o-mini does NOT support web_search with filters param, so exclude it # Ordered by cost-efficiency: mini models handle structured extraction equally well.
MODEL_FALLBACK_ORDER = ["gpt-4.1", "gpt-4o"] # Note: gpt-4o-mini does NOT support web_search with filters — excluded.
MODEL_FALLBACK_ORDER = ["gpt-5-mini", "gpt-4.1-mini", "gpt-4.1", "gpt-4o"]
def _log_error(msg: str): def _log_error(msg: str):
+6 -1
View File
@@ -130,7 +130,12 @@ def _extract_domain_queries(topic: str, events: List[Dict]) -> List[str]:
def _search_single_query(query: str, page: int = 1) -> Dict[str, Any]: def _search_single_query(query: str, page: int = 1) -> Dict[str, Any]:
"""Run a single search query against Gamma API.""" """Run a single search query against Gamma API."""
params = {"q": query, "page": str(page)} params = {
"q": query,
"page": str(page),
"events_status": "active",
"keep_closed_markets": "0",
}
url = f"{GAMMA_SEARCH_URL}?{urlencode(params)}" url = f"{GAMMA_SEARCH_URL}?{urlencode(params)}"
try: try:
+111
View File
@@ -0,0 +1,111 @@
"""Query type detection for source selection and scoring adjustments."""
import re
from typing import Literal
QueryType = Literal["product", "concept", "opinion", "how_to", "comparison", "breaking_news", "prediction"]
# Pattern-based classification (no LLM, no external deps)
_PRODUCT_PATTERNS = re.compile(
r"\b(price|pricing|cost|buy|purchase|deal|discount|subscription|plan|tier|free tier|alternative)\b", re.I
)
_CONCEPT_PATTERNS = re.compile(
r"\b(what is|what are|explain|definition|how does|how do|overview|introduction|guide to|primer)\b", re.I
)
_OPINION_PATTERNS = re.compile(
r"\b(worth it|thoughts on|opinion|review|experience with|recommend|should i|pros and cons|good or bad)\b", re.I
)
_HOWTO_PATTERNS = re.compile(
r"\b(how to|tutorial|step by step|setup|install|configure|deploy|migrate|implement|build a|create a|prompting|prompts?|best practices|tips|examples|animation|animations)\b",
re.I,
)
_COMPARISON_PATTERNS = re.compile(
r"\b(vs\.?|versus|compared to|comparison|better than|difference between|switch from)\b", re.I
)
_BREAKING_PATTERNS = re.compile(
r"\b(latest|breaking|just announced|launched|released|new|update|news|happened|today|this week)\b", re.I
)
_PREDICTION_PATTERNS = re.compile(
r"\b(predict|forecast|odds|chance|probability|election|outcome|bet on|market for)\b", re.I
)
def detect_query_type(topic: str) -> QueryType:
"""Classify a query into a type using pattern matching.
Returns the first match in priority order:
comparison > how_to > product > opinion > prediction > concept > breaking_news.
"""
# Most specific first
if _COMPARISON_PATTERNS.search(topic):
return "comparison"
if _HOWTO_PATTERNS.search(topic):
return "how_to"
if _PRODUCT_PATTERNS.search(topic):
return "product"
if _OPINION_PATTERNS.search(topic):
return "opinion"
if _PREDICTION_PATTERNS.search(topic):
return "prediction"
if _CONCEPT_PATTERNS.search(topic):
return "concept"
if _BREAKING_PATTERNS.search(topic):
return "breaking_news"
# Default: treat as breaking news (most common use case for "last 30 days")
return "breaking_news"
# Source tiering by query type.
# Tier 1: always run. Tier 2: run if available. Tier 3: opt-in only.
# Sources not listed are implicitly tier 3 (opt-in).
SOURCE_TIERS = {
"product": {"tier1": {"reddit", "x", "youtube"}, "tier2": {"web", "tiktok"}},
"concept": {"tier1": {"reddit", "hn", "web"}, "tier2": {"youtube", "x"}},
"opinion": {"tier1": {"reddit", "x"}, "tier2": {"youtube", "bluesky"}},
"how_to": {"tier1": {"youtube", "reddit", "hn"}, "tier2": {"web", "x"}},
"comparison": {"tier1": {"reddit", "hn", "youtube"}, "tier2": {"x", "web"}},
"breaking_news": {"tier1": {"x", "reddit", "web"}, "tier2": {"hn", "bluesky", "youtube"}},
"prediction": {"tier1": {"polymarket", "x", "reddit"}, "tier2": {"web", "hn", "youtube"}},
}
# WebSearch penalty adjustment by query type.
# Points subtracted from websearch score (0-100 scale). 0 = no penalty, 15 = full penalty.
# Concept/how_to queries benefit from authoritative web sources.
WEBSEARCH_PENALTY_BY_TYPE = {
"product": 15, # default: social discussion > blog posts
"concept": 0, # web docs are the best source
"opinion": 15, # social discussion > blog posts
"how_to": 5, # tutorials on web are valuable
"comparison": 10, # mix of social and web
"breaking_news": 10, # news sites are valuable
"prediction": 15, # social/market data > web articles
}
# Tiebreaker priority overrides by query type.
# Maps source type name to priority (lower = higher priority).
TIEBREAKER_BY_TYPE = {
"product": {"reddit": 0, "x": 1, "youtube": 2, "tiktok": 3, "instagram": 4, "hn": 5, "web": 6, "polymarket": 7},
"concept": {"hn": 0, "reddit": 1, "web": 2, "youtube": 3, "x": 4, "tiktok": 5, "instagram": 6, "polymarket": 7},
"opinion": {"reddit": 0, "x": 1, "bluesky": 2, "youtube": 3, "hn": 4, "tiktok": 5, "web": 6, "polymarket": 7},
"how_to": {"youtube": 0, "reddit": 1, "hn": 2, "web": 3, "x": 4, "tiktok": 5, "instagram": 6, "polymarket": 7},
"comparison": {"reddit": 0, "hn": 1, "youtube": 2, "x": 3, "web": 4, "tiktok": 5, "instagram": 6, "polymarket": 7},
"breaking_news": {"x": 0, "reddit": 1, "web": 2, "hn": 3, "bluesky": 4, "tiktok": 5, "youtube": 6, "polymarket": 7},
"prediction": {"polymarket": 0, "x": 1, "reddit": 2, "web": 3, "hn": 4, "bluesky": 5, "youtube": 6, "tiktok": 7},
}
def is_source_enabled(source: str, query_type: QueryType, explicitly_requested: bool = False) -> bool:
"""Check if a source should run for a given query type.
Tier 1 and Tier 2 sources are enabled. Tier 3 (unlisted) sources only run
if explicitly requested via --search flag. Truth Social is always opt-in.
"""
if source == "truthsocial":
return explicitly_requested
if explicitly_requested:
return True
tiers = SOURCE_TIERS.get(query_type, SOURCE_TIERS["breaking_news"])
return source in tiers["tier1"] or source in tiers["tier2"]
+40 -31
View File
@@ -4,16 +4,19 @@ import math
from typing import List, Optional, Union from typing import List, Optional, Union
from . import dates, schema from . import dates, schema
from .query_type import QueryType, WEBSEARCH_PENALTY_BY_TYPE, TIEBREAKER_BY_TYPE
# Score weights for Reddit/X (has engagement) # Score weights for Reddit/X (has engagement)
WEIGHT_RELEVANCE = 0.45 WEIGHT_RELEVANCE = 0.45
WEIGHT_RECENCY = 0.25 WEIGHT_RECENCY = 0.25
WEIGHT_ENGAGEMENT = 0.30 WEIGHT_ENGAGEMENT = 0.30
# WebSearch weights (no engagement, reweighted to 100%) # WebSearch weights (no engagement data available)
WEBSEARCH_WEIGHT_RELEVANCE = 0.55 WEBSEARCH_WEIGHT_RELEVANCE = 0.55
WEBSEARCH_WEIGHT_RECENCY = 0.45 WEBSEARCH_WEIGHT_RECENCY = 0.45
WEBSEARCH_SOURCE_PENALTY = 15 # Points deducted for lacking engagement # Default web search penalty (fallback when query_type is not provided).
# Per-type penalties in query_type.WEBSEARCH_PENALTY_BY_TYPE.
WEBSEARCH_SOURCE_PENALTY = 15
# WebSearch date confidence adjustments # WebSearch date confidence adjustments
WEBSEARCH_VERIFIED_BONUS = 10 # Bonus for URL-verified recent date (high confidence) WEBSEARCH_VERIFIED_BONUS = 10 # Bonus for URL-verified recent date (high confidence)
@@ -642,19 +645,17 @@ def score_polymarket_items(items: List[schema.PolymarketItem]) -> List[schema.Po
return items return items
def score_websearch_items(items: List[schema.WebSearchItem]) -> List[schema.WebSearchItem]: def score_websearch_items(items: List[schema.WebSearchItem], query_type: QueryType = None) -> List[schema.WebSearchItem]:
"""Compute scores for WebSearch items WITHOUT engagement metrics. """Compute scores for WebSearch items WITHOUT engagement metrics.
Uses reweighted formula: 55% relevance + 45% recency - 15pt source penalty. Uses reweighted formula: 55% relevance + 45% recency - penalty.
This ensures WebSearch items rank below comparable Reddit/X items. Penalty varies by query type: concept queries get 0 penalty (web docs
are authoritative), while product/opinion queries get full 15pt penalty
Date confidence adjustments: (social discussion is more valuable).
- High confidence (URL-verified date): +10 bonus
- Med confidence (snippet-extracted date): no change
- Low confidence (no date signals): -20 penalty
Args: Args:
items: List of WebSearch items items: List of WebSearch items
query_type: Query classification for penalty adjustment
Returns: Returns:
Items with updated scores Items with updated scores
@@ -682,8 +683,9 @@ def score_websearch_items(items: List[schema.WebSearchItem]) -> List[schema.WebS
WEBSEARCH_WEIGHT_RECENCY * rec_score WEBSEARCH_WEIGHT_RECENCY * rec_score
) )
# Apply source penalty (WebSearch < Reddit/X for same relevance/recency) # Apply source penalty (varies by query type)
overall -= WEBSEARCH_SOURCE_PENALTY penalty = WEBSEARCH_PENALTY_BY_TYPE.get(query_type, WEBSEARCH_SOURCE_PENALTY) if query_type else WEBSEARCH_SOURCE_PENALTY
overall -= penalty
# Apply date confidence adjustments # Apply date confidence adjustments
# High confidence (URL-verified): reward with bonus # High confidence (URL-verified): reward with bonus
@@ -699,15 +701,36 @@ def score_websearch_items(items: List[schema.WebSearchItem]) -> List[schema.WebS
return items return items
def sort_items(items: List[Union[schema.RedditItem, schema.XItem, schema.WebSearchItem, schema.YouTubeItem, schema.TikTokItem, schema.InstagramItem, schema.HackerNewsItem, schema.PolymarketItem]]) -> List: _ITEM_SOURCE_MAP = {
"""Sort items by score (descending), then date, then source priority. schema.RedditItem: "reddit",
schema.XItem: "x",
schema.YouTubeItem: "youtube",
schema.TikTokItem: "tiktok",
schema.InstagramItem: "instagram",
schema.HackerNewsItem: "hn",
schema.BlueskyItem: "bluesky",
schema.TruthSocialItem: "truthsocial",
schema.PolymarketItem: "polymarket",
}
_DEFAULT_TIEBREAKER = {"reddit": 0, "x": 1, "youtube": 2, "tiktok": 3, "instagram": 4, "hn": 5, "bluesky": 6, "truthsocial": 7, "polymarket": 8, "web": 9}
def sort_items(items: List[Union[schema.RedditItem, schema.XItem, schema.WebSearchItem, schema.YouTubeItem, schema.TikTokItem, schema.InstagramItem, schema.HackerNewsItem, schema.PolymarketItem]], query_type: QueryType = None) -> List:
"""Sort items by score (descending), then date, then source tiebreaker.
Tiebreaker (tertiary sort key, after score and date): source priority
varies by query type. YouTube ranks first for how_to, X ranks first
for breaking_news, Polymarket ranks first for prediction.
Args: Args:
items: List of items to sort items: List of items to sort
query_type: Query classification for tiebreaker adjustment
Returns: Returns:
Sorted items Sorted items
""" """
tiebreaker = TIEBREAKER_BY_TYPE.get(query_type, _DEFAULT_TIEBREAKER) if query_type else _DEFAULT_TIEBREAKER
def sort_key(item): def sort_key(item):
# Primary: score descending (negate for descending) # Primary: score descending (negate for descending)
score = -item.score score = -item.score
@@ -716,23 +739,9 @@ def sort_items(items: List[Union[schema.RedditItem, schema.XItem, schema.WebSear
date = item.date or "0000-00-00" date = item.date or "0000-00-00"
date_key = -int(date.replace("-", "")) date_key = -int(date.replace("-", ""))
# Tertiary: source priority (Reddit > X > YouTube > TikTok > HN > Polymarket > WebSearch) # Tertiary: query-type-aware source priority
if isinstance(item, schema.RedditItem): source_name = _ITEM_SOURCE_MAP.get(type(item), "web")
source_priority = 0 source_priority = tiebreaker.get(source_name, 99)
elif isinstance(item, schema.XItem):
source_priority = 1
elif isinstance(item, schema.YouTubeItem):
source_priority = 2
elif isinstance(item, schema.TikTokItem):
source_priority = 3
elif isinstance(item, schema.InstagramItem):
source_priority = 4
elif isinstance(item, schema.HackerNewsItem):
source_priority = 5
elif isinstance(item, schema.PolymarketItem):
source_priority = 6
else: # WebSearchItem
source_priority = 7
# Quaternary: title/text for stability # Quaternary: title/text for stability
text = getattr(item, "title", "") or getattr(item, "text", "") text = getattr(item, "title", "") or getattr(item, "text", "")
+197
View File
@@ -0,0 +1,197 @@
"""Tests for Brave Search module, including LLM Context endpoint."""
import sys
import os
import unittest
# Ensure scripts/ is on path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'scripts'))
from lib.brave_search import (
_normalize_results,
_normalize_llm_context,
_days_between,
_brave_freshness,
_parse_brave_date,
EXCLUDED_DOMAINS,
)
class TestDaysBetween(unittest.TestCase):
def test_same_day(self):
self.assertEqual(_days_between("2026-03-01", "2026-03-01"), 1)
def test_one_week(self):
self.assertEqual(_days_between("2026-03-01", "2026-03-08"), 7)
def test_invalid_dates(self):
self.assertEqual(_days_between("bad", "dates"), 30)
class TestBraveFreshness(unittest.TestCase):
def test_one_day(self):
self.assertEqual(_brave_freshness(1), "pd")
def test_one_week(self):
self.assertEqual(_brave_freshness(7), "pw")
def test_one_month(self):
self.assertEqual(_brave_freshness(31), "pm")
def test_longer_returns_range(self):
result = _brave_freshness(60)
self.assertIn("to", result)
def test_none(self):
self.assertIsNone(_brave_freshness(None))
class TestParseBraveDate(unittest.TestCase):
def test_hours_ago(self):
result = _parse_brave_date("3 hours ago", None)
self.assertIsNotNone(result)
self.assertRegex(result, r"\d{4}-\d{2}-\d{2}")
def test_days_ago(self):
result = _parse_brave_date("5 days ago", None)
self.assertIsNotNone(result)
def test_weeks_ago(self):
result = _parse_brave_date("2 weeks ago", None)
self.assertIsNotNone(result)
def test_iso_date(self):
self.assertEqual(_parse_brave_date("2026-03-10T12:00:00", None), "2026-03-10")
def test_none(self):
self.assertIsNone(_parse_brave_date(None, None))
class TestNormalizeResults(unittest.TestCase):
def test_merges_news_and_web(self):
response = {
"news": {"results": [
{"url": "https://news.example.com/a", "title": "News A", "description": "News desc"},
]},
"web": {"results": [
{"url": "https://blog.example.com/b", "title": "Blog B", "description": "Blog desc"},
]},
}
items = _normalize_results(response, "2026-03-01", "2026-03-10")
self.assertEqual(len(items), 2)
self.assertEqual(items[0]["title"], "News A")
self.assertEqual(items[1]["title"], "Blog B")
def test_excludes_reddit_and_x(self):
response = {
"web": {"results": [
{"url": "https://www.reddit.com/r/test/123", "title": "Reddit", "description": "text"},
{"url": "https://x.com/user/status/1", "title": "X post", "description": "text"},
{"url": "https://example.com/ok", "title": "OK", "description": "text"},
]},
}
items = _normalize_results(response, "2026-03-01", "2026-03-10")
self.assertEqual(len(items), 1)
self.assertEqual(items[0]["title"], "OK")
def test_default_relevance(self):
response = {"web": {"results": [
{"url": "https://a.com", "title": "A", "description": "desc"},
]}}
items = _normalize_results(response, "2026-03-01", "2026-03-10")
self.assertEqual(items[0]["relevance"], 0.6)
class TestNormalizeLlmContext(unittest.TestCase):
def _make_response(self, generic=None, sources=None):
return {
"grounding": {"generic": generic or []},
"sources": sources or {},
}
def test_basic_result(self):
resp = self._make_response(
generic=[{
"url": "https://docs.example.com/page",
"title": "Example Page",
"snippets": ["First chunk of text.", "Second chunk of text."],
}],
sources={
"https://docs.example.com/page": {
"title": "Example Page",
"hostname": "docs.example.com",
"age": ["2026-03-05", "5 days ago"],
}
},
)
items = _normalize_llm_context(resp)
self.assertEqual(len(items), 1)
item = items[0]
self.assertEqual(item["title"], "Example Page")
self.assertEqual(item["url"], "https://docs.example.com/page")
self.assertIn("First chunk", item["snippet"])
self.assertIn("Second chunk", item["snippet"])
self.assertEqual(item["date"], "2026-03-05")
self.assertEqual(item["date_confidence"], "med")
self.assertEqual(item["relevance"], 0.7)
self.assertEqual(item["source_domain"], "docs.example.com")
def test_excludes_reddit(self):
resp = self._make_response(
generic=[
{"url": "https://www.reddit.com/r/test", "title": "Reddit", "snippets": ["text"]},
{"url": "https://example.com", "title": "OK", "snippets": ["text"]},
],
)
items = _normalize_llm_context(resp)
self.assertEqual(len(items), 1)
self.assertEqual(items[0]["title"], "OK")
def test_empty_grounding(self):
resp = self._make_response()
items = _normalize_llm_context(resp)
self.assertEqual(items, [])
def test_snippet_truncation(self):
long_snippet = "x" * 2000
resp = self._make_response(
generic=[{"url": "https://a.com", "title": "A", "snippets": [long_snippet]}],
)
items = _normalize_llm_context(resp)
self.assertLessEqual(len(items[0]["snippet"]), 1500)
def test_no_date_gives_low_confidence(self):
resp = self._make_response(
generic=[{"url": "https://a.com", "title": "A", "snippets": ["text"]}],
sources={"https://a.com": {"hostname": "a.com", "age": None}},
)
items = _normalize_llm_context(resp)
self.assertIsNone(items[0]["date"])
self.assertEqual(items[0]["date_confidence"], "low")
def test_multiple_age_entries_picks_first_valid(self):
resp = self._make_response(
generic=[{"url": "https://a.com", "title": "A", "snippets": ["text"]}],
sources={"https://a.com": {
"hostname": "a.com",
"age": ["Monday, March 10, 2026", "2026-03-10", "1 day ago"],
}},
)
items = _normalize_llm_context(resp)
self.assertEqual(items[0]["date"], "2026-03-10")
def test_ids_are_sequential(self):
resp = self._make_response(
generic=[
{"url": "https://a.com", "title": "A", "snippets": ["a"]},
{"url": "https://b.com", "title": "B", "snippets": ["b"]},
{"url": "https://c.com", "title": "C", "snippets": ["c"]},
],
)
items = _normalize_llm_context(resp)
ids = [item["id"] for item in items]
self.assertEqual(ids, ["W1", "W2", "W3"])
if __name__ == "__main__":
unittest.main()
+95 -29
View File
@@ -28,21 +28,47 @@ class TestParseVersion(unittest.TestCase):
self.assertIsNone(result) self.assertIsNone(result)
class TestIsMainlineOpenAIModel(unittest.TestCase): class TestIsSearchCapableModel(unittest.TestCase):
def test_gpt5_is_mainline(self): def test_gpt5_is_capable(self):
self.assertTrue(models.is_search_capable_model("gpt-5"))
def test_gpt52_is_capable(self):
self.assertTrue(models.is_search_capable_model("gpt-5.2"))
def test_gpt5_mini_is_capable(self):
self.assertTrue(models.is_search_capable_model("gpt-5-mini"))
def test_gpt41_mini_is_capable(self):
self.assertTrue(models.is_search_capable_model("gpt-4.1-mini"))
def test_gpt4o_is_capable(self):
self.assertTrue(models.is_search_capable_model("gpt-4o"))
def test_gpt4o_mini_not_capable(self):
"""gpt-4o-mini does not support web_search with domain filtering."""
self.assertFalse(models.is_search_capable_model("gpt-4o-mini"))
def test_nano_not_capable(self):
"""nano models don't support web_search."""
self.assertFalse(models.is_search_capable_model("gpt-4.1-nano"))
self.assertFalse(models.is_search_capable_model("gpt-5-nano"))
def test_gpt4_not_capable(self):
self.assertFalse(models.is_search_capable_model("gpt-4"))
def test_codex_not_capable(self):
self.assertFalse(models.is_search_capable_model("gpt-5.1-codex"))
def test_backward_compat_alias(self):
"""is_mainline_openai_model still works as alias."""
self.assertTrue(models.is_mainline_openai_model("gpt-5")) self.assertTrue(models.is_mainline_openai_model("gpt-5"))
def test_gpt52_is_mainline(self):
self.assertTrue(models.is_mainline_openai_model("gpt-5.2"))
def test_gpt5_mini_is_not_mainline(self):
self.assertFalse(models.is_mainline_openai_model("gpt-5-mini"))
def test_gpt4_is_not_mainline(self):
self.assertFalse(models.is_mainline_openai_model("gpt-4"))
class TestSelectOpenAIModel(unittest.TestCase): class TestSelectOpenAIModel(unittest.TestCase):
def setUp(self):
from lib import cache
cache.MODEL_CACHE_FILE.unlink(missing_ok=True)
def test_pinned_policy(self): def test_pinned_policy(self):
result = models.select_openai_model( result = models.select_openai_model(
"fake-key", "fake-key",
@@ -51,20 +77,8 @@ class TestSelectOpenAIModel(unittest.TestCase):
) )
self.assertEqual(result, "gpt-5.1") self.assertEqual(result, "gpt-5.1")
def test_auto_with_mock_models(self): def test_prefers_mini_over_mainline(self):
mock_models = [ """Mini models should be preferred for cost-efficiency."""
{"id": "gpt-5.2", "created": 1704067200},
{"id": "gpt-5.1", "created": 1701388800},
{"id": "gpt-5", "created": 1698710400},
]
result = models.select_openai_model(
"fake-key",
policy="auto",
mock_models=mock_models
)
self.assertEqual(result, "gpt-5.2")
def test_auto_filters_variants(self):
mock_models = [ mock_models = [
{"id": "gpt-5.2", "created": 1704067200}, {"id": "gpt-5.2", "created": 1704067200},
{"id": "gpt-5-mini", "created": 1704067200}, {"id": "gpt-5-mini", "created": 1704067200},
@@ -75,8 +89,50 @@ class TestSelectOpenAIModel(unittest.TestCase):
policy="auto", policy="auto",
mock_models=mock_models mock_models=mock_models
) )
self.assertEqual(result, "gpt-5-mini")
def test_prefers_newer_generation_mini(self):
"""gpt-5-mini should beat gpt-4.1-mini (newer generation)."""
mock_models = [
{"id": "gpt-4.1-mini", "created": 1701388800},
{"id": "gpt-5-mini", "created": 1704067200},
{"id": "gpt-4.1", "created": 1698710400},
]
result = models.select_openai_model(
"fake-key",
policy="auto",
mock_models=mock_models
)
self.assertEqual(result, "gpt-5-mini")
def test_falls_back_to_mainline_when_no_mini(self):
"""Without mini models, mainline models are selected."""
mock_models = [
{"id": "gpt-5.2", "created": 1704067200},
{"id": "gpt-4.1", "created": 1698710400},
]
result = models.select_openai_model(
"fake-key",
policy="auto",
mock_models=mock_models
)
self.assertEqual(result, "gpt-5.2") self.assertEqual(result, "gpt-5.2")
def test_filters_unsupported_variants(self):
"""Nano, codex, preview models should be excluded."""
mock_models = [
{"id": "gpt-5-nano", "created": 1704067200},
{"id": "gpt-5.1-codex", "created": 1704067200},
{"id": "gpt-4o-mini", "created": 1704067200},
{"id": "gpt-4.1-mini", "created": 1698710400},
]
result = models.select_openai_model(
"fake-key",
policy="auto",
mock_models=mock_models
)
self.assertEqual(result, "gpt-4.1-mini")
class TestSelectXAIModel(unittest.TestCase): class TestSelectXAIModel(unittest.TestCase):
def test_latest_policy(self): def test_latest_policy(self):
@@ -106,6 +162,10 @@ class TestSelectXAIModel(unittest.TestCase):
class TestGetModels(unittest.TestCase): class TestGetModels(unittest.TestCase):
def setUp(self):
from lib import cache
cache.MODEL_CACHE_FILE.unlink(missing_ok=True)
def test_no_keys_returns_none(self): def test_no_keys_returns_none(self):
config = {} config = {}
result = models.get_models(config) result = models.get_models(config)
@@ -114,9 +174,12 @@ class TestGetModels(unittest.TestCase):
def test_openai_key_only(self): def test_openai_key_only(self):
config = {"OPENAI_API_KEY": "sk-test"} config = {"OPENAI_API_KEY": "sk-test"}
mock_models = [{"id": "gpt-5.2", "created": 1704067200}] mock_models = [
{"id": "gpt-5.2", "created": 1704067200},
{"id": "gpt-5-mini", "created": 1704067200},
]
result = models.get_models(config, mock_openai_models=mock_models) result = models.get_models(config, mock_openai_models=mock_models)
self.assertEqual(result["openai"], "gpt-5.2") self.assertEqual(result["openai"], "gpt-5-mini")
self.assertIsNone(result["xai"]) self.assertIsNone(result["xai"])
def test_both_keys(self): def test_both_keys(self):
@@ -124,10 +187,13 @@ class TestGetModels(unittest.TestCase):
"OPENAI_API_KEY": "sk-test", "OPENAI_API_KEY": "sk-test",
"XAI_API_KEY": "xai-test", "XAI_API_KEY": "xai-test",
} }
mock_openai = [{"id": "gpt-5.2", "created": 1704067200}] mock_openai = [
{"id": "gpt-5.2", "created": 1704067200},
{"id": "gpt-5-mini", "created": 1704067200},
]
mock_xai = [{"id": "grok-4-1-fast-non-reasoning", "created": 1704067200}] mock_xai = [{"id": "grok-4-1-fast-non-reasoning", "created": 1704067200}]
result = models.get_models(config, mock_openai, mock_xai) result = models.get_models(config, mock_openai, mock_xai)
self.assertEqual(result["openai"], "gpt-5.2") self.assertEqual(result["openai"], "gpt-5-mini")
self.assertEqual(result["xai"], "grok-4-1-fast-non-reasoning") self.assertEqual(result["xai"], "grok-4-1-fast-non-reasoning")
+10 -5
View File
@@ -64,13 +64,18 @@ class TestIsModelAccessError(unittest.TestCase):
class TestModelFallbackOrder(unittest.TestCase): class TestModelFallbackOrder(unittest.TestCase):
"""Tests for MODEL_FALLBACK_ORDER constant.""" """Tests for MODEL_FALLBACK_ORDER constant."""
def test_contains_gpt4o(self): def test_mini_first(self):
"""Fallback list should include gpt-4o.""" """Mini models should come first (cost-efficient for structured extraction)."""
self.assertEqual(MODEL_FALLBACK_ORDER[0], "gpt-5-mini")
def test_contains_mainline_fallbacks(self):
"""Fallback list should include mainline models as last resort."""
self.assertIn("gpt-4.1", MODEL_FALLBACK_ORDER)
self.assertIn("gpt-4o", MODEL_FALLBACK_ORDER) self.assertIn("gpt-4o", MODEL_FALLBACK_ORDER)
def test_gpt41_is_first(self): def test_no_gpt4o_mini(self):
"""gpt-4.1 should be the first fallback option.""" """gpt-4o-mini should NOT be in fallback (no domain filtering support)."""
self.assertEqual(MODEL_FALLBACK_ORDER[0], "gpt-4.1") self.assertNotIn("gpt-4o-mini", MODEL_FALLBACK_ORDER)
if __name__ == "__main__": if __name__ == "__main__":
+156
View File
@@ -0,0 +1,156 @@
"""Tests for query type detection and source tiering."""
import sys
import unittest
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent / "scripts"))
from lib.query_type import (
detect_query_type,
is_source_enabled,
WEBSEARCH_PENALTY_BY_TYPE,
TIEBREAKER_BY_TYPE,
SOURCE_TIERS,
)
class TestDetectQueryType(unittest.TestCase):
def test_product_queries(self):
self.assertEqual(detect_query_type("cursor IDE pricing"), "product")
self.assertEqual(detect_query_type("is Claude Pro worth the cost"), "product")
self.assertEqual(detect_query_type("best free tier LLM API"), "product")
def test_concept_queries(self):
self.assertEqual(detect_query_type("what is WebTransport"), "concept")
self.assertEqual(detect_query_type("explain React Server Components"), "concept")
self.assertEqual(detect_query_type("how does MCP protocol work"), "concept")
def test_opinion_queries(self):
self.assertEqual(detect_query_type("is cursor worth it"), "opinion")
self.assertEqual(detect_query_type("thoughts on Claude Code"), "opinion")
self.assertEqual(detect_query_type("should i switch to Neovim"), "opinion")
def test_howto_queries(self):
self.assertEqual(detect_query_type("how to deploy on Vercel"), "how_to")
self.assertEqual(detect_query_type("tutorial for building MCP servers"), "how_to")
self.assertEqual(detect_query_type("step by step Kubernetes setup"), "how_to")
self.assertEqual(detect_query_type("nano banana pro prompting"), "how_to")
self.assertEqual(detect_query_type("remotion animations for Claude Code"), "how_to")
def test_comparison_queries(self):
self.assertEqual(detect_query_type("cursor vs windsurf"), "comparison")
self.assertEqual(detect_query_type("Claude compared to GPT-5"), "comparison")
self.assertEqual(detect_query_type("difference between React and Vue"), "comparison")
def test_breaking_news_queries(self):
self.assertEqual(detect_query_type("latest AI funding rounds"), "breaking_news")
self.assertEqual(detect_query_type("OpenAI just announced GPT-6"), "breaking_news")
def test_prediction_queries(self):
self.assertEqual(detect_query_type("odds of Fed rate cut"), "prediction")
self.assertEqual(detect_query_type("predict the next recession"), "prediction")
self.assertEqual(detect_query_type("election outcome 2028"), "prediction")
def test_default_is_breaking_news(self):
self.assertEqual(detect_query_type("tariffs"), "breaking_news")
self.assertEqual(detect_query_type("AI agents"), "breaking_news")
def test_comparison_beats_product(self):
"""Comparison is more specific than product."""
self.assertEqual(detect_query_type("cursor vs windsurf pricing"), "comparison")
def test_howto_beats_concept(self):
"""How-to is more specific than concept."""
self.assertEqual(detect_query_type("how to explain transformers"), "how_to")
def test_will_alone_not_prediction(self):
"""Bare 'will' should not trigger prediction classification."""
self.assertNotEqual(detect_query_type("Will React 19 support concurrent mode"), "prediction")
def test_or_for_not_comparison(self):
"""'or X for Y' should not trigger comparison classification."""
self.assertNotEqual(detect_query_type("best tools or libraries for Python"), "comparison")
class TestIsSourceEnabled(unittest.TestCase):
def test_truthsocial_always_opt_in(self):
for qt in ["product", "concept", "opinion", "breaking_news", "prediction"]:
self.assertFalse(is_source_enabled("truthsocial", qt))
self.assertTrue(is_source_enabled("truthsocial", "breaking_news", explicitly_requested=True))
def test_tier1_sources_enabled(self):
self.assertTrue(is_source_enabled("reddit", "product"))
self.assertTrue(is_source_enabled("youtube", "how_to"))
self.assertTrue(is_source_enabled("polymarket", "prediction"))
self.assertTrue(is_source_enabled("x", "breaking_news"))
def test_tier2_sources_enabled(self):
self.assertTrue(is_source_enabled("web", "product"))
self.assertTrue(is_source_enabled("bluesky", "opinion"))
self.assertTrue(is_source_enabled("x", "how_to"))
self.assertTrue(is_source_enabled("youtube", "breaking_news"))
self.assertTrue(is_source_enabled("hn", "prediction"))
def test_tier3_sources_disabled_by_default(self):
self.assertFalse(is_source_enabled("instagram", "concept"))
self.assertFalse(is_source_enabled("tiktok", "comparison"))
self.assertFalse(is_source_enabled("bluesky", "product"))
def test_explicit_request_overrides_tier(self):
self.assertTrue(is_source_enabled("instagram", "concept", explicitly_requested=True))
self.assertTrue(is_source_enabled("tiktok", "comparison", explicitly_requested=True))
class TestWebSearchPenalty(unittest.TestCase):
def test_concept_has_zero_penalty(self):
self.assertEqual(WEBSEARCH_PENALTY_BY_TYPE["concept"], 0)
def test_product_has_full_penalty(self):
self.assertEqual(WEBSEARCH_PENALTY_BY_TYPE["product"], 15)
def test_howto_has_reduced_penalty(self):
self.assertLess(WEBSEARCH_PENALTY_BY_TYPE["how_to"], WEBSEARCH_PENALTY_BY_TYPE["product"])
def test_all_query_types_have_penalty(self):
for qt in ["product", "concept", "opinion", "how_to", "comparison", "breaking_news", "prediction"]:
self.assertIn(qt, WEBSEARCH_PENALTY_BY_TYPE)
class TestTiebreakerPriority(unittest.TestCase):
def test_youtube_highest_for_howto(self):
self.assertEqual(TIEBREAKER_BY_TYPE["how_to"]["youtube"], 0)
def test_x_highest_for_breaking_news(self):
self.assertEqual(TIEBREAKER_BY_TYPE["breaking_news"]["x"], 0)
def test_polymarket_highest_for_prediction(self):
self.assertEqual(TIEBREAKER_BY_TYPE["prediction"]["polymarket"], 0)
def test_hn_highest_for_concept(self):
self.assertEqual(TIEBREAKER_BY_TYPE["concept"]["hn"], 0)
def test_all_query_types_have_tiebreakers(self):
for qt in ["product", "concept", "opinion", "how_to", "comparison", "breaking_news", "prediction"]:
self.assertIn(qt, TIEBREAKER_BY_TYPE)
class TestSourceTiers(unittest.TestCase):
def test_all_query_types_have_tiers(self):
for qt in ["product", "concept", "opinion", "how_to", "comparison", "breaking_news", "prediction"]:
self.assertIn(qt, SOURCE_TIERS)
self.assertIn("tier1", SOURCE_TIERS[qt])
self.assertIn("tier2", SOURCE_TIERS[qt])
def test_truthsocial_not_in_any_tier(self):
for qt, tiers in SOURCE_TIERS.items():
self.assertNotIn("truthsocial", tiers["tier1"], f"truthsocial in tier1 for {qt}")
self.assertNotIn("truthsocial", tiers["tier2"], f"truthsocial in tier2 for {qt}")
if __name__ == "__main__":
unittest.main()