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,
xai_x,
youtube_yt,
query_type as qt,
)
@@ -631,8 +632,10 @@ def _search_web(
topic, from_date, to_date, config["PARALLEL_API_KEY"], depth=depth,
)
elif backend == "brave":
use_llm_ctx = os.environ.get("BRAVE_LLM_CONTEXT", "").strip() == "1"
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":
raw_results = openrouter_search.search_web(
@@ -1692,14 +1695,19 @@ def main():
else:
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
search_do_hackernews = True
search_do_bluesky = has_bluesky
search_do_truthsocial = has_truthsocial
search_do_polymarket = True
search_run_youtube = has_ytdlp
search_run_tiktok = has_tiktok
search_run_instagram = has_instagram
# Source defaults are query-type-aware (Truth Social always opt-in,
# Bluesky only for query types where it adds signal)
search_do_hackernews = qt.is_source_enabled("hn", query_type) if not args.search else True
search_do_bluesky = has_bluesky and qt.is_source_enabled("bluesky", query_type)
search_do_truthsocial = False # Always opt-in (requires --search truthsocial)
search_do_polymarket = qt.is_source_enabled("polymarket", query_type)
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
if 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_ts = score.score_truthsocial_items(filtered_ts) if filtered_ts 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
sorted_reddit = score.sort_items(scored_reddit)
sorted_x = score.sort_items(scored_x)
sorted_youtube = score.sort_items(scored_youtube) if scored_youtube else []
sorted_tiktok = score.sort_items(scored_tiktok) if scored_tiktok else []
sorted_ig = score.sort_items(scored_ig) if scored_ig else []
sorted_hn = score.sort_items(scored_hn) if scored_hn else []
sorted_bsky = score.sort_items(scored_bsky) if scored_bsky else []
sorted_ts = score.sort_items(scored_ts) if scored_ts else []
sorted_pm = score.sort_items(scored_pm) if scored_pm else []
sorted_web = score.sort_items(scored_web) if scored_web else []
# Sort items (query-type-aware tiebreaker ordering)
sorted_reddit = score.sort_items(scored_reddit, query_type=query_type)
sorted_x = score.sort_items(scored_x, query_type=query_type)
sorted_youtube = score.sort_items(scored_youtube, query_type=query_type) if scored_youtube else []
sorted_tiktok = score.sort_items(scored_tiktok, query_type=query_type) if scored_tiktok else []
sorted_ig = score.sort_items(scored_ig, query_type=query_type) if scored_ig else []
sorted_hn = score.sort_items(scored_hn, query_type=query_type) if scored_hn else []
sorted_bsky = score.sort_items(scored_bsky, query_type=query_type) if scored_bsky else []
sorted_ts = score.sort_items(scored_ts, query_type=query_type) if scored_ts else []
sorted_pm = score.sort_items(scored_pm, query_type=query_type) if scored_pm else []
sorted_web = score.sort_items(scored_web, query_type=query_type) if scored_web else []
# Dedupe items
deduped_reddit = dedupe.dedupe_reddit(sorted_reddit)
+119 -3
View File
@@ -1,7 +1,12 @@
"""Brave Search web search for last30days skill.
Uses the Brave Search API as a fallback web search backend.
Simple, cheap (free tier: 2,000 queries/month), widely available.
Uses the Brave Search API as a web search backend.
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
"""
@@ -16,6 +21,7 @@ from urllib.parse import urlencode, urlparse
from . import http
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_MAP = {1: "pd", 7: "pw", 31: "pm"}
@@ -33,6 +39,7 @@ def search_web(
to_date: str,
api_key: str,
depth: str = "default",
use_llm_context: bool = False,
) -> List[Dict[str, Any]]:
"""Search the web via Brave Search API.
@@ -42,6 +49,7 @@ def search_web(
to_date: End date (YYYY-MM-DD)
api_key: Brave Search API key
depth: 'quick', 'default', or 'deep'
use_llm_context: Use LLM Context endpoint for pre-extracted content
Returns:
List of result dicts with keys: url, title, snippet, source_domain, date, relevance
@@ -49,6 +57,9 @@ def search_web(
Raises:
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)
# Calculate days for freshness filter
@@ -81,6 +92,48 @@ def search_web(
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:
"""Calculate days between two YYYY-MM-DD dates."""
try:
@@ -138,7 +191,7 @@ def _normalize_results(
continue
if domain.startswith("www."):
domain = domain[4:]
except Exception:
except (ValueError, TypeError):
domain = ""
title = _clean_html(str(result.get("title", "")).strip())
@@ -169,6 +222,69 @@ def _normalize_results(
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:
"""Remove HTML tags and decode entities."""
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 sys
from typing import Dict, List, Optional, Tuple
from . import cache, http, env
# OpenAI API
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"]
# 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_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",
}
@@ -32,30 +47,45 @@ def parse_version(model_id: str) -> Optional[Tuple[int, ...]]:
return None
def is_mainline_openai_model(model_id: str) -> bool:
"""Check if model is a mainline GPT model (not mini/nano/chat/codex/pro)."""
def is_search_capable_model(model_id: str) -> bool:
"""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()
# Must be gpt-4o, gpt-4.1+, or gpt-5+ series (mainline, not mini/nano/etc)
if not re.match(r'^gpt-(?:4o|4\.1|5)(\.\d+)*$', model_lower):
# gpt-4o-mini does NOT support web_search with filters — exclude it
if model_lower.startswith("gpt-4o-mini"):
return False
# Exclude variants
excludes = ['mini', 'nano', 'chat', 'codex', 'pro', 'preview', 'turbo']
for exc in excludes:
# Must be gpt-4o, gpt-4.1[-mini], or gpt-5[-mini] series
if not re.match(r'^gpt-(?:4o|4\.1|5)(\.\d+)*(-mini)?$', model_lower):
return False
# Exclude unsupported variants
for exc in ['nano', 'chat', 'codex', 'pro', 'preview', 'turbo', 'search']:
if exc in model_lower:
return False
return True
# Backward compat alias
is_mainline_openai_model = is_search_capable_model
def select_openai_model(
api_key: str,
policy: str = "auto",
pin: Optional[str] = None,
mock_models: Optional[List[Dict]] = None,
) -> 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:
api_key: OpenAI API key
@@ -82,27 +112,29 @@ def select_openai_model(
headers = {"Authorization": f"Bearer {api_key}"}
response = http.get(OPENAI_MODELS_URL, headers=headers)
models = response.get("data", [])
except http.HTTPError:
# Fall back to known models
except http.HTTPError as e:
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]
# Filter to mainline models
candidates = [m for m in models if is_mainline_openai_model(m.get("id", ""))]
candidates = [m for m in models if is_search_capable_model(m.get("id", ""))]
if not candidates:
# No gpt-5 models found, use fallback
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):
version = parse_version(m.get("id", "")) or (0,)
created = m.get("created", 0)
return (version, created)
model_id = m.get("id", "")
version = parse_version(model_id) or (0,)
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)
selected = candidates[0]["id"]
# Cache the selection
cache.set_cached_model("openai", selected)
return selected
+4 -3
View File
@@ -7,9 +7,10 @@ from typing import Any, Dict, List, Optional
from . import http, env
# Fallback models when the selected model isn't accessible (e.g., org not verified for GPT-5)
# Note: gpt-4o-mini does NOT support web_search with filters param, so exclude it
MODEL_FALLBACK_ORDER = ["gpt-4.1", "gpt-4o"]
# Fallback models when the selected model isn't accessible (e.g., org not verified).
# Ordered by cost-efficiency: mini models handle structured extraction equally well.
# 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):
+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]:
"""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)}"
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 . import dates, schema
from .query_type import QueryType, WEBSEARCH_PENALTY_BY_TYPE, TIEBREAKER_BY_TYPE
# Score weights for Reddit/X (has engagement)
WEIGHT_RELEVANCE = 0.45
WEIGHT_RECENCY = 0.25
WEIGHT_ENGAGEMENT = 0.30
# WebSearch weights (no engagement, reweighted to 100%)
# WebSearch weights (no engagement data available)
WEBSEARCH_WEIGHT_RELEVANCE = 0.55
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_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
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.
Uses reweighted formula: 55% relevance + 45% recency - 15pt source penalty.
This ensures WebSearch items rank below comparable Reddit/X items.
Date confidence adjustments:
- High confidence (URL-verified date): +10 bonus
- Med confidence (snippet-extracted date): no change
- Low confidence (no date signals): -20 penalty
Uses reweighted formula: 55% relevance + 45% recency - penalty.
Penalty varies by query type: concept queries get 0 penalty (web docs
are authoritative), while product/opinion queries get full 15pt penalty
(social discussion is more valuable).
Args:
items: List of WebSearch items
query_type: Query classification for penalty adjustment
Returns:
Items with updated scores
@@ -682,8 +683,9 @@ def score_websearch_items(items: List[schema.WebSearchItem]) -> List[schema.WebS
WEBSEARCH_WEIGHT_RECENCY * rec_score
)
# Apply source penalty (WebSearch < Reddit/X for same relevance/recency)
overall -= WEBSEARCH_SOURCE_PENALTY
# Apply source penalty (varies by query type)
penalty = WEBSEARCH_PENALTY_BY_TYPE.get(query_type, WEBSEARCH_SOURCE_PENALTY) if query_type else WEBSEARCH_SOURCE_PENALTY
overall -= penalty
# Apply date confidence adjustments
# High confidence (URL-verified): reward with bonus
@@ -699,15 +701,36 @@ def score_websearch_items(items: List[schema.WebSearchItem]) -> List[schema.WebS
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:
"""Sort items by score (descending), then date, then source priority.
_ITEM_SOURCE_MAP = {
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:
items: List of items to sort
query_type: Query classification for tiebreaker adjustment
Returns:
Sorted items
"""
tiebreaker = TIEBREAKER_BY_TYPE.get(query_type, _DEFAULT_TIEBREAKER) if query_type else _DEFAULT_TIEBREAKER
def sort_key(item):
# Primary: score descending (negate for descending)
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_key = -int(date.replace("-", ""))
# Tertiary: source priority (Reddit > X > YouTube > TikTok > HN > Polymarket > WebSearch)
if isinstance(item, schema.RedditItem):
source_priority = 0
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
# Tertiary: query-type-aware source priority
source_name = _ITEM_SOURCE_MAP.get(type(item), "web")
source_priority = tiebreaker.get(source_name, 99)
# Quaternary: title/text for stability
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)
class TestIsMainlineOpenAIModel(unittest.TestCase):
def test_gpt5_is_mainline(self):
class TestIsSearchCapableModel(unittest.TestCase):
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"))
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):
def setUp(self):
from lib import cache
cache.MODEL_CACHE_FILE.unlink(missing_ok=True)
def test_pinned_policy(self):
result = models.select_openai_model(
"fake-key",
@@ -51,20 +77,8 @@ class TestSelectOpenAIModel(unittest.TestCase):
)
self.assertEqual(result, "gpt-5.1")
def test_auto_with_mock_models(self):
mock_models = [
{"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):
def test_prefers_mini_over_mainline(self):
"""Mini models should be preferred for cost-efficiency."""
mock_models = [
{"id": "gpt-5.2", "created": 1704067200},
{"id": "gpt-5-mini", "created": 1704067200},
@@ -75,8 +89,50 @@ class TestSelectOpenAIModel(unittest.TestCase):
policy="auto",
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")
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):
def test_latest_policy(self):
@@ -106,6 +162,10 @@ class TestSelectXAIModel(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):
config = {}
result = models.get_models(config)
@@ -114,9 +174,12 @@ class TestGetModels(unittest.TestCase):
def test_openai_key_only(self):
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)
self.assertEqual(result["openai"], "gpt-5.2")
self.assertEqual(result["openai"], "gpt-5-mini")
self.assertIsNone(result["xai"])
def test_both_keys(self):
@@ -124,10 +187,13 @@ class TestGetModels(unittest.TestCase):
"OPENAI_API_KEY": "sk-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}]
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")
+10 -5
View File
@@ -64,13 +64,18 @@ class TestIsModelAccessError(unittest.TestCase):
class TestModelFallbackOrder(unittest.TestCase):
"""Tests for MODEL_FALLBACK_ORDER constant."""
def test_contains_gpt4o(self):
"""Fallback list should include gpt-4o."""
def test_mini_first(self):
"""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)
def test_gpt41_is_first(self):
"""gpt-4.1 should be the first fallback option."""
self.assertEqual(MODEL_FALLBACK_ORDER[0], "gpt-4.1")
def test_no_gpt4o_mini(self):
"""gpt-4o-mini should NOT be in fallback (no domain filtering support)."""
self.assertNotIn("gpt-4o-mini", MODEL_FALLBACK_ORDER)
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()