+585

5pkgs/kagi-search/kagi_search.py

View file Reviewed

··· 1 1 + #!/usr/bin/env python3 2 2 + """ 3 3 + Kagi search script using session token authentication. 4 4 + Scrapes search results from Kagi web interface. 5 5 + """ 6 6 + 7 7 + import json 8 8 + import subprocess 9 9 + import sys 10 10 + import os 11 11 + from pathlib import Path 12 12 + from typing import List, Dict, Optional, Any 13 13 + from urllib.request import Request, HTTPCookieProcessor, build_opener 14 14 + from urllib.parse import urlencode 15 15 + from urllib.error import URLError, HTTPError 16 16 + from http.cookiejar import CookieJar 17 17 + import gzip 18 18 + from bs4 import BeautifulSoup, Tag 19 19 + import time 20 20 + from dataclasses import dataclass 21 21 + import argparse 22 22 + import logging 23 23 + 24 24 + # Module-level logger 25 25 + logger = logging.getLogger(__name__) 26 26 + 27 27 + 28 28 + def colorize(text: str, color: str = "", bold: bool = False, dim: bool = False) -> str: 29 29 + """ 30 30 + Colorize text if colors are enabled (TTY detected and NO_COLOR not set). 31 31 + 32 32 + Args: 33 33 + text: The text to colorize 34 34 + color: Color name (red, green, yellow, blue, magenta, cyan, white) 35 35 + bold: Whether to make text bold 36 36 + dim: Whether to make text dim/faint 37 37 + 38 38 + Returns: 39 39 + Colorized text or plain text based on environment 40 40 + """ 41 41 + # Check if we should use colors 42 42 + if not sys.stdout.isatty() or os.environ.get("NO_COLOR"): 43 43 + return text 44 44 + 45 45 + # Color codes 46 46 + colors = { 47 47 + "red": "\033[91m", 48 48 + "green": "\033[92m", 49 49 + "yellow": "\033[93m", 50 50 + "blue": "\033[94m", 51 51 + "magenta": "\033[95m", 52 52 + "cyan": "\033[96m", 53 53 + "white": "\033[97m", 54 54 + } 55 55 + 56 56 + # Start with empty escape sequence 57 57 + escape = "" 58 58 + 59 59 + # Add color if specified 60 60 + if color and color.lower() in colors: 61 61 + escape = colors[color.lower()] 62 62 + 63 63 + # Add bold if requested 64 64 + if bold: 65 65 + escape = "\033[1m" + escape 66 66 + 67 67 + # Add dim if requested 68 68 + if dim: 69 69 + escape = "\033[2m" + escape 70 70 + 71 71 + # Return plain text if no formatting requested 72 72 + if not escape: 73 73 + return text 74 74 + 75 75 + # Return formatted text 76 76 + return f"{escape}{text}\033[0m" 77 77 + 78 78 + 79 79 + def hyperlink(url: str, text: str = "") -> str: 80 80 + """ 81 81 + Create a terminal hyperlink if supported (TTY detected and NO_COLOR not set). 82 82 + 83 83 + Args: 84 84 + url: The URL to link to 85 85 + text: The text to display (defaults to URL if not provided) 86 86 + 87 87 + Returns: 88 88 + Hyperlinked text or plain text based on environment 89 89 + """ 90 90 + # Check if we should use hyperlinks (same conditions as colors) 91 91 + if not sys.stdout.isatty() or os.environ.get("NO_COLOR"): 92 92 + return text or url 93 93 + 94 94 + # Use the URL as text if no text provided 95 95 + display_text = text or url 96 96 + 97 97 + # OSC 8 hyperlink format: ESC]8;;URL ESC\TEXT ESC]8;; ESC\ 98 98 + return f"\033]8;;{url}\033\\{display_text}\033]8;;\033\\" 99 99 + 100 100 + 101 101 + @dataclass 102 102 + class SearchResult: 103 103 + """Represents a single search result from Kagi.""" 104 104 + 105 105 + title: str 106 106 + url: str 107 107 + snippet: str 108 108 + 109 109 + 110 110 + @dataclass 111 111 + class QuickAnswer: 112 112 + """Represents a Kagi Quick Answer response.""" 113 113 + 114 114 + html: str 115 115 + markdown: str 116 116 + raw_text: str 117 117 + references: List[Dict[str, Any]] 118 118 + 119 119 + 120 120 + class KagiSearch: 121 121 + """Kagi search client using session token authentication.""" 122 122 + 123 123 + def __init__(self, session_token: Optional[str] = None, config_path: Optional[str] = None): 124 124 + """ 125 125 + Initialize Kagi search client. 126 126 + 127 127 + Args: 128 128 + session_token: Kagi session token (if not provided, will load from config) 129 129 + config_path: Path to config file (defaults to ~/.config/kagi/config.json) 130 130 + """ 131 131 + self.base_url = "https://kagi.com" 132 132 + self.user_agent = "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36" 133 133 + 134 134 + # Load config and get session token 135 135 + if not session_token: 136 136 + config = self._load_config(config_path) 137 137 + session_token = self._get_session_token(config) 138 138 + 139 139 + self.session_token = session_token 140 140 + 141 141 + # Set up cookie handling 142 142 + self.cookie_jar = CookieJar() 143 143 + self.opener = build_opener(HTTPCookieProcessor(self.cookie_jar)) 144 144 + 145 145 + # Authenticate with token 146 146 + self._authenticate() 147 147 + 148 148 + def _load_config(self, config_path: Optional[str] = None) -> Dict[str, Any]: 149 149 + """Load configuration from file.""" 150 150 + if not config_path: 151 151 + config_file = Path.home() / ".config" / "kagi" / "config.json" 152 152 + else: 153 153 + config_file = Path(config_path) 154 154 + 155 155 + if not config_file.exists(): 156 156 + # Create default config 157 157 + config_file.parent.mkdir(parents=True, exist_ok=True) 158 158 + default_config = { 159 159 + "password_command": "rbw get kagi-session-link", 160 160 + "timeout": 30, 161 161 + "max_retries": 5, 162 162 + } 163 163 + with open(config_file, "w") as f: 164 164 + json.dump(default_config, f, indent=2) 165 165 + msg = colorize(f"Created default config at {config_file}", color="green") 166 166 + print(msg, file=sys.stderr) 167 167 + return default_config 168 168 + 169 169 + with open(config_file, "r") as f: 170 170 + data: Dict[str, Any] = json.load(f) 171 171 + return data 172 172 + 173 173 + def _get_session_token(self, config: Dict[str, Any]) -> str: 174 174 + """Get session token using password command from config.""" 175 175 + password_command = config.get("password_command", "rbw get kagi-session-link") 176 176 + 177 177 + try: 178 178 + # Execute the password command 179 179 + result = subprocess.run( 180 180 + password_command.split(), capture_output=True, text=True, check=True 181 181 + ) 182 182 + session_link = result.stdout.strip() 183 183 + 184 184 + # Extract token from session link 185 185 + if "token=" in session_link: 186 186 + return session_link.split("token=")[1].split("&")[0] 187 187 + else: 188 188 + # Assume the entire output is the token 189 189 + return session_link 190 190 + 191 191 + except subprocess.CalledProcessError as e: 192 192 + error_msg = colorize(f"Error executing password command: {e}", color="red", bold=True) 193 193 + print(error_msg, file=sys.stderr) 194 194 + if e.stderr: 195 195 + stderr_msg = colorize(f"stderr: {e.stderr}", color="red") 196 196 + print(stderr_msg, file=sys.stderr) 197 197 + sys.exit(1) 198 198 + except Exception as e: 199 199 + error_msg = colorize(f"Error getting session token: {e}", color="red", bold=True) 200 200 + print(error_msg, file=sys.stderr) 201 201 + sys.exit(1) 202 202 + 203 203 + def _authenticate(self) -> None: 204 204 + """Authenticate with Kagi using session token.""" 205 205 + token_url = f"{self.base_url}/html/search?token={self.session_token}" 206 206 + request = Request(token_url) 207 207 + request.add_header("User-Agent", self.user_agent) 208 208 + 209 209 + try: 210 210 + response = self.opener.open(request, timeout=30) 211 211 + final_url = response.geturl() 212 212 + 213 213 + # Check if we were redirected to home page or html search (successful auth) 214 214 + if final_url in [f"{self.base_url}/", self.base_url, f"{self.base_url}/html/search"]: 215 215 + return # Success 216 216 + elif "/signin" in final_url or "/welcome" in final_url: 217 217 + raise Exception(f"Authentication failed - redirected to {final_url}") 218 218 + except Exception as e: 219 219 + raise Exception(f"Failed to authenticate with token: {e}") 220 220 + 221 221 + def search(self, query: str, limit: int = 10) -> List[SearchResult]: 222 222 + """ 223 223 + Search Kagi and return results. 224 224 + 225 225 + Args: 226 226 + query: Search query 227 227 + limit: Maximum number of results to return 228 228 + 229 229 + Returns: 230 230 + List of SearchResult objects 231 231 + """ 232 232 + # Construct search URL for HTML version (no token needed, we use cookies) 233 233 + params = {"q": query} 234 234 + search_url = f"{self.base_url}/html/search?{urlencode(params)}" 235 235 + 236 236 + max_retries = 5 237 237 + retry_delay = 1.0 238 238 + 239 239 + for attempt in range(max_retries): 240 240 + try: 241 241 + # Create request with headers 242 242 + request = Request(search_url) 243 243 + request.add_header("User-Agent", self.user_agent) 244 244 + request.add_header( 245 245 + "Accept", 246 246 + "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,*/*;q=0.8", 247 247 + ) 248 248 + request.add_header("Accept-Language", "en-US,en;q=0.5") 249 249 + request.add_header("DNT", "1") 250 250 + request.add_header("Connection", "keep-alive") 251 251 + request.add_header("Upgrade-Insecure-Requests", "1") 252 252 + 253 253 + # Make request using opener (with cookies) 254 254 + response = self.opener.open(request, timeout=30) 255 255 + 256 256 + # Check if we're redirected to sign in 257 257 + final_url = response.geturl() 258 258 + if "/signin" in final_url or "/welcome" in final_url: 259 259 + raise Exception(f"Authentication failed - redirected to {final_url}") 260 260 + 261 261 + # Read response 262 262 + content = response.read() 263 263 + 264 264 + # Handle gzip encoding if present 265 265 + if response.headers.get("Content-Encoding") == "gzip": 266 266 + content = gzip.decompress(content) 267 267 + 268 268 + html_content = content.decode("utf-8") 269 269 + 270 270 + # Parse results 271 271 + soup = BeautifulSoup(html_content, "html.parser") 272 272 + 273 273 + # Find results container 274 274 + results_box = soup.find(class_="results-box") 275 275 + if not results_box or isinstance(results_box, str): 276 276 + if attempt < max_retries - 1: 277 277 + time.sleep(retry_delay) 278 278 + continue 279 279 + raise Exception("No results box found on page") 280 280 + 281 281 + # Extract search results 282 282 + results = [] 283 283 + if not isinstance(results_box, Tag): 284 284 + raise Exception("Results box is not a Tag element") 285 285 + search_results = results_box.find_all(class_="search-result") 286 286 + 287 287 + for result in search_results[:limit]: 288 288 + # Extract title 289 289 + title_elem = result.find(class_="__sri-title") 290 290 + title = "" 291 291 + if title_elem: 292 292 + # Get only the text content, not the UI elements 293 293 + # Remove "More results from this site", "Remove results from this site", etc. 294 294 + title_text = title_elem.get_text(separator=" ", strip=True) 295 295 + # Split by common Kagi UI elements and take the first part 296 296 + for separator in [ 297 297 + "More results from", 298 298 + "Remove results from", 299 299 + "Open page in", 300 300 + ]: 301 301 + if separator in title_text: 302 302 + title_text = title_text.split(separator)[0] 303 303 + title = title_text.strip() 304 304 + 305 305 + # Extract URL 306 306 + url_box = result.find(class_="__sri-url-box") 307 307 + url = "" 308 308 + if url_box: 309 309 + link = url_box.find("a", href=True) 310 310 + if link: 311 311 + href = link.get("href", "") 312 312 + url = str(href) if href else "" 313 313 + 314 314 + # Extract snippet 315 315 + desc_elem = result.find(class_="__sri-desc") 316 316 + snippet = desc_elem.get_text(strip=True) if desc_elem else "" 317 317 + 318 318 + if title and url: 319 319 + results.append(SearchResult(title=title, url=url, snippet=snippet)) 320 320 + 321 321 + if results: 322 322 + return results 323 323 + elif attempt < max_retries - 1: 324 324 + time.sleep(retry_delay) 325 325 + continue 326 326 + 327 327 + except (URLError, HTTPError) as e: 328 328 + if attempt < max_retries - 1: 329 329 + time.sleep(retry_delay) 330 330 + continue 331 331 + raise Exception(f"Request failed: {e}") 332 332 + 333 333 + return [] 334 334 + 335 335 + def _get_session_cookie(self) -> str: 336 336 + """Extract kagi_session cookie value for X-Kagi-Authorization header.""" 337 337 + for cookie in self.cookie_jar: 338 338 + if cookie.name == "kagi_session": 339 339 + return str(cookie.value) if cookie.value else "" 340 340 + return "" 341 341 + 342 342 + def get_quick_answer(self, query: str) -> Optional[QuickAnswer]: 343 343 + """ 344 344 + Get Kagi Quick Answer for a query. 345 345 + 346 346 + Args: 347 347 + query: Search query 348 348 + 349 349 + Returns: 350 350 + QuickAnswer object or None if no answer available 351 351 + """ 352 352 + # Construct Quick Answer URL - POST request with query in URL 353 353 + params = {"q": query} 354 354 + quick_answer_url = f"{self.base_url}/mother/context?{urlencode(params)}" 355 355 + logger.debug(f"Quick Answer URL: {quick_answer_url}") 356 356 + logger.info("Fetching Quick Answer...") 357 357 + 358 358 + try: 359 359 + # Get session cookie for authorization header 360 360 + session_cookie = self._get_session_cookie() 361 361 + 362 362 + # Create POST request with headers matching the working curl command 363 363 + request = Request(quick_answer_url, data=b"", method="POST") 364 364 + request.add_header("User-Agent", self.user_agent) 365 365 + request.add_header("Accept", "application/vnd.kagi.stream") 366 366 + request.add_header("Accept-Language", "en-US,en;q=0.5") 367 367 + request.add_header("Accept-Encoding", "gzip, deflate") 368 368 + request.add_header("Referer", f"{self.base_url}/search?{urlencode(params)}") 369 369 + request.add_header("Origin", self.base_url) 370 370 + request.add_header("Connection", "keep-alive") 371 371 + request.add_header("Content-Length", "0") 372 372 + if session_cookie: 373 373 + request.add_header("X-Kagi-Authorization", session_cookie) 374 374 + 375 375 + # Make request using opener (with cookies) 376 376 + logger.debug("Making Quick Answer POST request...") 377 377 + response = self.opener.open(request, timeout=30) 378 378 + logger.debug(f"Response status: {response.getcode()}") 379 379 + 380 380 + # Check if we're redirected to sign in 381 381 + final_url = response.geturl() 382 382 + logger.debug(f"Final URL: {final_url}") 383 383 + if "/signin" in final_url or "/welcome" in final_url: 384 384 + raise Exception(f"Authentication failed - redirected to {final_url}") 385 385 + 386 386 + # Read response 387 387 + content = response.read() 388 388 + logger.debug(f"Response content length: {len(content)} bytes") 389 389 + 390 390 + # Handle gzip encoding if present 391 391 + if response.headers.get("Content-Encoding") == "gzip": 392 392 + logger.debug("Decompressing gzip content") 393 393 + content = gzip.decompress(content) 394 394 + 395 395 + # Parse streaming response - lines with prefixes like "hi:", "tokens.json:", "new_message.json:" 396 396 + content_str = content.decode("utf-8") 397 397 + lines = content_str.strip().split("\n") 398 398 + logger.debug(f"Response has {len(lines)} lines") 399 399 + 400 400 + final_data = None 401 401 + 402 402 + for line in lines: 403 403 + line = line.strip() 404 404 + if not line: 405 405 + continue 406 406 + logger.debug(f"Processing line: {line[:100]}...") # Log first 100 chars 407 407 + 408 408 + # Look for the new_message.json line which contains the final answer 409 409 + if line.startswith("new_message.json:"): 410 410 + json_str = line[len("new_message.json:") :] 411 411 + logger.debug(f"Parsing new_message JSON: {json_str[:200]}...") 412 412 + try: 413 413 + # Use JSONDecoder to parse only the first JSON object 414 414 + # This handles cases where there's trailing data after the JSON 415 415 + decoder = json.JSONDecoder() 416 416 + final_data, _ = decoder.raw_decode(json_str) 417 417 + logger.debug("Successfully parsed new_message data") 418 418 + except json.JSONDecodeError as e: 419 419 + logger.error(f"Failed to parse new_message JSON: {e}") 420 420 + logger.debug(f"Full line: {line}") 421 421 + 422 422 + if not final_data: 423 423 + logger.debug("No new_message data found in response") 424 424 + return None 425 425 + 426 426 + # New response format has fields: md, reply, references_md, etc. 427 427 + markdown = final_data.get("md", "") 428 428 + html = final_data.get("reply", "") 429 429 + references_md = final_data.get("references_md", "") 430 430 + 431 431 + # Parse references from references_md (format: [^1]: [Title](URL) (percent%)) 432 432 + references = [] 433 433 + if references_md: 434 434 + import re 435 435 + 436 436 + # Match pattern like: [^1]: [Title](URL) (22%) 437 437 + # Use non-greedy match for URL with lookahead to handle URLs containing parentheses 438 438 + ref_pattern = r"\[\^\d+\]:\s*\[([^\]]+)\]\((.+?)\)\s*\((\d+)%\)" 439 439 + for match in re.finditer(ref_pattern, references_md): 440 440 + references.append( 441 441 + { 442 442 + "title": match.group(1), 443 443 + "url": match.group(2), 444 444 + "contribution": f"{match.group(3)}%", 445 445 + } 446 446 + ) 447 447 + 448 448 + # If no content, return None 449 449 + if not html and not markdown: 450 450 + logger.debug("No content found in Quick Answer") 451 451 + return None 452 452 + 453 453 + logger.debug(f"Quick Answer found with {len(references)} references") 454 454 + 455 455 + return QuickAnswer( 456 456 + html=html, markdown=markdown, raw_text=markdown, references=references 457 457 + ) 458 458 + 459 459 + except Exception as e: 460 460 + # Quick Answer might not be available for all queries 461 461 + logger.debug(f"Quick Answer error: {type(e).__name__}: {e}") 462 462 + return None 463 463 + 464 464 + 465 465 + def main() -> None: 466 466 + """Main entry point for command line usage.""" 467 467 + parser = argparse.ArgumentParser(description="Search Kagi using session token") 468 468 + parser.add_argument("query", nargs="?", help="Search query") 469 469 + parser.add_argument( 470 470 + "-n", "--num-results", type=int, default=3, help="Number of link results (default: 3)" 471 471 + ) 472 472 + parser.add_argument( 473 473 + "-l", "--links", action="store_true", help="Show search result links (hidden by default)" 474 474 + ) 475 475 + parser.add_argument("-t", "--token", help="Session token (overrides config)") 476 476 + parser.add_argument("-c", "--config", help="Config file path") 477 477 + parser.add_argument("-j", "--json", action="store_true", help="Output as JSON") 478 478 + parser.add_argument("-d", "--debug", action="store_true", help="Enable debug logging to stderr") 479 479 + 480 480 + args = parser.parse_args() 481 481 + 482 482 + # Configure logging - always to stderr 483 483 + log_level = logging.DEBUG if args.debug else logging.WARNING 484 484 + logging.basicConfig( 485 485 + level=log_level, 486 486 + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", 487 487 + stream=sys.stderr, 488 488 + ) 489 489 + 490 490 + # If no query provided, read from stdin 491 491 + if not args.query: 492 492 + if sys.stdin.isatty(): 493 493 + parser.error("No query provided") 494 494 + args.query = sys.stdin.read().strip() 495 495 + 496 496 + # Initialize client 497 497 + try: 498 498 + client = KagiSearch(session_token=args.token, config_path=args.config) 499 499 + except Exception as e: 500 500 + error_msg = colorize(f"Error initializing Kagi client: {e}", color="red", bold=True) 501 501 + print(error_msg, file=sys.stderr) 502 502 + sys.exit(1) 503 503 + 504 504 + # Perform search 505 505 + try: 506 506 + results = client.search(args.query, limit=args.num_results) if args.links else [] 507 507 + quick_answer = client.get_quick_answer(args.query) 508 508 + except Exception as e: 509 509 + error_msg = colorize(f"Search failed: {e}", color="red", bold=True) 510 510 + print(error_msg, file=sys.stderr) 511 511 + sys.exit(1) 512 512 + 513 513 + # Output results 514 514 + if args.json: 515 515 + output: Dict[str, Any] = { 516 516 + "results": [{"title": r.title, "url": r.url, "snippet": r.snippet} for r in results] 517 517 + } 518 518 + if quick_answer: 519 519 + output["quick_answer"] = { 520 520 + "markdown": quick_answer.markdown, 521 521 + "raw_text": quick_answer.raw_text, 522 522 + "references": quick_answer.references, 523 523 + } 524 524 + print(json.dumps(output, indent=2)) 525 525 + else: 526 526 + # Display Quick Answer if available 527 527 + if quick_answer: 528 528 + qa_title = colorize("Quick Answer", color="cyan", bold=True) 529 529 + print(f"\n{qa_title}") 530 530 + print(colorize("─" * 80, color="cyan", dim=True)) 531 531 + 532 532 + # Display raw text for terminal (cleaner than HTML) 533 533 + if quick_answer.raw_text: 534 534 + print(quick_answer.raw_text) 535 535 + elif quick_answer.markdown: 536 536 + # Fallback to markdown if no raw text 537 537 + print(quick_answer.markdown) 538 538 + 539 539 + # Display references 540 540 + if quick_answer.references: 541 541 + print() 542 542 + refs_title = colorize("References:", color="cyan", dim=True) 543 543 + print(refs_title) 544 544 + for i, ref in enumerate(quick_answer.references[:5], 1): # Limit to 5 refs 545 545 + ref_num = colorize(f"[{i}]", color="cyan", dim=True) 546 546 + ref_title = ref.get("title", "") 547 547 + ref_url = ref.get("url", "") 548 548 + if ref_title and ref_url: 549 549 + ref_link = hyperlink(ref_url, colorize(ref_title, color="blue")) 550 550 + print(f" {ref_num} {ref_link}") 551 551 + url_text = colorize(ref_url, color="green", dim=True) 552 552 + url_with_link = hyperlink(ref_url, url_text) 553 553 + print(f" {url_with_link}") 554 554 + 555 555 + print(colorize("─" * 80, color="cyan", dim=True)) 556 556 + 557 557 + # Display search result links only when --links is passed 558 558 + if results: 559 559 + if quick_answer: 560 560 + print() # Extra newline after Quick Answer 561 561 + for i, result in enumerate(results, 1): 562 562 + # Result number in yellow/bold 563 563 + number = colorize(f"{i}.", color="yellow", bold=True) 564 564 + # Title in blue/bold with hyperlink 565 565 + title = colorize(result.title, color="blue", bold=True) 566 566 + title_with_link = hyperlink(result.url, title) 567 567 + print(f"\n{number} {title_with_link}") 568 568 + 569 569 + # URL in green with hyperlink 570 570 + url_text = colorize(result.url, color="green") 571 571 + url_with_link = hyperlink(result.url, url_text) 572 572 + print(f" {url_with_link}") 573 573 + 574 574 + # Snippet in default color but dim 575 575 + if result.snippet: 576 576 + snippet = colorize(result.snippet, dim=True) 577 577 + print(f" {snippet}") 578 578 + 579 579 + if not quick_answer and not results and not args.json: 580 580 + error_msg = colorize("No results found", color="red") 581 581 + print(error_msg, file=sys.stderr) 582 582 + 583 583 + 584 584 + if __name__ == "__main__": 585 585 + main()

