1+ from __future__ import annotations
2+
3+ import argparse
4+ import os
5+ import sys
6+ from dataclasses import dataclass
7+ from pathlib import Path
8+ from typing import Any , Final
9+ from urllib .parse import urlparse
10+
11+ from dotenv import load_dotenv
12+ from upstash_search import Search
13+
14+
15+ DEFAULT_INDEX_NAME : Final [str ] = 'default'
16+ DEFAULT_LIMIT : Final [int ] = 10
17+
18+
19+ @dataclass (slots = True )
20+ class ArticleRelevancyRow :
21+ path : str
22+ title : str
23+ score : float
24+
25+
26+ def load_environment () -> None :
27+ repo_root = Path (__file__ ).resolve ().parents [1 ]
28+ env_path = repo_root / '.env.development'
29+ if not env_path .exists ():
30+ raise ValueError ('Missing .env.development. Create it locally before using this script.' )
31+
32+ load_dotenv (dotenv_path = env_path )
33+
34+
35+ def resolve_upstash_credentials () -> tuple [str , str , str ]:
36+ index_name = (os .environ .get ('UPSTASH_SEARCH_INDEX_NAME' ) or '' ).strip () or DEFAULT_INDEX_NAME
37+
38+ url = (
39+ (os .environ .get ('UPSTASH_SEARCH_REST_URL' ) or '' ).strip ()
40+ or (os .environ .get ('PUBLIC_UPSTASH_SEARCH_REST_URL' ) or '' ).strip ()
41+ )
42+
43+ token = (
44+ (os .environ .get ('PUBLIC_UPSTASH_SEARCH_READONLY_TOKEN' ) or '' ).strip ()
45+ or (os .environ .get ('UPSTASH_SEARCH_REST_TOKEN' ) or '' ).strip ()
46+ or (os .environ .get ('UPSTASH_SEARCH_WRITE_TOKEN' ) or '' ).strip ()
47+ or (os .environ .get ('UPSTASH_SEARCH_REST_WRITE_TOKEN' ) or '' ).strip ()
48+ )
49+
50+ if not url :
51+ raise ValueError ('Missing UPSTASH_SEARCH_REST_URL or PUBLIC_UPSTASH_SEARCH_REST_URL' )
52+ if not token :
53+ raise ValueError ('Missing a usable Upstash Search token' )
54+
55+ return url , token , index_name
56+
57+
58+ def get_value (source : Any , key : str ) -> Any :
59+ if isinstance (source , dict ):
60+ return source .get (key )
61+ return getattr (source , key , None )
62+
63+
64+ def normalize_path (value : Any ) -> str :
65+ if not isinstance (value , str ):
66+ return ''
67+
68+ trimmed = value .strip ()
69+ if not trimmed :
70+ return ''
71+
72+ if trimmed .startswith ('http://' ) or trimmed .startswith ('https://' ):
73+ parsed = urlparse (trimmed )
74+ path = parsed .path or ''
75+ suffix = ''
76+ if parsed .query :
77+ suffix = f'?{ parsed .query }
78+ if parsed .fragment :
79+ suffix = f'{ suffix } { parsed .fragment }
80+ return f'{ path } { suffix }
81+
82+ return trimmed
83+
84+
85+ def get_base_article_path (path : str ) -> str :
86+ without_query = path .split ('?' , 1 )[0 ]
87+ without_fragment = without_query .split ('#' , 1 )[0 ]
88+ return without_fragment .rstrip ('/' )
89+
90+
91+ def is_article_path (path : str ) -> bool :
92+ return path .startswith ('/articles/' ) and len (path .split ('/' )) >= 3
93+
94+
95+ def slug_to_title (slug : str ) -> str :
96+ words = [word for word in slug .replace ('_' , '-' ).split ('-' ) if word ]
97+ return ' ' .join (word .capitalize () for word in words )
98+
99+
100+ def get_article_title (result : Any , article_path : str ) -> str :
101+ content = get_value (result , 'content' ) or {}
102+ title = get_value (content , 'title' )
103+ if isinstance (title , str ) and title .strip ():
104+ result_path = normalize_path (
105+ get_value (get_value (result , 'metadata' ) or {}, 'path' )
106+ or get_value (get_value (result , 'metadata' ) or {}, 'url' )
107+ or get_value (content , 'path' )
108+ or get_value (content , 'url' )
109+ )
110+ if '#' not in result_path :
111+ return title .strip ()
112+
113+ slug = article_path .rsplit ('/' , 1 )[- 1 ]
114+ return slug_to_title (slug )
115+
116+
117+ def collect_article_relevancy_rows (results : list [Any ]) -> list [ArticleRelevancyRow ]:
118+ grouped : dict [str , ArticleRelevancyRow ] = {}
119+
120+ for result in results :
121+ metadata = get_value (result , 'metadata' ) or {}
122+ content = get_value (result , 'content' ) or {}
123+ path = normalize_path (
124+ get_value (metadata , 'path' )
125+ or get_value (metadata , 'url' )
126+ or get_value (content , 'path' )
127+ or get_value (content , 'url' )
128+ )
129+
130+ article_path = get_base_article_path (path )
131+ if not is_article_path (article_path ):
132+ continue
133+
134+ raw_score = get_value (result , 'score' )
135+ if not isinstance (raw_score , int | float ):
136+ continue
137+
138+ candidate = ArticleRelevancyRow (
139+ path = article_path ,
140+ title = get_article_title (result , article_path ),
141+ score = float (raw_score ),
142+ )
143+
144+ existing = grouped .get (article_path )
145+ if existing is None :
146+ grouped [article_path ] = candidate
147+ continue
148+
149+ if candidate .score > existing .score :
150+ title = existing .title
151+ if '#' not in path and candidate .title .strip ():
152+ title = candidate .title
153+ grouped [article_path ] = ArticleRelevancyRow (path = article_path , title = title , score = candidate .score )
154+ continue
155+
156+ if '#' not in path and candidate .title .strip () and existing .title == slug_to_title (article_path .rsplit ('/' , 1 )[- 1 ]):
157+ grouped [article_path ] = ArticleRelevancyRow (path = article_path , title = candidate .title , score = existing .score )
158+
159+ return sorted (grouped .values (), key = lambda row : row .score , reverse = True )
160+
161+
162+ def format_score (score : float ) -> str :
163+ return f'{ score :.6f} .rstrip ('0' ).rstrip ('.' )
164+
165+
166+ def format_results_table (rows : list [ArticleRelevancyRow ]) -> str :
167+ title_header = 'Article Title'
168+ score_header = 'Relevancy Score'
169+
170+ title_width = max ([len (title_header ), * (len (row .title ) for row in rows )])
171+ score_values = [format_score (row .score ) for row in rows ]
172+ score_width = max ([len (score_header ), * (len (value ) for value in score_values )])
173+
174+ separator = f'+-{ '-' * title_width } { '-' * score_width }
175+ header = f'| { title_header .ljust (title_width )} { score_header .rjust (score_width )}
176+
177+ body = [
178+ f'| { row .title .ljust (title_width )} { format_score (row .score ).rjust (score_width )}
179+ for row in rows
180+ ]
181+
182+ return '\n ' .join ([separator , header , separator , * body , separator ])
183+
184+
185+ def run_search (* , query : str , limit : int , index_name : str | None = None ) -> list [ArticleRelevancyRow ]:
186+ url , token , default_index_name = resolve_upstash_credentials ()
187+ client = Search (url = url , token = token )
188+ raw_results = client .index (index_name or default_index_name ).search (query , limit = limit , reranking = True )
189+ return collect_article_relevancy_rows (raw_results )
190+
191+
192+ def parse_args (argv : list [str ]) -> argparse .Namespace :
193+ parser = argparse .ArgumentParser (
194+ description = 'Query Upstash Search and print article relevancy scores in a table.'
195+ )
196+ parser .add_argument ('query' , help = 'Search query. Wrap multi-word queries in quotes.' )
197+ parser .add_argument ('--limit' , type = int , default = DEFAULT_LIMIT , help = 'Maximum Upstash results to request.' )
198+ parser .add_argument ('--index-name' , default = None , help = 'Override the Upstash index name.' )
199+ return parser .parse_args (argv )
200+
201+
202+ def main (argv : list [str ] | None = None ) -> int :
203+ args = parse_args (argv or sys .argv [1 :])
204+
205+ try :
206+ load_environment ()
207+ rows = run_search (query = args .query , limit = args .limit , index_name = args .index_name )
208+ except Exception as exc : # noqa: BLE001
209+ print (f'[search:relevancy] { exc } , file = sys .stderr )
210+ return 1
211+
212+ if not rows :
213+ print (f'No article results found for query: { args .query } )
214+ return 0
215+
216+ print (format_results_table (rows ))
217+ return 0
218+
219+
220+ if __name__ == '__main__' :
221+ raise SystemExit (main ())
0 commit comments