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 abc import ABC, abstractmethod20from datetime import datetime21from pathlib import Path22from typing import TYPE_CHECKING2324import pandas as pd2526from src.utils.constants import PROJECT_ROOT27from src.utils.logging import get_logger2829if TYPE_CHECKING:30    from src.embeddings.manager import EmbeddingManager31    from src.search.semantic import SemanticSearch32    from src.storage.database import DatabaseManager3334logger = get_logger("research.base")353637class ResearchModule(ABC):3839    def __init__(40        self,41        db: DatabaseManager,42        embedding_manager: EmbeddingManager | None = None,43        search: SemanticSearch | None = None,44    ) -> None:45        self.db = db46        self.emb = embedding_manager47        self.search = search4849    @property50    @abstractmethod51    def name(self) -> str: ...5253    @property54    @abstractmethod55    def description(self) -> str: ...5657    @abstractmethod58    def compute(59        self,60        tickers: list[str] | None = None,61        start_year: int | None = None,62        end_year: int | None = None,63    ) -> pd.DataFrame: ...6465    def save_results(66        self,67        results: pd.DataFrame,68        output_dir: Path | None = None,69    ) -> None:70        if output_dir is None:71            output_dir = PROJECT_ROOT / "outputs" / self.name7273        output_dir.mkdir(parents=True, exist_ok=True)74        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")7576        parquet_path = output_dir / f"{timestamp}.parquet"77        csv_path = output_dir / f"{timestamp}.csv"7879        results.to_parquet(parquet_path, index=False, compression="zstd")80        results.to_csv(csv_path, index=False)8182        logger.info(83            "Saved %d rows for module '%s' to %s",84            len(results), self.name, output_dir,85        )8687    def load_results(self, output_dir: Path | None = None) -> pd.DataFrame:88        if output_dir is None:89            output_dir = PROJECT_ROOT / "outputs" / self.name9091        parquet_files = sorted(output_dir.glob("*.parquet"))92        if not parquet_files:93            raise FileNotFoundError(94                f"No results found for module '{self.name}' in {output_dir}"95            )9697        latest = parquet_files[-1]98        logger.info("Loading results from %s", latest)99        return pd.read_parquet(latest)100