spb/uqo-research-db

1# ============================================================2# Project: EDGAR Research Database3# Description: SEC EDGAR Research Infrastructure4#5# Authors:6# - Manel Kammoun — manel.kammoun@uqo.ca7# - Djerrycharli Tandjambianda — djerrycharli.tandjambianda@uqo.ca8# - Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca9#10# Institution:11# Université du Québec en Outaouais (UQO)12#13# License:14# Research / Academic Use15# ============================================================1617from __future__ import annotations1819import argparse20import sys21from pathlib import Path2223sys.path.insert(0, str(Path(__file__).resolve().parent.parent))2425from src.embeddings import EmbeddingManager26from src.search import SemanticSearch27from src.storage.database import DatabaseManager28from src.storage.parquet_manager import ParquetManager29from src.utils.logging import setup_logging3031logger = setup_logging("search_filings")3233MAX_TEXT_DISPLAY = 200343536def parse_args() -> argparse.Namespace:37    parser = argparse.ArgumentParser(38        description="Semantic search across SEC filing content"39    )40    parser.add_argument(41        "--query",42        type=str,43        required=True,44        help="Search query text",45    )46    parser.add_argument(47        "--top-k",48        type=int,49        default=20,50        help="Number of results to return (default: 20)",51    )52    parser.add_argument(53        "--form-type",54        type=str,55        default=None,56        help="Filter by form type (e.g., 10-K)",57    )58    parser.add_argument(59        "--ticker",60        type=str,61        default=None,62        help="Filter by ticker symbol",63    )64    parser.add_argument(65        "--start-date",66        type=str,67        default=None,68        help="Filter by filing start date (YYYY-MM-DD)",69    )70    parser.add_argument(71        "--end-date",72        type=str,73        default=None,74        help="Filter by filing end date (YYYY-MM-DD)",75    )76    return parser.parse_args()777879def format_results(results) -> str:80    if results.empty:81        return "No results found."8283    lines: list[str] = []84    lines.append("")85    lines.append(f"{'Rank':<6} {'Score':<8} {'Ticker':<8} {'Form':<8} {'Date':<12} {'Text Preview'}")86    lines.append("-" * 100)8788    for i, row in results.iterrows():89        text = str(row.get("text", ""))[:MAX_TEXT_DISPLAY].replace("\n", " ")90        if len(str(row.get("text", ""))) > MAX_TEXT_DISPLAY:91            text += "..."9293        lines.append(94            f"{i + 1:<6} "95            f"{row.get('similarity_score', 0):<8.4f} "96            f"{str(row.get('ticker', '')):<8} "97            f"{str(row.get('form_type', '')):<8} "98            f"{str(row.get('filing_date', '')):<12} "99            f"{text}"100        )101102    lines.append("")103    lines.append(f"Total results: {len(results)}")104    return "\n".join(lines)105106107def main() -> None:108    args = parse_args()109110    logger.info("Query: '%s'", args.query)111    logger.info("Top-K: %d", args.top_k)112113    db = DatabaseManager()114    db.initialize()115    parquet = ParquetManager()116    emb = EmbeddingManager(db, parquet)117    search = SemanticSearch(db, emb)118119    try:120        results = search.search(121            query=args.query,122            top_k=args.top_k,123            form_type=args.form_type,124            ticker=args.ticker,125            start_date=args.start_date,126            end_date=args.end_date,127        )128    except NotImplementedError as e:129        logger.error("Search failed: %s", e)130        db.close()131        sys.exit(1)132    except Exception:133        logger.error("Search failed", exc_info=True)134        db.close()135        sys.exit(1)136137    print(format_results(results))138139    db.close()140141142if __name__ == "__main__":143    main()144