+31

5pkgs/kagi-search/package.nix

View file Reviewed

··· 1 1 + { 2 2 + pkgs, 3 3 + ... 4 4 + }: 5 5 + let 6 6 + python = pkgs.python3.withPackages ( 7 7 + ps: with ps; [ 8 8 + beautifulsoup4 9 9 + ] 10 10 + ); 11 11 + in 12 12 + pkgs.writeShellApplication { 13 13 + name = "kagi-search"; 14 14 + runtimeInputs = [ 15 15 + python 16 16 + pkgs.rbw 17 17 + ]; 18 18 + text = '' 19 19 + if ! rbw unlocked 2>/dev/null; then 20 20 + if [ -t 0 ]; then 21 21 + rbw unlock 22 22 + else 23 23 + saved_pinentry=$(rbw config show | ${pkgs.jq}/bin/jq -r '.pinentry') 24 24 + rbw config set pinentry ${pkgs.pinentry-rofi}/bin/pinentry-rofi 25 25 + trap 'rbw config set pinentry "$saved_pinentry"' EXIT 26 26 + rbw unlock 27 27 + fi 28 28 + fi 29 29 + ${python}/bin/python ${./kagi_search.py} "$@" 30 30 + ''; 31 31 + }

+6

tools/pi/flake-module.nix

