diff --git a/scripts/last30days.py b/scripts/last30days.py index b456009..3d6d1b3 100644 --- a/scripts/last30days.py +++ b/scripts/last30days.py @@ -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) diff --git a/scripts/lib/brave_search.py b/scripts/lib/brave_search.py index 33b64a5..7efd83d 100644 --- a/scripts/lib/brave_search.py +++ b/scripts/lib/brave_search.py @@ -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) diff --git a/scripts/lib/models.py b/scripts/lib/models.py index 2d5ec7f..dd25ad0 100644 --- a/scripts/lib/models.py +++ b/scripts/lib/models.py @@ -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 diff --git a/scripts/lib/openai_reddit.py b/scripts/lib/openai_reddit.py index f98e8e4..afe4b61 100644 --- a/scripts/lib/openai_reddit.py +++ b/scripts/lib/openai_reddit.py @@ -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): diff --git a/scripts/lib/polymarket.py b/scripts/lib/polymarket.py index 7ce87e0..27c93da 100644 --- a/scripts/lib/polymarket.py +++ b/scripts/lib/polymarket.py @@ -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: diff --git a/scripts/lib/query_type.py b/scripts/lib/query_type.py new file mode 100644 index 0000000..bc48c85 --- /dev/null +++ b/scripts/lib/query_type.py @@ -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"] diff --git a/scripts/lib/score.py b/scripts/lib/score.py index f52bb04..ffd108d 100644 --- a/scripts/lib/score.py +++ b/scripts/lib/score.py @@ -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", "") diff --git a/tests/test_brave_search.py b/tests/test_brave_search.py new file mode 100644 index 0000000..e8d9255 --- /dev/null +++ b/tests/test_brave_search.py @@ -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() diff --git a/tests/test_models.py b/tests/test_models.py index 493a4d7..36db570 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -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") diff --git a/tests/test_openai_reddit.py b/tests/test_openai_reddit.py index 607be97..74cc4dc 100644 --- a/tests/test_openai_reddit.py +++ b/tests/test_openai_reddit.py @@ -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__": diff --git a/tests/test_query_type.py b/tests/test_query_type.py new file mode 100644 index 0000000..bd36410 --- /dev/null +++ b/tests/test_query_type.py @@ -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()