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 annotations1819from typing import TYPE_CHECKING2021import numpy as np22import pandas as pd2324from src.utils.logging import get_logger2526if TYPE_CHECKING:27    from src.embeddings.manager import EmbeddingManager28    from src.storage.database import DatabaseManager2930logger = get_logger("search.semantic")313233class SemanticSearch:3435    def __init__(self, db: DatabaseManager, embedding_manager: EmbeddingManager) -> None:36        self.db = db37        self.emb = embedding_manager3839    def search(40        self,41        query: str,42        top_k: int = 20,43        form_type: str | None = None,44        ticker: str | None = None,45        start_date: str | None = None,46        end_date: str | None = None,47        source_type: str = "chunk",48    ) -> pd.DataFrame:49        model_name = self.emb.model_name50        query_vector = self.emb.registry.encode([query], model_name=model_name, show_progress=False)51        query_vector = query_vector[0]5253        conditions: list[str] = ["em.source_type = ?"]54        params: list = [source_type]5556        if form_type:57            conditions.append("f.form_type = ?")58            params.append(form_type)59        if ticker:60            conditions.append("f.ticker = ?")61            params.append(ticker)62        if start_date:63            conditions.append("f.filing_date >= ?")64            params.append(start_date)65        if end_date:66            conditions.append("f.filing_date <= ?")67            params.append(end_date)6869        conditions.append("em.model_name = ?")70        params.append(model_name)7172        where_clause = " AND ".join(conditions)7374        if source_type == "chunk":75            sql = f"""76                SELECT77                    em.source_id,78                    c.chunk_text AS text,79                    c.accession_number,80                    f.ticker,81                    f.company_name,82                    f.filing_date,83                    f.form_type,84                    s.section_name85                FROM embeddings_metadata em86                JOIN filing_chunks c ON em.source_id = c.chunk_id87                JOIN filings f ON c.accession_number = f.accession_number88                LEFT JOIN filing_sections s ON c.section_id = s.section_id89                WHERE {where_clause}90            """91        else:92            sql = f"""93                SELECT94                    em.source_id,95                    sec.text,96                    sec.accession_number,97                    f.ticker,98                    f.company_name,99                    f.filing_date,100                    f.form_type,101                    sec.section_name102                FROM embeddings_metadata em103                JOIN filing_sections sec ON em.source_id = sec.section_id104                JOIN filings f ON sec.accession_number = f.accession_number105                WHERE {where_clause}106            """107108        candidates_df = self.db.fetch_df(sql, params)109110        if candidates_df.empty:111            logger.info("No candidates found for query")112            return pd.DataFrame()113114        source_ids = candidates_df["source_id"].tolist()115        vectors = self.emb.get_vectors(source_ids, model_name=model_name)116117        if vectors.size == 0:118            logger.warning("No embedding vectors found for candidates")119            return pd.DataFrame()120121        similarities = self._cosine_similarity(query_vector, vectors)122        candidates_df = candidates_df.copy()123        candidates_df["similarity_score"] = similarities124125        results = (126            candidates_df127            .sort_values("similarity_score", ascending=False)128            .head(top_k)129            .reset_index(drop=True)130        )131132        logger.info(133            "Search returned %d results (top score=%.4f)",134            len(results),135            results["similarity_score"].iloc[0] if not results.empty else 0.0,136        )137        return results138139    def search_similar(self, source_id: str, top_k: int = 10) -> pd.DataFrame:140        meta_sql = "SELECT source_type, model_name FROM embeddings_metadata WHERE source_id = ? LIMIT 1"141        meta_df = self.db.fetch_df(meta_sql, [source_id])142143        if meta_df.empty:144            logger.warning("Source ID '%s' not found in embeddings metadata", source_id)145            return pd.DataFrame()146147        source_type = meta_df["source_type"].iloc[0]148        model_name = meta_df["model_name"].iloc[0]149150        source_vector = self.emb.get_vectors([source_id], model_name=model_name)151        if source_vector.size == 0:152            return pd.DataFrame()153        source_vector = source_vector[0]154155        all_meta_sql = (156            "SELECT source_id FROM embeddings_metadata "157            "WHERE source_type = ? AND model_name = ? AND source_id != ?"158        )159        all_meta_df = self.db.fetch_df(all_meta_sql, [source_type, model_name, source_id])160161        if all_meta_df.empty:162            return pd.DataFrame()163164        candidate_ids = all_meta_df["source_id"].tolist()165        vectors = self.emb.get_vectors(candidate_ids, model_name=model_name)166167        if vectors.size == 0:168            return pd.DataFrame()169170        similarities = self._cosine_similarity(source_vector, vectors)171172        if source_type == "chunk":173            detail_sql = """174                SELECT c.chunk_id AS source_id, c.chunk_text AS text, c.accession_number,175                       f.ticker, f.company_name, f.filing_date, f.form_type176                FROM filing_chunks c177                JOIN filings f ON c.accession_number = f.accession_number178                WHERE c.chunk_id IN ({})179            """.format(", ".join(["?"] * len(candidate_ids)))180        else:181            detail_sql = """182                SELECT s.section_id AS source_id, s.text, s.accession_number,183                       f.ticker, f.company_name, f.filing_date, f.form_type184                FROM filing_sections s185                JOIN filings f ON s.accession_number = f.accession_number186                WHERE s.section_id IN ({})187            """.format(", ".join(["?"] * len(candidate_ids)))188189        detail_df = self.db.fetch_df(detail_sql, candidate_ids)190191        sim_df = pd.DataFrame({"source_id": candidate_ids, "similarity_score": similarities})192        results = detail_df.merge(sim_df, on="source_id")193        results = (194            results195            .sort_values("similarity_score", ascending=False)196            .head(top_k)197            .reset_index(drop=True)198        )199        return results200201    def search_companies(self, query: str, top_k: int = 5) -> pd.DataFrame:202        results = self.search(query, top_k=200)203204        if results.empty:205            return pd.DataFrame()206207        company_agg = (208            results209            .groupby(["ticker", "company_name"])210            .agg(211                avg_similarity=("similarity_score", "mean"),212                max_similarity=("similarity_score", "max"),213                match_count=("similarity_score", "count"),214            )215            .reset_index()216            .sort_values("avg_similarity", ascending=False)217            .head(top_k)218            .reset_index(drop=True)219        )220        return company_agg221222    @staticmethod223    def _cosine_similarity(query_vec: np.ndarray, matrix: np.ndarray) -> np.ndarray:224        query_norm = np.linalg.norm(query_vec)225        if query_norm == 0:226            return np.zeros(matrix.shape[0])227        matrix_norms = np.linalg.norm(matrix, axis=1)228        matrix_norms[matrix_norms == 0] = 1.0229        return np.dot(matrix, query_vec) / (matrix_norms * query_norm)230