View file Reviewed

··· 48 48 extensions = [ 49 49 ./questionnaire.ts 50 50 ./auto-theme.ts 51 51 + ./websearch.ts 51 52 ]; 52 53 defaultProvider = "anthropic"; 53 54 defaultModel = "claude-opus-4-6"; ··· 60 61 "nix fmt *" 61 62 ]; 62 63 }; 64 64 + 65 65 + extraPackages = [ 66 66 + self.packages.${pkgs.system}.kagi-search 67 67 + pkgs.curl 68 68 + ]; 63 69 64 70 pluginOverrides = '' 65 71 # Fix keybinding conflicts in extension source

+223

tools/pi/websearch.ts

View file Reviewed

··· 1 1 + /** 2 2 + * Web Search Tool - Search the web via Kagi and fetch page contents 3 3 + * 4 4 + * Provides two tools: 5 5 + * - websearch: Search the web via Kagi (uses kagi-search CLI with rbw for auth) 6 6 + * - webfetch: Fetch and extract text content from a URL 7 7 + */ 8 8 + 9 9 + import type { ExtensionAPI } from "@mariozechner/pi-coding-agent"; 10 10 + import { Text } from "@mariozechner/pi-tui"; 11 11 + import { Type } from "@sinclair/typebox"; 12 12 + 13 13 + interface SearchResult { 14 14 + title: string; 15 15 + url: string; 16 16 + snippet: string; 17 17 + } 18 18 + 19 19 + interface KagiOutput { 20 20 + results: SearchResult[]; 21 21 + quick_answer?: { 22 22 + markdown: string; 23 23 + raw_text: string; 24 24 + references: { title: string; url: string; contribution: string }[]; 25 25 + }; 26 26 + } 27 27 + 28 28 + export default function websearch(pi: ExtensionAPI) { 29 29 + pi.registerTool({ 30 30 + name: "websearch", 31 31 + label: "Web Search", 32 32 + description: 33 33 + "Search the web using Kagi. Returns a quick answer (AI summary) and search result links. Use this when you need to find current information online.", 34 34 + promptSnippet: "Search the web via Kagi for current information", 35 35 + parameters: Type.Object({ 36 36 + query: Type.String({ description: "Search query" }), 37 37 + num_results: Type.Optional( 38 38 + Type.Number({ 39 39 + description: "Number of link results to return (default: 5)", 40 40 + }), 41 41 + ), 42 42 + }), 43 43 + 44 44 + async execute(_toolCallId, params, signal) { 45 45 + const numResults = params.num_results ?? 5; 46 46 + const result = await pi.exec( 47 47 + "kagi-search", 48 48 + [ 49 49 + "--json", 50 50 + "--links", 51 51 + "--num-results", 52 52 + String(numResults), 53 53 + params.query, 54 54 + ], 55 55 + { signal, timeout: 30000 }, 56 56 + ); 57 57 + 58 58 + if (result.code !== 0) { 59 59 + throw new Error( 60 60 + `kagi-search failed (exit ${result.code}): ${result.stderr}`, 61 61 + ); 62 62 + } 63 63 + 64 64 + let output: KagiOutput; 65 65 + try { 66 66 + output = JSON.parse(result.stdout); 67 67 + } catch { 68 68 + throw new Error( 69 69 + `Failed to parse kagi-search output: ${result.stdout.slice(0, 200)}`, 70 70 + ); 71 71 + } 72 72 + 73 73 + const parts: string[] = []; 74 74 + 75 75 + if (output.quick_answer) { 76 76 + parts.push("## Quick Answer\n"); 77 77 + parts.push( 78 78 + output.quick_answer.markdown || output.quick_answer.raw_text, 79 79 + ); 80 80 + if (output.quick_answer.references?.length) { 81 81 + parts.push("\n\n### References"); 82 82 + for (const ref of output.quick_answer.references) { 83 83 + parts.push(`- [${ref.title}](${ref.url}) (${ref.contribution})`); 84 84 + } 85 85 + } 86 86 + } 87 87 + 88 88 + if (output.results?.length) { 89 89 + parts.push("\n\n## Search Results\n"); 90 90 + for (const [i, r] of output.results.entries()) { 91 91 + parts.push(`${i + 1}. ${r.title}\n ${r.url}\n ${r.snippet}`); 92 92 + } 93 93 + } 94 94 + 95 95 + if ( 96 96 + !output.quick_answer && 97 97 + (!output.results || output.results.length === 0) 98 98 + ) { 99 99 + return { 100 100 + content: [ 101 101 + { type: "text", text: "No results found for: " + params.query }, 102 102 + ], 103 103 + details: { query: params.query, results: [], quick_answer: null }, 104 104 + }; 105 105 + } 106 106 + 107 107 + return { 108 108 + content: [{ type: "text", text: parts.join("\n") }], 109 109 + details: { query: params.query, ...output }, 110 110 + }; 111 111 + }, 112 112 + 113 113 + renderCall(args, theme) { 114 114 + let text = theme.fg("toolTitle", theme.bold("websearch ")); 115 115 + text += theme.fg("accent", `"${args.query}"`); 116 116 + if (args.num_results) { 117 117 + text += theme.fg("dim", ` (${args.num_results} results)`); 118 118 + } 119 119 + return new Text(text, 0, 0); 120 120 + }, 121 121 + 122 122 + renderResult(result, _options, theme) { 123 123 + const details = result.details as 124 124 + | (KagiOutput & { query: string }) 125 125 + | undefined; 126 126 + if (!details) { 127 127 + const text = result.content[0]; 128 128 + return new Text(text?.type === "text" ? text.text : "", 0, 0); 129 129 + } 130 130 + 131 131 + const lines: string[] = []; 132 132 + 133 133 + if (details.quick_answer) { 134 134 + lines.push(theme.fg("accent", theme.bold("Quick Answer"))); 135 135 + const preview = 136 136 + details.quick_answer.markdown?.slice(0, 300) || 137 137 + details.quick_answer.raw_text?.slice(0, 300) || 138 138 + ""; 139 139 + lines.push( 140 140 + theme.fg("text", preview + (preview.length >= 300 ? "..." : "")), 141 141 + ); 142 142 + } 143 143 + 144 144 + if (details.results?.length) { 145 145 + if (lines.length) lines.push(""); 146 146 + for (const r of details.results) { 147 147 + lines.push( 148 148 + `${theme.fg("accent", r.title)}\n ${theme.fg("dim", r.url)}\n ${theme.fg("muted", r.snippet)}`, 149 149 + ); 150 150 + } 151 151 + } 152 152 + 153 153 + if (!lines.length) { 154 154 + return new Text(theme.fg("warning", "No results found"), 0, 0); 155 155 + } 156 156 + 157 157 + return new Text(lines.join("\n"), 0, 0); 158 158 + }, 159 159 + }); 160 160 + 161 161 + pi.registerTool({ 162 162 + name: "webfetch", 163 163 + label: "Web Fetch", 164 164 + description: 165 165 + "Fetch a web page and extract its text content. Use this to read the contents of a specific URL found via websearch or provided by the user.", 166 166 + promptSnippet: "Fetch and read the text content of a web page URL", 167 167 + parameters: Type.Object({ 168 168 + url: Type.String({ description: "URL to fetch" }), 169 169 + }), 170 170 + 171 171 + async execute(_toolCallId, params, signal) { 172 172 + const result = await pi.exec( 173 173 + "bash", 174 174 + [ 175 175 + "-c", 176 176 + `curl -sL --max-time 15 --max-filesize 2097152 -H 'User-Agent: Mozilla/5.0' ${JSON.stringify(params.url)} | sed -e 's/<script[^>]*>.*<\\/script>//g' -e 's/<style[^>]*>.*<\\/style>//g' -e 's/<[^>]*>//g' -e '/^[[:space:]]*$/d' | head -c 50000`, 177 177 + ], 178 178 + { signal, timeout: 20000 }, 179 179 + ); 180 180 + 181 181 + if (result.code !== 0) { 182 182 + throw new Error( 183 183 + `Failed to fetch URL (exit ${result.code}): ${result.stderr}`, 184 184 + ); 185 185 + } 186 186 + 187 187 + const content = result.stdout.trim(); 188 188 + if (!content) { 189 189 + throw new Error("No content extracted from URL: " + params.url); 190 190 + } 191 191 + 192 192 + return { 193 193 + content: [{ type: "text", text: content }], 194 194 + details: { url: params.url, length: content.length }, 195 195 + }; 196 196 + }, 197 197 + 198 198 + renderCall(args, theme) { 199 199 + let text = theme.fg("toolTitle", theme.bold("webfetch ")); 200 200 + text += theme.fg("accent", args.url as string); 201 201 + return new Text(text, 0, 0); 202 202 + }, 203 203 + 204 204 + renderResult(result, _options, theme) { 205 205 + const details = result.details as 206 206 + | { url: string; length: number } 207 207 + | undefined; 208 208 + if (!details) { 209 209 + const text = result.content[0]; 210 210 + return new Text(text?.type === "text" ? text.text : "", 0, 0); 211 211 + } 212 212 + const preview = 213 213 + result.content[0]?.type === "text" 214 214 + ? result.content[0].text.slice(0, 200) + "..." 215 215 + : ""; 216 216 + return new Text( 217 217 + `${theme.fg("success", "✓")} ${theme.fg("dim", details.url)} ${theme.fg("muted", `(${details.length} chars)`)}\n${theme.fg("muted", preview)}`, 218 218 + 0, 219 219 + 0, 220 220 + ); 221 221 + }, 222 222 + }); 223 223 + }